From 1cbb100e631dfc8607998ab7c5ff9bd603347788 Mon Sep 17 00:00:00 2001 From: Cam Quilici Date: Sun, 14 Jun 2026 00:23:33 -0500 Subject: [PATCH 1/3] [Bugfix][ROCm] Fix MiniMax-M3 FP8 KV cache dtype Signed-off-by: Cam Quilici --- tests/kernels/attention/test_minimax_m3.py | 38 +++++++++++++++++++ .../minimax_m3/common/ops/sparse_attn.py | 9 ++++- .../minimax_m3/common/sparse_attention.py | 4 +- 3 files changed, 48 insertions(+), 3 deletions(-) diff --git a/tests/kernels/attention/test_minimax_m3.py b/tests/kernels/attention/test_minimax_m3.py index 32b4bc97ede7..faff3c476857 100644 --- a/tests/kernels/attention/test_minimax_m3.py +++ b/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,42 @@ 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_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, + ], +) +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, diff --git a/vllm/models/minimax_m3/common/ops/sparse_attn.py b/vllm/models/minimax_m3/common/ops/sparse_attn.py index 40287e166b2a..2b82c5db2ee2 100644 --- a/vllm/models/minimax_m3/common/ops/sparse_attn.py +++ b/vllm/models/minimax_m3/common/ops/sparse_attn.py @@ -24,6 +24,11 @@ # One sparse block == one KV page. SPARSE_BLOCK_SIZE = 128 +_FP8_DTYPES = ( + torch.float8_e4m3fn, + torch.float8_e4m3fnuz, + torch.float8_e5m2, +) _SPARSE_ATTN_NUM_STAGES_KWARG: dict | None = None @@ -456,7 +461,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 +518,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 diff --git a/vllm/models/minimax_m3/common/sparse_attention.py b/vllm/models/minimax_m3/common/sparse_attention.py index 8cca0e8e2998..9fcf0b8bb25d 100644 --- a/vllm/models/minimax_m3/common/sparse_attention.py +++ b/vllm/models/minimax_m3/common/sparse_attention.py @@ -292,7 +292,9 @@ def __init__( 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 + torch.float8_e5m2 + if "e5m2" in kv_cache_dtype + else current_platform.fp8_dtype() ) # Sparse selection parameters (block_size == page size == SPARSE_BLOCK_SIZE). self.topk_blocks = topk_blocks From 3e9faf5d02adf420fe3b69058f8ea328411053d4 Mon Sep 17 00:00:00 2001 From: Cameron Quilici Date: Sun, 14 Jun 2026 11:25:06 -0500 Subject: [PATCH 2/3] Update sparse_attn.py Co-authored-by: Hongxia Yang <62075498+hongxiayang@users.noreply.github.com> Signed-off-by: Cameron Quilici --- vllm/models/minimax_m3/common/ops/sparse_attn.py | 1 + 1 file changed, 1 insertion(+) diff --git a/vllm/models/minimax_m3/common/ops/sparse_attn.py b/vllm/models/minimax_m3/common/ops/sparse_attn.py index 2b82c5db2ee2..7b6fb73cba9a 100644 --- a/vllm/models/minimax_m3/common/ops/sparse_attn.py +++ b/vllm/models/minimax_m3/common/ops/sparse_attn.py @@ -28,6 +28,7 @@ torch.float8_e4m3fn, torch.float8_e4m3fnuz, torch.float8_e5m2, + torch.float8_e5m2fnuz, ) _SPARSE_ATTN_NUM_STAGES_KWARG: dict | None = None From 7e7d9df4ec1e3aef7eb309410a8034326589c517 Mon Sep 17 00:00:00 2001 From: Cam Quilici Date: Sun, 14 Jun 2026 15:35:31 -0500 Subject: [PATCH 3/3] fix(rocm): handle E5M2 FNUZ KV cache Signed-off-by: Cam Quilici --- tests/kernels/attention/test_minimax_m3.py | 8 +++++++- vllm/models/minimax_m3/common/sparse_attention.py | 13 ++++++++----- 2 files changed, 15 insertions(+), 6 deletions(-) diff --git a/tests/kernels/attention/test_minimax_m3.py b/tests/kernels/attention/test_minimax_m3.py index faff3c476857..a69f85f7e08c 100644 --- a/tests/kernels/attention/test_minimax_m3.py +++ b/tests/kernels/attention/test_minimax_m3.py @@ -86,7 +86,12 @@ def _allocate_main_kv_via_contract( [ ("fp8", current_platform.fp8_dtype()), ("fp8_e4m3", current_platform.fp8_dtype()), - ("fp8_e5m2", torch.float8_e5m2), + ( + "fp8_e5m2", + torch.float8_e5m2fnuz + if current_platform.is_fp8_fnuz() + else torch.float8_e5m2, + ), ], ) def test_sparse_impl_uses_platform_fp8_dtype( @@ -111,6 +116,7 @@ def test_sparse_impl_uses_platform_fp8_dtype( torch.float8_e4m3fn, torch.float8_e4m3fnuz, torch.float8_e5m2, + torch.float8_e5m2fnuz, ], ) def test_sparse_kernels_recognize_fp8_dtypes(dtype: torch.dtype): diff --git a/vllm/models/minimax_m3/common/sparse_attention.py b/vllm/models/minimax_m3/common/sparse_attention.py index 9fcf0b8bb25d..ab8f24186e6b 100644 --- a/vllm/models/minimax_m3/common/sparse_attention.py +++ b/vllm/models/minimax_m3/common/sparse_attention.py @@ -291,11 +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 current_platform.fp8_dtype() - ) + 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