From 6fbbcf2151c34cd9313c223b8cbf19afa77f2ea6 Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Tue, 28 Jul 2026 18:09:41 -0700 Subject: [PATCH] [BugFix] Stop dummy runs from writing mamba state through stale block-table rows (#49757) Signed-off-by: Nick Hill Signed-off-by: Jeff Ma Co-authored-by: Jeff Ma --- tests/v1/worker/test_gpu_block_table.py | 63 +++++++++++++++++++++++++ vllm/v1/worker/block_table.py | 5 ++ vllm/v1/worker/gpu/block_table.py | 8 +++- vllm/v1/worker/gpu/model_runner.py | 4 ++ vllm/v1/worker/gpu_model_runner.py | 8 +++- 5 files changed, 86 insertions(+), 2 deletions(-) diff --git a/tests/v1/worker/test_gpu_block_table.py b/tests/v1/worker/test_gpu_block_table.py index 31acd475ad..365e365a57 100644 --- a/tests/v1/worker/test_gpu_block_table.py +++ b/tests/v1/worker/test_gpu_block_table.py @@ -130,3 +130,66 @@ def test_block_tables_apply_staged_writes_single_group(): block_tables.block_tables[0].gpu[0, :2], torch.tensor([1, 2], dtype=torch.int32, device=device), ) + + +def test_v1_block_table_move_row_clears_vacated_row(): + """condense() moves the last row into a freed slot; the vacated row must + not keep stale block ids. Padded dummy-run batches dereference stale rows + as mamba state slots (bypassing the NULL_BLOCK_ID fill of real decode + padding) and write state in place there — corrupting the blocks' new + owner once they are reallocated, e.g. to an in-flight NIXL load.""" + from vllm.v1.worker.block_table import BlockTable + + block_table = BlockTable( + block_size=16, + max_num_reqs=4, + max_num_blocks_per_req=8, + max_num_batched_tokens=64, + pin_memory=False, + device=torch.device("cuda"), + kernel_block_size=16, + cp_kv_cache_interleave_size=1, + ) + block_table.add_row([7, 8, 9], row_idx=0) + block_table.add_row([4, 5], row_idx=1) + + block_table.move_row(1, 0) + + assert block_table.block_table.np[0, :2].tolist() == [4, 5] + assert block_table.num_blocks_per_row[0] == 2 + # The vacated source row routes to the reserved null block. + assert block_table.num_blocks_per_row[1] == 0 + assert (block_table.block_table.np[1] == 0).all() + + +def test_get_dummy_block_tables_returns_zeroed_rows(): + """Dummy runs bypass the gather, so the persistent input_block_tables + hold the previous real step's rows. Mamba/GDN metadata routes in-place + state writes through block_table[:, 0] (dummy slot mappings are + PAD-filled, state indices are not), so stale rows would direct dummy + state writes at freed — possibly reallocated — blocks. + get_dummy_block_tables must hand out zeroed (null block) rows while + preserving the persistent storage address for CUDA graphs.""" + device = torch.device("cuda") + block_tables = BlockTables( + block_sizes=[16], + max_num_reqs=4, + max_num_batched_tokens=64, + max_num_blocks_per_group=[8], + device=device, + kernel_block_sizes=[16], + ) + # Simulate a real step: stage a request's blocks and gather them into + # the persistent input block tables. + block_tables.append_block_ids(req_index=0, new_block_ids=([1, 2],), overwrite=True) + block_tables.apply_staged_writes() + idx_mapping = torch.zeros(1, dtype=torch.int32, device=device) + block_tables.gather_block_tables(idx_mapping, num_reqs_padded=1) + torch.accelerator.synchronize() + assert block_tables.input_block_tables[0][0, 0].item() == 1 + + dummy = block_tables.get_dummy_block_tables(num_reqs=1) + torch.accelerator.synchronize() + assert (dummy[0] == 0).all() + # CUDA graph invariant: same persistent tensor, not a fresh allocation. + assert dummy[0].data_ptr() == block_tables.input_block_tables[0].data_ptr() diff --git a/vllm/v1/worker/block_table.py b/vllm/v1/worker/block_table.py index 332eda4cbd..4d9f256ee0 100644 --- a/vllm/v1/worker/block_table.py +++ b/vllm/v1/worker/block_table.py @@ -144,6 +144,11 @@ class BlockTable: block_table_np = self.block_table.np block_table_np[tgt, :num_blocks] = block_table_np[src, :num_blocks] self.num_blocks_per_row[tgt] = num_blocks + # Clear the vacated source row: dummy-run batches dereference stale + # rows as mamba state slots and write state in place there, possibly + # after the blocks have been freed and reallocated. + block_table_np[src, :num_blocks] = 0 + self.num_blocks_per_row[src] = 0 def swap_row(self, src: int, tgt: int) -> None: src_tgt, tgt_src = [src, tgt], [tgt, src] diff --git a/vllm/v1/worker/gpu/block_table.py b/vllm/v1/worker/gpu/block_table.py index df5d2be629..5d17868b17 100644 --- a/vllm/v1/worker/gpu/block_table.py +++ b/vllm/v1/worker/gpu/block_table.py @@ -163,7 +163,13 @@ class BlockTables: # Therefore, this method must return the persistent tensor # with the same memory address as that used during the model's forward pass, # rather than allocating a new tensor. - return tuple(block_table[:num_reqs] for block_table in self.input_block_tables) + # + # Zero the rows so dummy runs write mamba state to the reserved null + # block rather than through the previous real step's (stale) block + # ids, which may point at blocks since freed and reallocated. + return tuple( + block_table[:num_reqs].zero_() for block_table in self.input_block_tables + ) def compute_slot_mappings( self, diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 9182ada439..b16e75bf45 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1290,6 +1290,10 @@ class GPUModelRunner(LoRAModelRunnerMixin): slot_mappings, self.attn_groups, self.kv_cache_config, + # FULL replay reads capture-time metadata buffers. Re-stage them + # from the zeroed dummy block tables instead of retaining state + # indices from the previous real batch. + for_capture=dummy_run and batch_desc.cg_mode == CUDAGraphMode.FULL, ) input_ids = input_batch.input_ids diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 0859b6dacd..3019b7750a 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -6064,7 +6064,13 @@ class GPUModelRunner( num_reqs=num_reqs_padded, max_query_len=max_query_len, ubatch_slices=(ubatch_slices_padded if pad_attn else ubatch_slices), - for_cudagraph_capture=is_graph_capturing, + # FULL replay reads capture-time metadata buffers. Re-stage them + # from the zeroed dummy block tables instead of retaining state + # indices from the previous real batch. + for_cudagraph_capture=( + is_graph_capturing + or cudagraph_runtime_mode == CUDAGraphMode.FULL + ), slot_mappings=slot_mappings_by_group, use_spec_decode=self.speculative_config is not None, )