[Bugfix][Spec Decode] Clear stale shared block-table rows (#61174)

Signed-off-by: Cheng Rui <286040359@qq.com>
Co-authored-by: Codex <noreply@openai.com>
Co-authored-by: Giancarlo Delfin <32987265+TheEpicDolphin@users.noreply.github.com>

tests/v1/worker/test_gpu_block_table.py

@@ -305,3 +305,104 @@ def test_dummy_request_slot_mapping_is_pad():
         num_tokens_padded=3,
     )
     assert dummy[0].tolist() == [PAD_SLOT_ID] * 3
+
+
+@pytest.fixture
+def populated_block_tables():
+    tables = BlockTables([16, 32], 8, 64, [3, 5], torch.device("cuda"))
+    for req_index in range(8):
+        tables.append_block_ids(
+            req_index, ([req_index + 1] * 3, [req_index + 1] * 5), overwrite=True
+        )
+    tables.apply_staged_writes()
+    tables.gather_block_tables(torch.arange(8, device="cuda"), num_reqs_padded=8)
+    return tables
+
+
+@pytest.mark.parametrize("previous_dummy", [False, True])
+@pytest.mark.parametrize("next_dummy", [False, True])
+def test_shared_block_tables_clear_rows_from_larger_batch(
+    populated_block_tables, previous_dummy, next_dummy
+):
+    """A smaller target batch must leave no stale rows for larger draft graphs."""
+    tables = populated_block_tables
+    if previous_dummy:
+        for table in tables.get_dummy_block_tables(8):
+            table.fill_(99)  # Profiling supplies valid context/state slots.
+    ptrs = [table.data_ptr() for table in tables.input_block_tables]
+    idx = torch.tensor([7, 2, 0, 5, 1], device="cuda")
+    for num_reqs in (3, 5, 0):
+        if next_dummy:
+            out = tables.get_dummy_block_tables(num_reqs)
+            num_rows = num_reqs
+        else:
+            num_rows = num_reqs + 1
+            out = tables.gather_block_tables(idx[:num_reqs], num_rows)
+        for table, view, ptr in zip(tables.input_block_tables, out, ptrs):
+            assert view.shape == (num_rows, table.shape[1])
+            assert table.data_ptr() == ptr
+            if num_rows:
+                assert view.data_ptr() == ptr
+            assert torch.count_nonzero(table[num_reqs:]).item() == 0
+            if next_dummy:
+                assert torch.count_nonzero(view).item() == 0
+                view.fill_(99)
+            else:
+                expected = (idx[:num_reqs, None] + 1).expand(-1, table.shape[1])
+                torch.testing.assert_close(view[:num_reqs], expected.to(torch.int32))
+
+
+def test_external_block_table_output_preserves_shared_rows(populated_block_tables):
+    """PCP-local gathers must neither clear shared rows nor reset their tracking."""
+    tables = populated_block_tables
+    previous = [table.clone() for table in tables.input_block_tables]
+    out = tuple(torch.full_like(table[:3], -1) for table in previous)
+    out_ptrs = torch.tensor(
+        [table.data_ptr() for table in out], dtype=torch.uint64, device="cuda"
+    )
+    idx = torch.tensor([7, 2], device="cuda")
+    result = tables.gather_block_tables(idx, 3, out=out, out_ptrs=out_ptrs)
+    for table, before, view in zip(tables.input_block_tables, previous, result):
+        torch.testing.assert_close(table, before)
+        assert view[:, 0].tolist() == [8, 3, 0]
+    tables.gather_block_tables(idx[:1], num_reqs_padded=1)
+    for table in tables.input_block_tables:
+        assert torch.count_nonzero(table[1:]).item() == 0
+
+
+@pytest.mark.parametrize("use_cuda_graph", [False, True])
+def test_larger_draft_batch_uses_null_kpool_slots(
+    populated_block_tables, use_cuda_graph
+):
+    """Draft padding must not resolve to another request's physical KPool slots."""
+    from vllm.v1.attention.backends.mla.indexer import compute_kpool_tail_slot_mapping
+
+    tables = populated_block_tables
+    idx = torch.tensor([2, 0, 4, 1, 3], device="cuda")
+    query_start_loc = torch.arange(9, dtype=torch.int32, device="cuda").clamp_max_(5)
+    slot_mapping = torch.full((8,), -1, dtype=torch.int64, device="cuda")
+    positions = torch.arange(8, dtype=torch.int64, device="cuda")
+
+    def compute_slots():
+        return compute_kpool_tail_slot_mapping(
+            slot_mapping,
+            tables.input_block_tables[0],
+            query_start_loc,
+            positions,
+            num_actual_tokens=8,
+            num_reqs=8,
+            kpool=8,
+        )
+
+    if use_cuda_graph:
+        compute_slots()
+        torch.accelerator.synchronize()
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            slots = compute_slots()
+    tables.gather_block_tables(idx, num_reqs_padded=5)
+    if use_cuda_graph:
+        graph.replay()
+    else:
+        slots = compute_slots()
+    assert slots.tolist() == [24, 9, 42, 19, 36, -1, -1, -1]

vllm/v1/worker/gpu/block_table.py

@@ -74,6 +74,7 @@ class BlockTables:
         self.input_block_tables: list[torch.Tensor] = [
             torch.zeros_like(b.gpu) for b in self.block_tables
         ]
+        self._num_dirty_rows = 0
 
         self.slot_mappings = torch.zeros(
             self.num_kv_cache_groups,
@@ -151,15 +152,18 @@ class BlockTables:
     ) -> tuple[torch.Tensor, ...]:
         if self.num_kv_cache_groups == 0:
             return ()
+        num_reqs = idx_mapping.shape[0]
+        num_rows = num_reqs_padded
         if out is None:
             out = tuple(self.input_block_tables)
             out_ptrs = self.input_block_table_ptrs
+            num_rows = max(num_rows, self._num_dirty_rows)
+            self._num_dirty_rows = num_reqs
         else:
             assert out_ptrs is not None
             assert len(out) == self.num_kv_cache_groups
-        num_reqs = idx_mapping.shape[0]
-        # Launch kernel with num_reqs_padded to fuse zeroing of padded rows.
-        _gather_block_tables_kernel[(self.num_kv_cache_groups, num_reqs_padded)](
+        # Also clear rows left by a larger batch in the shared buffer.
+        _gather_block_tables_kernel[(self.num_kv_cache_groups, num_rows)](
             idx_mapping,
             self.block_table_ptrs,
             out_ptrs,
@@ -180,9 +184,12 @@ class BlockTables:
         # 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
-        )
+        num_rows = max(num_reqs, self._num_dirty_rows)
+        for block_table in self.input_block_tables:
+            block_table[:num_rows].zero_()
+        # Profiling may populate the returned rows with valid block IDs.
+        self._num_dirty_rows = num_reqs
+        return tuple(block_table[:num_reqs] for block_table in self.input_block_tables)
 
     def compute_slot_mappings(
         self,