[BugFix] Stop dummy runs from writing mamba state through stale block-table rows (#49757)

Signed-off-by: Nick Hill <[email protected]>
Signed-off-by: Jeff Ma <[email protected]>
Co-authored-by: Jeff Ma <[email protected]>

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()

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]

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,

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

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,
                 )