61bb590877
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>
409 lines
15 KiB
Python
409 lines
15 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm.platforms import current_platform
|
|
from vllm.v1.attention.backends.utils import PAD_SLOT_ID
|
|
from vllm.v1.worker.gpu.block_table import BlockTables
|
|
|
|
pytestmark = pytest.mark.skipif(
|
|
not current_platform.is_cuda(),
|
|
reason="requires CUDA",
|
|
)
|
|
|
|
|
|
def test_block_tables_apply_staged_writes_fuses_kv_groups(monkeypatch):
|
|
device = torch.device("cuda")
|
|
block_tables = BlockTables(
|
|
block_sizes=[16, 32, 8],
|
|
max_num_reqs=4,
|
|
max_num_batched_tokens=64,
|
|
max_num_blocks_per_group=[8, 8, 8],
|
|
device=device,
|
|
)
|
|
|
|
def fail_if_apply_write_called():
|
|
pytest.fail("multi-group writes should use the fused apply kernel")
|
|
|
|
for block_table in block_tables.block_tables:
|
|
monkeypatch.setattr(block_table, "apply_write", fail_if_apply_write_called)
|
|
|
|
block_tables.append_block_ids(
|
|
req_index=0,
|
|
new_block_ids=([1, 2], [10, 11], []),
|
|
overwrite=True,
|
|
)
|
|
block_tables.append_block_ids(
|
|
req_index=1,
|
|
new_block_ids=([3], [12], [5, 6]),
|
|
overwrite=True,
|
|
)
|
|
block_tables.apply_staged_writes()
|
|
torch.accelerator.synchronize()
|
|
|
|
assert torch.equal(
|
|
block_tables.block_tables[0].gpu[0, :2],
|
|
torch.tensor([1, 2], dtype=torch.int32, device=device),
|
|
)
|
|
assert torch.equal(
|
|
block_tables.block_tables[1].gpu[0, :2],
|
|
torch.tensor([10, 11], dtype=torch.int32, device=device),
|
|
)
|
|
assert torch.equal(
|
|
block_tables.block_tables[0].gpu[1, :1],
|
|
torch.tensor([3], dtype=torch.int32, device=device),
|
|
)
|
|
assert torch.equal(
|
|
block_tables.block_tables[1].gpu[1, :1],
|
|
torch.tensor([12], dtype=torch.int32, device=device),
|
|
)
|
|
assert torch.equal(
|
|
block_tables.block_tables[2].gpu[1, :2],
|
|
torch.tensor([5, 6], dtype=torch.int32, device=device),
|
|
)
|
|
assert block_tables.num_blocks.np[0, 0] == 2
|
|
assert block_tables.num_blocks.np[1, 0] == 2
|
|
assert block_tables.num_blocks.np[2, 0] == 0
|
|
assert block_tables.num_blocks.np[0, 1] == 1
|
|
assert block_tables.num_blocks.np[1, 1] == 1
|
|
assert block_tables.num_blocks.np[2, 1] == 2
|
|
assert torch.equal(
|
|
block_tables.num_blocks.gpu[:, :2],
|
|
torch.tensor([[2, 1], [2, 1], [0, 2]], dtype=torch.int32, device=device),
|
|
)
|
|
|
|
for block_table in block_tables.block_tables:
|
|
assert not block_table._staged_write_indices
|
|
assert not block_table._staged_write_starts
|
|
assert not block_table._staged_write_contents
|
|
assert not block_table._staged_write_cu_lens
|
|
|
|
block_tables.append_block_ids(
|
|
req_index=0,
|
|
new_block_ids=([7], [13], [8]),
|
|
overwrite=False,
|
|
)
|
|
block_tables.apply_staged_writes()
|
|
torch.accelerator.synchronize()
|
|
|
|
assert torch.equal(
|
|
block_tables.block_tables[0].gpu[0, :3],
|
|
torch.tensor([1, 2, 7], dtype=torch.int32, device=device),
|
|
)
|
|
assert torch.equal(
|
|
block_tables.block_tables[1].gpu[0, :3],
|
|
torch.tensor([10, 11, 13], dtype=torch.int32, device=device),
|
|
)
|
|
assert torch.equal(
|
|
block_tables.block_tables[2].gpu[0, :1],
|
|
torch.tensor([8], dtype=torch.int32, device=device),
|
|
)
|
|
assert block_tables.num_blocks.np[0, 0] == 3
|
|
assert block_tables.num_blocks.np[1, 0] == 3
|
|
assert block_tables.num_blocks.np[2, 0] == 1
|
|
|
|
|
|
def test_block_tables_apply_staged_writes_single_group():
|
|
device = torch.device("cuda")
|
|
block_tables = BlockTables(
|
|
block_sizes=[16],
|
|
max_num_reqs=2,
|
|
max_num_batched_tokens=16,
|
|
max_num_blocks_per_group=[4],
|
|
device=device,
|
|
)
|
|
|
|
block_tables.append_block_ids(
|
|
req_index=0,
|
|
new_block_ids=([1, 2],),
|
|
overwrite=True,
|
|
)
|
|
block_tables.apply_staged_writes()
|
|
torch.accelerator.synchronize()
|
|
|
|
assert torch.equal(
|
|
block_tables.block_tables[0].gpu[0, :2],
|
|
torch.tensor([1, 2], dtype=torch.int32, device=device),
|
|
)
|
|
|
|
|
|
def test_block_tables_skip_custom_slot_mapping_groups():
|
|
device = torch.device("cuda")
|
|
block_tables = BlockTables(
|
|
block_sizes=[8, 262144],
|
|
max_num_reqs=1,
|
|
max_num_batched_tokens=4,
|
|
max_num_blocks_per_group=[1, 1],
|
|
device=device,
|
|
slot_mapping_enabled=[False, True],
|
|
)
|
|
block_tables.append_block_ids(
|
|
req_index=0,
|
|
new_block_ids=([7], [12]),
|
|
overwrite=True,
|
|
)
|
|
block_tables.apply_staged_writes()
|
|
|
|
idx_mapping = torch.tensor([0], dtype=torch.int32, device=device)
|
|
query_start_loc = torch.tensor([0, 2], dtype=torch.int32, device=device)
|
|
positions = torch.tensor([153797, 165757], dtype=torch.int64, device=device)
|
|
slot_mappings = block_tables.compute_slot_mappings(
|
|
idx_mapping,
|
|
query_start_loc,
|
|
positions,
|
|
num_tokens_padded=2,
|
|
)
|
|
torch.accelerator.synchronize()
|
|
|
|
assert slot_mappings[0].tolist() == [-1, -1]
|
|
assert slot_mappings[1].tolist() == [
|
|
12 * 262144 + 153797,
|
|
12 * 262144 + 165757,
|
|
]
|
|
|
|
|
|
@pytest.mark.parametrize("cp_rank", range(4))
|
|
def test_dcp_slot_mapping_only_shards_sharded_groups(cp_rank: int):
|
|
"""Only sharded groups use DCP interleave in logical-block coordinates."""
|
|
device = torch.device("cuda")
|
|
block_tables = BlockTables(
|
|
block_sizes=[128, 128],
|
|
max_num_reqs=1,
|
|
max_num_batched_tokens=1024,
|
|
max_num_blocks_per_group=[2, 8],
|
|
device=device,
|
|
dcp_sharded=[True, False],
|
|
cp_size=4,
|
|
cp_rank=cp_rank,
|
|
cp_interleave=128,
|
|
)
|
|
block_tables.append_block_ids(
|
|
req_index=0,
|
|
new_block_ids=([5, 9], list(range(10, 18))),
|
|
overwrite=True,
|
|
)
|
|
block_tables.apply_staged_writes()
|
|
|
|
idx_mapping = torch.zeros(1, dtype=torch.int32, device=device)
|
|
query_start_loc = torch.tensor([0, 1024], dtype=torch.int32, device=device)
|
|
positions = torch.arange(1024, dtype=torch.int64, device=device)
|
|
actual = block_tables.compute_slot_mappings(
|
|
idx_mapping,
|
|
query_start_loc,
|
|
positions,
|
|
num_tokens_padded=1024,
|
|
)
|
|
|
|
expected = torch.full((1024,), -1, dtype=torch.int64, device=device)
|
|
first_start = cp_rank * 128
|
|
second_start = 512 + first_start
|
|
expected[first_start : first_start + 128] = torch.arange(
|
|
5 * 128, 6 * 128, dtype=torch.int64, device=device
|
|
)
|
|
expected[second_start : second_start + 128] = torch.arange(
|
|
9 * 128, 10 * 128, dtype=torch.int64, device=device
|
|
)
|
|
assert torch.equal(actual[0], expected)
|
|
assert torch.equal(actual[1], positions + 10 * 128)
|
|
|
|
|
|
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"),
|
|
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,
|
|
)
|
|
# 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()
|
|
|
|
|
|
def test_dummy_request_slot_mapping_is_pad():
|
|
"""idx_mapping == -1 marks a dummy (or CUDA-graph padding) request.
|
|
|
|
A dummy draft decode step must not resolve slots through the persistent
|
|
block-table row of a real request slot, which may be stale and point at
|
|
blocks that are now in the prefix cache.
|
|
"""
|
|
device = torch.device("cuda")
|
|
block_tables = BlockTables(
|
|
block_sizes=[4],
|
|
max_num_reqs=2,
|
|
max_num_batched_tokens=16,
|
|
max_num_blocks_per_group=[4],
|
|
device=device,
|
|
)
|
|
block_tables.append_block_ids(req_index=0, new_block_ids=([7, 8],), overwrite=True)
|
|
block_tables.apply_staged_writes()
|
|
query_start_loc = torch.tensor([0, 3], dtype=torch.int32, device=device)
|
|
positions = torch.tensor([1, 2, 3], dtype=torch.int64, device=device)
|
|
|
|
real = block_tables.compute_slot_mappings(
|
|
torch.tensor([0], dtype=torch.int32, device=device),
|
|
query_start_loc,
|
|
positions,
|
|
num_tokens_padded=3,
|
|
)
|
|
assert real[0].tolist() == [7 * 4 + 1, 7 * 4 + 2, 7 * 4 + 3]
|
|
|
|
dummy = block_tables.compute_slot_mappings(
|
|
torch.tensor([-1], dtype=torch.int32, device=device),
|
|
query_start_loc,
|
|
positions,
|
|
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]
|