A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 3, 03:00 UTC
[Bugfix] Zero stale is_prefilling in padded CUDA graph rows for Mamba (#41873)
Signed-off-by: Lanze Liu <lanzetech@gmail.com>
Lanze Liu committed
May 21, 2026 at 15:42 UTC
39d5fa96a7c687f9ed7e14a5a52064965356cede
2 files changed
+68
tests/v1/attention/test_attention_splitting.py
+65
@@ -379,3 +379,68 @@ def test_prefill_split_across_ubatches(
379
# Map to original request index
380
orig_idx = split_req_idx + j
381
assert int(second_meta.seq_lens[j]) == seq_lens[orig_idx]
382
+
383
+
384
+def test_build_attention_metadata_zeros_stale_is_prefilling():
385
+ """_build_attention_metadata zeroes is_prefilling for padded rows."""
386
+ from unittest.mock import MagicMock, patch
387
+
388
+ from vllm.v1.attention.backend import CommonAttentionMetadata
389
+ from vllm.v1.worker.gpu_model_runner import GPUModelRunner
390
+
391
+ num_reqs = 3
392
+ num_reqs_padded = 5
393
+
394
+ # Real rows [0-2] have known computed/prompt values; padded rows [3-4]
395
+ # carry stale data from a prior prefill (num_computed < num_prompt → True).
396
+ num_computed = torch.tensor([50, 100, 200, 10, 20], dtype=torch.int32)
397
+ num_prompt = torch.tensor([50, 200, 200, 100, 200], dtype=torch.int32)
398
+
399
+ runner = MagicMock()
400
+ runner.kv_cache_config.kv_cache_groups = [
401
+ MagicMock()
402
+ ] # non-empty: skip early return
403
+ runner.attn_groups = [[]] # empty inner list: inner loop never runs
404
+ runner.input_batch.num_computed_tokens_cpu_tensor = num_computed
405
+ runner.input_batch.num_prompt_tokens_cpu_tensor = num_prompt
406
+ runner.optimistic_seq_lens_cpu = torch.tensor([100, 200, 300, 0, 0])
407
+ runner.query_start_loc.gpu = torch.zeros(num_reqs_padded + 1, dtype=torch.int32)
408
+ runner.query_start_loc.cpu = torch.zeros(num_reqs_padded + 1, dtype=torch.int32)
409
+ runner.seq_lens = torch.zeros(num_reqs_padded, dtype=torch.int32)
410
+ runner.positions = torch.zeros(num_reqs_padded, dtype=torch.int64)
411
+ runner.routed_experts_initialized = False
412
+ runner.use_async_spec_decode = False
413
+ runner.dcp_world_size = 1
414
+ runner.speculative_config = None
415
+ runner.is_mm_prefix_lm = False
416
+ runner._get_encoder_seq_lens.return_value = (None, None)
417
+
418
+ # Intercept CommonAttentionMetadata construction to capture is_prefilling.
419
+ # With speculative_config=None the constructor is called exactly once (for
420
+ # cm_base), so captured reflects what the fix produced before storage.
421
+ captured_is_prefilling = None
422
+ original_init = CommonAttentionMetadata.__init__
423
+
424
+ def capturing_init(self, *args, **kwargs):
425
+ nonlocal captured_is_prefilling
426
+ if "is_prefilling" in kwargs:
427
+ captured_is_prefilling = kwargs["is_prefilling"]
428
+ original_init(self, *args, **kwargs)
429
+
430
+ with patch.object(CommonAttentionMetadata, "__init__", capturing_init):
431
+ GPUModelRunner._build_attention_metadata(
432
+ runner,
433
+ num_tokens=num_reqs,
434
+ num_reqs=num_reqs,
435
+ max_query_len=1,
436
+ num_tokens_padded=num_reqs_padded,
437
+ num_reqs_padded=num_reqs_padded,
438
+ slot_mappings={0: torch.zeros(num_reqs_padded, dtype=torch.int64)},
439
+ )
440
+
441
+ assert captured_is_prefilling is not None
442
+ assert not captured_is_prefilling[0] # decode (50 >= 50)
443
+ assert captured_is_prefilling[1] # prefill (100 < 200)
444
+ assert not captured_is_prefilling[2] # decode (200 >= 200)
445
+ assert not captured_is_prefilling[3] # stale data (10 < 100) zeroed
446
+ assert not captured_is_prefilling[4] # stale data (20 < 200) zeroed
vllm/v1/worker/gpu_model_runner.py
+3
@@ -2268,6 +2268,9 @@ class GPUModelRunner(
2268
# Used by mamba backends to distinguish actual decodes from
2269
# short extends.
2270
is_prefilling = num_computed_tokens_cpu < num_prompt_tokens_cpu
2271
+ # Zero out padded rows so stale data from condense() doesn't
2272
+ # misclassify padding as prefill in CUDA graph mode.
2273
+ is_prefilling[num_reqs:] = False
2274
2275
if self.use_async_spec_decode:
2276
# GPU tensors are authoritative in async mode.