[Minimax M3] Enable fp8 indexer cache on triton indexer for non-SM100 architectures. (#59081)
Signed-off-by: Jarrel Seah <jarrel@harrison.ai> Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com> Signed-off-by: Rajas Mhaskar <rmhaskar@nvidia.com> Co-authored-by: Jarrel Seah <jarrel@harrison.ai> Co-authored-by: Tyler Michael Smith <tlrmchlsmth@gmail.com> Co-authored-by: OpenAI Codex <codex@openai.com>
rmhaskar committed
Sep 28, 2026 at 16:21 UTC
e03875f4d283bc003b08af351a0b316e18cd798b
3 files changed
+133
-9
tests/kernels/attention/test_minimax_m3_fp8_triton_indexer.py
new
+109
@@ -0,0 +1,109 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""fp8 query and index-K cache on the Triton indexer must agree with bf16.
4
+
5
+The MiniMax M3 lightning-indexer scores only feed a top-k block ranking (the
6
+attention itself reads the main KV cache), so an fp8 (e4m3) side cache trades
7
+~2% relative score error for half the per-step index-K read bandwidth — the
8
+dominant linear-in-context decode cost at long context. #45892 enabled fp8
9
+index caches on the SM100 MSA path; this covers the Triton impl used
10
+everywhere else (e.g. SM120), for both the decode and prefill scorers.
11
+"""
12
+
13
+import pytest
14
+import torch
15
+
16
+from vllm.models.minimax_m3.common.ops.index_topk import (
17
+ SPARSE_BLOCK_SIZE,
18
+ minimax_m3_index_decode,
19
+ minimax_m3_index_score,
20
+ minimax_m3_index_topk,
21
+)
22
+from vllm.platforms import current_platform
23
+
24
+FP8_CUDA_SUPPORTED = current_platform.is_cuda() and current_platform.supports_fp8()
25
+PAGE = SPARSE_BLOCK_SIZE
26
+HEAD_DIM = 128
27
+HEADS = 1
28
+TOPK = 16
29
+NUM_PAGES = 100
30
+SEQ_LEN = 96 * PAGE + 37
31
+PLANTED = (7, 40, 88)
32
+
33
+
34
+def _make_inputs(device: str) -> tuple[torch.Tensor, torch.Tensor]:
35
+ torch.manual_seed(3)
36
+ keys = torch.randn(NUM_PAGES * PAGE, HEAD_DIM, dtype=torch.bfloat16, device=device)
37
+ # RMS-normalize like the real (post index_k_norm) keys.
38
+ keys = keys * keys.float().pow(2).mean(-1, keepdim=True).add(1e-6).rsqrt().to(
39
+ torch.bfloat16
40
+ )
41
+ q = torch.randn(1, HEADS, HEAD_DIM, dtype=torch.bfloat16, device=device)
42
+ for b in PLANTED:
43
+ keys[b * PAGE + 5] = (q[0, 0] * 3).to(torch.bfloat16)
44
+ return q, keys
45
+
46
+
47
+def _check_selection(top_bf16: set[int], top_fp8: set[int]) -> None:
48
+ # Planted (clearly relevant) blocks must never be lost.
49
+ assert all(b in top_fp8 for b in PLANTED)
50
+ # Selection may differ only at the noise floor (near-tied filler blocks).
51
+ assert len(top_bf16 & top_fp8) >= TOPK - 2
52
+
53
+
54
+@pytest.mark.skipif(not FP8_CUDA_SUPPORTED, reason="CUDA FP8 support required")
55
+@torch.inference_mode()
56
+def test_fp8_index_query_and_cache_decode_topk_matches_bf16() -> None:
57
+ dev = "cuda"
58
+ q, keys = _make_inputs(dev)
59
+ block_table = torch.arange(NUM_PAGES, dtype=torch.int32, device=dev).unsqueeze(0)
60
+ seq_lens = torch.tensor([SEQ_LEN], dtype=torch.int32, device=dev)
61
+
62
+ def select(cache: torch.Tensor) -> set[int]:
63
+ query = q.to(cache.dtype)
64
+ idx = minimax_m3_index_decode(
65
+ query,
66
+ cache.view(NUM_PAGES, PAGE, HEAD_DIM).contiguous(),
67
+ block_table,
68
+ seq_lens,
69
+ SEQ_LEN,
70
+ TOPK,
71
+ 0,
72
+ 0,
73
+ HEADS,
74
+ 1,
75
+ 1,
76
+ )
77
+ return {int(b) for b in idx[0, 0].tolist() if b >= 0}
78
+
79
+ _check_selection(select(keys), select(keys.to(torch.float8_e4m3fn)))
80
+
81
+
82
+@pytest.mark.skipif(not FP8_CUDA_SUPPORTED, reason="CUDA FP8 support required")
83
+@torch.inference_mode()
84
+def test_fp8_index_query_and_cache_prefill_topk_matches_bf16() -> None:
85
+ dev = "cuda"
86
+ q, keys = _make_inputs(dev)
87
+ block_table = torch.arange(NUM_PAGES, dtype=torch.int32, device=dev).unsqueeze(0)
88
+ seq_lens = torch.tensor([SEQ_LEN], dtype=torch.int32, device=dev)
89
+ # single decode-like query token at the end of the sequence
90
+ cu_seqlens_q = torch.tensor([0, 1], dtype=torch.int32, device=dev)
91
+ prefix_lens = torch.tensor([SEQ_LEN - 1], dtype=torch.int32, device=dev)
92
+
93
+ def select(cache: torch.Tensor) -> set[int]:
94
+ query = q.to(cache.dtype)
95
+ score = minimax_m3_index_score(
96
+ query,
97
+ cache.view(NUM_PAGES, PAGE, HEAD_DIM).contiguous(),
98
+ block_table,
99
+ cu_seqlens_q,
100
+ seq_lens,
101
+ prefix_lens,
102
+ 1,
103
+ SEQ_LEN,
104
+ HEADS,
105
+ )
106
+ idx = minimax_m3_index_topk(score, cu_seqlens_q, prefix_lens, 1, TOPK, 0, 0)
107
+ return {int(b) for b in idx[0, 0].tolist() if b >= 0}
108
+
109
+ _check_selection(select(keys), select(keys.to(torch.float8_e4m3fn)))
vllm/models/minimax_m3/common/indexer.py
+10
-5
@@ -502,16 +502,21 @@ def select_indexer_impl_cls(
502
On Blackwell (SM100) with ``topk_blocks == 16`` (the only width fmha_sm100's
503
``sparse_topk_select`` kernel supports), the fmha_sm100 score + top-k path is
504
used for both bf16 and fp8 index caches. Everything else falls back to the
505
- Triton indexer (bf16 only).
505
+ Triton indexer, with fp8 restricted to CUDA platforms that advertise fp8
506
+ support.
507
"""
508
if indexer_kv_dtype in ("mxfp4", "nvfp4"):
509
raise NotImplementedError(
510
f"indexer_kv_dtype={indexer_kv_dtype!r} needs the (not-yet-added) "
511
"CuteDSL indexer impl."
512
)
512
- is_sm100 = (
513
- current_platform.is_cuda() and current_platform.is_device_capability_family(100)
514
- )
513
+ is_cuda = current_platform.is_cuda()
514
+ is_fp8 = indexer_kv_dtype in ("fp8", "fp8_e4m3")
515
+ if is_fp8 and not (is_cuda and current_platform.supports_fp8()):
516
+ raise NotImplementedError(
517
+ f"indexer_kv_dtype={indexer_kv_dtype!r} requires CUDA fp8 support."
518
+ )
519
+ is_sm100 = is_cuda and current_platform.is_device_capability_family(100)
520
use_msa = (
521
is_sm100
522
and topk_blocks == 16
@@ -530,7 +535,7 @@ def select_indexer_impl_cls(
535
indexer_kv_dtype,
536
)
537
return MiniMaxM3IndexerMSAImpl
533
- if indexer_kv_dtype != "bf16":
538
+ if indexer_kv_dtype not in ("bf16", "fp8", "fp8_e4m3"):
539
raise NotImplementedError(
540
f"indexer_kv_dtype={indexer_kv_dtype!r} is not supported by the "
541
"Triton indexer impl."
vllm/models/minimax_m3/common/ops/index_topk.py
+14
-4
@@ -123,6 +123,10 @@ def _index_block_score_kernel(
123
order=(1, 0),
124
)
125
q = tl.load(q_ptrs, boundary_check=(0,), padding_option="zero")
126
+ # FP8 has no mixed-dtype dot with bf16/fp32. Leave bf16 and fp32 loads
127
+ # in their stored dtype so fp32 pipeline tests keep full precision.
128
+ if q.dtype.is_fp8():
129
+ q = q.to(tl.bfloat16)
130
q_start = prefix_len + pid_q * BLOCK_SIZE_Q
131
132
off_q = tl.arange(0, BLOCK_SIZE_Q) + pid_q * BLOCK_SIZE_Q + prefix_len
@@ -146,7 +150,9 @@ def _index_block_score_kernel(
150
+ off_k[None, :] * stride_ik_pos
151
+ off_d[:, None] * stride_ik_d,
152
)
149
- qk = tl.dot(q, k)
153
+ if k.dtype.is_fp8():
154
+ k = k.to(tl.bfloat16)
155
+ qk = tl.dot(q, k, out_dtype=tl.float32)
156
# apply causal mask as needed
157
if q_start < i + BLOCK_SIZE_K:
158
qk = tl.where(off_q[:, None] >= pos[None, :], qk, float("-inf"))
@@ -360,6 +366,8 @@ def _decode_index_score_kernel(
366
mask=q_mask[None, :],
367
other=0.0,
368
) # [D,HQ]
369
+ if q.dtype.is_fp8():
370
+ q = q.to(tl.bfloat16)
371
for blk in tl.range(chunk_start_block, chunk_end_block):
372
page = tl.load(bt_row + blk).to(tl.int64)
373
pos = blk * BLOCK_SIZE_K + off_k
@@ -373,9 +381,11 @@ def _decode_index_score_kernel(
381
+ off_k[:, None] * stride_ik_pos
382
+ off_d * stride_ik_d,
383
) # [N,D]
376
- # fp32 accumulation is required for the fp8 (e4m3) index cache: q/k are
377
- # loaded in their stored dtype (bf16 or e4m3) and the MMA accumulates in
378
- # fp32 so the per-block max score is exact for the fp8 indexer too.
384
+ # Upcast only FP8 cache loads. BF16 operands keep FP8 dot
385
+ # instructions out of the fallback path; FP32 accumulation
386
+ # preserves score accuracy. BF16/FP32 loads stay as stored.
387
+ if k.dtype.is_fp8():
388
+ k = k.to(tl.bfloat16)
389
kq = tl.dot(k, q, out_dtype=tl.float32) # [N,HQ]
390
kq = tl.where(pos_mask & q_mask[None, :], kq, float("-inf"))
391
score = tl.max(kq, axis=0) # [HQ]