[BugFix] Stop dummy runs from writing mamba state through stale block-table rows (#49757)
Signed-off-by: Nick Hill <nickhill123@gmail.com> Signed-off-by: Jeff Ma <jeffjma@umich.edu> Co-authored-by: Jeff Ma <jeffjma@umich.edu>
Nick Hill committed
Jul 28, 2026 at 18:09 UTC
6fbbcf2151c34cd9313c223b8cbf19afa77f2ea6
5 files changed
+86
-2
tests/v1/worker/test_gpu_block_table.py
+63
@@ -130,3 +130,66 @@ def test_block_tables_apply_staged_writes_single_group():
130
block_tables.block_tables[0].gpu[0, :2],
131
torch.tensor([1, 2], dtype=torch.int32, device=device),
132
)
133
+
134
+
135
+def test_v1_block_table_move_row_clears_vacated_row():
136
+ """condense() moves the last row into a freed slot; the vacated row must
137
+ not keep stale block ids. Padded dummy-run batches dereference stale rows
138
+ as mamba state slots (bypassing the NULL_BLOCK_ID fill of real decode
139
+ padding) and write state in place there — corrupting the blocks' new
140
+ owner once they are reallocated, e.g. to an in-flight NIXL load."""
141
+ from vllm.v1.worker.block_table import BlockTable
142
+
143
+ block_table = BlockTable(
144
+ block_size=16,
145
+ max_num_reqs=4,
146
+ max_num_blocks_per_req=8,
147
+ max_num_batched_tokens=64,
148
+ pin_memory=False,
149
+ device=torch.device("cuda"),
150
+ kernel_block_size=16,
151
+ cp_kv_cache_interleave_size=1,
152
+ )
153
+ block_table.add_row([7, 8, 9], row_idx=0)
154
+ block_table.add_row([4, 5], row_idx=1)
155
+
156
+ block_table.move_row(1, 0)
157
+
158
+ assert block_table.block_table.np[0, :2].tolist() == [4, 5]
159
+ assert block_table.num_blocks_per_row[0] == 2
160
+ # The vacated source row routes to the reserved null block.
161
+ assert block_table.num_blocks_per_row[1] == 0
162
+ assert (block_table.block_table.np[1] == 0).all()
163
+
164
+
165
+def test_get_dummy_block_tables_returns_zeroed_rows():
166
+ """Dummy runs bypass the gather, so the persistent input_block_tables
167
+ hold the previous real step's rows. Mamba/GDN metadata routes in-place
168
+ state writes through block_table[:, 0] (dummy slot mappings are
169
+ PAD-filled, state indices are not), so stale rows would direct dummy
170
+ state writes at freed — possibly reallocated — blocks.
171
+ get_dummy_block_tables must hand out zeroed (null block) rows while
172
+ preserving the persistent storage address for CUDA graphs."""
173
+ device = torch.device("cuda")
174
+ block_tables = BlockTables(
175
+ block_sizes=[16],
176
+ max_num_reqs=4,
177
+ max_num_batched_tokens=64,
178
+ max_num_blocks_per_group=[8],
179
+ device=device,
180
+ kernel_block_sizes=[16],
181
+ )
182
+ # Simulate a real step: stage a request's blocks and gather them into
183
+ # the persistent input block tables.
184
+ block_tables.append_block_ids(req_index=0, new_block_ids=([1, 2],), overwrite=True)
185
+ block_tables.apply_staged_writes()
186
+ idx_mapping = torch.zeros(1, dtype=torch.int32, device=device)
187
+ block_tables.gather_block_tables(idx_mapping, num_reqs_padded=1)
188
+ torch.accelerator.synchronize()
189
+ assert block_tables.input_block_tables[0][0, 0].item() == 1
190
+
191
+ dummy = block_tables.get_dummy_block_tables(num_reqs=1)
192
+ torch.accelerator.synchronize()
193
+ assert (dummy[0] == 0).all()
194
+ # CUDA graph invariant: same persistent tensor, not a fresh allocation.
195
+ assert dummy[0].data_ptr() == block_tables.input_block_tables[0].data_ptr()
vllm/v1/worker/block_table.py
+5
@@ -144,6 +144,11 @@ class BlockTable:
144
block_table_np = self.block_table.np
145
block_table_np[tgt, :num_blocks] = block_table_np[src, :num_blocks]
146
self.num_blocks_per_row[tgt] = num_blocks
147
+ # Clear the vacated source row: dummy-run batches dereference stale
148
+ # rows as mamba state slots and write state in place there, possibly
149
+ # after the blocks have been freed and reallocated.
150
+ block_table_np[src, :num_blocks] = 0
151
+ self.num_blocks_per_row[src] = 0
152
153
def swap_row(self, src: int, tgt: int) -> None:
154
src_tgt, tgt_src = [src, tgt], [tgt, src]
vllm/v1/worker/gpu/block_table.py
+7
-1
@@ -163,7 +163,13 @@ class BlockTables:
163
# Therefore, this method must return the persistent tensor
164
# with the same memory address as that used during the model's forward pass,
165
# rather than allocating a new tensor.
166
- return tuple(block_table[:num_reqs] for block_table in self.input_block_tables)
166
+ #
167
+ # Zero the rows so dummy runs write mamba state to the reserved null
168
+ # block rather than through the previous real step's (stale) block
169
+ # ids, which may point at blocks since freed and reallocated.
170
+ return tuple(
171
+ block_table[:num_reqs].zero_() for block_table in self.input_block_tables
172
+ )
173
174
def compute_slot_mappings(
175
self,
vllm/v1/worker/gpu/model_runner.py
+4
@@ -1290,6 +1290,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1290
slot_mappings,
1291
self.attn_groups,
1292
self.kv_cache_config,
1293
+ # FULL replay reads capture-time metadata buffers. Re-stage them
1294
+ # from the zeroed dummy block tables instead of retaining state
1295
+ # indices from the previous real batch.
1296
+ for_capture=dummy_run and batch_desc.cg_mode == CUDAGraphMode.FULL,
1297
)
1298
1299
input_ids = input_batch.input_ids
vllm/v1/worker/gpu_model_runner.py
+7
-1
@@ -6064,7 +6064,13 @@ class GPUModelRunner(
6064
num_reqs=num_reqs_padded,
6065
max_query_len=max_query_len,
6066
ubatch_slices=(ubatch_slices_padded if pad_attn else ubatch_slices),
6067
- for_cudagraph_capture=is_graph_capturing,
6067
+ # FULL replay reads capture-time metadata buffers. Re-stage them
6068
+ # from the zeroed dummy block tables instead of retaining state
6069
+ # indices from the previous real batch.
6070
+ for_cudagraph_capture=(
6071
+ is_graph_capturing
6072
+ or cudagraph_runtime_mode == CUDAGraphMode.FULL
6073
+ ),
6074
slot_mappings=slot_mappings_by_group,
6075
use_spec_decode=self.speculative_config is not None,
6076
)