[Perf][GLM] Fuse the kpool tail slot mapping into one Triton kernel

tests/v1/attention/test_kpool_tail_slot_mapping.py

@@ -391,3 +391,66 @@ def test_interleaved_decode_pollution_legacy_vs_circular():
 
     # The circular mapping keeps the rings isolated under interleaving.
     torch.testing.assert_close(circular, ground_truth)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device")
+@pytest.mark.parametrize(
+    "per_req,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),
+    ],
+)
+def test_triton_mapping_matches_cpu(per_req, 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."""
+    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)))
+
+    ref = circular_tail_slots(slot_mapping, bt, qsl, positions, num_actual, num_reqs)
+    got = circular_tail_slots(
+        slot_mapping.cuda(),
+        bt.cuda(),
+        qsl.cuda().to(torch.int32),
+        positions.cuda(),
+        num_actual,
+        num_reqs,
+    )
+    torch.testing.assert_close(got.cpu(), ref)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device")
+def test_triton_mapping_reads_strided_block_table():
+    """The kernel must address the block table through its real stride(0),
+    not a dense assumption (the #57477 class of bug): a tail block table that
+    is a column view of a wider allocation must still read column 0."""
+    per_req = [list(range(10)), list(range(12))]
+    positions, qsl, slot_mapping, num_actual, num_reqs = make_batch(per_req)
+    own_blocks = [11, 22]
+    wide = torch.zeros(num_reqs, 8, dtype=torch.int32)
+    bt = wide[:, 1:4]  # non-contiguous view, stride(0) == 8
+    bt[:, 0] = torch.tensor(own_blocks, dtype=torch.int32)
+
+    ref = circular_tail_slots(
+        slot_mapping, bt.contiguous(), qsl, positions, num_actual, num_reqs
+    )
+    got = circular_tail_slots(
+        slot_mapping.cuda(),
+        bt.cuda(),
+        qsl.cuda().to(torch.int32),
+        positions.cuda(),
+        num_actual,
+        num_reqs,
+    )
+    torch.testing.assert_close(got.cpu(), ref)
+    for req, blk in enumerate(own_blocks):
+        start, end = int(qsl[req]), int(qsl[req + 1])
+        assert (got[start:end] // KPOOL == blk).all()

vllm/v1/attention/backends/mla/indexer.py

@@ -634,6 +634,38 @@ class DeepseekV32IndexerMetadata:
     prefill: DeepseekV32IndexerPrefillMetadata | None = None
 
 
+@triton.jit(do_not_specialize=["num_reqs", "num_actual_tokens", "num_tokens"])
+def _kpool_tail_slot_mapping_kernel(
+    slot_mapping_ptr,
+    block_table_ptr,
+    block_table_stride,
+    query_start_loc_ptr,
+    positions_ptr,
+    out_ptr,
+    num_reqs,
+    num_actual_tokens,
+    num_tokens,
+    kpool,
+    BLOCK: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    if pid < num_reqs:
+        start = tl.load(query_start_loc_ptr + pid)
+        # Tokens past the last request's boundary (if any) also map to it.
+        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)
+        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)
+    else:
+        offs = num_actual_tokens + (pid - num_reqs) * BLOCK + tl.arange(0, BLOCK)
+        mask = offs < num_tokens
+        tl.store(out_ptr + offs, tl.load(slot_mapping_ptr + offs, mask=mask), mask=mask)
+
+
 def compute_kpool_tail_slot_mapping(
     slot_mapping: torch.Tensor,
     block_table: torch.Tensor,
@@ -646,10 +678,33 @@ def compute_kpool_tail_slot_mapping(
 ) -> torch.Tensor:
     """Map every token to its request's one circular tail block."""
     if out is None:
-        out = slot_mapping.clone()
+        out = torch.empty_like(slot_mapping)
     else:
         assert out.shape == slot_mapping.shape
-        out.copy_(slot_mapping)
+    if slot_mapping.is_cuda and slot_mapping.dim() == 1 and num_reqs > 0:
+        block = 256
+        num_tokens = slot_mapping.shape[0]
+        num_actual_tokens = min(num_actual_tokens, num_tokens)
+        grid = (num_reqs + triton.cdiv(num_tokens - num_actual_tokens, block),)
+        _kpool_tail_slot_mapping_kernel[grid](
+            slot_mapping,
+            block_table,
+            block_table.stride(0),
+            query_start_loc,
+            positions,
+            out,
+            num_reqs,
+            num_actual_tokens,
+            num_tokens,
+            kpool,
+            BLOCK=block,
+            num_warps=4,
+        )
+        return out
+    # Torch fallback: CPU tensors (the CPU unit tests), non-1D or empty inputs.
+    # Production always takes the Triton path above — spec-decode tokens arrive
+    # flattened token-major, so slot_mapping is 1D there too.
+    out.copy_(slot_mapping)
     if num_actual_tokens == 0:
         return out
     tokens = torch.arange(num_actual_tokens, device=slot_mapping.device)