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

[M3] Improve indexer for long-context decode (sm100) (#48582)

Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>

Thien Tran committed Jul 17, 2026 at 09:12 UTC fe784ff22e630a31fd798f392b01e0a75c18f047

7 files changed +756 -63

tests/kernels/attention/test_minimax_m3.py

+118 -8

@@ -185,6 +185,43 @@ def _assert_topk_indices_equal_unordered(

185 assert set(actual_row) == set(expected_row)

186

187

188 +def _reference_decode_index_score(

189 + idx_q: torch.Tensor,

190 + index_kv_cache: torch.Tensor,

191 + block_table: torch.Tensor,

192 + seq_lens: torch.Tensor,

193 + decode_query_len: int,

194 + score_block_stride: int,

195 +) -> torch.Tensor:

196 + total_q, num_idx_heads, _ = idx_q.shape

197 + out = torch.full(

198 + (num_idx_heads, total_q, score_block_stride),

199 + -float("inf"),

200 + device=idx_q.device,

201 + dtype=torch.float32,

202 + )

203 + for req_id, seq_len in enumerate(seq_lens.tolist()):

204 + num_blocks = (seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE

205 + token_start = req_id * decode_query_len

206 + q = idx_q[token_start : token_start + decode_query_len].float()

207 + pages = block_table[req_id, :num_blocks]

208 + k = index_kv_cache[pages].reshape(num_blocks * BLOCK_SIZE, -1).float()

209 + score = torch.einsum("qhd,kd->hqk", q, k)

210 + q_pos = (

211 + seq_len

212 + - decode_query_len

213 + + torch.arange(decode_query_len, device=idx_q.device)

214 + )

215 + k_pos = torch.arange(k.shape[0], device=idx_q.device)

216 + score.masked_fill_(k_pos[None, :] > q_pos[:, None], -float("inf"))

217 + out[:, token_start : token_start + decode_query_len, :num_blocks] = (

218 + score.reshape(num_idx_heads, decode_query_len, num_blocks, BLOCK_SIZE)

219 + .max(dim=3)

220 + .values

221 + )

222 + return out

223 +

224 +

225 def test_prefill_index_topk_correctness():

226 topk = 6

227 init_blocks = 0

@@ -383,8 +420,8 @@ def test_fmha_sm100_indexer_matches_reference(q_lens, prefix_lens, index_dtype):

420 _assert_topk_indices_equal_unordered(actual, expected)

421

422

386 -# Full impl-level parity: drive both MiniMaxM3IndexerMSAImpl (fmha_sm100 score +

387 -# Triton top-k) and MiniMaxM3IndexerTritonImpl through their real metadata

423 +# Full impl-level parity: drive both MiniMaxM3IndexerMSAImpl (fmha/CuteDSL score

424 +# + unified top-k) and MiniMaxM3IndexerTritonImpl through their real metadata

425 # builders on the SAME CommonAttentionMetadata + index cache, and assert the

426 # selected blocks agree. This exercises all the metadata the impl/kernels consume

427 # (decode/prefill split, cu_seqlens_q rebasing, prefix_lens, kv_indices gather,

@@ -395,7 +432,8 @@ def test_fmha_sm100_indexer_matches_reference(q_lens, prefix_lens, index_dtype):

432 reason="fmha_sm100 indexer requires SM100 (Blackwell).",

433 )

434 @pytest.mark.parametrize("topk", [8, 16])

398 -def test_msa_indexer_impl_matches_triton(topk, monkeypatch):

435 +@pytest.mark.parametrize("index_dtype", [torch.bfloat16, torch.float8_e4m3fn])

436 +def test_msa_indexer_impl_matches_triton(topk, index_dtype, monkeypatch):

437 import vllm.models.minimax_m3.common.indexer as indexer_mod

438 from tests.v1.attention.utils import (

439 BatchSpec,

@@ -439,13 +477,13 @@ def test_msa_indexer_impl_matches_triton(topk, monkeypatch):

477 block_table = common.block_table_tensor

478 num_pages = int(block_table.max().item()) + 1

479 index_cache = torch.zeros(

442 - num_pages, BLOCK_SIZE, head_dim, device=device, dtype=DTYPE

480 + num_pages, BLOCK_SIZE, head_dim, device=device, dtype=index_dtype

481 )

482 for r, seq_len in enumerate(batch.seq_lens):

483 for b in range((seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE):

484 index_cache[block_table[r, b]] = float(b + 1)

485 index_q = torch.ones(

448 - num_tokens, num_idx_heads * head_dim, device=device, dtype=DTYPE

486 + num_tokens, num_idx_heads * head_dim, device=device, dtype=index_dtype

487 )

488

489 spec = MLAAttentionSpec(

@@ -588,9 +626,8 @@ def test_decode_index_topk_correctness(

626 )

627 @pytest.mark.parametrize("num_idx_heads", [1, 4])

628 def test_decode_index_topk_fp8(num_idx_heads: int):

591 - """The fp8 (e4m3) indexer cache feeds the Triton decode kernel on the MSA

592 - path. The kernel must score in fp32 (no scaling) so its top-k matches a

593 - reference computed from the dequantized fp8 values."""

629 + """The standalone Triton path must score FP8 inputs in FP32 so its top-k

630 + matches a reference computed from the dequantized FP8 values."""

631 torch.manual_seed(0)

632 topk, init_blocks, local_blocks, head_dim = 8, 0, 1, 128

633 decode_query_len = 1

@@ -641,6 +678,79 @@ def test_decode_index_topk_fp8(num_idx_heads: int):

678 _assert_topk_indices_equal_unordered(actual, expected)

679

680

681 +@pytest.mark.skipif(

682 + not current_platform.is_device_capability_family(100),

683 + reason="CuteDSL index decode score requires Blackwell.",

684 +)

685 +@pytest.mark.parametrize(

686 + ("dtype", "decode_query_len", "max_decode_query_len"),

687 + [

688 + (torch.bfloat16, 1, 1),

689 + (torch.bfloat16, 3, 8),

690 + (torch.float8_e4m3fn, 1, 1),

691 + (torch.float8_e4m3fn, 3, 8),

692 + (torch.float8_e4m3fn, 8, 8),

693 + ],

694 +)

695 +def test_decode_index_score_cutedsl_correctness(

696 + dtype: torch.dtype,

697 + decode_query_len: int,

698 + max_decode_query_len: int,

699 +):

700 + pytest.importorskip("cutlass")

701 + from vllm.models.minimax_m3.nvidia.ops import (

702 + minimax_m3_index_decode_score_cutedsl,

703 + )

704 +

705 + torch.manual_seed(0)

706 + init_blocks, local_blocks = 0, 0

707 + num_idx_heads, head_dim = 4, 128

708 + active_seq_lens = torch.tensor((1025, 4097), device="cuda", dtype=torch.int32)

709 + batch = active_seq_lens.numel()

710 + total_q = batch * decode_query_len

711 + max_seq_len = int(active_seq_lens.max())

712 + max_blocks = (max_seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE

713 + score_block_stride = ((max_blocks + 15) // 16) * 16

714 + num_pages = batch * max_blocks

715 + block_table = torch.randperm(num_pages, device="cuda", dtype=torch.int32).reshape(

716 + batch, max_blocks

717 + )

718 + idx_q = torch.randn(total_q, num_idx_heads, head_dim, device="cuda").to(dtype)

719 + index_kv_cache = torch.randn(num_pages, BLOCK_SIZE, head_dim, device="cuda").to(

720 + dtype

721 + )

722 + unified_score = torch.full(

723 + (total_q, num_idx_heads, score_block_stride),

724 + -float("inf"),

725 + device="cuda",

726 + dtype=torch.float32,

727 + )

728 + score = unified_score.transpose(0, 1)

729 +

730 + minimax_m3_index_decode_score_cutedsl(

731 + idx_q,

732 + index_kv_cache,

733 + block_table,

734 + active_seq_lens,

735 + max_seq_len=max_seq_len,

736 + init_blocks=init_blocks,

737 + local_blocks=local_blocks,

738 + num_kv_heads=num_idx_heads,

739 + decode_query_len=decode_query_len,

740 + max_decode_query_len=max_decode_query_len,

741 + score_out=score,

742 + )

743 + expected = _reference_decode_index_score(

744 + idx_q,

745 + index_kv_cache,

746 + block_table,

747 + active_seq_lens,

748 + decode_query_len,

749 + score_block_stride,

750 + )

751 + torch.testing.assert_close(score, expected)

752 +

753 +

754 # Sparse attention kernels.

755 def _reference_sparse_attn(

756 q: torch.Tensor,

vllm/cute_utils/__init__.py

+48 -13

@@ -1,11 +1,33 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -from cutlass import BFloat16, Float32, Int64, Uint32, cute

3 +import torch

4 +from cutlass import (

5 + BFloat16,

6 + Float8E4M3FN,

7 + Float16,

8 + Float32,

9 + Int32,

10 + Int64,

11 + Uint32,

12 + cute,

13 +)

14 from cutlass._mlir import ir

15 from cutlass._mlir.dialects import llvm, vector

16 from cutlass.cute.nvgpu import cpasync

17 from cutlass.cutlass_dsl import T, dsl_user_op

19 +_TORCH_TO_CUTE_DTYPE = {

20 + torch.bfloat16: BFloat16,

21 + torch.float8_e4m3fn: Float8E4M3FN,

22 +}

23 +

24 +_CUTE_TO_PTX_DTYPE = {

25 + BFloat16: "bf16",

26 + Float16: "f16",

27 + Float8E4M3FN: "e4m3",

28 + Float32: "f32",

29 +}

30 +

31 # https://github.com/NVIDIA/cutlass/blob/v4.3.2/include/cute/arch/copy_sm90_desc.hpp#L193-L197

32 EVICT_NORMAL = Int64(0x1000000000000000)

33 EVICT_FIRST = Int64(0x12F0000000000000)

@@ -63,22 +85,35 @@ def fence_before_tma_store(*, loc=None, ip=None):

87 @dsl_user_op

66 -def mma_bf16(

67 - a: cute.TensorSSA, b: cute.TensorSSA, c: cute.TensorSSA, *, loc=None, ip=None

68 -):

69 - if a.element_type == BFloat16:

70 - a = cute.recast_tensor(a, Uint32)

71 - if b.element_type == BFloat16:

72 - b = cute.recast_tensor(b, Uint32)

73 -

74 - mlir_ty = Float32.mlir_type

88 +def mma_sync(a, b, c: cute.Tensor, *, loc=None, ip=None):

89 + a_ty = _CUTE_TO_PTX_DTYPE[a.element_type]

90 + b_ty = _CUTE_TO_PTX_DTYPE[b.element_type]

91 + c_ty = _CUTE_TO_PTX_DTYPE[c.element_type]

92 + mlir_ty = c.element_type.mlir_type

93 + K = 256 // a.element_type.width # 32B

94 +

95 + # Recast expects tensor-backed fragments; materialize SSA fragments here so

96 + # callsites can pass converted FP8 fragments directly.

97 + if isinstance(a, cute.TensorSSA):

98 + a_ = cute.make_rmem_tensor_like(a)

99 + a_.store(a, loc=loc, ip=ip)

100 + a = a_

101 + if isinstance(b, cute.TensorSSA):

102 + b_ = cute.make_rmem_tensor_like(b)

103 + b_.store(b, loc=loc, ip=ip)

104 + b = b_

105 +

106 + a = cute.recast_tensor(a, Int32, loc=loc, ip=ip)

107 + b = cute.recast_tensor(b, Int32, loc=loc, ip=ip)

108 out = llvm.inline_asm(

109 llvm.StructType.get_literal([mlir_ty] * 4),

110 [a[i].ir_value(loc=loc, ip=ip) for i in range(4)]

111 + [b[i].ir_value(loc=loc, ip=ip) for i in range(2)]

112 + [c[i].ir_value(loc=loc, ip=ip) for i in range(4)],

80 - "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "

81 - "{$0, $1, $2, $3}, {$4, $5, $6, $7}, {$8, $9}, "

113 + f"mma.sync.aligned.m16n8k{K}.row.col.{c_ty}.{a_ty}.{b_ty}.{c_ty} "

114 + "{$0, $1, $2, $3}, "

115 + "{$4, $5, $6, $7}, "

116 + "{$8, $9}, "

117 "{$10, $11, $12, $13};",

118 "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f",

119 has_side_effects=False,

@@ -92,7 +127,7 @@ def mma_bf16(

127 loc=loc,

128 ip=ip,

129 )

95 - return cute.TensorSSA(vec, 4, Float32)

130 + return cute.TensorSSA(vec, 4, c.element_type)

131

132

133 def _bf16x2_unary(asm: str, a: Uint32, *, loc=None, ip=None) -> Uint32:

vllm/cute_utils/cvt.py

+26

@@ -84,6 +84,32 @@ def fp8x4_to_bf16x4(x: Uint32, *, loc=None, ip=None) -> cute.TensorSSA:

84 return cute.TensorSSA(vec, 2, Uint32)

87 +@dsl_user_op

88 +def fp8x4_to_fp16x4(x: Uint32, *, loc=None, ip=None) -> cute.TensorSSA:

89 + out = llvm.inline_asm(

90 + llvm.StructType.get_literal([T.i32()] * 2),

91 + [x.ir_value(loc=loc, ip=ip)],

92 + "{\n\t"

93 + ".reg .b16 lo, hi;\n\t"

94 + "mov.b32 {lo, hi}, $2;\n\t"

95 + "cvt.rn.f16x2.e4m3x2 $0, lo;\n\t"

96 + "cvt.rn.f16x2.e4m3x2 $1, hi;\n\t"

97 + "}\n",

98 + "=r,=r,r",

99 + has_side_effects=False,

100 + is_align_stack=False,

101 + loc=loc,

102 + ip=ip,

103 + )

104 + vec = vector.from_elements(

105 + ir.VectorType.get([2], T.i32(), loc=loc),

106 + [llvm.extractvalue(T.i32(), out, [i], loc=loc, ip=ip) for i in range(2)],

107 + loc=loc,

108 + ip=ip,

109 + )

110 + return cute.TensorSSA(vec, 2, Uint32)

111 +

112 +

113 @dsl_user_op

114 def fp32x4_to_fp8x4(

115 a0: Float32,

vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/kernel_kkt_inv_uw.py

+63 -36

@@ -16,7 +16,7 @@ from vllm.cute_utils import (

16 _tcgen05,

17 cvt,

18 fence_before_tma_store,

19 - mma_bf16,

19 + mma_sync,

20 simple_tma_copy,

21 )

@@ -536,14 +536,16 @@ class Sm100ChunkUWKernel:

536 zeros_f32.fill(0.0)

537

538 Ai_bf16 = cute.make_rmem_tensor(8, BFloat16)

539 - mma_B_bf16 = cute.make_rmem_tensor(8, BFloat16)

540 - M_bf16 = cute.make_rmem_tensor(8, BFloat16)

539 + mma_B_bf16 = cute.make_rmem_tensor((4, 2), BFloat16)

540 + M_bf16 = cute.make_rmem_tensor((4, 2), BFloat16)

541 acc = cute.make_rmem_tensor((4, 2), Float32)

542

543 - # share the same storage

543 + # ldmatrix copies require rank-1 destination views.

544 + mma_B_bf16_ldsm = cute.group_modes(mma_B_bf16, 0, 2)

545 +

546 + # Packed aliases for BF16x2 arithmetic.

547 Ai = cute.recast_tensor(Ai_bf16, Uint32)

545 - mma_B = cute.logical_divide(cute.recast_tensor(mma_B_bf16, Uint32), 2)

546 - M = cute.logical_divide(cute.recast_tensor(M_bf16, Uint32), 2)

548 + M = cute.recast_tensor(M_bf16, Uint32)

549

550 # construct rmem-backed identity matrix

551 eye = cute.make_rmem_tensor(4, Uint32)

@@ -563,7 +565,11 @@ class Sm100ChunkUWKernel:

565 Ai_f32 = cute.logical_divide(cvt.bf16x2_to_fp32x2(Ai), 4)

566

567 # M is holding -(I+A), stay constant throughout the iterations

566 - cute.copy(ldsm_trans_atom, sA_ldsm[None, (warp_id_, warp_id_)], M_bf16)

568 + cute.copy(

569 + ldsm_trans_atom,

570 + sA_ldsm[None, (warp_id_, warp_id_)],

571 + cute.group_modes(M_bf16, 0, 2),

572 + )

573 for i in cutlass.range_constexpr(4):

574 M[i] = _bf16x2_sub(_bf16x2_neg(eye[i]), M[i])

575

@@ -572,8 +578,8 @@ class Sm100ChunkUWKernel:

578 # First MMA: -AiM = Ai @ (-M)

579 cute.copy(stsm_atom, Ai_bf16, sA_ldsm[None, (warp_id_, warp_id_)])

580 cute.arch.sync_warp()

575 - acc[None, 0] = mma_bf16(Ai, M[None, 0], zeros_f32)

576 - acc[None, 1] = mma_bf16(Ai, M[None, 1], zeros_f32)

581 + acc[None, 0] = mma_sync(Ai_bf16, M_bf16[None, 0], zeros_f32)

582 + acc[None, 1] = mma_sync(Ai_bf16, M_bf16[None, 1], zeros_f32)

583 Ai_bf16.store(acc.load().to(BFloat16))

584

585 # Second MMA: Ai_new = 2Ai + (-AiM) @ Ai

@@ -582,10 +588,14 @@ class Sm100ChunkUWKernel:

588 cute.copy(

589 ldsm_trans_atom,

590 sA_ldsm[None, (warp_id_, warp_id_)],

585 - mma_B_bf16,

591 + mma_B_bf16_ldsm,

592 + )

593 + Ai_f32[None, 0] = mma_sync(

594 + Ai_bf16, mma_B_bf16[None, 0], Ai_f32[None, 0]

595 + )

596 + Ai_f32[None, 1] = mma_sync(

597 + Ai_bf16, mma_B_bf16[None, 1], Ai_f32[None, 1]

598 )

587 - Ai_f32[None, 0] = mma_bf16(Ai, mma_B[None, 0], Ai_f32[None, 0])

588 - Ai_f32[None, 1] = mma_bf16(Ai, mma_B[None, 1], Ai_f32[None, 1])

599 Ai_bf16.store(Ai_f32.load().to(BFloat16))

600

601 cute.copy(stsm_atom, Ai_bf16, sAi_ldsm[None, (warp_id_, warp_id_)])

@@ -613,23 +623,24 @@ class Sm100ChunkUWKernel:

623 neg_Ai = cute.make_rmem_tensor(4, Uint32)

624 for i in cutlass.range_constexpr(4):

625 neg_Ai[i] = _bf16x2_neg(Ai[i])

626 + neg_Ai_bf16 = cute.recast_tensor(neg_Ai, BFloat16)

627

628 cute.copy(

629 ldsm_trans_atom,

630 sA_ldsm[None, (warp_id_, warp_id_ - 1)],

620 - mma_B_bf16,

631 + mma_B_bf16_ldsm,

632 )

622 - acc[None, 0] = mma_bf16(neg_Ai, mma_B[None, 0], zeros_f32)

623 - acc[None, 1] = mma_bf16(neg_Ai, mma_B[None, 1], zeros_f32)

633 + acc[None, 0] = mma_sync(neg_Ai_bf16, mma_B_bf16[None, 0], zeros_f32)

634 + acc[None, 1] = mma_sync(neg_Ai_bf16, mma_B_bf16[None, 1], zeros_f32)

635 Ai_bf16.store(acc.load().to(BFloat16))

636

637 cute.copy(

638 ldsm_trans_atom,

639 sAi_ldsm[None, (warp_id_ - 1, warp_id_ - 1)],

629 - mma_B_bf16,

640 + mma_B_bf16_ldsm,

641 )

631 - acc[None, 0] = mma_bf16(Ai, mma_B[None, 0], zeros_f32)

632 - acc[None, 1] = mma_bf16(Ai, mma_B[None, 1], zeros_f32)

642 + acc[None, 0] = mma_sync(Ai_bf16, mma_B_bf16[None, 0], zeros_f32)

643 + acc[None, 1] = mma_sync(Ai_bf16, mma_B_bf16[None, 1], zeros_f32)

644 Ai_bf16.store(acc.load().to(BFloat16))

645 store_ab_abg(

646 acc,

@@ -660,10 +671,10 @@ class Sm100ChunkUWKernel:

671 cute.copy(

672 ldsm_trans_atom,

673 sAi_ldsm[None, (tile_col, tile_col)],

663 - mma_B_bf16,

674 + mma_B_bf16_ldsm,

675 )

665 - acc[None, 0] = mma_bf16(Ai, mma_B[None, 0], zeros_f32)

666 - acc[None, 1] = mma_bf16(Ai, mma_B[None, 1], zeros_f32)

676 + acc[None, 0] = mma_sync(Ai_bf16, mma_B_bf16[None, 0], zeros_f32)

677 + acc[None, 1] = mma_sync(Ai_bf16, mma_B_bf16[None, 1], zeros_f32)

678

679 cute.copy(

680 ldsm_atom, sA_ldsm[None, (warp_id_, tile_col + 1)], Ai_bf16

@@ -671,10 +682,10 @@ class Sm100ChunkUWKernel:

682 cute.copy(

683 ldsm_trans_atom,

684 sAi_ldsm[None, (tile_col + 1, tile_col)],

674 - mma_B_bf16,

685 + mma_B_bf16_ldsm,

686 )

676 - acc[None, 0] = mma_bf16(Ai, mma_B[None, 0], acc[None, 0])

677 - acc[None, 1] = mma_bf16(Ai, mma_B[None, 1], acc[None, 1])

687 + acc[None, 0] = mma_sync(Ai_bf16, mma_B_bf16[None, 0], acc[None, 0])

688 + acc[None, 1] = mma_sync(Ai_bf16, mma_B_bf16[None, 1], acc[None, 1])

689

690 tmp = cute.make_rmem_tensor(8, BFloat16)

691 tmp.store(acc.load().to(BFloat16))

@@ -687,10 +698,10 @@ class Sm100ChunkUWKernel:

698 cute.copy(

699 ldsm_trans_atom,

700 sAi_ldsm[None, (warp_id_, tile_col)],

690 - mma_B_bf16,

701 + mma_B_bf16_ldsm,

702 )

692 - acc[None, 0] = mma_bf16(Ai, mma_B[None, 0], zeros_f32)

693 - acc[None, 1] = mma_bf16(Ai, mma_B[None, 1], zeros_f32)

703 + acc[None, 0] = mma_sync(Ai_bf16, mma_B_bf16[None, 0], zeros_f32)

704 + acc[None, 1] = mma_sync(Ai_bf16, mma_B_bf16[None, 1], zeros_f32)

705 tmp.store(acc.load().to(BFloat16))

706 cute.copy(stsm_atom, tmp, sAi_ldsm[None, (warp_id_, tile_col)])

707 store_ab_abg(

@@ -708,15 +719,27 @@ class Sm100ChunkUWKernel:

719 # warp3: Ai30 = -Ai33 @ (A30 @ Ai00 + A31 @ Ai10 + A32 @ Ai20)

720 if warp_id_ == 3:

721 cute.copy(ldsm_atom, sA_ldsm[None, (3, 0)], Ai_bf16)

711 - cute.copy(ldsm_trans_atom, sAi_ldsm[None, (0, 0)], mma_B_bf16)

712 - acc[None, 0] = mma_bf16(Ai, mma_B[None, 0], zeros_f32)

713 - acc[None, 1] = mma_bf16(Ai, mma_B[None, 1], zeros_f32)

722 + cute.copy(

723 + ldsm_trans_atom,

724 + sAi_ldsm[None, (0, 0)],

725 + mma_B_bf16_ldsm,

726 + )

727 + acc[None, 0] = mma_sync(Ai_bf16, mma_B_bf16[None, 0], zeros_f32)

728 + acc[None, 1] = mma_sync(Ai_bf16, mma_B_bf16[None, 1], zeros_f32)

729

730 for i in cutlass.range_constexpr(1, 3):

731 cute.copy(ldsm_atom, sA_ldsm[None, (3, i)], Ai_bf16)

717 - cute.copy(ldsm_trans_atom, sAi_ldsm[None, (i, 0)], mma_B_bf16)

718 - acc[None, 0] = mma_bf16(Ai, mma_B[None, 0], acc[None, 0])

719 - acc[None, 1] = mma_bf16(Ai, mma_B[None, 1], acc[None, 1])

732 + cute.copy(

733 + ldsm_trans_atom,

734 + sAi_ldsm[None, (i, 0)],

735 + mma_B_bf16_ldsm,

736 + )

737 + acc[None, 0] = mma_sync(

738 + Ai_bf16, mma_B_bf16[None, 0], acc[None, 0]

739 + )

740 + acc[None, 1] = mma_sync(

741 + Ai_bf16, mma_B_bf16[None, 1], acc[None, 1]

742 + )

743

744 tmp = cute.make_rmem_tensor(8, BFloat16)

745 tmp.store(acc.load().to(BFloat16))

@@ -726,9 +749,13 @@ class Sm100ChunkUWKernel:

749 cute.copy(ldsm_atom, sAi_ldsm[None, (3, 3)], Ai_bf16)

750 for i in cutlass.range_constexpr(4):

751 Ai[i] = _bf16x2_neg(Ai[i])

729 - cute.copy(ldsm_trans_atom, sAi_ldsm[None, (3, 0)], mma_B_bf16)

730 - acc[None, 0] = mma_bf16(Ai, mma_B[None, 0], zeros_f32)

731 - acc[None, 1] = mma_bf16(Ai, mma_B[None, 1], zeros_f32)

752 + cute.copy(

753 + ldsm_trans_atom,

754 + sAi_ldsm[None, (3, 0)],

755 + mma_B_bf16_ldsm,

756 + )

757 + acc[None, 0] = mma_sync(Ai_bf16, mma_B_bf16[None, 0], zeros_f32)

758 + acc[None, 1] = mma_sync(Ai_bf16, mma_B_bf16[None, 1], zeros_f32)

759 tmp.store(acc.load().to(BFloat16))

760 cute.copy(stsm_atom, tmp, sAi_ldsm[None, (3, 0)])

761 store_ab_abg(

vllm/models/minimax_m3/nvidia/indexer_msa.py

+13 -6

@@ -13,10 +13,9 @@ Prefill scores with ``fmha_sm100``'s score-only (``OnlyScore``) path (much faste

13 than Triton for the wide prefill score, benchmarked ~3-5x), writing its

14 ``max_score`` straight into the buffer's prefill region (stride-aware, no copy).

16 -Decode scores with the Triton split-K ``minimax_m3_index_decode_score`` (a

17 -purpose-built vector x matrix score, no wasted tensor-core tiles, 256-way

18 -split-K, cudagraph-safe by shape-constant grids), writing into the decode

19 -region. Its tuning heuristics are kept; only the top-k is shared with prefill.

16 +Decode scores with CuteDSL when the flattened query tile is supported and fall

17 +back to Triton otherwise, writing into the decode region. Only the top-k is

18 +shared with prefill.

20 ``fmha_sm100`` imports are function-local so this module is import-safe on

21 AMD / non-SM100.

@@ -39,6 +38,7 @@ from vllm.models.minimax_m3.common.indexer import (

38 from vllm.models.minimax_m3.common.ops.index_topk import (

39 minimax_m3_index_decode_score,

40 )

41 +from vllm.models.minimax_m3.nvidia.ops import minimax_m3_index_decode_score_cutedsl

42 from vllm.v1.attention.backend import (

43 AttentionBackend,

44 AttentionCGSupport,

@@ -257,7 +257,7 @@ class MiniMaxM3IndexerMSAMetadataBuilder(MiniMaxM3IndexerMetadataBuilder):

257

258

259 class MiniMaxM3IndexerMSAImpl(MiniMaxM3IndexerImpl):

260 - """Decode: Triton fused score+top-k. Prefill: fmha_sm100 OnlyScore + top-k."""

260 + """Decode: CuteDSL/Triton score. Prefill: fmha_sm100 OnlyScore + top-k."""

261

262 indexer_backend_cls: ClassVar[type[AttentionBackend]] = MiniMaxM3IndexerMSABackend

263

@@ -296,7 +296,14 @@ class MiniMaxM3IndexerMSAImpl(MiniMaxM3IndexerImpl):

296 # writes by strides). Top-k is deferred to the single unified call below.

297 if md.decode is not None:

298 d = md.decode

299 - minimax_m3_index_decode_score(

299 + # max_decode_query_len avoids recompiles across runtime decode sizes.

300 + # Fall back when the flattened Q tile gets too wide for this kernel.

301 + decode_score = (

302 + minimax_m3_index_decode_score_cutedsl

303 + if self.num_index_heads * d.max_decode_query_len <= 32

304 + else minimax_m3_index_decode_score

305 + )

306 + decode_score(

307 index_q[:nd],

308 kv,

309 d.block_table,

vllm/models/minimax_m3/nvidia/ops/__init__.py new

@@ -0,0 +1,5 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +from .index_decode_score import minimax_m3_index_decode_score_cutedsl

4 +

5 +__all__ = ["minimax_m3_index_decode_score_cutedsl"]

vllm/models/minimax_m3/nvidia/ops/index_decode_score.py new

+483

@@ -0,0 +1,483 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""CuteDSL MiniMax M3 index decode score kernel.

4 +

5 +The kernel computes decode-time index block scores with TMA + ``mma.sync``.

6 +We use ``mma.sync`` instead of tcgen05 because this score GEMM has a very small

7 +N dimension and benefits more from higher CTA occupancy than from a deeper

8 +single-CTA tcgen05 pipeline.

9 +

10 +The implementation should be portable to SM90/SM120 in principle, but it is

11 +currently validated only for SM100.

12 +"""

13 +

14 +from functools import cache

15 +

16 +import cutlass

17 +import torch

18 +from cuda.bindings.driver import CUstream

19 +from cutlass import Float8E4M3FN, Float16, Float32, Int32, Int64, Uint32, cute

20 +from cutlass.cute.nvgpu import cpasync, warp

21 +from quack.compile_utils import make_fake_tensor

22 +

23 +from vllm.cute_utils import (

24 + _TORCH_TO_CUTE_DTYPE,

25 + EVICT_FIRST,

26 + cvt,

27 + mma_sync,

28 + simple_tma_copy,

29 +)

30 +

31 +

32 +@cute.jit

33 +def _fp8_to_f16_mma_fragments(src: cute.Tensor):

34 + src_elems = cute.size(src)

35 + src_u32 = cute.recast_tensor(src, Uint32)

36 + src_f16 = cute.make_rmem_tensor(src_elems, Float16)

37 + src_f16_u32 = cute.recast_tensor(src_f16, Uint32)

38 + # This packed conversion is faster and emits fewer SASS instructions than

39 + # src.load().to(Float16).

40 + for i in cutlass.range_constexpr(src_elems // 4):

41 + converted = cvt.fp8x4_to_fp16x4(src_u32[i])

42 + src_f16_u32[i * 2] = converted[0]

43 + src_f16_u32[i * 2 + 1] = converted[1]

44 + lower = cute.make_rmem_tensor(src_elems // 2, Float16)

45 + upper = cute.make_rmem_tensor(src_elems // 2, Float16)

46 +

47 + # FP8 ldmatrix gives four consecutive values along K. Split each group

48 + # into the lower two and upper two values for two FP16 MMA k-fragments.

49 + for i in cutlass.range_constexpr(src_elems // 2):

50 + lower[i] = src_f16[(i // 2) * 4 + i % 2]

51 + upper[i] = src_f16[(i // 2) * 4 + 2 + i % 2]

52 + return lower, upper

53 +

54 +

55 +class IndexDecodeScoreKernel:

56 + BLOCK_K = 128

57 + BAR_MMA = 1

58 + num_stages = 2

59 +

60 + def __init__(

61 + self,

62 + dtype: type[cutlass.Numeric],

63 + num_heads: int,

64 + max_decode_query_len: int,

65 + split_k: int,

66 + head_dim: int = 128,

67 + ):

68 + self.dtype = dtype

69 + self.num_heads = num_heads

70 + self.max_decode_query_len = max_decode_query_len

71 + self.split_k = split_k

72 + self.head_dim = head_dim

73 +

74 + @cute.jit

75 + def __call__(

76 + self,

77 + gQ: cute.Tensor, # [bs * runtime_decode_query_len, num_heads, head_dim]

78 + gK_cache: cute.Tensor, # [num_pages, page_size, head_dim]

79 + block_table: cute.Tensor, # [bs, max_pages]

80 + score: cute.Tensor, # [num_heads, bs * runtime_decode_query_len, max_pages]

81 + seq_lens: cute.Tensor, # [bs]

82 + stream: CUstream,

83 + ):

84 + dtype = self.dtype

85 + num_heads = self.num_heads

86 + head_dim = self.head_dim

87 + BLOCK_K = self.BLOCK_K

88 + num_stages = self.num_stages

89 + MAX_DQL = self.max_decode_query_len

90 + BLOCK_Q = num_heads * MAX_DQL

91 + assert BLOCK_Q <= 32

92 +

93 + batch = seq_lens.shape[0]

94 + decode_query_len = gQ.shape[0] // batch

95 + grid = (batch, self.split_k, 1)

96 + block = (32 * 5, 1, 1)

97 +

98 + tma_g2s = cpasync.CopyBulkTensorTileG2SOp()

99 + swizzle_128B = cute.make_swizzle(3, 4, 3)

100 + elems = 128 * 8 // dtype.width

101 +

102 + sQ_layout = cute.make_layout(

103 + (MAX_DQL, num_heads, (elems, head_dim // elems)),

104 + stride=(elems, MAX_DQL * elems, (1, BLOCK_Q * elems)),

105 + )

106 + sQ_layout = cute.make_composed_layout(swizzle_128B, 0, sQ_layout)

107 + Q_tma = cpasync.make_tiled_tma_atom(

108 + tma_g2s,

109 + cute.logical_divide(gQ, (None, None, elems)),

110 + sQ_layout,

111 + cta_tiler=(MAX_DQL, num_heads, head_dim),

112 + )

113 +

114 + sK_layout = cute.make_layout(

115 + (1, BLOCK_K, (elems, head_dim // elems), num_stages),

116 + stride=(0, elems, (1, BLOCK_K * elems), BLOCK_K * head_dim),

117 + )

118 + sK_layout = cute.make_composed_layout(swizzle_128B, 0, sK_layout)

119 + K_tma = cpasync.make_tiled_tma_atom(

120 + tma_g2s,

121 + cute.logical_divide(gK_cache, (None, None, elems)),

122 + sK_layout,

123 + cta_tiler=(1, BLOCK_K, head_dim),

124 + )

125 +

126 + self.kernel(

127 + Q_tma,

128 + K_tma,

129 + block_table,

130 + score,

131 + seq_lens,

132 + decode_query_len,

133 + ).launch(grid=grid, block=block, stream=stream, use_pdl=True)

134 +

135 + @cute.kernel

136 + def kernel(

137 + self,

138 + Q_tma: cpasync.TmaInfo,

139 + K_tma: cpasync.TmaInfo,

140 + block_table: cute.Tensor,

141 + score: cute.Tensor,

142 + seq_lens: cute.Tensor,

143 + decode_query_len,

144 + ):

145 + tid, _, _ = cute.arch.thread_idx()

146 + batch_id, split_id, _ = cute.arch.block_idx()

147 + _, split_k, _ = cute.arch.grid_dim()

148 + warp_id = cute.arch.make_warp_uniform(cute.arch.warp_idx())

149 + lane_id = cute.arch.lane_idx()

150 +

151 + NUM_HEADS = self.num_heads

152 + MAX_DQL = self.max_decode_query_len

153 + BLOCK_Q = NUM_HEADS * MAX_DQL

154 + BLOCK_K = self.BLOCK_K

155 + head_dim = self.head_dim

156 + dtype = self.dtype

157 + MMA_N = 8

158 + num_stages = self.num_stages

159 + Q_TILES = cute.ceil_div(BLOCK_Q, MMA_N)

160 + EPI_Q = Q_TILES * MMA_N

161 +

162 + smem = cutlass.utils.SmemAllocator()

163 + sK = smem.allocate_tensor(

164 + dtype,

165 + K_tma.smem_layout.outer,

166 + byte_alignment=128,

167 + swizzle=K_tma.smem_layout.inner,

168 + )[0, None, None, None]

169 + # alias sQ with the 1st stage of sK

170 + sQ_tma = cute.make_tensor(

171 + sK[None, None, 0].iterator, layout=Q_tma.smem_layout.outer

172 + )

173 + # TMA sees Q as (query, head, dim), while ldmatrix consumes a

174 + # flattened Q column mode. The target profile keeps the rank-2 view even

175 + # for degenerate shapes like DQL1.

176 + q_tma_elems = 128 * 8 // dtype.width

177 + sQ = cute.coalesce(

178 + cute.group_modes(sQ_tma, 0, 2),

179 + target_profile=(BLOCK_Q, (q_tma_elems, head_dim // q_tma_elems)),

180 + )

181 + epi_buffer = smem.allocate_tensor(Float32, cute.make_layout((EPI_Q, 4)))

182 +

183 + tma_full_mbar = smem.allocate_array(Int64, num_stages)

184 + tma_empty_mbar = smem.allocate_array(Int64, num_stages)

185 +

186 + seqlen = seq_lens[batch_id]

187 + num_blocks = cute.ceil_div(seqlen, BLOCK_K)

188 +

189 + if split_id < num_blocks:

190 + if warp_id == 0:

191 + with cute.arch.elect_one():

192 + for i in cutlass.range_constexpr(num_stages):

193 + cute.arch.mbarrier_init(tma_full_mbar + i, 1)

194 + cute.arch.mbarrier_init(tma_empty_mbar + i, 128)

195 + cute.arch.mbarrier_init_fence()

196 + elif warp_id == 1:

197 + cpasync.prefetch_descriptor(Q_tma.atom)

198 + cpasync.prefetch_descriptor(K_tma.atom)

199 + cute.arch.sync_threads()

200 +

201 + cute.arch.griddepcontrol_wait()

202 + cute.arch.griddepcontrol_launch_dependents()

203 +

204 + if warp_id == 4:

205 + # TMA warp

206 + tma_stage = 0

207 + tma_parity = 1

208 +

209 + gQ_tile = cute.local_tile(

210 + cute.domain_offset(

211 + (batch_id * decode_query_len, 0, 0),

212 + Q_tma.tma_tensor,

213 + ),

214 + tiler=(MAX_DQL, NUM_HEADS, head_dim),

215 + coord=(0, 0, 0),

216 + )

217 + cute.arch.mbarrier_wait(tma_empty_mbar, tma_parity)

218 + with cute.arch.elect_one():

219 + Q_size = BLOCK_Q * head_dim * (dtype.width // 8)

220 + cute.arch.mbarrier_arrive_and_expect_tx(tma_full_mbar, Q_size)

221 + # TMA bounds-checks rows when runtime decode_query_len is smaller

222 + # than MAX_DQL; padded Q columns are masked before global stores.

223 + simple_tma_copy(Q_tma.atom, gQ_tile, sQ_tma, tma_full_mbar)

224 +

225 + tma_stage = (tma_stage + 1) % num_stages

226 + if tma_stage == 0:

227 + tma_parity ^= 1

228 +

229 + for block_id in range(split_id, num_blocks, split_k):

230 + page_id = block_table[batch_id, block_id]

231 + gK_tile = K_tma.tma_tensor[page_id, None, None]

232 + k_mbar = tma_full_mbar + tma_stage

233 +

234 + cute.arch.mbarrier_wait(tma_empty_mbar + tma_stage, tma_parity)

235 + with cute.arch.elect_one():

236 + K_size = BLOCK_K * head_dim * (dtype.width // 8)

237 + cute.arch.mbarrier_arrive_and_expect_tx(k_mbar, K_size)

238 + simple_tma_copy(

239 + K_tma.atom,

240 + gK_tile,

241 + sK[None, None, tma_stage],

242 + k_mbar,

243 + cache_policy=EVICT_FIRST,

244 + )

245 +

246 + tma_stage = (tma_stage + 1) % num_stages

247 + if tma_stage == 0:

248 + tma_parity ^= 1

249 +

250 + else:

251 + # MMA warps

252 + # each warp handles K[32, head_dim] @ Q[BLOCK_Q, head_dim].T

253 + sK_warp = cute.local_tile(

254 + sK, (32, head_dim, num_stages), (warp_id, 0, 0)

255 + )

256 + q_start = seqlen - decode_query_len

257 +

258 + elems = 128 // dtype.width # 16B

259 + MMA_K = 32 * 8 // dtype.width # 32B

260 +

261 + # Pre-compute ldmatrix address.

262 + # sK loads a [16 x 16B] tile:

263 + # ((16, (16B, 2), 1), (32 / 16, head_dim / 32B, num_stages))

264 + # sQ loads an [8 x 32B] tile:

265 + # ((8, (16B, 4)), (BLOCK_Q / MMA_N, head_dim / 64B))

266 + sK_ldsm = cute.zipped_divide(

267 + sK_warp, (16, cute.make_layout((elems, 2)), 1)

268 + )

269 + sQ_ldsm = cute.zipped_divide(sQ, (MMA_N, cute.make_layout((elems, 4))))

270 +

271 + # sK: (16B, (32 / 16, head_dim / 32B, num_stages))

272 + # sQ: (16B, (BLOCK_Q / MMA_N, head_dim / 64B))

273 + sK_ldsm = sK_ldsm[(lane_id % 16, (None, lane_id // 16), 0), None]

274 + sQ_ldsm = sQ_ldsm[(lane_id % MMA_N, (None, lane_id // 8)), None]

275 +

276 + ldsm_op = warp.LdMatrix8x8x16bOp(num_matrices=4)

277 + ldsm_atom = cute.make_copy_atom(ldsm_op, dtype)

278 +

279 + rQ = cute.make_rmem_tensor(

280 + ((elems // 2, 2), head_dim // (MMA_K * 2), Q_TILES), dtype

281 + )

282 + rK = cute.make_rmem_tensor((elems, 2, head_dim // MMA_K), dtype)

283 + rC = cute.make_rmem_tensor((4, 2, Q_TILES), Float32)

284 +

285 + if warp_id == 0:

286 + cute.arch.mbarrier_wait(tma_full_mbar, 0)

287 + cute.arch.barrier(barrier_id=self.BAR_MMA, number_of_threads=128)

288 + for q in cutlass.range_constexpr(Q_TILES):

289 + cute.copy(ldsm_atom, sQ_ldsm[None, (q, None)], rQ[None, None, q])

290 + cute.arch.mbarrier_arrive(tma_empty_mbar)

291 +

292 + tma_stage = 1 % self.num_stages

293 + tma_parity = 0

294 + if tma_stage == 0:

295 + tma_parity ^= 1

296 +

297 + # sm100 doesn't have native mma.sync.f8. ptxas lowers mma.sync.f8

298 + # to F2FP.F16.E4M3 + HMMA; doing the conversion explicitly gives

299 + # better codegen while keeping the two FP16 k-fragments visible.

300 + if cutlass.const_expr(dtype is Float8E4M3FN):

301 + rQ_f16 = cute.make_rmem_tensor(

302 + (4, head_dim // MMA_K, Q_TILES, 2), Float16

303 + )

304 + q_lower, q_upper = _fp8_to_f16_mma_fragments(rQ)

305 + rQ_f16[None, None, None, 0].store(q_lower.load())

306 + rQ_f16[None, None, None, 1].store(q_upper.load())

307 +

308 + for block_id in range(split_id, num_blocks, split_k):

309 + rC.fill(0.0)

310 +

311 + if warp_id == 0:

312 + cute.arch.mbarrier_wait(tma_full_mbar + tma_stage, tma_parity)

313 + cute.arch.barrier(barrier_id=self.BAR_MMA, number_of_threads=128)

314 +

315 + for k in cutlass.range_constexpr(head_dim // MMA_K):

316 + cute.copy(

317 + ldsm_atom,

318 + sK_ldsm[None, (None, k, tma_stage)],

319 + rK[None, None, k],

320 + )

321 + for m in cutlass.range_constexpr(2):

322 + if cutlass.const_expr(dtype is Float8E4M3FN):

323 + rK_lower, rK_upper = _fp8_to_f16_mma_fragments(

324 + rK[None, m, k]

325 + )

326 + for n in cutlass.range_constexpr(Q_TILES):

327 + rC[None, m, n] = mma_sync(

328 + rK_lower,

329 + rQ_f16[None, k, n, 0],

330 + rC[None, m, n],

331 + )

332 + rC[None, m, n] = mma_sync(

333 + rK_upper,

334 + rQ_f16[None, k, n, 1],

335 + rC[None, m, n],

336 + )

337 + else:

338 + for n in cutlass.range_constexpr(Q_TILES):

339 + rC[None, m, n] = mma_sync(

340 + rK[None, m, k],

341 + rQ[(None, k % 2), k // 2, n],

342 + rC[None, m, n],

343 + )

344 +

345 + cute.arch.mbarrier_arrive(tma_empty_mbar + tma_stage)

346 +

347 + k_start = block_id * BLOCK_K + warp_id * 32

348 +

349 + # causal mask

350 + for q in cutlass.range_constexpr(Q_TILES):

351 + for i in cutlass.range_constexpr(4):

352 + for j in cutlass.range_constexpr(2):

353 + col = q * 8 + (lane_id % 4) * 2 + j

354 + q_local_pos = col % MAX_DQL

355 + q_pos = q_start + q_local_pos

356 + k_pos = k_start + i * 8 + lane_id // 4

357 + rC[q * 8 + i * 2 + j] = (

358 + rC[q * 8 + i * 2 + j]

359 + if q_pos >= k_pos

360 + else float("-inf")

361 + )

362 +

363 + for q in cutlass.range_constexpr(Q_TILES):

364 + # thread-reduction along BLOCK_K dim

365 + rScore = cute.make_rmem_tensor(2, Float32)

366 + rScore.fill(float("-inf"))

367 + for i in cutlass.range_constexpr(4):

368 + rScore[0] = cute.arch.fmax(rScore[0], rC[i * 2 + 0 + q * 8])

369 + rScore[1] = cute.arch.fmax(rScore[1], rC[i * 2 + 1 + q * 8])

370 +

371 + # warp-reduction among lanes 0,4,8,12,...

372 + for i in cutlass.range_constexpr(3):

373 + offset = 4 << i

374 + other0 = cute.arch.shuffle_sync_bfly(

375 + rScore[0], offset=offset, mask=-1, mask_and_clamp=31

376 + )

377 + other1 = cute.arch.shuffle_sync_bfly(

378 + rScore[1], offset=offset, mask=-1, mask_and_clamp=31

379 + )

380 + rScore[0] = cute.arch.fmax(rScore[0], other0)

381 + rScore[1] = cute.arch.fmax(rScore[1], other1)

382 +

383 + # store to smem for 4-warp reduction

384 + if lane_id * 2 < MMA_N:

385 + epi_buffer[q * MMA_N + lane_id * 2 + 0, warp_id] = rScore[0]

386 + epi_buffer[q * MMA_N + lane_id * 2 + 1, warp_id] = rScore[1]

387 + cute.arch.barrier(barrier_id=self.BAR_MMA, number_of_threads=128)

388 +

389 + head_id = lane_id // MAX_DQL

390 + q_local_pos = lane_id - head_id * MAX_DQL

391 + valid_q = head_id < NUM_HEADS and q_local_pos < decode_query_len

392 + if lane_id < BLOCK_Q and valid_q:

393 + final_score = epi_buffer[lane_id, 0]

394 + for i in cutlass.range_constexpr(1, 4):

395 + final_score = cute.arch.fmax(

396 + final_score, epi_buffer[lane_id, i]

397 + )

398 +

399 + t = batch_id * decode_query_len + q_local_pos

400 + score[head_id, t, block_id] = final_score

401 +

402 + tma_stage = (tma_stage + 1) % self.num_stages

403 + if tma_stage == 0:

404 + tma_parity ^= 1

405 +

406 + @cache

407 + @staticmethod

408 + def compile(

409 + dtype: type[cutlass.Numeric],

410 + num_heads: int,

411 + max_decode_query_len: int,

412 + split_k: int,

413 + head_dim: int = 128,

414 + ):

415 + bs = cute.sym_int()

416 + total_tokens = cute.sym_int()

417 + BLOCK_K = IndexDecodeScoreKernel.BLOCK_K

418 +

419 + q = make_fake_tensor(

420 + dtype, (total_tokens, num_heads, head_dim), divisibility=16

421 + )

422 + k_cache = make_fake_tensor(

423 + dtype, (cute.sym_int(), BLOCK_K, head_dim), divisibility=16

424 + )

425 + block_table = make_fake_tensor(Int32, (bs, cute.sym_int()), divisibility=1)

426 + score = make_fake_tensor(

427 + Float32, (num_heads, total_tokens, cute.sym_int()), divisibility=4

428 + )

429 + seq_lens = make_fake_tensor(Int32, (bs,), divisibility=1)

430 + kernel = IndexDecodeScoreKernel(

431 + dtype,

432 + num_heads,

433 + max_decode_query_len,

434 + split_k,

435 + head_dim,

436 + )

437 + stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)

438 + return cute.compile(

439 + kernel,

440 + q,

441 + k_cache,

442 + block_table,

443 + score,

444 + seq_lens,

445 + stream,

446 + options="--enable-tvm-ffi",

447 + )

448 +

449 +

450 +def minimax_m3_index_decode_score_cutedsl(

451 + idx_q: torch.Tensor,

452 + index_kv_cache: torch.Tensor,

453 + block_table: torch.Tensor,

454 + seq_lens: torch.Tensor,

455 + max_seq_len: int,

456 + init_blocks: int,

457 + local_blocks: int,

458 + num_kv_heads: int,

459 + decode_query_len: int,

460 + max_decode_query_len: int,

461 + score_out: torch.Tensor,

462 +) -> torch.Tensor:

463 + if idx_q.dtype not in (torch.bfloat16, torch.float8_e4m3fn):

464 + raise TypeError("CuteDSL index decode score supports BF16 and FP8 E4M3 only")

465 + total_tokens, num_heads, head_dim = idx_q.shape

466 + batch = block_table.shape[0]

467 + assert index_kv_cache.shape[1] == IndexDecodeScoreKernel.BLOCK_K

468 + assert total_tokens == batch * decode_query_len

469 + assert 1 <= decode_query_len <= max_decode_query_len

470 + assert num_heads * max_decode_query_len <= 32

471 + dtype = _TORCH_TO_CUTE_DTYPE[idx_q.dtype]

472 + del max_seq_len, init_blocks, local_blocks, num_kv_heads

473 + score = score_out

474 + split_k = 256

475 + kernel = IndexDecodeScoreKernel.compile(

476 + dtype,

477 + num_heads,

478 + max_decode_query_len,

479 + split_k,

480 + head_dim,

481 + )

482 + kernel(idx_q, index_kv_cache, block_table, score, seq_lens)

483 + return score