Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions tests/models/qwen4_exp/test_qsa_pre_indexer_amd.py
Original file line number Diff line number Diff line change
Expand Up @@ -55,9 +55,10 @@ def _reference_gemma_rmsnorm(


def _make_block_table(block_counts):
num_blocks = sum(block_counts)
# Block 0 is the null block, which no request owns.
num_blocks = sum(block_counts) + 1
table = torch.full((len(block_counts), max(block_counts)), -1, dtype=torch.int32)
physical_blocks = torch.randperm(num_blocks)
physical_blocks = torch.randperm(sum(block_counts)) + 1
offset = 0
for request, count in enumerate(block_counts):
table[request, :count] = physical_blocks[offset : offset + count]
Expand Down
18 changes: 11 additions & 7 deletions tests/models/qwen4_exp/test_qsa_reference.py
Original file line number Diff line number Diff line change
Expand Up @@ -315,7 +315,7 @@ def test_qsa_circular_buffer_metadata_keeps_only_each_requests_suffix() -> None:
builder.k_work_metadata_buffer = torch.empty(0, 2, dtype=torch.int32, device=device)
query_start_loc = torch.tensor([0, 7, 13, 13], dtype=torch.int32, device=device)
token_to_req = torch.tensor([0] * 7 + [1] * 6 + [0] * 3, device=device)
block_table = torch.tensor([[1], [0], [2]], dtype=torch.int32, device=device)
block_table = torch.tensor([[1], [3], [2]], dtype=torch.int32, device=device)
common = SimpleNamespace(
num_actual_tokens=16,
num_reqs=3,
Expand All @@ -340,17 +340,21 @@ def test_qsa_circular_buffer_metadata_keeps_only_each_requests_suffix() -> None:
4,
-1,
-1,
3,
0,
1,
2,
15,
12,
13,
14,
-1,
-1,
-1,
]

assert metadata.slot_mapping.tolist() == expected

# A dummy batch puts every request on the null block, which owns no ring.
block_table.zero_()
assert builder.build(0, common).slot_mapping.tolist() == [-1] * 16


@pytest.mark.parametrize("chunk_start", list(range(8)))
def test_qsa_circular_buffer_survives_one_speculative_step(chunk_start: int) -> None:
Expand All @@ -367,15 +371,15 @@ def test_qsa_circular_buffer_survives_one_speculative_step(chunk_start: int) ->
query_len = num_spec + 1

slots = qsa_cache.circular_qsa_slot_mapping(
torch.tensor([[0]], dtype=torch.int32),
torch.tensor([[1]], dtype=torch.int32),
torch.zeros(query_len, dtype=torch.int32),
torch.arange(chunk_start, chunk_start + query_len),
capacity,
query_start_loc=torch.tensor([0, query_len], dtype=torch.int32),
)

committed = torch.arange(chunk_start - chunk_start % compress_ratio, chunk_start)
assert set(slots.tolist()).isdisjoint((committed % capacity).tolist())
assert set((slots % capacity).tolist()).isdisjoint((committed % capacity).tolist())


def _qsa_key_cache(
Expand Down
30 changes: 19 additions & 11 deletions tests/v1/attention/test_kpool_tail_slot_mapping.py
Original file line number Diff line number Diff line change
Expand Up @@ -227,8 +227,10 @@ def test_circular_mapping_matches_generic_for_short_requests(prompt_len):


def test_circular_mapping_preserves_padding_and_empty_batch():
own_blocks = [5, 9]
per_req = [list(range(10)), list(range(12))]
# The last request is a padding request on the null block, which owns no
# tail block and must never be written.
own_blocks = [5, 9, 0]
per_req = [list(range(10)), list(range(12)), [0, 1]]
padded_len = sum(len(p) for p in per_req) + 8
positions, qsl, slot_mapping, num_actual, num_reqs = make_batch(
per_req, padded_len=padded_len
Expand All @@ -237,7 +239,8 @@ def test_circular_mapping_preserves_padding_and_empty_batch():

out = circular_tail_slots(slot_mapping, bt, qsl, positions, num_actual, num_reqs)
assert out.shape == slot_mapping.shape
assert torch.equal(out[num_actual:], torch.full_like(out[num_actual:], -1))
pad_start = int(qsl[-2])
assert torch.equal(out[pad_start:], torch.full_like(out[pad_start:], -1))

empty = circular_tail_slots(slot_mapping, bt, qsl, positions[:0], 0, num_reqs)
assert torch.equal(empty, slot_mapping)
Expand Down Expand Up @@ -430,25 +433,27 @@ def test_interleaved_decode_pollution_legacy_vs_circular():

@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device")
@pytest.mark.parametrize(
"per_req,num_actual,padded_len",
"per_req,own_blocks,num_actual,padded_len",
[
([list(range(10)), list(range(12))], 22, 22),
([list(range(10)), list(range(12))], 22, 30),
([[3, 4], [0], [7, 8, 9]], 6, 8),
([[5]], 1, 1),
([list(range(10)), list(range(12))], [5, 6], 22, 22),
([list(range(10)), list(range(12))], [5, 6], 22, 30),
([[3, 4], [0], [7, 8, 9]], [5, 0, 7], 6, 8),
([[5]], [5], 1, 1),
([[0], [0]], [0, 0], 2, 4), # dummy run: every row on the null block
],
)
def test_triton_mapping_matches_cpu(per_req, num_actual, padded_len):
def test_triton_mapping_matches_cpu(per_req, own_blocks, num_actual, padded_len):
"""The CUDA (Triton) path must match the CPU torch reference, including
tokens between the last request boundary and num_actual_tokens (mapped to
the last request) and untouched padding beyond num_actual."""
the last request), untouched padding beyond num_actual, and PAD for
requests on the null block."""
positions, qsl, slot_mapping, _, num_reqs = make_batch(
per_req, padded_len=padded_len
)
# Replace the all--1 placeholder slots with sentinel values to check the
# padding range is copied through untouched.
slot_mapping = torch.arange(padded_len, dtype=torch.int64) + 1000
bt = make_tail_block_table(list(range(5, 5 + num_reqs)))
bt = make_tail_block_table(own_blocks)

ref = circular_tail_slots(slot_mapping, bt, qsl, positions, num_actual, num_reqs)
got = circular_tail_slots(
Expand All @@ -460,6 +465,9 @@ def test_triton_mapping_matches_cpu(per_req, num_actual, padded_len):
num_reqs,
)
torch.testing.assert_close(got.cpu(), ref)
for req, blk in enumerate(own_blocks):
if blk == 0:
assert (got[int(qsl[req]) : int(qsl[req + 1])] == -1).all()


@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device")
Expand Down
7 changes: 5 additions & 2 deletions vllm/models/qwen4_exp/common/qsa_cache.py
Original file line number Diff line number Diff line change
Expand Up @@ -127,7 +127,9 @@ def circular_qsa_slot_mapping(
valid = (requests >= 0) & (requests < block_table.shape[0]) & (positions >= 0)
safe_requests = requests.clamp(0, block_table.shape[0] - 1)
physical_blocks = block_table[safe_requests, 0].long()
valid &= physical_blocks >= 0
# Dummy and padding requests own no ring and sit on the null block,
# which must never be written.
valid &= physical_blocks > 0
slots = physical_blocks * compressor_state_size + positions.remainder(
compressor_state_size
)
Expand Down Expand Up @@ -288,7 +290,8 @@ def _build_qsa_metadata_kernel(
mask=valid,
other=-1,
)
valid &= physical_block >= 0
# Dummy and padding requests own no ring and sit on the null block.
valid &= physical_block > 0
slot = physical_block * circular_buffer_size + (
logical_position % circular_buffer_size
)
Expand Down
9 changes: 7 additions & 2 deletions vllm/v1/attention/backends/mla/indexer.py
Original file line number Diff line number Diff line change
Expand Up @@ -702,11 +702,15 @@ def _kpool_tail_slot_mapping_kernel(
end = tl.load(query_start_loc_ptr + pid + 1)
end = tl.where(pid == num_reqs - 1, num_actual_tokens, end)
own_block = tl.load(block_table_ptr + pid * block_table_stride).to(tl.int64)
# Every scheduled request owns a tail block; the runners give dummy and
# padding requests the null block, which must never be written.
has_block = own_block != 0
for i in range(start, end, BLOCK):
offs = i + tl.arange(0, BLOCK)
mask = offs < end
pos = tl.load(positions_ptr + offs, mask=mask, other=0).to(tl.int64)
tl.store(out_ptr + offs, own_block * kpool + pos % kpool, mask=mask)
slot = tl.where(has_block, own_block * kpool + pos % kpool, -1)
tl.store(out_ptr + offs, slot, mask=mask)
else:
offs = num_actual_tokens + (pid - num_reqs) * BLOCK + tl.arange(0, BLOCK)
mask = offs < num_tokens
Expand Down Expand Up @@ -759,7 +763,8 @@ def compute_kpool_tail_slot_mapping(
req = req.clamp_(min=0, max=num_reqs - 1)
own_block = block_table[:num_reqs, 0].index_select(0, req).to(torch.int64)
pos = positions[:num_actual_tokens].to(torch.int64)
out[:num_actual_tokens] = own_block * kpool + torch.remainder(pos, kpool)
slots = own_block * kpool + torch.remainder(pos, kpool)
out[:num_actual_tokens] = torch.where(own_block != 0, slots, -1)
return out


Expand Down
Loading