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 6, 02:00 UTC
Fix Mamba state corruption from referencing stale block table entries (#37728) (#37728) (#37728)
Ming Yang committed
Mar 24, 2026 at 10:29 UTC
c07e2ca6e0ace42aaeae96b5ce6b8525e592e2ae
3 files changed
+17
vllm/v1/worker/block_table.py
+10
@@ -121,6 +121,12 @@ class BlockTable:
121
self.num_blocks_per_row[row_idx] = 0
122
self.append_row(block_ids, row_idx)
123
124
+ def clear_row(self, row_idx: int) -> None:
125
+ num_blocks = self.num_blocks_per_row[row_idx]
126
+ if num_blocks > 0:
127
+ self.block_table.np[row_idx, :num_blocks] = 0
128
+ self.num_blocks_per_row[row_idx] = 0
129
+
130
def move_row(self, src: int, tgt: int) -> None:
131
num_blocks = self.num_blocks_per_row[src]
132
block_table_np = self.block_table.np
@@ -275,6 +281,10 @@ class MultiGroupBlockTable:
281
for i, block_table in enumerate(self.block_tables):
282
block_table.add_row(block_ids[i], row_idx)
283
284
+ def clear_row(self, row_idx: int) -> None:
285
+ for block_table in self.block_tables:
286
+ block_table.clear_row(row_idx)
287
+
288
def move_row(self, src: int, tgt: int) -> None:
289
for block_table in self.block_tables:
290
block_table.move_row(src, tgt)
vllm/v1/worker/gpu_input_batch.py
+1
@@ -496,6 +496,7 @@ class InputBatch:
496
self._req_ids[req_index] = None
497
self.req_output_token_ids[req_index] = None
498
self.spec_token_ids[req_index].clear()
499
+ self.block_table.clear_row(req_index)
500
501
# LoRA
502
lora_id = self.request_lora_mapping[req_index]
vllm/v1/worker/gpu_model_runner.py
+6
@@ -5376,6 +5376,12 @@ class GPUModelRunner(
5376
self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens
5377
self.query_start_loc.copy_to_gpu()
5378
5379
+ # Sync block table CPU->GPU so cleared rows from
5380
+ # remove_request() are visible to the attention metadata
5381
+ # builder. Without this, stale block IDs from finished
5382
+ # requests can corrupt Mamba state.
5383
+ self.input_batch.block_table.commit_block_table(num_reqs_padded)
5384
+
5385
pad_attn = cudagraph_runtime_mode == CUDAGraphMode.FULL
5386
attn_metadata, _ = self._build_attention_metadata(
5387
num_tokens=num_tokens_unpadded,