[Bugfix] Fix double-counted max_q_seqlen in decode delta kv_seqlens

create_model_inputs_delta / create_model_inputs_delta_valid_only build
kv_seqlens as [seq.num_all_ids + max_q_seqlen]. num_all_ids can be one decode
step stale here -- EngineLoop prefetches the next inputs before
_finish_forward_output() advances the sequence -- so the +max_q_seqlen recovers
this forward's kv length. But the reductions then added max_q_seqlen a SECOND
time and used batch_size = len(self.running_seqs), which counts
scheduler-dropped invalid seqs:

    sum_kv_seqlen = sum(kv_seqlens) + batch_size * max_q_seqlen
    max_kv_seqlen = max(kv_seqlens) + max_q_seqlen

so max_kv_seqlen / sum_kv_seqlen were over-counted (max by max_q_seqlen,
scaling with spec/MTP num_decode_tokens), over-allocating the attention grid +
kv-cache resources.

Reduce over kv_seqlens directly; the +max_q_seqlen is already applied once in
the comprehension.

Fixes #4024

Co-authored-by: Claude <noreply@anthropic.com>

Author: waynehacking8

lmdeploy/pytorch/engine/inputs_maker.py

@@ -603,9 +603,15 @@ def create_model_inputs_delta(self):
         else:
             num_ignored_history = torch.zeros(len(valid_seqs), dtype=torch.long)
 
+        # num_all_ids can be one decode step stale here: EngineLoop prefetches
+        # the next inputs before _finish_forward_output() advances the sequence,
+        # so +max_q_seqlen recovers this forward's kv length. The bug was adding
+        # max_q_seqlen AGAIN in the reductions, plus using batch_size (which
+        # counts scheduler-dropped invalid seqs) instead of reducing over the
+        # valid seqs only (#4024).
         kv_seqlens = [seq.num_all_ids + max_q_seqlen for seq in valid_seqs]
-        sum_kv_seqlen = sum(kv_seqlens) + batch_size * max_q_seqlen
-        max_kv_seqlen = max(kv_seqlens) + max_q_seqlen
+        sum_kv_seqlen = sum(kv_seqlens)
+        max_kv_seqlen = max(kv_seqlens)
 
         output = ModelInputsDelta(
             indices=None,
@@ -650,13 +656,15 @@ def create_model_inputs_delta_valid_only(self):
 
         num_decode_tokens = self.engine_strategy.get_num_decode_tokens()
         max_q_seqlen = num_decode_tokens
+        # Keep +max_q_seqlen (num_all_ids may be one decode step stale), but do
+        # not add it a second time in the reductions or use batch_size (#4024).
         kv_seqlens = [seq.num_all_ids + max_q_seqlen for seq in valid_seqs]
         if len(kv_seqlens) == 0:
             sum_kv_seqlen = 0
             max_kv_seqlen = 0
         else:
-            sum_kv_seqlen = sum(kv_seqlens) + batch_size * max_q_seqlen
-            max_kv_seqlen = max(kv_seqlens) + max_q_seqlen
+            sum_kv_seqlen = sum(kv_seqlens)
+            max_kv_seqlen = max(kv_seqlens)
 
         output = ModelInputsDelta(
             indices=None,

tests/pytorch/engine/test_inputs_maker.py

@@ -333,3 +333,26 @@ def test_state_prefix_cache_save_offsets_are_compact():
 
     assert src_offsets == (5, 6)
     assert dst_offsets == (21, 22)
+
+
+def test_create_model_inputs_delta_valid_only_kv_seqlen_no_double_count():
+    # Regression for #4024: kv_seqlens = num_all_ids + max_q_seqlen is correct
+    # (num_all_ids can be one decode step stale due to EngineLoop prefetch), but
+    # the old code added max_q_seqlen a SECOND time in the reductions and used
+    # batch_size (which counts scheduler-dropped invalid seqs), over-inflating
+    # max_kv_seqlen / sum_kv_seqlen.
+    maker = InputsMakerAsync.__new__(InputsMakerAsync)
+    maker.engine_strategy = SimpleNamespace(get_num_decode_tokens=lambda: 4)
+    maker.running_seqs = [
+        SimpleNamespace(status=MessageStatus.RUNNING, num_all_ids=100),
+        SimpleNamespace(status=MessageStatus.RUNNING, num_all_ids=250),
+        SimpleNamespace(status=MessageStatus.STOPPED, num_all_ids=70),  # dropped
+    ]
+
+    output, valid_seqs, invalid_seqs = maker.create_model_inputs_delta_valid_only()
+
+    assert [seq.num_all_ids for seq in valid_seqs] == [100, 250]
+    assert len(invalid_seqs) == 1
+    # kv_seqlens = [104, 254]; reduce over the valid seqs, add max_q_seqlen once
+    assert output.max_kv_seqlen == 254  # old (buggy): 254 + 4 = 258
+    assert output.sum_kv_seqlen == 358  # old (buggy): 358 + 3 * 4 = 370