[Bugfix] Keep the GLM-5.3 kpool tail and Qwen4 QSA ring out of the null block (#59528)

Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

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]

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(

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")

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
         )

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