@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 1257512305

[DSA] Wire DeepGEMM sparse MQA logits into the DeepSeek V4.1 indexer (#56254)

Signed-off-by: Jared Wen <w13431838023@gmail.com> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>

Jared Wen committed Sep 16, 2026 at 05:07 UTC 12575123059c2142e39c9af4e217a4792c3c8176
13 files changed +1830 -26
tests/kernels/test_fused_indexer_q_rope_quant.py
+52
@@ -560,3 +560,55 @@ def test_indexer_k_store_roundtrips_through_rocm_gather(block_size, compress_rat
560 torch.testing.assert_close(
561 k_scale[emitted], expected_scale[emitted], rtol=0, atol=0
562 )
563 +
564 +
565 +@pytest.mark.parametrize("use_cutedsl", [False, True])
566 +@pytest.mark.skipif(
567 + not (
568 + current_platform.is_cuda() and current_platform.is_device_capability_family(100)
569 + ),
570 + reason="MXFP4 indexer cache requires an SM100-family GPU",
571 +)
572 +@torch.inference_mode()
573 +def test_fused_indexer_q_rope_quant_writes_bf16_weights(use_cutedsl):
574 + """The MXFP4 path can emit the per-head weights in bf16 (what DeepGEMM's
575 + sparse MQA-logits kernels take) instead of fp32; the values are the fp32
576 + result rounded once, so the scoring kernel needs no cast."""
577 + if use_cutedsl and not has_cutedsl():
578 + pytest.skip("cutedsl (cutlass) not installed")
579 +
580 + device = "cuda"
581 + torch.manual_seed(0)
582 + num_tokens, n_head = 257, 32
583 + q = torch.randn(num_tokens, n_head, HEAD_DIM, dtype=torch.bfloat16, device=device)
584 + positions = torch.randint(
585 + 0, MAX_POS, (num_tokens,), dtype=torch.int64, device=device
586 + )
587 + cos_sin_cache = torch.randn(MAX_POS, ROPE_DIM, dtype=torch.float32, device=device)
588 + weights = torch.randn(num_tokens, n_head, dtype=torch.bfloat16, device=device)
589 + softmax_scale, head_scale = HEAD_DIM**-0.5, 0.125 # pow2 head_scale, see above
590 +
591 + _, weights_ref = _reference(
592 + positions, q, cos_sin_cache, weights, softmax_scale, head_scale, n_head, True
593 + )
594 + cutedsl_patch = (
595 + mock.patch(
596 + "vllm.models.deepseek_v4.common.ops.fused_indexer_q.has_cutedsl",
597 + return_value=False,
598 + )
599 + if not use_cutedsl
600 + else contextlib.nullcontext()
601 + )
602 + with cutedsl_patch:
603 + _, weights_fused = fused_indexer_q_rope_quant(
604 + positions,
605 + q.clone(),
606 + cos_sin_cache,
607 + weights,
608 + softmax_scale,
609 + head_scale,
610 + use_fp4=True,
611 + weights_out_dtype=torch.bfloat16,
612 + )
613 + assert weights_fused.dtype == torch.bfloat16
614 + assert torch.equal(weights_fused, weights_ref.to(torch.bfloat16))
tests/model_executor/layers/test_mla_short_prefill_indexer.py
+544
@@ -450,3 +450,547 @@ def test_candidate_kernels_preserve_packed_bounds_and_padding(
450 keep = valid & (block >= 0) & ((cols - ks[:, None]) // block_size == block)
451 reference = torch.where(keep, 3.0, -torch.inf)
452 torch.testing.assert_close(logits, reference, rtol=0, atol=0)
453 +
454 +
455 +# --- DeepGEMM sparse-MQA indexer path (V4.1 two-level candidate filtering) ---
456 +
457 +
458 +def _make_candidates(lens: torch.Tensor, num_candidates: int, cbk: int):
459 + """Production-faithful candidates: top min(avail, K) blocks per row —
460 + full coverage for lens <= K*cbk, a fixed-size subset beyond that."""
461 + avail = ((lens.cpu() + cbk - 1) // cbk).tolist()
462 + rows = []
463 + for a in avail:
464 + n = min(a, num_candidates)
465 + picks = torch.randperm(a)[:n].sort().values[:num_candidates].tolist()
466 + rows.append(picks + [-1] * (num_candidates - len(picks)))
467 + return torch.tensor(rows, dtype=torch.int32)
468 +
469 +
470 +def _candidate_token_mask(
471 + candidates: torch.Tensor, cbk: int, total: int, ks: torch.Tensor | None = None
472 +):
473 + """[rows, total] bool mask of the candidate-covered token positions.
474 +
475 + ``ks`` converts request-local candidate blocks to absolute positions
476 + (prefill's packed workspace); pass None when already request-local.
477 + """
478 + rows = candidates.shape[0]
479 + device = candidates.device
480 + # Sentinel block keeps -1 padding from scattering onto real columns.
481 + width = max(total, (candidates.shape[1] + 1) * cbk)
482 + if ks is not None:
483 + width += int(ks.max()) + cbk
484 + in_cand = torch.zeros(rows, width, dtype=torch.bool, device=device)
485 + block_id = torch.where(candidates >= 0, candidates.long(), candidates.shape[1])
486 + tok = block_id.unsqueeze(2) * cbk + torch.arange(cbk, device=device)
487 + if ks is not None:
488 + tok = tok + ks.unsqueeze(1).unsqueeze(2)
489 + keep = (candidates >= 0).unsqueeze(2).expand(-1, -1, cbk)
490 + in_cand.scatter_(1, tok.flatten(1), keep.flatten(1))
491 + return in_cand[:, :total]
492 +
493 +
494 +def _reference_topk(masked: torch.Tensor, topk: int) -> torch.Tensor:
495 + """Plain ``torch.topk`` over -inf-masked logits; slots past a row's valid
496 + count are -1 (the indexer's "no token" sentinel)."""
497 + values, indices = masked.topk(topk, dim=1)
498 + return torch.where(values > float("-inf"), indices.int(), -1)
499 +
500 +
501 +def _assert_same_selection(
502 + logits: torch.Tensor,
503 + got: torch.Tensor,
504 + expected: torch.Tensor,
505 + row_starts: torch.Tensor | None = None,
506 +) -> None:
507 + """Both top-k index sets must pick the same multiset of logit values per
508 + row (-1 padding must match too). Indices themselves may differ: the
509 + top-k kernels break ties among equal values in nondeterministic order."""
510 + for r in range(logits.shape[0]):
511 + base = int(row_starts[r]) if row_starts is not None else 0
512 + g, e = got[r][got[r] >= 0], expected[r][expected[r] >= 0]
513 + assert g.numel() == e.numel(), f"row {r}: {g.numel()} vs {e.numel()} selected"
514 + vg = logits[r][(base + g).long()].sort().values
515 + ve = logits[r][(base + e).long()].sort().values
516 + assert torch.equal(vg, ve), f"row {r}: selected values differ"
517 +
518 +
519 +def _sparse_scratch(rows: int, num_sparse: int, topk: int) -> dict:
520 + """Caller-owned buffers the sparse path writes into (the metadata builder
521 + owns these in production)."""
522 + i32 = dict(dtype=torch.int32, device="cuda")
523 + return dict(
524 + sparse_indices=torch.zeros(rows, num_sparse, **i32),
525 + end=torch.zeros(rows, **i32),
526 + col_indices=torch.zeros(rows, topk, **i32),
527 + )
528 +
529 +
530 +def _expand_candidates_reference(
531 + candidates: torch.Tensor, ks: torch.Tensor, ke: torch.Tensor, cbk: int, sbk: int
532 +) -> tuple[torch.Tensor, torch.Tensor]:
533 + """PyTorch reference for `candidate_blocks_to_sparse_indices`: expand
534 + each candidate block into ``cbk // sbk`` sparse blocks, keep those inside
535 + ``[0, ke - ks)``, sort ascending, anchor at ``ks // sbk``, pad by
536 + repeating the last valid block, and count the valid sparse columns."""
537 + ratio = cbk // sbk
538 + rows, num_candidates = candidates.shape
539 + num_sparse = num_candidates * ratio
540 + big = torch.iinfo(torch.int32).max
541 + offsets = torch.arange(ratio, device=candidates.device, dtype=torch.int32)
542 + blocks = (candidates.unsqueeze(-1) * ratio + offsets).flatten(1)
543 + max_block = ((ke - ks).clamp(min=0) + sbk - 1) // sbk
544 + valid = (blocks >= 0) & (blocks < max_block.unsqueeze(1))
545 + sorted_blocks, _ = torch.where(valid, blocks, big).sort(dim=1)
546 + n_valid = (sorted_blocks != big).sum(dim=1, dtype=torch.int32)
547 + base = (ks // sbk).unsqueeze(1)
548 + kernel_blocks = torch.where(sorted_blocks != big, sorted_blocks, 0) + base
549 + last_slot = (n_valid - 1).clamp(min=0).unsqueeze(1).long()
550 + last_block = kernel_blocks.gather(1, last_slot)
551 + pad = torch.arange(num_sparse, device=candidates.device) >= n_valid.unsqueeze(1)
552 + sparse_indices = torch.where(pad, last_block, kernel_blocks)
553 + last_start = last_block.squeeze(1) * sbk + ks % sbk
554 + last_fill = (ke - last_start).clamp(min=1, max=sbk)
555 + end = torch.where(n_valid > 0, (n_valid - 1) * sbk + last_fill, 0)
556 + return sparse_indices.int(), end.int()
557 +
558 +
559 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Triton kernel")
560 +def test_candidate_blocks_to_sparse_indices_math():
561 + """Candidate->sparse-block expansion: ks anchoring, range filtering,
562 + repeat-last padding, and the per-row valid column count (``end``)."""
563 + from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import (
564 + candidate_blocks_to_sparse_indices,
565 + )
566 +
567 + # Rows cover: exactly one block, an unaligned ks with a partial tail
568 + # block, an empty row, duplicates / -1 padding / out-of-range garbage
569 + # (warmup runs touch uninitialized buffers), and int32-overflow junk.
570 + ks = torch.tensor([64, 13, 100, 8, 0], dtype=torch.int32, device="cuda")
571 + ke = torch.tensor([72, 30, 100, 40, 1000], dtype=torch.int32, device="cuda")
572 + candidates = torch.tensor(
573 + [
574 + [0, -1, -1, -1],
575 + [2, 0, -1, -1],
576 + [3, -1, -1, -1],
577 + [1, 1, 7, 0],
578 + [2**30, -(2**30), 0, -1],
579 + ],
580 + dtype=torch.int32,
581 + device="cuda",
582 + )
583 + indices, end = candidate_blocks_to_sparse_indices(
584 + candidates, ks, ke, candidate_block_size=8, sparse_block_kv=8
585 + )
586 + assert indices.tolist() == [
587 + [8, 8, 8, 8],
588 + [1, 3, 3, 3],
589 + [12, 12, 12, 12],
590 + [1, 2, 2, 2],
591 + [0, 0, 0, 0],
592 + ]
593 + # end: row 1's last block covers [29, 37) but ke = 30 -> one valid token;
594 + # row 3 keeps the duplicated block (production candidates are unique).
595 + assert end.tolist() == [8, 9, 0, 24, 8]
596 +
597 + # Wider candidate blocks expand into multiple sparse blocks.
598 + idx2, end2 = candidate_blocks_to_sparse_indices(
599 + torch.tensor([[1, 0]], dtype=torch.int32, device="cuda"),
600 + torch.tensor([0], dtype=torch.int32, device="cuda"),
601 + torch.tensor([48], dtype=torch.int32, device="cuda"),
602 + candidate_block_size=16,
603 + sparse_block_kv=8,
604 + )
605 + assert idx2.tolist() == [[0, 1, 2, 3]]
606 + assert end2.tolist() == [32]
607 +
608 +
609 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Triton kernel")
610 +@pytest.mark.parametrize("cbk,sbk", [(8, 8), (16, 8), (16, 16)])
611 +def test_candidate_blocks_to_sparse_indices_matches_reference(cbk: int, sbk: int):
612 + """The Triton expansion equals the PyTorch reference on production-shaped
613 + random inputs, including strided (column-sliced) candidate rows and
614 + strided output buffers."""
615 + from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import (
616 + candidate_blocks_to_sparse_indices,
617 + )
618 +
619 + torch.manual_seed(0)
620 + rows, num_candidates = 41, 2048
621 + ratio = cbk // sbk
622 + lens = torch.randint(1, 40000, (rows,), dtype=torch.int32)
623 + lens[0], lens[1] = 1, 3
624 + ks = torch.zeros(rows, dtype=torch.int32)
625 + ks[1:] = lens.cumsum(0)[:-1]
626 + ks, ke = ks.cuda(), (ks + lens).cuda()
627 + candidates = _make_candidates(lens, num_candidates, cbk).cuda()
628 + # Row 2 gets out-of-range garbage past the padding, as a CUDA-graph
629 + # warm-up on uninitialized buffers would.
630 + candidates[2, -3:] = torch.tensor([2**30, -(2**29), 10**6], dtype=torch.int32)
631 +
632 + ref_indices, ref_end = _expand_candidates_reference(candidates, ks, ke, cbk, sbk)
633 +
634 + cand_buf = torch.full((rows, num_candidates + 64), -7, dtype=torch.int32).cuda()
635 + cand_buf[:, :num_candidates] = candidates
636 + si_buf = torch.zeros(rows, num_candidates * ratio + 32, dtype=torch.int32).cuda()
637 + end_buf = torch.zeros(rows * 2, dtype=torch.int32, device="cuda")
638 + ks_buf = torch.zeros(rows * 2, dtype=torch.int32, device="cuda")
639 + ks_buf[::2] = ks
640 + got_indices, got_end = candidate_blocks_to_sparse_indices(
641 + cand_buf[:, :num_candidates],
642 + ks_buf[::2],
643 + ke,
644 + cbk,
645 + sbk,
646 + out=(si_buf[:, : num_candidates * ratio], end_buf[::2]),
647 + )
648 + assert torch.equal(got_indices, ref_indices)
649 + assert torch.equal(got_end, ref_end)
650 + assert torch.equal(
651 + cand_buf[:, num_candidates:], torch.full_like(cand_buf[:, num_candidates:], -7)
652 + )
653 + assert not si_buf[:, num_candidates * ratio :].any()
654 + assert not end_buf[1::2].any()
655 +
656 +
657 +def _skip_unless_sm100_sparse_kernels():
658 + if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10:
659 + pytest.skip("SM100 required")
660 + deep_gemm = pytest.importorskip("deep_gemm")
661 + if not hasattr(deep_gemm, "fp8_fp4_sparse_mqa_logits"):
662 + pytest.skip("DeepGEMM >= 2.8 required for sparse MQA logits")
663 + from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import (
664 + has_deep_select,
665 + )
666 +
667 + if not has_deep_select():
668 + pytest.skip("DeepSelect extension (vllm._deepselect_C) required")
669 + return deep_gemm
670 +
671 +
672 +def _quant_fp4(x: torch.Tensor):
673 + from deep_gemm.utils import per_token_cast_to_fp4
674 +
675 + return per_token_cast_to_fp4(x, use_ue8m0=True, gran_k=32, use_packed_ue8m0=True)
676 +
677 +
678 +@pytest.mark.parametrize("capture", [False, True])
679 +def test_sparse_mqa_logits_prefill_matches_dense_masked(capture: bool):
680 + """The sparse-MQA candidate path selects exactly the same logit values as
681 + ``torch.topk`` over dense bf16 logits masked to the candidates (the two
682 + DeepGEMM kernels agree bitwise on candidate columns), incl. under CUDA
683 + graph capture, and a later layer reusing the step's expansion/schedule
684 + selects the same values."""
685 + deep_gemm = _skip_unless_sm100_sparse_kernels()
686 + from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import (
687 + sparse_mqa_logits_prefill_chunk,
688 + )
689 +
690 + torch.manual_seed(0)
691 + rows, num_heads, head_dim = 37, 32, 128
692 + cbk, topk, num_candidates = 8, 512, 2048 # V4.1-Flash production values
693 + lens = torch.randint(1, 30000, (rows,), dtype=torch.int32)
694 + lens[0], lens[1], lens[2] = 1, 5, 2048 # single token, sub-block, long
695 + ks = torch.zeros(rows, dtype=torch.int32)
696 + ks[1:] = lens.cumsum(0)[:-1]
697 + ks, ke, lens = ks.cuda(), (ks + lens).cuda(), lens.cuda()
698 + total = int(lens.sum())
699 + q_fp, q_sf = _quant_fp4(
700 + torch.randn(rows * num_heads, head_dim, device="cuda", dtype=torch.bfloat16)
701 + )
702 + k_fp, k_sf = _quant_fp4(
703 + torch.randn(total, head_dim, device="cuda", dtype=torch.bfloat16)
704 + )
705 + weights = torch.randn(rows, num_heads, device="cuda", dtype=torch.bfloat16)
706 + candidates = _make_candidates(lens, num_candidates, cbk).cuda()
707 +
708 + # Reference: dense bf16 logits masked to the candidates, then torch.topk.
709 + dense = deep_gemm.fp8_fp4_mqa_logits(
710 + (
711 + q_fp.view(torch.int8).view(rows, num_heads, head_dim // 2),
712 + q_sf.view(rows, num_heads),
713 + ),
714 + (k_fp.view(torch.int8), k_sf.view(total)),
715 + weights,
716 + ks,
717 + ke,
718 + clean_logits=False,
719 + logits_dtype=torch.bfloat16,
720 + )
721 + cols = torch.arange(total, device="cuda")
722 + in_range = (cols - ks.unsqueeze(1) >= 0) & (
723 + cols - ks.unsqueeze(1) < lens.unsqueeze(1)
724 + )
725 + masked = dense.masked_fill(
726 + ~(in_range & _candidate_token_mask(candidates, cbk, total, ks)),
727 + float("-inf"),
728 + ).float()
729 + ref_abs = _reference_topk(masked, topk) # packed-workspace positions
730 + ref = torch.where(ref_abs >= 0, ref_abs - ks.unsqueeze(1), -1)
731 +
732 + out = torch.full((rows, topk), -1, dtype=torch.int32, device="cuda")
733 + scratch = _sparse_scratch(rows, num_candidates, topk)
734 +
735 + def run_sparse(kernel_metadata=None):
736 + return sparse_mqa_logits_prefill_chunk(
737 + q_fp.view(torch.int8).view(rows, num_heads, head_dim // 2),
738 + q_sf.view(rows, num_heads),
739 + k_fp.view(torch.int8),
740 + k_sf.view(total),
741 + weights,
742 + ks,
743 + ke,
744 + candidates,
745 + cbk,
746 + cbk,
747 + topk,
748 + out,
749 + kernel_metadata=kernel_metadata,
750 + **scratch,
751 + )
752 +
753 + if capture:
754 + stream = torch.cuda.Stream()
755 + stream.wait_stream(torch.cuda.current_stream())
756 + with torch.cuda.stream(stream):
757 + run_sparse() # warm up JIT + the DeepGEMM workspace pre-capture
758 + graph = torch.cuda.CUDAGraph()
759 + with torch.cuda.graph(graph, stream=stream):
760 + run_sparse()
761 + out.fill_(-1)
762 + graph.replay()
763 + else:
764 + kernel_metadata = run_sparse()
765 + _assert_same_selection(dense, out, ref, row_starts=ks)
766 + out.fill_(-1)
767 + run_sparse(kernel_metadata)
768 +
769 + _assert_same_selection(dense, out, ref, row_starts=ks)
770 +
771 +
772 +def _build_fp4_paged_cache(req_lens: torch.Tensor, page_kv: int, head_dim: int):
773 + """MXFP4 paged cache in DeepGEMM's fused layout: per page, all tokens'
774 + fp4 data first, then their packed UE8M0 scales."""
775 + device = req_lens.device
776 + pages_per = (req_lens + page_kv - 1) // page_kv
777 + num_pages, max_pages = int(pages_per.sum()), int(pages_per.max())
778 + kv_cache = torch.zeros(
779 + num_pages, page_kv, 1, head_dim // 2 + 4, dtype=torch.uint8, device=device
780 + )
781 + block_table = torch.zeros(
782 + len(req_lens), max_pages, dtype=torch.int32, device=device
783 + )
784 + page_pool = torch.randperm(num_pages, device=device)
785 + offset = 0
786 + for b in range(len(req_lens)):
787 + n = int(pages_per[b])
788 + block_table[b, :n] = page_pool[offset : offset + n].to(torch.int32)
789 + offset += n
790 + k_fp, k_sf = _quant_fp4(
791 + torch.randn(int(req_lens[b]), head_dim, device=device, dtype=torch.bfloat16)
792 + )
793 + rows68 = torch.zeros(n * page_kv, 68, dtype=torch.uint8, device=device)
794 + rows68[: int(req_lens[b]), : head_dim // 2] = k_fp.view(torch.uint8)
795 + rows68[: int(req_lens[b]), head_dim // 2 :] = k_sf.view(torch.uint8).view(-1, 4)
796 + by_page = rows68.view(n, page_kv, 68)
797 + split = torch.cat(
798 + [by_page[..., :64].reshape(n, -1), by_page[..., 64:].reshape(n, -1)], dim=1
799 + ).view(n, page_kv, 1, head_dim // 2 + 4)
800 + kv_cache[block_table[b, :n].long()] = split
801 + return kv_cache, block_table
802 +
803 +
804 +@pytest.mark.parametrize("next_n", [1, 2])
805 +@pytest.mark.parametrize("capture", [False, True])
806 +def test_sparse_mqa_logits_paged_matches_dense_masked(next_n: int, capture: bool):
807 + """Paged sparse-MQA decode selects exactly the values ``torch.topk`` picks
808 + from the dense bf16 masked logits, incl. MTP rows (next_n > 1 flattened to
809 + one query row each), CUDA graph capture/replay and schedule reuse by a
810 + later layer."""
811 + deep_gemm = _skip_unless_sm100_sparse_kernels()
812 + from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import (
813 + sparse_mqa_logits_paged_decode,
814 + )
815 +
816 + torch.manual_seed(0)
817 + batch, num_heads, head_dim, page_kv = 9, 32, 128, 128
818 + cbk, topk, num_candidates = 8, 512, 2048
819 + rows = batch * next_n
820 + req_lens = torch.randint(1, 30000, (batch,), dtype=torch.int32, device="cuda")
821 + req_lens[0], req_lens[1] = 1, 30000 # tiny row + a filtered (subset) row
822 + # Row (b, j) sees L_b - next_n + j + 1 tokens (native MTP layout).
823 + row_lens = (
824 + (req_lens.unsqueeze(1) - next_n + torch.arange(1, next_n + 1, device="cuda"))
825 + .clamp(min=0)
826 + .to(torch.int32)
827 + )
828 + kv_cache, block_table = _build_fp4_paged_cache(req_lens, page_kv, head_dim)
829 + assert kv_cache.stride(0) % 512 == 0 # the paged sparse kernel requires it
830 +
831 + q_fp, q_sf = _quant_fp4(
832 + torch.randn(rows * num_heads, head_dim, device="cuda", dtype=torch.bfloat16)
833 + )
834 + q = q_fp.view(torch.int8).view(batch, next_n, num_heads, head_dim // 2)
835 + q_scale = q_sf.view(batch, next_n, num_heads)
836 + weights = torch.randn(rows, num_heads, device="cuda", dtype=torch.bfloat16)
837 + candidates = _make_candidates(row_lens.flatten(), num_candidates, cbk).cuda()
838 +
839 + def reference_topk():
840 + context_lens_2d = row_lens.view(batch, next_n)
841 + sched = deep_gemm.get_paged_mqa_logits_metadata(
842 + context_lens_2d, page_kv, deep_gemm.get_num_sms(), None
843 + )
844 + logits = deep_gemm.fp8_fp4_paged_mqa_logits(
845 + (q, q_scale),
846 + kv_cache,
847 + weights,
848 + context_lens_2d,
849 + block_table,
850 + sched,
851 + int(req_lens.max()),
852 + False,
853 + torch.bfloat16,
854 + None,
855 + )
856 + lens_flat = row_lens.flatten()
857 + cols = torch.arange(logits.shape[1], device="cuda")
858 + masked = logits.masked_fill(
859 + ~(
860 + (cols.unsqueeze(0) < lens_flat.unsqueeze(1))
861 + & _candidate_token_mask(candidates, cbk, logits.shape[1])
862 + ),
863 + float("-inf"),
864 + ).float()
865 + return _reference_topk(masked, topk), logits
866 +
867 + expected, dense_logits = reference_topk()
868 + out = torch.full((rows, topk), -1, dtype=torch.int32, device="cuda")
869 +
870 + # The metadata builder flattens every query to one row: per-row context
871 + # lengths, an expanded block table and a row -> request map.
872 + context_lens = row_lens.flatten().contiguous()
873 + row_block_table = block_table.repeat_interleave(next_n, dim=0)
874 + row_indices = torch.arange(batch, dtype=torch.int32, device="cuda")
875 + row_indices = row_indices.repeat_interleave(next_n)
876 + row_ks = torch.zeros(rows, dtype=torch.int32, device="cuda")
877 + scratch = _sparse_scratch(rows, num_candidates, topk)
878 +
879 + def run_sparse(o: torch.Tensor, kernel_metadata=None):
880 + return sparse_mqa_logits_paged_decode(
881 + q.view(rows, 1, num_heads, head_dim // 2),
882 + q_scale.view(rows, 1, num_heads),
883 + kv_cache,
884 + weights,
885 + context_lens,
886 + row_block_table,
887 + row_indices,
888 + candidates,
889 + cbk,
890 + cbk,
891 + topk,
892 + o,
893 + row_ks=row_ks,
894 + kernel_metadata=kernel_metadata,
895 + **scratch,
896 + )
897 +
898 + if capture:
899 + stream = torch.cuda.Stream()
900 + stream.wait_stream(torch.cuda.current_stream())
901 + with torch.cuda.stream(stream):
902 + run_sparse(out)
903 + graph = torch.cuda.CUDAGraph()
904 + with torch.cuda.graph(graph, stream=stream):
905 + run_sparse(out)
906 + out.fill_(-1)
907 + graph.replay()
908 + else:
909 + kernel_metadata = run_sparse(out)
910 + _assert_same_selection(dense_logits, out, expected)
911 + out.fill_(-1)
912 + run_sparse(out, kernel_metadata)
913 +
914 + _assert_same_selection(dense_logits, out, expected)
915 +
916 +
917 +def test_sparse_mqa_logits_accuracy_vs_torch():
918 + """Sparse-kernel numerics vs a PyTorch fp32 reference on the dequantized
919 + MXFP4 inputs."""
920 + _skip_unless_sm100_sparse_kernels()
921 + from deep_gemm.utils import cast_back_from_fp4
922 +
923 + from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import (
924 + candidate_blocks_to_sparse_indices,
925 + )
926 + from vllm.utils.deep_gemm import (
927 + fp8_fp4_sparse_mqa_logits,
928 + get_sparse_mqa_logits_metadata,
929 + )
930 +
931 + torch.manual_seed(0)
932 + rows, num_heads, head_dim = 17, 32, 128
933 + sbk, num_candidates = 8, 2048
934 + lens = torch.randint(1, 1500, (rows,), dtype=torch.int32)
935 + lens[0] = 1
936 + ks = torch.zeros(rows, dtype=torch.int32)
937 + ks[1:] = lens.cumsum(0)[:-1]
938 + ks, ke, lens = ks.cuda(), (ks + lens).cuda(), lens.cuda()
939 + total = int(lens.sum())
940 + q_fp, q_sf = _quant_fp4(
941 + torch.randn(rows * num_heads, head_dim, device="cuda", dtype=torch.bfloat16)
942 + )
943 + k_fp, k_sf = _quant_fp4(
944 + torch.randn(total, head_dim, device="cuda", dtype=torch.bfloat16)
945 + )
946 + weights = torch.randn(rows, num_heads, device="cuda", dtype=torch.bfloat16)
947 + candidates = _make_candidates(lens, num_candidates, sbk).cuda()
948 +
949 + sparse_idx, end = candidate_blocks_to_sparse_indices(candidates, ks, ke, sbk, sbk)
950 + metadata = get_sparse_mqa_logits_metadata(
951 + ks, ke, total, sparse_idx, torch.int8, sbk
952 + )
953 + logits = fp8_fp4_sparse_mqa_logits(
954 + (
955 + q_fp.view(torch.int8).view(rows, num_heads, head_dim // 2),
956 + q_sf.view(rows, num_heads),
957 + ),
958 + (k_fp.view(torch.int8), k_sf.view(total)),
959 + weights,
960 + metadata,
961 + sparse_idx.shape[1],
962 + sbk,
963 + )
964 +
965 + # PyTorch fp32 reference on the dequantized inputs.
966 + q_deq = cast_back_from_fp4(q_fp, q_sf, gran_k=32, use_packed_ue8m0=True)
967 + k_deq = cast_back_from_fp4(k_fp, k_sf, gran_k=32, use_packed_ue8m0=True)
968 + dots = torch.einsum(
969 + "rhd,td->rht",
970 + q_deq.view(rows, num_heads, head_dim).float(),
971 + k_deq.float(),
972 + )
973 + ref = (dots.relu() * weights.float().unsqueeze(-1)).sum(1) # [rows, total]
974 +
975 + # Compare on valid sparse columns (absolute packed-workspace positions).
976 + width = sparse_idx.shape[1] * sbk
977 + cols = torch.arange(width, device="cuda")
978 + valid = cols.unsqueeze(0) < end.unsqueeze(1).long()
979 + pos = (
980 + sparse_idx.repeat_interleave(sbk, dim=1).long() * sbk
981 + + (ks % sbk).unsqueeze(1)
982 + + (cols % sbk).unsqueeze(0)
983 + ).clamp(min=0, max=total - 1)
984 + abs_err = (logits.float() - ref.gather(1, pos)).abs()[valid]
985 + row_scale = ref.abs().amax(dim=1)
986 + norm_err = ((logits.float() - ref.gather(1, pos)).abs() / row_scale.unsqueeze(1))[
987 + valid
988 + ]
989 + print(
990 + f"sparse vs torch fp32: max_abs={abs_err.max().item():.4f}; "
991 + f"rowmax-normalized p50={norm_err.median().item():.5f} "
992 + f"p99={norm_err.quantile(0.99).item():.5f} max={norm_err.max().item():.5f}"
993 + )
994 + # bf16-level numerics: far below the value gaps that change top-k selection.
995 + assert norm_err.quantile(0.99).item() < 1e-2
996 + assert norm_err.max().item() < 2e-2
vllm/config/attention.py
+9
@@ -86,6 +86,15 @@ class AttentionConfig:
86 indexer). Quantized formats (fp8, mxfp4, nvfp4) require indexer kernel
87 support in the backend."""
88
89 + indexer_sparse_logits: bool = False
90 + """DeepSeek V4.1 two-level indexer: score only the candidate blocks with
91 + DeepGEMM's sparse MQA-logits kernels instead of computing dense logits over
92 + the whole context and masking them. Requires `indexer_kv_dtype="mxfp4"`,
93 + an SM100-class GPU, DeepGEMM >= 2.8 and the DeepSelect top-k extension
94 + (the top-k runs on the kernels' bf16 logits). The sparse path costs
95 + O(candidate blocks) per query regardless of context length, so it pays off
96 + for long contexts (roughly 32K tokens and beyond) and is slower below."""
97 +
98 hisparse_config: HiSparseConfig | None = None
99 """HiSparse host-resident KV configuration. Setting this enables experimental
100 Model Runner V2-only HiSparse sparse-MLA decode hot-buffering. It is inferred
vllm/model_executor/kernels/attention/dsa/sparse_mqa_logits.py new
+469
@@ -0,0 +1,469 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""DeepGEMM sparse MQA-logits path for the DeepSeek V4.1 two-level indexer.
4 +
5 +The candidate-source indexer publishes the top candidate blocks (in units of
6 +``candidate_block_size`` compressed positions). Instead of scoring every KV
7 +position and masking the dense logits down to those blocks, this path scores
8 +only the candidate tokens with DeepGEMM's sparse MQA logits kernels
9 +(``fp8_fp4_(paged_)sparse_mqa_logits``, DeepGEMM >= 2.8, SM100 only), runs
10 +DeepSelect's top-k directly on the bf16 logits they produce, and remaps the
11 +selected columns back to request-local positions.
12 +
13 +Only the MXFP4 indexer cache is supported: the kernels require UE8M0-packed
14 +(granularity-32) Q/KV scales, which is exactly the MXFP4 cache layout. The
15 +FP8 indexer cache stores one fp32 scale per quant block instead.
16 +
17 +The functions here take caller-owned output buffers (``sparse_indices``,
18 +``end``, ``col_indices``) so the hot path allocates nothing; the indexer
19 +metadata builder owns them.
20 +"""
21 +
22 +import importlib.util
23 +
24 +import torch
25 +
26 +from vllm.model_executor.layers.indexer_topk import (
27 + deep_select_topk,
28 + get_deep_select_stride_requirement,
29 +)
30 +from vllm.platforms import current_platform
31 +from vllm.triton_utils import tl, triton
32 +from vllm.utils.deep_gemm import (
33 + fp8_fp4_paged_sparse_mqa_logits,
34 + fp8_fp4_sparse_mqa_logits,
35 + get_paged_sparse_mqa_logits_metadata,
36 + get_sparse_mqa_logits_metadata,
37 +)
38 +
39 +# Sparse-block sizes (in tokens) supported by the DeepGEMM kernels.
40 +SPARSE_BLOCK_KV_CHOICES = (16, 8)
41 +
42 +
43 +def has_deep_select() -> bool:
44 + """Whether the DeepSelect top-k extension is built and this GPU runs it."""
45 + return (
46 + importlib.util.find_spec("vllm._deepselect_C") is not None
47 + and current_platform.is_cuda()
48 + and current_platform.is_device_capability_family(100)
49 + )
50 +
51 +
52 +def check_deep_select_layout(num_sparse_cols: int, topk_tokens: int) -> None:
53 + """Raise unless DeepSelect accepts the sparse logits layout.
54 +
55 + The sparse kernels return contiguous bf16 ``[rows, num_sparse_cols]``
56 + logits and the top-k writes int32 ``[rows, topk_tokens]`` indices;
57 + DeepSelect needs both row strides aligned (1024B / 32B on SM100).
58 + """
59 + in_align, out_align = get_deep_select_stride_requirement()
60 + logits_row_bytes = num_sparse_cols * torch.bfloat16.itemsize
61 + if num_sparse_cols >= 2**23 or logits_row_bytes % in_align != 0:
62 + raise ValueError(
63 + f"DeepSelect needs the sparse logits row ({num_sparse_cols} bf16 "
64 + f"columns, {logits_row_bytes} B) to be {in_align}B-aligned."
65 + )
66 + if topk_tokens > 4096 or (topk_tokens * torch.int32.itemsize) % out_align != 0:
67 + raise ValueError(
68 + f"DeepSelect needs topk <= 4096 with a {out_align}B-aligned int32 "
69 + f"output row, got topk={topk_tokens}."
70 + )
71 +
72 +
73 +_INT32_BIG = tl.constexpr(2**31 - 1)
74 +
75 +
76 +@triton.jit
77 +def _expand_candidates_kernel(
78 + cand_ptr,
79 + cand_stride,
80 + ks_ptr,
81 + ks_stride,
82 + ke_ptr,
83 + ke_stride,
84 + si_ptr,
85 + si_stride,
86 + end_ptr,
87 + end_stride,
88 + K: tl.constexpr,
89 + S: tl.constexpr,
90 + RATIO: tl.constexpr,
91 + SBK: tl.constexpr,
92 +):
93 + """One program per row: expand + filter + sort candidate blocks.
94 +
95 + Each row's valid sparse blocks come out sorted ascending (no dedup pass:
96 + production candidates are unique, and the kernel tolerates repeats),
97 + padded by repeating the last valid block — DeepGEMM's padding convention.
98 + ``end`` counts valid sparse-token columns; valid columns form a prefix
99 + because block ids are sorted.
100 + """
101 + row = tl.program_id(0)
102 + ks = tl.load(ks_ptr + row * ks_stride)
103 + ke = tl.load(ke_ptr + row * ke_stride)
104 + cols = tl.arange(0, K)
105 + cand = tl.load(cand_ptr + row * cand_stride + cols) # K is pow2, no mask
106 + if RATIO == 2:
107 + cand = tl.interleave(cand * RATIO, cand * RATIO + 1)
108 +
109 + # Keep only blocks inside the row's KV range [0, ke - ks); the lower
110 + # bound is free since candidates are non-negative.
111 + kv_len = tl.maximum(ke - ks, 0)
112 + max_block = (kv_len + SBK - 1) // SBK
113 + valid = (cand >= 0) & (cand < max_block)
114 + vals = tl.where(valid, cand, _INT32_BIG)
115 + vals = tl.sort(vals)
116 +
117 + # Anchor at the row's ks rounded down to a sparse-block boundary
118 + # (identity for the paged path, where ks == 0).
119 + base = ks // SBK
120 + si = tl.where(vals != _INT32_BIG, vals, 0) + base
121 + n_valid = tl.sum(valid.to(tl.int32))
122 + last_rel = tl.max(tl.where(valid, cand, -1))
123 + pad_val = tl.where(n_valid > 0, last_rel + base, base)
124 + tl.store(
125 + si_ptr + row * si_stride + tl.arange(0, S),
126 + tl.where(tl.arange(0, S) < n_valid, si, pad_val),
127 + )
128 +
129 + # Valid sparse-token column count; only the last block can be partial.
130 + # Absolute start of the last block: (base + last_rel) * SBK + ks % SBK,
131 + # which is exactly last_rel * SBK + ks.
132 + last_start = last_rel * SBK + ks
133 + fill = tl.minimum(tl.maximum(ke - last_start, 1), SBK)
134 + end = tl.where(n_valid > 0, (n_valid - 1) * SBK + fill, 0)
135 + tl.store(end_ptr + row * end_stride, end)
136 +
137 +
138 +@triton.jit
139 +def _sparse_topk_remap_kernel(
140 + col_ptr,
141 + col_stride,
142 + si_ptr,
143 + si_stride,
144 + ks_ptr,
145 + ks_stride,
146 + out_ptr,
147 + out_stride,
148 + k,
149 + SBK: tl.constexpr,
150 + K_POW2: tl.constexpr,
151 +):
152 + """Remap sparse top-k columns to request-local positions.
153 +
154 + Column ``j * SBK + o`` scores the token at
155 + ``(si[row, j] - ks // SBK) * SBK + o``; -1 columns stay -1.
156 + """
157 + row = tl.program_id(0)
158 + ks_base = tl.load(ks_ptr + row * ks_stride) // SBK
159 + cols = tl.arange(0, K_POW2)
160 + c = tl.load(col_ptr + row * col_stride + cols, mask=cols < k, other=-1)
161 + valid = c >= 0
162 + cc = tl.where(valid, c, 0)
163 + blocks = tl.load(si_ptr + row * si_stride + cc // SBK)
164 + pos = (blocks - ks_base) * SBK + cc % SBK
165 + tl.store(out_ptr + row * out_stride + cols, tl.where(valid, pos, -1), mask=cols < k)
166 +
167 +
168 +def pick_sparse_block_kv(candidate_block_size: int) -> int:
169 + """Largest sparse-block size (tokens) dividing the candidate block size."""
170 + for sbk in SPARSE_BLOCK_KV_CHOICES:
171 + if candidate_block_size % sbk == 0:
172 + return sbk
173 + raise ValueError(
174 + f"candidate_block_size={candidate_block_size} is not a multiple of "
175 + f"any supported sparse_block_kv {SPARSE_BLOCK_KV_CHOICES}."
176 + )
177 +
178 +
179 +def candidate_blocks_to_sparse_indices(
180 + candidate_blocks: torch.Tensor,
181 + row_ks: torch.Tensor,
182 + row_ke: torch.Tensor,
183 + candidate_block_size: int,
184 + sparse_block_kv: int,
185 + out: tuple[torch.Tensor, torch.Tensor] | None = None,
186 +) -> tuple[torch.Tensor, torch.Tensor]:
187 + """Expand request-local candidate blocks into DeepGEMM sparse block ids.
188 +
189 + Args:
190 + candidate_blocks: [rows, K] int32 request-local candidate block ids
191 + (-1 padded), in units of ``candidate_block_size`` positions. K
192 + must be a power of two (the in-kernel sort); rows may be strided.
193 + row_ks/row_ke: [rows] int32 per-row K range; bounds are in the same
194 + (packed-workspace) coordinates the sparse kernel iterates over.
195 + Pass zeros for the paged path, whose blocks are context-relative.
196 + candidate_block_size: Positions per candidate block.
197 + sparse_block_kv: Positions per sparse block (8 or 16).
198 + out: Optional ``(sparse_indices, end)`` buffers to write into,
199 + shaped [rows, K * ratio] and [rows], both int32.
200 +
201 + Returns:
202 + sparse_indices: [rows, K * ratio] int32 block ids for DeepGEMM,
203 + sorted ascending per row with the valid prefix first (unique in
204 + production, where candidates are unique) and the remaining slots
205 + repeating the last valid block (the kernel's padding convention).
206 + With the unaligned-ks variant, block ``i`` starts at
207 + ``i * sparse_block_kv + ks % sparse_block_kv``.
208 + end: [rows] int32 count of valid sparse-token columns; valid columns
209 + always form a prefix because the block ids are sorted.
210 + """
211 + rows, num_candidates = candidate_blocks.shape
212 + ratio = candidate_block_size // sparse_block_kv
213 + assert ratio * sparse_block_kv == candidate_block_size and ratio in (1, 2), (
214 + candidate_block_size,
215 + sparse_block_kv,
216 + )
217 + assert num_candidates & (num_candidates - 1) == 0, num_candidates
218 + num_sparse = num_candidates * ratio
219 + if out is None:
220 + device = candidate_blocks.device
221 + out = (
222 + torch.empty(rows, num_sparse, dtype=torch.int32, device=device),
223 + torch.empty(rows, dtype=torch.int32, device=device),
224 + )
225 + sparse_indices, end = out
226 + assert sparse_indices.shape == (rows, num_sparse) and end.shape == (rows,)
227 + assert candidate_blocks.stride(1) == 1 and sparse_indices.stride(1) == 1
228 + for t in (candidate_blocks, row_ks, row_ke, sparse_indices, end):
229 + assert t.dtype == torch.int32, t.dtype
230 +
231 + _expand_candidates_kernel[(rows,)](
232 + candidate_blocks,
233 + candidate_blocks.stride(0),
234 + row_ks,
235 + row_ks.stride(0),
236 + row_ke,
237 + row_ke.stride(0),
238 + sparse_indices,
239 + sparse_indices.stride(0),
240 + end,
241 + end.stride(0),
242 + K=num_candidates,
243 + S=num_sparse,
244 + RATIO=ratio,
245 + SBK=sparse_block_kv,
246 + num_warps=16,
247 + )
248 + return sparse_indices, end
249 +
250 +
251 +def sparse_topk_remap(
252 + logits: torch.Tensor,
253 + sparse_indices: torch.Tensor,
254 + end: torch.Tensor,
255 + row_ks: torch.Tensor,
256 + sparse_block_kv: int,
257 + topk_tokens: int,
258 + topk_indices: torch.Tensor,
259 + *,
260 + col_indices: torch.Tensor,
261 +) -> None:
262 + """Row top-k over sparse logits, remapped to request-local positions.
263 +
264 + Column ``j * sparse_block_kv + o`` of ``logits`` scores the token at
265 + request-local position
266 + ``(sparse_indices[row, j] - ks // sparse_block_kv) * sparse_block_kv + o``.
267 +
268 + DeepSelect consumes the bf16 logits as they come out of the sparse
269 + kernels; the valid columns form a prefix (block ids are sorted), so the
270 + per-row ``end`` bounds drive it directly. Slots beyond a row's valid
271 + count come back as -1.
272 +
273 + Args:
274 + col_indices: [rows, topk_tokens] int32 scratch for the sparse-column
275 + top-k result (row stride 32B-aligned for DeepSelect).
276 + """
277 + rows, width = logits.shape
278 + assert logits.dtype == torch.bfloat16 and width >= topk_tokens
279 + assert end.dtype == torch.int32 and end.shape == (rows,)
280 + cols = col_indices[:rows]
281 + deep_select_topk(logits, topk_tokens, end=end, output_idx=cols)
282 + _sparse_topk_remap_kernel[(rows,)](
283 + cols,
284 + cols.stride(0),
285 + sparse_indices,
286 + sparse_indices.stride(0),
287 + row_ks,
288 + row_ks.stride(0),
289 + topk_indices,
290 + topk_indices.stride(0),
291 + topk_tokens,
292 + SBK=sparse_block_kv,
293 + K_POW2=triton.next_power_of_2(topk_tokens),
294 + )
295 +
296 +
297 +def sparse_mqa_logits_prefill_chunk(
298 + q: torch.Tensor,
299 + q_scale: torch.Tensor,
300 + k_quant: torch.Tensor,
301 + k_scale: torch.Tensor,
302 + weights: torch.Tensor,
303 + cu_seqlen_ks: torch.Tensor,
304 + cu_seqlen_ke: torch.Tensor,
305 + candidate_blocks: torch.Tensor,
306 + candidate_block_size: int,
307 + sparse_block_kv: int,
308 + topk_tokens: int,
309 + topk_indices: torch.Tensor,
310 + *,
311 + sparse_indices: torch.Tensor,
312 + end: torch.Tensor,
313 + col_indices: torch.Tensor,
314 + kernel_metadata: torch.Tensor | None = None,
315 +) -> torch.Tensor:
316 + """Sparse-MQA logits + top-k for one prefill chunk (packed KV workspace).
317 +
318 + Args:
319 + q: [rows, H, D] packed Q values (MXFP4 viewed as int8).
320 + q_scale: [rows, H] int32 packed UE8M0 Q scales.
321 + k_quant: [total_kv, D] packed K workspace (MXFP4 viewed as int8).
322 + k_scale: [total_kv] int32 packed UE8M0 K scales.
323 + weights: [rows, H] bf16 per-head weights; the sparse kernels take
324 + bf16 and do not fold the Q scale in.
325 + cu_seqlen_ks/cu_seqlen_ke: [rows] int32 per-token K bounds in the
326 + packed workspace.
327 + candidate_blocks: [rows, K] int32 request-local candidate block ids.
328 + topk_indices: [rows, topk_tokens] output buffer.
329 + sparse_indices/end/col_indices: Caller-owned scratch, see
330 + `candidate_blocks_to_sparse_indices` and `sparse_topk_remap`.
331 + kernel_metadata: DeepGEMM schedule from a previous call with the same
332 + candidates and bounds (i.e. another indexer layer in the same
333 + step). When given, the candidate expansion is skipped and
334 + ``sparse_indices``/``end`` are assumed to be up to date.
335 +
336 + Returns:
337 + The DeepGEMM schedule metadata used, for reuse by later layers.
338 + """
339 + assert weights.dtype == torch.bfloat16, weights.dtype
340 + if kernel_metadata is None:
341 + candidate_blocks_to_sparse_indices(
342 + candidate_blocks,
343 + cu_seqlen_ks,
344 + cu_seqlen_ke,
345 + candidate_block_size,
346 + sparse_block_kv,
347 + out=(sparse_indices, end),
348 + )
349 + kernel_metadata = get_sparse_mqa_logits_metadata(
350 + cu_seqlen_ks,
351 + cu_seqlen_ke,
352 + k_quant.shape[0],
353 + sparse_indices,
354 + q.dtype,
355 + sparse_block_kv,
356 + )
357 + logits = fp8_fp4_sparse_mqa_logits(
358 + (q, q_scale),
359 + (k_quant, k_scale),
360 + weights,
361 + kernel_metadata,
362 + sparse_indices.shape[1],
363 + sparse_block_kv,
364 + )
365 + sparse_topk_remap(
366 + logits,
367 + sparse_indices,
368 + end,
369 + cu_seqlen_ks,
370 + sparse_block_kv,
371 + topk_tokens,
372 + topk_indices,
373 + col_indices=col_indices,
374 + )
375 + return kernel_metadata
376 +
377 +
378 +def sparse_mqa_logits_paged_decode(
379 + q: torch.Tensor,
380 + q_scale: torch.Tensor,
381 + kv_cache: torch.Tensor,
382 + weights: torch.Tensor,
383 + context_lens: torch.Tensor,
384 + block_table: torch.Tensor,
385 + row_indices: torch.Tensor,
386 + candidate_blocks: torch.Tensor,
387 + candidate_block_size: int,
388 + sparse_block_kv: int,
389 + topk_tokens: int,
390 + topk_indices: torch.Tensor,
391 + *,
392 + row_ks: torch.Tensor,
393 + sparse_indices: torch.Tensor,
394 + end: torch.Tensor,
395 + col_indices: torch.Tensor,
396 + kernel_metadata: torch.Tensor | None = None,
397 +) -> torch.Tensor:
398 + """Sparse-MQA logits + top-k for paged decode rows.
399 +
400 + Every query is one row (spec-decode rows are flattened by the metadata
401 + builder), so ``block_table`` is already per row.
402 +
403 + Args:
404 + q: [rows, 1, H, D] packed Q values (MXFP4 viewed as int8).
405 + q_scale: [rows, 1, H] int32 packed UE8M0 Q scales.
406 + kv_cache: 4D [num_pages, page_kv, 1, head_bytes] uint8 fused cache
407 + view; the page stride must be 512B-aligned.
408 + weights: [rows, H] bf16 per-head weights.
409 + context_lens: [rows] int32 per-row (compressed) context lengths.
410 + block_table: [rows, P] int32 per-row page table, ``stride(-1) == 1``.
411 + row_indices: [rows] int32 row -> request map (pairing only affects
412 + scheduling).
413 + candidate_blocks: [rows, K] int32 candidate block ids.
414 + topk_indices: [rows, topk_tokens] output buffer.
415 + row_ks: [rows] int32 zeros (paged blocks are context-relative).
416 + sparse_indices/end/col_indices: Caller-owned scratch.
417 + kernel_metadata: See `sparse_mqa_logits_prefill_chunk`.
418 +
419 + Returns:
420 + The DeepGEMM schedule metadata used, for reuse by later layers.
421 + """
422 + rows = q.shape[0]
423 + page_kv = kv_cache.shape[1]
424 + assert weights.dtype == torch.bfloat16, weights.dtype
425 + assert kv_cache.stride(0) % 512 == 0, (
426 + "paged cache page stride must be 512B-aligned: "
427 + f"shape={tuple(kv_cache.shape)} stride={kv_cache.stride()}"
428 + )
429 + assert page_kv % sparse_block_kv == 0, (page_kv, sparse_block_kv)
430 + assert block_table.shape[0] == rows and block_table.stride(-1) == 1
431 + assert context_lens.shape == (rows,) and context_lens.is_contiguous()
432 +
433 + if kernel_metadata is None:
434 + candidate_blocks_to_sparse_indices(
435 + candidate_blocks,
436 + row_ks,
437 + context_lens,
438 + candidate_block_size,
439 + sparse_block_kv,
440 + out=(sparse_indices, end),
441 + )
442 + kernel_metadata = get_paged_sparse_mqa_logits_metadata(
443 + context_lens,
444 + block_table,
445 + row_indices,
446 + page_kv,
447 + sparse_indices,
448 + q.dtype,
449 + sparse_block_kv,
450 + )
451 + logits = fp8_fp4_paged_sparse_mqa_logits(
452 + (q, q_scale),
453 + kv_cache,
454 + weights,
455 + kernel_metadata,
456 + sparse_indices.shape[1],
457 + sparse_block_kv,
458 + )
459 + sparse_topk_remap(
460 + logits,
461 + sparse_indices,
462 + end,
463 + row_ks,
464 + sparse_block_kv,
465 + topk_tokens,
466 + topk_indices,
467 + col_indices=col_indices,
468 + )
469 + return kernel_metadata
vllm/model_executor/layers/sparse_mqa_indexer.py new
+209
@@ -0,0 +1,209 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Sparse Attention Indexer that scores only the candidate blocks.
4 +
5 +DeepSeek V4.1 two-level selection: the candidate-source indexer publishes
6 +the top candidate blocks and later indexers pick their top-k inside them.
7 +`SparseAttnIndexer` does that by computing dense logits over the whole
8 +context and masking; this layer instead calls DeepGEMM's sparse MQA-logits
9 +kernels on the candidate blocks only, so the work is O(candidate blocks)
10 +instead of O(context). It requires the `DeepseekV41SparseIndexerBackend`
11 +metadata (see ``AttentionConfig.indexer_sparse_logits``).
12 +"""
13 +
14 +import torch
15 +from torch import nn
16 +
17 +import vllm.envs as envs
18 +from vllm import _custom_ops as ops
19 +from vllm.forward_context import get_forward_context
20 +from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import (
21 + has_deep_select,
22 + sparse_mqa_logits_paged_decode,
23 + sparse_mqa_logits_prefill_chunk,
24 +)
25 +from vllm.model_executor.layers.sparse_attn_indexer import (
26 + _gather_workspace_shapes,
27 + kv_cache_as_quant_view,
28 +)
29 +from vllm.platforms import current_platform
30 +from vllm.utils.deep_gemm import has_deep_gemm_sparse_mqa
31 +from vllm.v1.attention.backends.mla.indexer import (
32 + DeepseekV32IndexerPrefillChunkMetadata,
33 +)
34 +from vllm.v1.attention.backends.mla.sparse_indexer import (
35 + DeepseekV41SparseIndexerMetadata,
36 +)
37 +from vllm.v1.worker.workspace import current_workspace_manager
38 +
39 +
40 +def _prefill_k_workspaces(
41 + total_seq_lens: int, head_dim: int
42 +) -> tuple[torch.Tensor, torch.Tensor]:
43 + """The packed MXFP4 K-gather workspace ``(values, scales)``, shared with
44 + the dense indexer layers of the same model."""
45 + values_spec, scales_spec = _gather_workspace_shapes(
46 + total_seq_lens, head_dim, current_platform.fp8_dtype(), use_fp4_cache=True
47 + )
48 + k_quant, k_scale = current_workspace_manager().get_simultaneous(
49 + values_spec, scales_spec
50 + )
51 + return k_quant, k_scale
52 +
53 +
54 +def _gather_prefill_chunk_k(
55 + kv_cache: torch.Tensor,
56 + k_quant_full: torch.Tensor,
57 + k_scale_full: torch.Tensor,
58 + chunk: DeepseekV32IndexerPrefillChunkMetadata,
59 +) -> tuple[torch.Tensor, torch.Tensor]:
60 + """Gather one prefill chunk's paged K into the packed workspace."""
61 + assert chunk.local_cu_seq_lens is not None
62 + k_quant = k_quant_full[: chunk.max_local_total_seq_lens]
63 + k_scale = k_scale_full[: chunk.max_local_total_seq_lens]
64 + if not chunk.skip_kv_gather and chunk.local_total_seq_lens > 0:
65 + ops.cp_gather_indexer_k_quant_cache(
66 + kv_cache,
67 + k_quant,
68 + k_scale,
69 + chunk.block_table,
70 + chunk.local_cu_seq_lens,
71 + )
72 + return k_quant, k_scale
73 +
74 +
75 +class SparseMQAIndexer(nn.Module):
76 + """Candidate-consuming indexer on DeepGEMM's sparse MQA-logits kernels.
77 +
78 + Only valid for indexer layers that read candidate blocks with the MXFP4
79 + indexer cache on SM100. The K cache is written by the model before this
80 + runs; ``forward`` takes the same arguments as `SparseAttnIndexer` so the
81 + attention layer can call either.
82 + """
83 +
84 + weights_dtype = torch.bfloat16
85 + """Per-head weights dtype the sparse kernels take. The fused Q RoPE-quant
86 + kernel writes it directly so no cast runs per step."""
87 +
88 + def __init__(
89 + self,
90 + k_cache,
91 + topk_tokens: int,
92 + head_dim: int,
93 + max_total_seq_len: int,
94 + topk_indices_buffer: torch.Tensor,
95 + candidate_blocks: torch.Tensor,
96 + candidate_block_size: int,
97 + ):
98 + super().__init__()
99 + if not (
100 + current_platform.is_cuda()
101 + and current_platform.is_device_capability_family(100)
102 + and has_deep_gemm_sparse_mqa()
103 + ):
104 + raise ValueError(
105 + "SparseMQAIndexer requires an SM100-class GPU and DeepGEMM >= 2.8."
106 + )
107 + if not has_deep_select():
108 + raise ValueError(
109 + "SparseMQAIndexer requires the DeepSelect top-k extension "
110 + "(vllm._deepselect_C)."
111 + )
112 + self.k_cache = k_cache
113 + self.topk_tokens = topk_tokens
114 + self.head_dim = head_dim
115 + self.max_total_seq_len = max_total_seq_len
116 + self.topk_indices_buffer = topk_indices_buffer
117 + self.candidate_blocks = candidate_blocks
118 + self.candidate_block_size = candidate_block_size
119 +
120 + def _reserve_workspaces(self, device: torch.device) -> None:
121 + """Profiling run: claim the K-gather workspace and the peak sparse
122 + logits allocation so the memory estimate covers them."""
123 + _prefill_k_workspaces(self.max_total_seq_len, self.head_dim)
124 + max_logits_bytes = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024
125 + torch.empty(max_logits_bytes, dtype=torch.uint8, device=device)
126 +
127 + def forward(
128 + self,
129 + hidden_states: torch.Tensor,
130 + q_quant: tuple[torch.Tensor, torch.Tensor],
131 + k: torch.Tensor | None,
132 + weights: torch.Tensor,
133 + ) -> torch.Tensor:
134 + attn_metadata = get_forward_context().attn_metadata
135 + if not isinstance(attn_metadata, dict):
136 + self._reserve_workspaces(hidden_states.device)
137 + return self.topk_indices_buffer
138 +
139 + metadata = attn_metadata[self.k_cache.prefix]
140 + assert isinstance(metadata, DeepseekV41SparseIndexerMetadata), (
141 + "SparseMQAIndexer needs DeepseekV41SparseIndexerBackend metadata"
142 + )
143 + assert k is None, "the model writes the indexer K cache"
144 + q_values, q_scale = q_quant
145 + kv_cache = self.k_cache.kv_cache
146 + topk_tokens = self.topk_tokens
147 + topk_indices_buffer = self.topk_indices_buffer
148 + topk_indices_buffer[: hidden_states.shape[0]] = -1
149 + sparse_block_kv = metadata.sparse_block_kv
150 +
151 + if metadata.num_prefills > 0:
152 + prefill = metadata.prefill
153 + sparse_chunks = metadata.sparse_prefill
154 + assert prefill is not None and sparse_chunks is not None
155 + k_quant_full, k_scale_full = _prefill_k_workspaces(
156 + self.max_total_seq_len, self.head_dim
157 + )
158 + for chunk, rows in zip(prefill.chunks, sparse_chunks):
159 + k_quant, k_scale = _gather_prefill_chunk_k(
160 + kv_cache, k_quant_full, k_scale_full, chunk
161 + )
162 + if chunk.local_total_seq_lens == 0:
163 + continue # buffer already holds -1
164 + start, end = chunk.token_start, chunk.token_end
165 + rows.kernel_metadata = sparse_mqa_logits_prefill_chunk(
166 + q_values[start:end].view(torch.int8),
167 + q_scale[start:end],
168 + k_quant.view(torch.int8),
169 + k_scale.view(torch.int32).squeeze(-1),
170 + weights[start:end],
171 + chunk.cu_seqlen_ks,
172 + chunk.cu_seqlen_ke,
173 + self.candidate_blocks[start:end],
174 + self.candidate_block_size,
175 + sparse_block_kv,
176 + topk_tokens,
177 + topk_indices_buffer[start:end, :topk_tokens],
178 + sparse_indices=rows.sparse_indices,
179 + end=rows.end,
180 + col_indices=rows.col_indices,
181 + kernel_metadata=rows.kernel_metadata,
182 + )
183 +
184 + if metadata.num_decodes > 0:
185 + decode_rows = metadata.sparse_decode
186 + assert decode_rows is not None and decode_rows.block_table is not None
187 + assert decode_rows.row_indices is not None
188 + num_rows = decode_rows.row_ke.shape[0]
189 + if num_rows > 0:
190 + decode_rows.kernel_metadata = sparse_mqa_logits_paged_decode(
191 + q_values[:num_rows].view(torch.int8).unsqueeze(1),
192 + q_scale[:num_rows].unsqueeze(1),
193 + kv_cache_as_quant_view(kv_cache, self.head_dim, use_fp4_cache=True),
194 + weights[:num_rows],
195 + decode_rows.row_ke,
196 + decode_rows.block_table,
197 + decode_rows.row_indices,
198 + self.candidate_blocks[:num_rows],
199 + self.candidate_block_size,
200 + sparse_block_kv,
201 + topk_tokens,
202 + topk_indices_buffer[:num_rows, :topk_tokens],
203 + row_ks=decode_rows.row_ks,
204 + sparse_indices=decode_rows.sparse_indices,
205 + end=decode_rows.end,
206 + col_indices=decode_rows.col_indices,
207 + kernel_metadata=decode_rows.kernel_metadata,
208 + )
209 + return topk_indices_buffer
vllm/models/deepseek_v4/common/ops/fused_indexer_q.py
+22 -2
@@ -309,12 +309,22 @@ class FusedIndexerQRopeQuantTritonKernel(
309 )
310
311
312 +def _indexer_weights_out_dtypes(vllm_config: Any) -> tuple[torch.dtype, ...]:
313 + """Weights dtypes the model's indexer layers ask for: fp32 for the dense
314 + scoring kernels, plus bf16 when the DeepSeek V4.1 sparse-logits indexer
315 + (`SparseMQAIndexer`) is enabled."""
316 + if vllm_config.attention_config.indexer_sparse_logits:
317 + return (torch.float32, torch.bfloat16)
318 + return (torch.float32,)
319 +
320 +
321 class FusedIndexerQRopeMxFp4TritonKernel(
322 VllmTritonJitKernel["FusedIndexerQRopeMxFp4TritonKernel.CompileKey"]
323 ):
324 @dataclass(frozen=True)
325 class CompileKey:
326 dtype: torch.dtype
327 + weights_out_dtype: torch.dtype
328 num_heads: int
329 index_q_half_rot_dim: int
330 index_q_head_dim: int
@@ -434,12 +444,14 @@ class FusedIndexerQRopeMxFp4TritonKernel(
444 self,
445 *,
446 dtype: torch.dtype,
447 + weights_out_dtype: torch.dtype,
448 num_heads: int,
449 head_dim: int,
450 rope_dim: int,
451 ) -> CompileKey:
452 return self.CompileKey(
453 dtype=dtype,
454 + weights_out_dtype=weights_out_dtype,
455 num_heads=num_heads,
456 index_q_half_rot_dim=rope_dim // 2,
457 index_q_head_dim=head_dim,
@@ -456,6 +468,7 @@ class FusedIndexerQRopeMxFp4TritonKernel(
468
469 return self._trace_dispatch(self.dispatch)(
470 dtype=vllm_config.model_config.dtype,
471 + weights_out_dtype=_indexer_weights_out_dtypes(vllm_config),
472 num_heads=num_heads,
473 head_dim=head_dim,
474 rope_dim=rope_dim,
@@ -491,7 +504,7 @@ class FusedIndexerQRopeMxFp4TritonKernel(
504 shape=(1, compile_key.num_heads, scale_head_dim),
505 ),
506 index_weights_out=TritonWarmupTensor(
494 - torch.float32,
507 + compile_key.weights_out_dtype,
508 shape=(1, compile_key.num_heads),
509 ),
510 )
@@ -540,12 +553,17 @@ def fused_indexer_q_rope_quant(
553 index_weights_softmax_scale: float,
554 index_weights_head_scale: float,
555 use_fp4: bool = False,
556 + weights_out_dtype: torch.dtype = torch.float32,
557 ) -> tuple[
558 torch.Tensor | tuple[torch.Tensor, torch.Tensor],
559 torch.Tensor,
560 ]:
561 """Fused RoPE + quantize Q for the sparse indexer.
562
563 + ``weights_out_dtype`` is the dtype the downstream scoring kernel takes:
564 + fp32 for the dense MQA-logits kernels, bf16 for DeepGEMM's sparse
565 + MQA-logits kernels (CUDA MXFP4 path only).
566 +
567 Weight-fold semantics (important — the two paths differ):
568
569 FP8 path (use_fp4=False, default):
@@ -577,7 +595,7 @@ def fused_indexer_q_rope_quant(
595 num_index_q_heads = index_q.shape[1]
596 index_q_head_dim = index_q.shape[2]
597
580 - index_weights_out = torch.empty_like(index_weights, dtype=torch.float32)
598 + index_weights_out = torch.empty_like(index_weights, dtype=weights_out_dtype)
599
600 if use_fp4:
601 assert index_q_head_dim % MXFP4_BLOCK_SIZE == 0, (
@@ -613,6 +631,7 @@ def fused_indexer_q_rope_quant(
631 weights_out=index_weights_out,
632 )
633 elif current_platform.is_xpu():
634 + assert weights_out_dtype == torch.float32, weights_out_dtype
635 torch.ops.vllm.xpu_deepseek_fused_indexer_q_rope_mxfp4(
636 index_q,
637 positions,
@@ -647,6 +666,7 @@ def fused_indexer_q_rope_quant(
666 index_q_scale.view(torch.int32).squeeze(-1),
667 ), index_weights_out
668
669 + assert weights_out_dtype == torch.float32, weights_out_dtype
670 fp8_dtype = current_platform.fp8_dtype()
671 use_fnuz = fp8_dtype == torch.float8_e4m3fnuz
672 fp8_max = 224.0 if use_fnuz else 448.0
vllm/models/deepseek_v4/nvidia/ops/fused_indexer_q_cutedsl.py
+15 -2
@@ -165,6 +165,7 @@ class IndexerQMxFp4Kernel(VllmCuTeDSLJitKernel["IndexerQMxFp4Kernel.CompileKey"]
165 rope_dim: int
166 num_heads: int
167 cos_sin_dtype: type[cutlass.Numeric]
168 + weights_out_dtype: type[cutlass.Numeric]
169 coarsen: int
170
171 @staticmethod
@@ -177,6 +178,7 @@ class IndexerQMxFp4Kernel(VllmCuTeDSLJitKernel["IndexerQMxFp4Kernel.CompileKey"]
178 # process multiple heads at the same time to armotize RoPE load costs
179 coarsen = compile_key.coarsen
180 tb_size = IndexerQMxFp4Kernel.tb_size
181 + weights_out_dtype = compile_key.weights_out_dtype
182 assert num_heads % coarsen == 0
183
184 # later we will use 32B load = 16 BF16 elems
@@ -287,7 +289,7 @@ class IndexerQMxFp4Kernel(VllmCuTeDSLJitKernel["IndexerQMxFp4Kernel.CompileKey"]
289 weight_head_id = global_tid % num_heads
290 weights_out[weight_token_id, weight_head_id] = (
291 weights[weight_token_id, weight_head_id].to(Float32) * scale
290 - )
292 + ).to(weights_out_dtype)
293
294 @cute.jit
295 def host_entrypoint(
@@ -320,11 +322,13 @@ class IndexerQMxFp4Kernel(VllmCuTeDSLJitKernel["IndexerQMxFp4Kernel.CompileKey"]
322 self,
323 *,
324 cos_sin_dtype: type[cutlass.Numeric],
325 + weights_out_dtype: type[cutlass.Numeric],
326 **compile_key_fields: int,
327 ) -> CompileKey:
328 return self.CompileKey(
329 **compile_key_fields,
330 cos_sin_dtype=cos_sin_dtype,
331 + weights_out_dtype=weights_out_dtype,
332 )
333
334 def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
@@ -340,11 +344,17 @@ class IndexerQMxFp4Kernel(VllmCuTeDSLJitKernel["IndexerQMxFp4Kernel.CompileKey"]
344 if num_heads <= 0 or head_dim <= 0 or rope_dim <= 0:
345 return []
346
347 + # The DeepSeek V4.1 sparse-logits indexer takes bf16 weights; the
348 + # dense scoring kernels take fp32.
349 + weights_out_dtypes: tuple[type[cutlass.Numeric], ...] = (Float32,)
350 + if vllm_config.attention_config.indexer_sparse_logits:
351 + weights_out_dtypes = (Float32, BFloat16)
352 return self._trace_dispatch(self.dispatch)(
353 head_dim=head_dim,
354 rope_dim=rope_dim,
355 num_heads=num_heads,
356 cos_sin_dtype=Float32,
357 + weights_out_dtype=weights_out_dtypes,
358 coarsen=(1, 4),
359 )
360
@@ -380,7 +390,9 @@ class IndexerQMxFp4Kernel(VllmCuTeDSLJitKernel["IndexerQMxFp4Kernel.CompileKey"]
390 divisibility=4,
391 )
392 weights_out = make_fake_tensor(
383 - Float32, (num_tokens, compile_key.num_heads), divisibility=4
393 + compile_key.weights_out_dtype,
394 + (num_tokens, compile_key.num_heads),
395 + divisibility=4,
396 )
397 return (
398 positions,
@@ -413,6 +425,7 @@ class IndexerQMxFp4Kernel(VllmCuTeDSLJitKernel["IndexerQMxFp4Kernel.CompileKey"]
425 rope_dim=cos_sin_cache.shape[-1],
426 num_heads=num_heads,
427 cos_sin_dtype=torch_to_cute_dtype(cos_sin_cache.dtype),
428 + weights_out_dtype=torch_to_cute_dtype(weights_out.dtype),
429 coarsen=1 if num_tokens < 512 else 4,
430 )
431 launch_args = (
vllm/models/deepseek_v41/attention.py
+67 -20
@@ -4,6 +4,7 @@
4 DeepseekV4 MLA Attention Layer
5 """
6
7 +import math
8 from abc import ABC, abstractmethod
9 from collections.abc import Callable
10 from functools import cached_property
@@ -25,6 +26,7 @@ from vllm.model_executor.layers.linear import (
26 RowParallelLinear,
27 )
28 from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer
29 +from vllm.model_executor.layers.sparse_mqa_indexer import SparseMQAIndexer
30 from vllm.models.common.ops import fused_q_kv_rmsnorm
31 from vllm.models.deepseek_v41.common.ops import (
32 MXFP4_BLOCK_SIZE,
@@ -63,6 +65,9 @@ from vllm.v1.attention.backends.mla.indexer import (
65 dsa_indexer_uses_fp4,
66 get_max_prefill_buffer_size,
67 )
68 +from vllm.v1.attention.backends.mla.sparse_indexer import (
69 + DeepseekV41SparseIndexerBackend,
70 +)
71 from vllm.v1.attention.backends.mla.sparse_swa import DeepseekV4SWACache
72 from vllm.v1.kv_cache_interface import (
73 KVCacheSpec,
@@ -997,7 +1002,9 @@ class DeepseekV4IndexerCache(torch.nn.Module, AttentionLayerBase):
1002 self.cache_config = cache_config
1003 self.dtype = dtype
1004 self.compress_ratio = compress_ratio
1000 - compilation_config = get_current_vllm_config().compilation_config
1005 + vllm_config = get_current_vllm_config()
1006 + self.sparse_logits = vllm_config.attention_config.indexer_sparse_logits
1007 + compilation_config = vllm_config.compilation_config
1008 if prefix in compilation_config.static_forward_context:
1009 raise ValueError(f"Duplicate layer name: {prefix}")
1010 compilation_config.static_forward_context[prefix] = self
@@ -1010,6 +1017,9 @@ class DeepseekV4IndexerCache(torch.nn.Module, AttentionLayerBase):
1017 # head_dim already carries the fp8 scale padding
1018 # tokens_per_state=1 for V3.2, >1 for DeepseekV4; same cache layout.
1019 uses_fp8_ds_mla_layout = vllm_config.cache_config.cache_dtype == "fp8_ds_mla"
1020 + page_alignment = (
1021 + 576 if uses_fp8_ds_mla_layout and not _use_v41_mxfp8_kv_record() else 512
1022 + )
1023 return MLAAttentionSpec(
1024 block_size=self.cache_config.block_size,
1025 num_kv_heads=1,
@@ -1020,16 +1030,21 @@ class DeepseekV4IndexerCache(torch.nn.Module, AttentionLayerBase):
1030 # (the block stride is the sum of every page in the group), so it
1031 # has to carry the same padding: the packed record's TMA stride,
1032 # else 512B for FlashInfer sparse (#44577).
1023 - alignment=(
1024 - 576
1025 - if uses_fp8_ds_mla_layout and not _use_v41_mxfp8_kv_record()
1026 - else 512
1033 + alignment=page_alignment,
1034 + # DeepGEMM's paged sparse MQA logits address pages by the block
1035 + # stride, which packs every layer's page in this layout, and need
1036 + # it 512B-aligned; the main KV record read from the same blocks
1037 + # needs its own TMA stride, so keep both.
1038 + block_stride_alignment=(
1039 + math.lcm(512, page_alignment) if self.sparse_logits else None
1040 ),
1041 )
1042
1043 def forward(self): ...
1044
1045 def get_attn_backend(self) -> type[AttentionBackend]:
1046 + if self.sparse_logits:
1047 + return DeepseekV41SparseIndexerBackend
1048 return DeepseekV41IndexerBackend
1049
1050
@@ -1129,21 +1144,52 @@ class DeepseekV4Indexer(nn.Module):
1144 self.k_norm = RMSNorm(self.head_dim, config.rms_norm_eps)
1145 self.k_cache = k_cache
1146
1132 - self.indexer_op = SparseAttnIndexer(
1133 - self.k_cache,
1134 - self.quant_block_size,
1135 - self.scale_fmt,
1136 - self.topk_tokens,
1137 - self.head_dim,
1138 - self.max_model_len,
1139 - self.max_total_seq_len,
1140 - self.topk_indices_buffer,
1141 - skip_k_cache_insert=True,
1142 - use_fp4_cache=self.use_fp4_kv,
1143 - compress_ratio=self.compress_ratio,
1144 - candidate_blocks=candidate_block_buffer,
1145 - candidate_block_size=candidate_block_size,
1146 - candidate_write=candidate_write,
1147 + # Candidate consumers can score only the candidate blocks (opt-in);
1148 + # the candidate source and pre-candidate indexers stay dense.
1149 + use_sparse_logits = (
1150 + vllm_config.attention_config.indexer_sparse_logits
1151 + and candidate_block_buffer is not None
1152 + and not candidate_write
1153 + )
1154 + self.indexer_op: SparseAttnIndexer | SparseMQAIndexer
1155 + if use_sparse_logits:
1156 + if not self.use_fp4_kv:
1157 + raise ValueError(
1158 + "attention_config.indexer_sparse_logits requires the MXFP4 "
1159 + "indexer cache (attention_config.indexer_kv_dtype='mxfp4')."
1160 + )
1161 + assert candidate_block_buffer is not None
1162 + assert topk_indices_buffer is not None
1163 + self.indexer_op = SparseMQAIndexer(
1164 + self.k_cache,
1165 + self.topk_tokens,
1166 + self.head_dim,
1167 + self.max_total_seq_len,
1168 + topk_indices_buffer,
1169 + candidate_block_buffer,
1170 + candidate_block_size,
1171 + )
1172 + else:
1173 + self.indexer_op = SparseAttnIndexer(
1174 + self.k_cache,
1175 + self.quant_block_size,
1176 + self.scale_fmt,
1177 + self.topk_tokens,
1178 + self.head_dim,
1179 + self.max_model_len,
1180 + self.max_total_seq_len,
1181 + self.topk_indices_buffer,
1182 + skip_k_cache_insert=True,
1183 + use_fp4_cache=self.use_fp4_kv,
1184 + compress_ratio=self.compress_ratio,
1185 + candidate_blocks=candidate_block_buffer,
1186 + candidate_block_size=candidate_block_size,
1187 + candidate_write=candidate_write,
1188 + )
1189 + # The fused Q kernel writes the per-head weights in the dtype the
1190 + # scoring kernels take, so no cast runs per step.
1191 + self.indexer_weights_dtype = (
1192 + SparseMQAIndexer.weights_dtype if use_sparse_logits else torch.float32
1193 )
1194
1195 def _produce_k(
@@ -1231,6 +1277,7 @@ class DeepseekV4Indexer(nn.Module):
1277 self.softmax_scale,
1278 self.n_head**-0.5,
1279 use_fp4=self.use_fp4_kv,
1280 + weights_out_dtype=self.indexer_weights_dtype,
1281 )
1282 if isinstance(q_quant, tuple):
1283 q, q_scale = q_quant
vllm/utils/deep_gemm.py
+183
@@ -158,6 +158,10 @@ _grouped_fp4_impl: Callable[..., Any] | None = None
158 _fp8_fp4_mqa_logits_impl: Callable[..., Any] | None = None
159 _fp8_fp4_paged_mqa_logits_impl: Callable[..., Any] | None = None
160 _get_paged_mqa_logits_metadata_impl: Callable[..., Any] | None = None
161 +_fp8_fp4_sparse_mqa_logits_impl: Callable[..., Any] | None = None
162 +_fp8_fp4_paged_sparse_mqa_logits_impl: Callable[..., Any] | None = None
163 +_get_sparse_mqa_logits_metadata_impl: Callable[..., Any] | None = None
164 +_get_paged_sparse_mqa_logits_metadata_impl: Callable[..., Any] | None = None
165 _tf32_hc_prenorm_gemm_impl: Callable[..., Any] | None = None
166 _mega_mhc_impl: Callable[..., Any] | None = None
167 _get_mn_major_tma_aligned_tensor_impl: Callable[..., Any] | None = None
@@ -232,6 +236,9 @@ def _lazy_init() -> None:
236 global _grouped_impl, _grouped_masked_impl, _grouped_fp4_impl
237 global _fp8_fp4_mqa_logits_impl, _fp8_fp4_paged_mqa_logits_impl
238 global _get_paged_mqa_logits_metadata_impl
239 + global _fp8_fp4_sparse_mqa_logits_impl, _fp8_fp4_paged_sparse_mqa_logits_impl
240 + global _get_sparse_mqa_logits_metadata_impl
241 + global _get_paged_sparse_mqa_logits_metadata_impl
242 global _tf32_hc_prenorm_gemm_impl, _mega_mhc_impl
243 global _get_mn_major_tma_aligned_tensor_impl
244 global _get_mk_alignment_for_contiguous_layout_impl
@@ -295,6 +302,17 @@ def _lazy_init() -> None:
302 _get_paged_mqa_logits_metadata_impl = getattr(
303 _dg, "get_paged_mqa_logits_metadata", None
304 )
305 + # Sparse-indexer kernels (DeepGEMM >= 2.8, SM100 only).
306 + _fp8_fp4_sparse_mqa_logits_impl = getattr(_dg, "fp8_fp4_sparse_mqa_logits", None)
307 + _fp8_fp4_paged_sparse_mqa_logits_impl = getattr(
308 + _dg, "fp8_fp4_paged_sparse_mqa_logits", None
309 + )
310 + _get_sparse_mqa_logits_metadata_impl = getattr(
311 + _dg, "get_sparse_mqa_logits_metadata", None
312 + )
313 + _get_paged_sparse_mqa_logits_metadata_impl = getattr(
314 + _dg, "get_paged_sparse_mqa_logits_metadata", None
315 + )
316 _tf32_hc_prenorm_gemm_impl = getattr(_dg, "tf32_hc_prenorm_gemm", None)
317 _mega_mhc_impl = getattr(_dg, "mega_mhc", None)
318 _get_mn_major_tma_aligned_tensor_impl = getattr(
@@ -681,6 +699,171 @@ def fp8_fp4_paged_mqa_logits(
699 )
700
701
702 +def has_deep_gemm_sparse_mqa() -> bool:
703 + """Whether the installed DeepGEMM provides the sparse-indexer kernels
704 + (``fp8_fp4_(paged_)sparse_mqa_logits``, added in DeepGEMM 2.8, SM100-only)."""
705 + _lazy_init()
706 + return (
707 + _fp8_fp4_sparse_mqa_logits_impl is not None
708 + and _fp8_fp4_paged_sparse_mqa_logits_impl is not None
709 + and _get_sparse_mqa_logits_metadata_impl is not None
710 + and _get_paged_sparse_mqa_logits_metadata_impl is not None
711 + )
712 +
713 +
714 +def get_sparse_mqa_logits_metadata(
715 + cu_seqlen_ks: torch.Tensor,
716 + cu_seqlen_ke: torch.Tensor,
717 + num_kv_tokens: int,
718 + sparse_kv_block_indices: torch.Tensor,
719 + qk_dtype: torch.dtype,
720 + sparse_block_kv: int,
721 +) -> torch.Tensor:
722 + """Build scheduling metadata for `fp8_fp4_sparse_mqa_logits`.
723 +
724 + Always uses the unaligned-ks variant, which degenerates to the aligned
725 + variant when ``cu_seqlen_ks % sparse_block_kv == 0``, so callers never
726 + need a host-side sync to pick one.
727 +
728 + Args:
729 + cu_seqlen_ks/cu_seqlen_ke: Per-row K range bounds in the packed KV
730 + workspace, shape [num_q_tokens], dtype int32.
731 + num_kv_tokens: Total KV tokens in the packed workspace.
732 + sparse_kv_block_indices: Per-row candidate block ids, shape
733 + [num_q_tokens, num_max_sparse_blocks], dtype int32. Each row's
734 + valid prefix must be sorted ascending (repeats are tolerated but
735 + wasteful); pad by repeating the last valid block. Block ``i``
736 + covers tokens
737 + ``[i * sparse_block_kv + ks % sparse_block_kv, ...)``.
738 + qk_dtype: dtype of the packed Q values (``torch.int8`` for MXFP4,
739 + ``torch.float8_e4m3fn`` for FP8).
740 + sparse_block_kv: Tokens per sparse block, 8 or 16.
741 + """
742 + _lazy_init()
743 + if _get_sparse_mqa_logits_metadata_impl is None:
744 + return _missing()
745 + return _get_sparse_mqa_logits_metadata_impl(
746 + cu_seqlen_ks,
747 + cu_seqlen_ke,
748 + num_kv_tokens,
749 + sparse_kv_block_indices,
750 + qk_dtype,
751 + sparse_block_kv,
752 + use_unaligned_ks=True,
753 + )
754 +
755 +
756 +def get_paged_sparse_mqa_logits_metadata(
757 + context_lens: torch.Tensor,
758 + block_table: torch.Tensor,
759 + indices: torch.Tensor,
760 + page_kv: int,
761 + sparse_kv_block_indices: torch.Tensor,
762 + qk_dtype: torch.dtype,
763 + sparse_block_kv: int,
764 +) -> torch.Tensor:
765 + """Build scheduling metadata for `fp8_fp4_paged_sparse_mqa_logits`.
766 +
767 + Args:
768 + context_lens: Per-row context length, shape [num_q_tokens], int32.
769 + block_table: Per-row page table, shape [num_q_tokens, max_pages],
770 + int32; rows of queries paired with the same request must be
771 + identical (expand it per query row before calling).
772 + indices: Request index for each query row, shape [num_q_tokens].
773 + page_kv: Tokens per KV page; must be a multiple of sparse_block_kv.
774 + sparse_kv_block_indices: Per-row logical (context-relative) candidate
775 + block ids, [num_q_tokens, num_max_sparse_blocks], int32.
776 + qk_dtype: dtype of the packed Q values.
777 + sparse_block_kv: Tokens per sparse block, 8 or 16.
778 + """
779 + _lazy_init()
780 + if _get_paged_sparse_mqa_logits_metadata_impl is None:
781 + return _missing()
782 + return _get_paged_sparse_mqa_logits_metadata_impl(
783 + context_lens,
784 + block_table,
785 + indices,
786 + page_kv,
787 + sparse_kv_block_indices,
788 + qk_dtype,
789 + sparse_block_kv,
790 + )
791 +
792 +
793 +def fp8_fp4_sparse_mqa_logits(
794 + q: tuple[torch.Tensor, torch.Tensor],
795 + kv: tuple[torch.Tensor, torch.Tensor],
796 + weights: torch.Tensor,
797 + metadata: torch.Tensor,
798 + num_max_sparse_blocks: int,
799 + sparse_block_kv: int,
800 +) -> torch.Tensor:
801 + """Compute MQA logits only at the candidate blocks (prefill path).
802 +
803 + Args:
804 + q: ``(q_values, q_scale)``. q_values is [M, H, D] (packed FP4 viewed
805 + as int8, or FP8); q_scale is the packed UE8M0 scale tensor [M, H]
806 + int32 (mandatory, unlike the dense kernels).
807 + kv: ``(kv_values, kv_scale)`` — the packed KV workspace with UE8M0
808 + int32 scales.
809 + weights: [M, H] ``torch.bfloat16`` (the Q scale is NOT folded in).
810 + metadata: From `get_sparse_mqa_logits_metadata`.
811 + num_max_sparse_blocks: Width of each sparse-index row.
812 + sparse_block_kv: Tokens per sparse block, 8 or 16.
813 +
814 + Returns:
815 + bf16 logits of shape [M, num_max_sparse_blocks * sparse_block_kv];
816 + column ``j * sparse_block_kv + o`` scores the token at
817 + ``sparse_kv_block_indices[row, j] * sparse_block_kv + ks % sbk + o``.
818 + """
819 + _lazy_init()
820 + if _fp8_fp4_sparse_mqa_logits_impl is None:
821 + return _missing()
822 + return _fp8_fp4_sparse_mqa_logits_impl(
823 + q,
824 + kv,
825 + weights,
826 + metadata,
827 + num_max_sparse_blocks,
828 + sparse_block_kv,
829 + use_unaligned_ks=True,
830 + )
831 +
832 +
833 +def fp8_fp4_paged_sparse_mqa_logits(
834 + q: tuple[torch.Tensor, torch.Tensor],
835 + kv_cache: torch.Tensor,
836 + weights: torch.Tensor,
837 + metadata: torch.Tensor,
838 + num_max_sparse_blocks: int,
839 + sparse_block_kv: int,
840 +) -> torch.Tensor:
841 + """Compute MQA logits only at the candidate blocks (paged decode path).
842 +
843 + Args:
844 + q: ``(q_values, q_scale)``; q_values is [num_q_tokens, 1, H, D].
845 + kv_cache: Fused paged cache [num_pages, page_kv, 1, D' + 4] uint8,
846 + page stride 512B-aligned.
847 + weights: [num_q_tokens, H] ``torch.bfloat16``.
848 + metadata: From `get_paged_sparse_mqa_logits_metadata`.
849 +
850 + Returns:
851 + bf16 logits of shape
852 + [num_q_tokens, num_max_sparse_blocks * sparse_block_kv].
853 + """
854 + _lazy_init()
855 + if _fp8_fp4_paged_sparse_mqa_logits_impl is None:
856 + return _missing()
857 + return _fp8_fp4_paged_sparse_mqa_logits_impl(
858 + q,
859 + kv_cache,
860 + weights,
861 + metadata,
862 + num_max_sparse_blocks,
863 + sparse_block_kv,
864 + )
865 +
866 +
867 def tf32_hc_prenorm_gemm(
868 x: torch.Tensor,
869 fn: torch.Tensor,
vllm/v1/attention/backends/mla/indexer.py
+6 -1
@@ -1116,6 +1116,11 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
1116 for request_slice, query_slice in chunk_specs
1117 ]
1118
1119 + def _prefill_split_seq_lens(self, seq_lens_cpu: torch.Tensor) -> torch.Tensor:
1120 + """Per-request KV lengths the prefill chunker budgets logits with;
1121 + subclasses whose logits rows are wider than the context override."""
1122 + return seq_lens_cpu
1123 +
1124 @staticmethod
1125 def _split_indexer_prefill_chunks(
1126 compressed_seq_lens_cpu: torch.Tensor,
@@ -1278,7 +1283,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
1283 )
1284 else:
1285 chunk_specs = self._split_indexer_prefill_chunks(
1281 - compressed_seq_lens_cpu[num_decodes:],
1286 + self._prefill_split_seq_lens(compressed_seq_lens_cpu[num_decodes:]),
1287 prefill_query_lens_cpu,
1288 self.max_prefill_buffer_size,
1289 max_logits_bytes,
vllm/v1/attention/backends/mla/sparse_indexer.py new
+235
@@ -0,0 +1,235 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""DeepSeek V4.1 indexer backend for the DeepGEMM sparse MQA-logits path.
4 +
5 +Candidate-consuming indexer layers score only the candidate blocks published
6 +by the candidate-source layer (``fp8_fp4_(paged_)sparse_mqa_logits``) instead
7 +of dense logits over the whole context. This backend extends the dense
8 +indexer metadata with the per-row scratch that path needs, so the layer op
9 +allocates nothing per step, and lets all consumer layers of a step share one
10 +candidate expansion and DeepGEMM schedule (they all read the same candidate
11 +blocks). Opt in with ``AttentionConfig.indexer_sparse_logits``.
12 +"""
13 +
14 +from dataclasses import dataclass, fields
15 +
16 +import torch
17 +
18 +from vllm.config import VllmConfig
19 +from vllm.logger import init_logger
20 +from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import (
21 + check_deep_select_layout,
22 + has_deep_select,
23 + pick_sparse_block_kv,
24 +)
25 +from vllm.platforms import current_platform
26 +from vllm.utils.deep_gemm import has_deep_gemm_sparse_mqa
27 +from vllm.v1.attention.backend import CommonAttentionMetadata
28 +from vllm.v1.attention.backends.mla.indexer import (
29 + DeepseekV32IndexerMetadata,
30 + DeepseekV32IndexerMetadataBuilder,
31 + DeepseekV41IndexerBackend,
32 + dsa_indexer_uses_fp4,
33 +)
34 +from vllm.v1.kv_cache_interface import AttentionSpec
35 +
36 +logger = init_logger(__name__)
37 +
38 +
39 +@dataclass
40 +class SparseMQARowsMetadata:
41 + """Per-row state for one sparse-logits call: the decode rows, or one
42 + prefill chunk. Buffers are builder-owned views sized to the rows.
43 +
44 + ``kernel_metadata`` (the DeepGEMM schedule) is filled by the first
45 + consumer indexer layer of the step together with ``sparse_indices`` /
46 + ``end``; later consumer layers reuse all three.
47 + """
48 +
49 + row_ks: torch.Tensor
50 + """[rows] int32 K-range start per row; zeros for paged decode."""
51 + row_ke: torch.Tensor
52 + """[rows] int32 K-range end (prefill) or compressed context length."""
53 + sparse_indices: torch.Tensor
54 + """[rows, num_sparse_blocks] int32 DeepGEMM sparse block ids."""
55 + end: torch.Tensor
56 + """[rows] int32 valid sparse-column count per row."""
57 + col_indices: torch.Tensor
58 + """[rows, topk] int32 scratch for the sparse-column top-k."""
59 + block_table: torch.Tensor | None = None
60 + """Decode only: [rows, pages] per-row page table."""
61 + row_indices: torch.Tensor | None = None
62 + """Decode only: [rows] row -> request index."""
63 + kernel_metadata: torch.Tensor | None = None
64 +
65 +
66 +@dataclass
67 +class DeepseekV41SparseIndexerMetadata(DeepseekV32IndexerMetadata):
68 + sparse_block_kv: int = 0
69 + sparse_decode: SparseMQARowsMetadata | None = None
70 + sparse_prefill: list[SparseMQARowsMetadata] | None = None
71 + """Parallel to ``prefill.chunks``."""
72 +
73 +
74 +class DeepseekV41SparseIndexerBackend(DeepseekV41IndexerBackend):
75 + @staticmethod
76 + def get_name() -> str:
77 + return "DEEPSEEK_V41_SPARSE_INDEXER"
78 +
79 + @staticmethod
80 + def get_builder_cls() -> type["DeepseekV41SparseIndexerMetadataBuilder"]:
81 + return DeepseekV41SparseIndexerMetadataBuilder
82 +
83 +
84 +class DeepseekV41SparseIndexerMetadataBuilder(DeepseekV32IndexerMetadataBuilder):
85 + """Dense indexer metadata plus the sparse-logits row scratch.
86 +
87 + Every requirement of the sparse kernels is checked here, at engine start,
88 + rather than falling back to the dense path per step.
89 + """
90 +
91 + def __init__(
92 + self,
93 + kv_cache_spec: AttentionSpec,
94 + layer_names: list[str],
95 + vllm_config: VllmConfig,
96 + device: torch.device,
97 + **kwargs,
98 + ) -> None:
99 + super().__init__(kv_cache_spec, layer_names, vllm_config, device, **kwargs)
100 + hf_config = vllm_config.model_config.hf_text_config
101 + candidate_block_size = getattr(hf_config, "candidate_block_size", 0)
102 + candidate_topk_blocks = getattr(hf_config, "candidate_topk_blocks", 0)
103 + topk_tokens = getattr(hf_config, "index_topk", 0)
104 + prefix = "attention_config.indexer_sparse_logits"
105 + if candidate_block_size <= 0 or candidate_topk_blocks <= 0:
106 + raise ValueError(
107 + f"{prefix} requires a two-level (candidate block) indexer model "
108 + "such as DeepSeek-V4.1."
109 + )
110 + if not dsa_indexer_uses_fp4(vllm_config):
111 + raise ValueError(
112 + f"{prefix} requires the MXFP4 indexer cache "
113 + "(attention_config.indexer_kv_dtype='mxfp4'): the sparse kernels "
114 + "take UE8M0-packed granularity-32 scales."
115 + )
116 + if not (
117 + current_platform.is_cuda()
118 + and current_platform.is_device_capability_family(100)
119 + ):
120 + raise ValueError(f"{prefix} requires an SM100-class GPU.")
121 + if not has_deep_gemm_sparse_mqa():
122 + raise ValueError(
123 + f"{prefix} requires DeepGEMM >= 2.8 "
124 + "(fp8_fp4_sparse_mqa_logits not found)."
125 + )
126 + if self.dcp_world_size > 1 or self.use_pcp:
127 + raise NotImplementedError(
128 + f"{prefix} is not supported with decode/prefill context parallel."
129 + )
130 + if not self.supports_varlen:
131 + raise ValueError(
132 + f"{prefix} requires the varlen paged MQA logits decode layout."
133 + )
134 + # The top-k runs on the kernels' bf16 logits with DeepSelect; there is
135 + # no fp32 fallback, so the other top-k backends cannot be honored.
136 + topk_backend = vllm_config.kernel_config.sparse_indexer_topk_backend
137 + if not has_deep_select():
138 + raise ValueError(
139 + f"{prefix} requires the DeepSelect top-k extension "
140 + "(vllm._deepselect_C) for this GPU."
141 + )
142 + if topk_backend not in ("auto", "deep_select"):
143 + raise ValueError(
144 + f"{prefix} uses DeepSelect for the top-k; "
145 + f"kernel_config.sparse_indexer_topk_backend='{topk_backend}' "
146 + "is not supported with it."
147 + )
148 + self.sparse_block_kv = pick_sparse_block_kv(candidate_block_size)
149 + if kv_cache_spec.block_size % self.sparse_block_kv != 0:
150 + raise ValueError(
151 + f"{prefix}: indexer cache pages ({kv_cache_spec.block_size} "
152 + f"tokens) must be a multiple of {self.sparse_block_kv} tokens."
153 + )
154 + self.num_sparse_blocks = candidate_topk_blocks * (
155 + candidate_block_size // self.sparse_block_kv
156 + )
157 + num_sparse_cols = self.num_sparse_blocks * self.sparse_block_kv
158 + check_deep_select_layout(num_sparse_cols, topk_tokens)
159 + # Sparse logits are bf16 [rows, num_sparse_cols]; make the prefill
160 + # chunker (which budgets fp32 [rows, seq_len]) size chunks as if every
161 + # row were this wide.
162 + self.min_split_seq_len = (num_sparse_cols * 2 + 3) // 4
163 +
164 + # Rows are indexed by batch token (decode rows first, then prefill
165 + # chunks by token range), like ``arange_buffer``.
166 + max_rows = self.arange_buffer.shape[0]
167 + int32 = dict(dtype=torch.int32, device=device)
168 + self.sparse_indices_buffer = torch.zeros(
169 + (max_rows, self.num_sparse_blocks), **int32
170 + )
171 + self.end_buffer = torch.zeros(max_rows, **int32)
172 + self.col_indices_buffer = torch.zeros((max_rows, topk_tokens), **int32)
173 + self.zero_row_ks = torch.zeros(max_rows, **int32)
174 + logger.info_once(
175 + "DeepSeek V4.1 indexer: candidate consumers score only the candidate "
176 + "blocks with DeepGEMM sparse MQA logits (%d sparse blocks of %d "
177 + "tokens per row).",
178 + self.num_sparse_blocks,
179 + self.sparse_block_kv,
180 + )
181 +
182 + def _prefill_split_seq_lens(self, seq_lens_cpu: torch.Tensor) -> torch.Tensor:
183 + return seq_lens_cpu.clamp(min=self.min_split_seq_len)
184 +
185 + def _rows(self, start: int, end: int, row_ks, row_ke, **kwargs):
186 + return SparseMQARowsMetadata(
187 + row_ks=row_ks,
188 + row_ke=row_ke,
189 + sparse_indices=self.sparse_indices_buffer[start:end],
190 + end=self.end_buffer[start:end],
191 + col_indices=self.col_indices_buffer[start:end],
192 + **kwargs,
193 + )
194 +
195 + def build(
196 + self,
197 + common_prefix_len: int,
198 + common_attn_metadata: CommonAttentionMetadata,
199 + fast_build: bool = False,
200 + ) -> DeepseekV41SparseIndexerMetadata:
201 + base = super().build(common_prefix_len, common_attn_metadata, fast_build)
202 + metadata = DeepseekV41SparseIndexerMetadata(
203 + **{f.name: getattr(base, f.name) for f in fields(base)},
204 + sparse_block_kv=self.sparse_block_kv,
205 + )
206 + if base.decode is not None:
207 + decode = base.decode
208 + rows = decode.seq_lens.shape[0]
209 + # The varlen layout flattens every query to one row, so the base
210 + # decode tensors are already per row.
211 + assert decode.seq_lens.shape[1] == 1 and not decode.requires_padding
212 + assert decode.block_table.shape[0] == rows
213 + metadata.sparse_decode = self._rows(
214 + 0,
215 + rows,
216 + row_ks=self.zero_row_ks[:rows],
217 + row_ke=decode.seq_lens.view(-1),
218 + block_table=decode.block_table,
219 + row_indices=(
220 + decode.indices
221 + if decode.indices is not None
222 + else self.arange_buffer[:rows]
223 + ),
224 + )
225 + if base.prefill is not None:
226 + metadata.sparse_prefill = [
227 + self._rows(
228 + chunk.token_start,
229 + chunk.token_end,
230 + row_ks=chunk.cu_seqlen_ks,
231 + row_ke=chunk.cu_seqlen_ke,
232 + )
233 + for chunk in base.prefill.chunks
234 + ]
235 + return metadata
vllm/v1/core/kv_cache_utils.py
+9
@@ -1608,6 +1608,15 @@ def _get_kv_cache_bytes_per_block(
1608 ]
1609 if hot_page_sizes:
1610 bytes_per_block = round_up(bytes_per_block, math.lcm(*hot_page_sizes))
1611 + stride_alignments = [
1612 + spec.block_stride_alignment
1613 + for group in kv_cache_groups
1614 + for layer_name in group.layer_names
1615 + if isinstance(spec := _get_per_layer_spec(group, layer_name), MLAAttentionSpec)
1616 + and spec.block_stride_alignment
1617 + ]
1618 + if stride_alignments:
1619 + bytes_per_block = round_up(bytes_per_block, math.lcm(*stride_alignments))
1620 return bytes_per_block
1621
1622
vllm/v1/kv_cache_interface.py
+10 -1
@@ -653,6 +653,12 @@ class MLAAttentionSpec(FullAttentionSpec):
653 is_index_group_leader: bool = False
654 storage_block_size: int | None = None
655 """Token width used to view storage when it differs from the kernel block."""
656 + block_stride_alignment: int | None = None
657 + """Required alignment, in bytes, of the distance between consecutive
658 + blocks of this cache. In block-major layouts that distance is the whole
659 + block (all layers' pages), so the allocator rounds the block up to it.
660 + DeepGEMM's paged sparse MQA-logits kernels address pages as
661 + ``base + page * stride`` and need it 512B-aligned."""
662 # Group capability enabled when any member flattens a non-causal query block
663 # into decode rows. Runtime metadata still selects causal vs. non-causal mode.
664 non_causal_multi_token_decode: bool = False
@@ -674,6 +680,7 @@ class MLAAttentionSpec(FullAttentionSpec):
680 cache_role_set = {spec.cache_role for spec in specs}
681 index_group_leader_set = {spec.is_index_group_leader for spec in specs}
682 storage_block_size_set = set(spec.storage_block_size for spec in specs)
683 + block_stride_alignment_set = {spec.block_stride_alignment for spec in specs}
684 assert (
685 len(cache_dtype_str_set) == 1
686 and len(tokens_per_state_set) == 1
@@ -681,10 +688,11 @@ class MLAAttentionSpec(FullAttentionSpec):
688 and len(cache_role_set) == 1
689 and len(index_group_leader_set) == 1
690 and len(storage_block_size_set) == 1
691 + and len(block_stride_alignment_set) == 1
692 ), (
693 "All attention layers in the same KV cache group must use the same "
694 "quantization method, tokens per state, model version, cache role, "
687 - "index-sharing role, and storage block size."
695 + "index-sharing role, storage block size and block stride alignment."
696 )
697 merged_spec = cls(
698 block_size=specs[0].block_size,
@@ -701,6 +709,7 @@ class MLAAttentionSpec(FullAttentionSpec):
709 cache_role=cache_role_set.pop(),
710 is_index_group_leader=index_group_leader_set.pop(),
711 storage_block_size=storage_block_size_set.pop(),
712 + block_stride_alignment=block_stride_alignment_set.pop(),
713 non_causal_multi_token_decode=any(
714 spec.non_causal_multi_token_decode for spec in specs
715 ),