[Bugfix][ROCm] Fix MiniMax-M3 FP8 KV cache dtype

tests/kernels/attention/test_minimax_m3.py

@@ -15,11 +15,13 @@
     minimax_m3_index_topk,
 )
 from vllm.models.minimax_m3.common.ops.sparse_attn import (
+    _FP8_DTYPES,
     minimax_m3_sparse_attn,
     minimax_m3_sparse_attn_decode,
 )
 from vllm.models.minimax_m3.common.sparse_attention import (
     MiniMaxM3SparseBackend,
+    MiniMaxM3SparseTritonImpl,
 )
 from vllm.platforms import current_platform
 from vllm.v1.attention.backends.utils import set_kv_cache_layout
@@ -79,6 +81,48 @@ def _allocate_main_kv_via_contract(
 TOPK = 16
 
 
+@pytest.mark.parametrize(
+    ("kv_cache_dtype", "expected_dtype"),
+    [
+        ("fp8", current_platform.fp8_dtype()),
+        ("fp8_e4m3", current_platform.fp8_dtype()),
+        (
+            "fp8_e5m2",
+            torch.float8_e5m2fnuz
+            if current_platform.is_fp8_fnuz()
+            else torch.float8_e5m2,
+        ),
+    ],
+)
+def test_sparse_impl_uses_platform_fp8_dtype(
+    kv_cache_dtype: str,
+    expected_dtype: torch.dtype,
+):
+    impl = MiniMaxM3SparseTritonImpl(
+        num_heads=NUM_Q_HEADS,
+        head_size=HEAD_DIM,
+        scale=SM_SCALE,
+        num_kv_heads=NUM_KV_HEADS,
+        kv_cache_dtype=kv_cache_dtype,
+        topk_blocks=TOPK,
+        sparse_block_size=BLOCK_SIZE,
+    )
+    assert impl.kv_cache_fp8_dtype == expected_dtype
+
+
+@pytest.mark.parametrize(
+    "dtype",
+    [
+        torch.float8_e4m3fn,
+        torch.float8_e4m3fnuz,
+        torch.float8_e5m2,
+        torch.float8_e5m2fnuz,
+    ],
+)
+def test_sparse_kernels_recognize_fp8_dtypes(dtype: torch.dtype):
+    assert dtype in _FP8_DTYPES
+
+
 # Index top-k kernels.
 def _reference_index_topk(
     idx_q: torch.Tensor,

vllm/models/minimax_m3/common/ops/sparse_attn.py

@@ -24,6 +24,12 @@
 # One sparse block == one KV page.
 SPARSE_BLOCK_SIZE = 128
 
+_FP8_DTYPES = (
+    torch.float8_e4m3fn,
+    torch.float8_e4m3fnuz,
+    torch.float8_e5m2,
+    torch.float8_e5m2fnuz,
+)
 
 _SPARSE_ATTN_NUM_STAGES_KWARG: dict | None = None
 
@@ -456,7 +462,7 @@ def minimax_m3_sparse_attn(
     batch = cu_seqlens_q.shape[0] - 1
     topk = topk_idx.shape[-1]
     gqa_group_size = num_heads // num_kv_heads
-    use_fp8 = kv_cache.dtype in (torch.float8_e4m3fn, torch.float8_e5m2)
+    use_fp8 = kv_cache.dtype in _FP8_DTYPES
     grid = (max_query_len, num_kv_heads, batch)
     _gqa_sparse_fwd_kernel[grid](
         q,
@@ -513,7 +519,7 @@ def minimax_m3_sparse_attn_decode(
     assert total_q == seq_lens.shape[0] * decode_query_len
     max_topk = topk_idx.shape[-1]
     gqa_group_size = num_heads // num_kv_heads
-    use_fp8 = kv_cache.dtype in (torch.float8_e4m3fn, torch.float8_e5m2)
+    use_fp8 = kv_cache.dtype in _FP8_DTYPES
     use_pdl = current_platform.is_arch_support_pdl()
     # `launch_pdl` is a Triton runtime kwarg only some backends accept (CUDA
     # SM9+); this ROCm Triton rejects it even when False ("Keyword argument

vllm/models/minimax_m3/common/sparse_attention.py

@@ -291,9 +291,14 @@ def __init__(
         self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads
         self.kv_cache_dtype = kv_cache_dtype
         self.use_fp8_kv = is_quantized_kv_cache(kv_cache_dtype)
-        self.kv_cache_fp8_dtype = (
-            torch.float8_e5m2 if "e5m2" in kv_cache_dtype else torch.float8_e4m3fn
-        )
+        if "e5m2" in kv_cache_dtype:
+            self.kv_cache_fp8_dtype = (
+                torch.float8_e5m2fnuz
+                if current_platform.is_fp8_fnuz()
+                else torch.float8_e5m2
+            )
+        else:
+            self.kv_cache_fp8_dtype = current_platform.fp8_dtype()
         # Sparse selection parameters (block_size == page size == SPARSE_BLOCK_SIZE).
         self.topk_blocks = topk_blocks
         self.block_size = sparse_block_size