diff --git a/tests/models/qwen4_exp/test_qsa_pre_indexer_amd.py b/tests/models/qwen4_exp/test_qsa_pre_indexer_amd.py index 8a39504e00..7698198923 100644 --- a/tests/models/qwen4_exp/test_qsa_pre_indexer_amd.py +++ b/tests/models/qwen4_exp/test_qsa_pre_indexer_amd.py @@ -55,9 +55,10 @@ MIXED_BATCH = ([260, 259, 138], [1, 1, 37], [8, 8, 8]) 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] diff --git a/tests/models/qwen4_exp/test_qsa_reference.py b/tests/models/qwen4_exp/test_qsa_reference.py index bdcc4e3423..a55cc028f9 100644 --- a/tests/models/qwen4_exp/test_qsa_reference.py +++ b/tests/models/qwen4_exp/test_qsa_reference.py @@ -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, @@ -340,10 +340,10 @@ 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, @@ -351,6 +351,10 @@ def test_qsa_circular_buffer_metadata_keeps_only_each_requests_suffix() -> None: 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: @@ -367,7 +371,7 @@ 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, @@ -375,7 +379,7 @@ def test_qsa_circular_buffer_survives_one_speculative_step(chunk_start: int) -> ) 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( diff --git a/tests/v1/attention/test_kpool_tail_slot_mapping.py b/tests/v1/attention/test_kpool_tail_slot_mapping.py index da25eaefb2..183c63ec2f 100644 --- a/tests/v1/attention/test_kpool_tail_slot_mapping.py +++ b/tests/v1/attention/test_kpool_tail_slot_mapping.py @@ -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 @@ -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) @@ -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( @@ -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") diff --git a/vllm/models/qwen4_exp/common/qsa_cache.py b/vllm/models/qwen4_exp/common/qsa_cache.py index 13228185ec..f960dea071 100644 --- a/vllm/models/qwen4_exp/common/qsa_cache.py +++ b/vllm/models/qwen4_exp/common/qsa_cache.py @@ -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 ) @@ -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 ) diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 8b186502f2..1005b8e424 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -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 @@ -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