[Bugfix] Preserve Mamba state for padded prompt tails (#55178)
Signed-off-by: Nathan Sala <natouda@gmail.com> Signed-off-by: Nathan Sala <33957625+natsala13@users.noreply.github.com>
Nathan Sala committed
Sep 5, 2026 at 02:44 UTC
874df9373dab532543a0229fb2f144f7c14093ae
2 files changed
+26
-4
vllm/v1/attention/backends/mamba2_attn.py
+1
@@ -141,6 +141,7 @@ class Mamba2AttentionMetadataBuilder(
141
common_attn_metadata,
142
num_accepted_tokens=kwargs.get("num_accepted_tokens"),
143
prev_last_scheduled_idx=kwargs.get("prev_last_scheduled_idx"),
144
+ num_decode_draft_tokens_cpu=kwargs.get("num_decode_draft_tokens_cpu"),
145
)
146
147
seq_idx_p = None
vllm/v1/attention/backends/mamba_attn.py
+25
-4
@@ -273,6 +273,7 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
273
*,
274
num_accepted_tokens: torch.Tensor | None = None,
275
prev_last_scheduled_idx: torch.Tensor | None = None,
276
+ num_decode_draft_tokens_cpu: torch.Tensor | None = None,
277
**kwargs: Any,
278
) -> M:
279
"""
@@ -283,6 +284,7 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
284
common_attn_metadata,
285
num_accepted_tokens=num_accepted_tokens,
286
prev_last_scheduled_idx=prev_last_scheduled_idx,
287
+ num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu,
288
)
289
290
def _compute_chunk_metadata(
@@ -459,6 +461,7 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
461
*,
462
num_accepted_tokens: torch.Tensor | None = None,
463
prev_last_scheduled_idx: torch.Tensor | None = None,
464
+ num_decode_draft_tokens_cpu: torch.Tensor | None = None,
465
) -> M:
466
"""
467
Compute metadata common to both Mamba1 and Mamba2.
@@ -484,10 +487,28 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
487
seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound
488
assert seq_lens_cpu is not None
489
query_lens_cpu = torch.diff(common_attn_metadata.query_start_loc_cpu)
487
- single_token_prefill_rows = is_prefilling & (query_lens_cpu == 1)
488
- # First-token prefills have no prior Mamba state and must stay prefills.
489
- has_prior_state = seq_lens_cpu > 1
490
- prefill_to_decode = single_token_prefill_rows & has_prior_state
490
+
491
+ # First prompt chunks have no prior Mamba state and must stay prefills.
492
+ has_prior_state = seq_lens_cpu > query_lens_cpu
493
+ stateful_prefill_rows = is_prefilling & has_prior_state
494
+
495
+ # One-token prefills with prior state can use the decode/update path.
496
+ prefill_to_decode = stateful_prefill_rows & (query_lens_cpu == 1)
497
+
498
+ # The scheduler may pad a one-token remote prompt tail with placeholder
499
+ # drafts to retain the uniform K+1 decode graph. This is a speculative
500
+ # decode transaction even though the real token is still in the prompt:
501
+ # the decode kernels keep h(N) in the running slot and h(N+i) in scratch
502
+ # slots, so normal acceptance rollback remains valid. The prefill kernels
503
+ # only return h(N+K) and cannot roll the placeholders back.
504
+ if num_decode_draft_tokens_cpu is not None:
505
+ padded_prompt_tail_rows = (
506
+ stateful_prefill_rows
507
+ & (num_decode_draft_tokens_cpu >= 0)
508
+ & (query_lens_cpu == num_decode_draft_tokens_cpu + 1)
509
+ )
510
+ prefill_to_decode |= padded_prompt_tail_rows
511
+
512
if torch.any(prefill_to_decode).item():
513
# ReplaySSM handles these rows as single-token flushes (see the
514
# write-position derivation below), same as the baseline decode path.