A high-throughput and memory-efficient inference and serving engine for LLMs
[Bugfix][ROCm] Fix MiniMax-M3 FP8 KV cache dtype (#45720)
Signed-off-by: Cam Quilici <cjquilici@gmail.com> Signed-off-by: Cameron Quilici <cjquilici@gmail.com> Co-authored-by: Hongxia Yang <62075498+hongxiayang@users.noreply.github.com> Co-authored-by: TJian <tunjian.tan@embeddedllm.com>
Cameron Quilici committed
Jun 16, 2026 at 22:14 UTC
efd15e192a1a8c5372a46643555d87c9b8a662cb
3 files changed
+60
-5
tests/kernels/attention/test_minimax_m3.py
+44
@@ -15,11 +15,13 @@ from vllm.models.minimax_m3.common.ops.index_topk import (
15
minimax_m3_index_topk,
16
)
17
from vllm.models.minimax_m3.common.ops.sparse_attn import (
18
+ _FP8_DTYPES,
19
minimax_m3_sparse_attn,
20
minimax_m3_sparse_attn_decode,
21
)
22
from vllm.models.minimax_m3.common.sparse_attention import (
23
MiniMaxM3SparseBackend,
24
+ MiniMaxM3SparseTritonImpl,
25
)
26
from vllm.platforms import current_platform
27
from vllm.v1.attention.backends.utils import set_kv_cache_layout
@@ -79,6 +81,48 @@ SM_SCALE = HEAD_DIM**-0.5
81
TOPK = 16
82
83
84
+@pytest.mark.parametrize(
85
+ ("kv_cache_dtype", "expected_dtype"),
86
+ [
87
+ ("fp8", current_platform.fp8_dtype()),
88
+ ("fp8_e4m3", current_platform.fp8_dtype()),
89
+ (
90
+ "fp8_e5m2",
91
+ torch.float8_e5m2fnuz
92
+ if current_platform.is_fp8_fnuz()
93
+ else torch.float8_e5m2,
94
+ ),
95
+ ],
96
+)
97
+def test_sparse_impl_uses_platform_fp8_dtype(
98
+ kv_cache_dtype: str,
99
+ expected_dtype: torch.dtype,
100
+):
101
+ impl = MiniMaxM3SparseTritonImpl(
102
+ num_heads=NUM_Q_HEADS,
103
+ head_size=HEAD_DIM,
104
+ scale=SM_SCALE,
105
+ num_kv_heads=NUM_KV_HEADS,
106
+ kv_cache_dtype=kv_cache_dtype,
107
+ topk_blocks=TOPK,
108
+ sparse_block_size=BLOCK_SIZE,
109
+ )
110
+ assert impl.kv_cache_fp8_dtype == expected_dtype
111
+
112
+
113
+@pytest.mark.parametrize(
114
+ "dtype",
115
+ [
116
+ torch.float8_e4m3fn,
117
+ torch.float8_e4m3fnuz,
118
+ torch.float8_e5m2,
119
+ torch.float8_e5m2fnuz,
120
+ ],
121
+)
122
+def test_sparse_kernels_recognize_fp8_dtypes(dtype: torch.dtype):
123
+ assert dtype in _FP8_DTYPES
124
+
125
+
126
# Index top-k kernels.
127
def _reference_index_topk(
128
idx_q: torch.Tensor,
vllm/models/minimax_m3/common/ops/sparse_attn.py
+8
-2
@@ -24,6 +24,12 @@ from vllm.triton_utils import tl, triton
24
# One sparse block == one KV page.
25
SPARSE_BLOCK_SIZE = 128
26
27
+_FP8_DTYPES = (
28
+ torch.float8_e4m3fn,
29
+ torch.float8_e4m3fnuz,
30
+ torch.float8_e5m2,
31
+ torch.float8_e5m2fnuz,
32
+)
33
34
_SPARSE_ATTN_NUM_STAGES_KWARG: dict | None = None
35
@@ -456,7 +462,7 @@ def minimax_m3_sparse_attn(
462
batch = cu_seqlens_q.shape[0] - 1
463
topk = topk_idx.shape[-1]
464
gqa_group_size = num_heads // num_kv_heads
459
- use_fp8 = kv_cache.dtype in (torch.float8_e4m3fn, torch.float8_e5m2)
465
+ use_fp8 = kv_cache.dtype in _FP8_DTYPES
466
grid = (max_query_len, num_kv_heads, batch)
467
_gqa_sparse_fwd_kernel[grid](
468
q,
@@ -513,7 +519,7 @@ def minimax_m3_sparse_attn_decode(
519
assert total_q == seq_lens.shape[0] * decode_query_len
520
max_topk = topk_idx.shape[-1]
521
gqa_group_size = num_heads // num_kv_heads
516
- use_fp8 = kv_cache.dtype in (torch.float8_e4m3fn, torch.float8_e5m2)
522
+ use_fp8 = kv_cache.dtype in _FP8_DTYPES
523
use_pdl = current_platform.is_arch_support_pdl()
524
# `launch_pdl` is a Triton runtime kwarg only some backends accept (CUDA
525
# SM9+); this ROCm Triton rejects it even when False ("Keyword argument
vllm/models/minimax_m3/common/sparse_attention.py
+8
-3
@@ -291,9 +291,14 @@ class MiniMaxM3SparseImpl(AttentionImplBase[MiniMaxM3SparseMetadata]):
291
self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads
292
self.kv_cache_dtype = kv_cache_dtype
293
self.use_fp8_kv = is_quantized_kv_cache(kv_cache_dtype)
294
- self.kv_cache_fp8_dtype = (
295
- torch.float8_e5m2 if "e5m2" in kv_cache_dtype else torch.float8_e4m3fn
296
- )
294
+ if "e5m2" in kv_cache_dtype:
295
+ self.kv_cache_fp8_dtype = (
296
+ torch.float8_e5m2fnuz
297
+ if current_platform.is_fp8_fnuz()
298
+ else torch.float8_e5m2
299
+ )
300
+ else:
301
+ self.kv_cache_fp8_dtype = current_platform.fp8_dtype()
302
# Sparse selection parameters (block_size == page size == SPARSE_BLOCK_SIZE).
303
self.topk_blocks = topk_blocks
304
self.block_size = sparse_block_size