[Perf][GLM] Fuse the kpool tail slot mapping into one Triton kernel (#57534)
Signed-off-by: Jared Wen <jaredwen@inferact.ai> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Jared Wen committed
Sep 20, 2026 at 18:20 UTC
1b9fa3eaa8ff23d1a583f2f77dd2b33c9f896ee8
2 files changed
+120
-2
tests/v1/attention/test_kpool_tail_slot_mapping.py
+63
@@ -391,3 +391,66 @@ def test_interleaved_decode_pollution_legacy_vs_circular():
391
392
# The circular mapping keeps the rings isolated under interleaving.
393
torch.testing.assert_close(circular, ground_truth)
394
+
395
+
396
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device")
397
+@pytest.mark.parametrize(
398
+ "per_req,num_actual,padded_len",
399
+ [
400
+ ([list(range(10)), list(range(12))], 22, 22),
401
+ ([list(range(10)), list(range(12))], 22, 30),
402
+ ([[3, 4], [0], [7, 8, 9]], 6, 8),
403
+ ([[5]], 1, 1),
404
+ ],
405
+)
406
+def test_triton_mapping_matches_cpu(per_req, num_actual, padded_len):
407
+ """The CUDA (Triton) path must match the CPU torch reference, including
408
+ tokens between the last request boundary and num_actual_tokens (mapped to
409
+ the last request) and untouched padding beyond num_actual."""
410
+ positions, qsl, slot_mapping, _, num_reqs = make_batch(
411
+ per_req, padded_len=padded_len
412
+ )
413
+ # Replace the all--1 placeholder slots with sentinel values to check the
414
+ # padding range is copied through untouched.
415
+ slot_mapping = torch.arange(padded_len, dtype=torch.int64) + 1000
416
+ bt = make_tail_block_table(list(range(5, 5 + num_reqs)))
417
+
418
+ ref = circular_tail_slots(slot_mapping, bt, qsl, positions, num_actual, num_reqs)
419
+ got = circular_tail_slots(
420
+ slot_mapping.cuda(),
421
+ bt.cuda(),
422
+ qsl.cuda().to(torch.int32),
423
+ positions.cuda(),
424
+ num_actual,
425
+ num_reqs,
426
+ )
427
+ torch.testing.assert_close(got.cpu(), ref)
428
+
429
+
430
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device")
431
+def test_triton_mapping_reads_strided_block_table():
432
+ """The kernel must address the block table through its real stride(0),
433
+ not a dense assumption (the #57477 class of bug): a tail block table that
434
+ is a column view of a wider allocation must still read column 0."""
435
+ per_req = [list(range(10)), list(range(12))]
436
+ positions, qsl, slot_mapping, num_actual, num_reqs = make_batch(per_req)
437
+ own_blocks = [11, 22]
438
+ wide = torch.zeros(num_reqs, 8, dtype=torch.int32)
439
+ bt = wide[:, 1:4] # non-contiguous view, stride(0) == 8
440
+ bt[:, 0] = torch.tensor(own_blocks, dtype=torch.int32)
441
+
442
+ ref = circular_tail_slots(
443
+ slot_mapping, bt.contiguous(), qsl, positions, num_actual, num_reqs
444
+ )
445
+ got = circular_tail_slots(
446
+ slot_mapping.cuda(),
447
+ bt.cuda(),
448
+ qsl.cuda().to(torch.int32),
449
+ positions.cuda(),
450
+ num_actual,
451
+ num_reqs,
452
+ )
453
+ torch.testing.assert_close(got.cpu(), ref)
454
+ for req, blk in enumerate(own_blocks):
455
+ start, end = int(qsl[req]), int(qsl[req + 1])
456
+ assert (got[start:end] // KPOOL == blk).all()
vllm/v1/attention/backends/mla/indexer.py
+57
-2
@@ -634,6 +634,38 @@ class DeepseekV32IndexerMetadata:
634
prefill: DeepseekV32IndexerPrefillMetadata | None = None
635
636
637
+@triton.jit(do_not_specialize=["num_reqs", "num_actual_tokens", "num_tokens"])
638
+def _kpool_tail_slot_mapping_kernel(
639
+ slot_mapping_ptr,
640
+ block_table_ptr,
641
+ block_table_stride,
642
+ query_start_loc_ptr,
643
+ positions_ptr,
644
+ out_ptr,
645
+ num_reqs,
646
+ num_actual_tokens,
647
+ num_tokens,
648
+ kpool,
649
+ BLOCK: tl.constexpr,
650
+):
651
+ pid = tl.program_id(0)
652
+ if pid < num_reqs:
653
+ start = tl.load(query_start_loc_ptr + pid)
654
+ # Tokens past the last request's boundary (if any) also map to it.
655
+ end = tl.load(query_start_loc_ptr + pid + 1)
656
+ end = tl.where(pid == num_reqs - 1, num_actual_tokens, end)
657
+ own_block = tl.load(block_table_ptr + pid * block_table_stride).to(tl.int64)
658
+ for i in range(start, end, BLOCK):
659
+ offs = i + tl.arange(0, BLOCK)
660
+ mask = offs < end
661
+ pos = tl.load(positions_ptr + offs, mask=mask, other=0).to(tl.int64)
662
+ tl.store(out_ptr + offs, own_block * kpool + pos % kpool, mask=mask)
663
+ else:
664
+ offs = num_actual_tokens + (pid - num_reqs) * BLOCK + tl.arange(0, BLOCK)
665
+ mask = offs < num_tokens
666
+ tl.store(out_ptr + offs, tl.load(slot_mapping_ptr + offs, mask=mask), mask=mask)
667
+
668
+
669
def compute_kpool_tail_slot_mapping(
670
slot_mapping: torch.Tensor,
671
block_table: torch.Tensor,
@@ -646,10 +678,33 @@ def compute_kpool_tail_slot_mapping(
678
) -> torch.Tensor:
679
"""Map every token to its request's one circular tail block."""
680
if out is None:
649
- out = slot_mapping.clone()
681
+ out = torch.empty_like(slot_mapping)
682
else:
683
assert out.shape == slot_mapping.shape
652
- out.copy_(slot_mapping)
684
+ if slot_mapping.is_cuda and slot_mapping.dim() == 1 and num_reqs > 0:
685
+ block = 256
686
+ num_tokens = slot_mapping.shape[0]
687
+ num_actual_tokens = min(num_actual_tokens, num_tokens)
688
+ grid = (num_reqs + triton.cdiv(num_tokens - num_actual_tokens, block),)
689
+ _kpool_tail_slot_mapping_kernel[grid](
690
+ slot_mapping,
691
+ block_table,
692
+ block_table.stride(0),
693
+ query_start_loc,
694
+ positions,
695
+ out,
696
+ num_reqs,
697
+ num_actual_tokens,
698
+ num_tokens,
699
+ kpool,
700
+ BLOCK=block,
701
+ num_warps=4,
702
+ )
703
+ return out
704
+ # Torch fallback: CPU tensors (the CPU unit tests), non-1D or empty inputs.
705
+ # Production always takes the Triton path above — spec-decode tokens arrive
706
+ # flattened token-major, so slot_mapping is 1D there too.
707
+ out.copy_(slot_mapping)
708
if num_actual_tokens == 0:
709
return out
710
tokens = torch.arange(num_actual_tokens, device=slot_mapping.device)