[ROCm] MiniMax-M3: allow the AMD monokernel decode on MiniMax-M3.1

recipes/models/MiniMaxAI/MiniMax-M3.1-Flash.yaml

@@ -124,15 +124,18 @@ features:
           blackwell: unsupported
   mono_decode:
     label: "Mono decode"
-    description: "AMD MI355X single node, MXFP4 (Quark) checkpoint, low concurrency: a decode step of up to 16 tokens (4 requests with DSpark 3) runs each layer as one fused kernel; larger steps and steps with prefill run on the normal path. Uses FP8 KV and indexer K caches (fewer KV tokens than NVFP4) and FULL_DECODE_ONLY graphs, so prefill and mixed steps run eager."
+    description: "AMD MI355X single node, MXFP4 (Quark) checkpoint, low concurrency: decode steps of up to 16 tokens run fused kernels; larger steps and steps with prefill run on the normal path. Mono fuses each whole layer and keeps DSpark, but needs FP8 KV and indexer K caches. Monokernel fuses everything after attention and keeps the NVFP4 caches, but refuses DSpark: turn Speculative decoding off with it."
     strategies: [single_node_tp]
+    default_mode: mono
     modes:
       # The kernels read the experts as AITER MXFP4 tables, which only the
       # Quark export loads as; they need the AITER sparse PA path, i.e. FP8 KV
       # and FP8 indexer K caches in the shuffled KV layout, and no piecewise
       # graphs.
       mono:
         variants: [mxfp4]
+        label: "Mono"
+        description: "Each layer as one fused kernel, for decode steps of up to 16 tokens (4 requests with DSpark 3). FP8 KV and indexer K caches (fewer KV tokens than NVFP4) and FULL_DECODE_ONLY graphs, so prefill and mixed steps run eager."
         args:
           - "--kv-cache-dtype"
           - "fp8"
@@ -152,6 +155,29 @@ features:
           mi300x: unsupported
           mi325x: unsupported
           mi350x: unsupported
+      # The fused launch reads the shared expert as a 129th AITER MXFP4 expert,
+      # so the shared expert is re-quantized from MXFP8 to MXFP4 at load and
+      # fused. The fused QKV must have 2560 rows, so indexer_decode_cp is off.
+      # The NVFP4 KV cache comes from the amd baseline.
+      monokernel:
+        variants: [mxfp4]
+        label: "Monokernel"
+        description: "Everything after attention (all-reduces, norm, router, shared and routed experts) as one persistent kernel, for decode steps of up to 16 tokens, on the V2 model runner. Keeps the NVFP4 KV and indexer K caches. Re-quantizes the shared expert to MXFP4. No DSpark: turn Speculative decoding off."
+        args:
+          - "--monokernel"
+          - "--attention-config"
+          - '{"indexer_kv_dtype": "nvfp4", "indexer_decode_cp": false}'
+          - "--quantization-config.targets"
+          - '{"*shared_experts*": "mxfp4"}'
+        env:
+          VLLM_USE_V2_MODEL_RUNNER: "1"
+          VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: "1"
+        hardware:
+          hopper: unsupported
+          blackwell: unsupported
+          mi300x: unsupported
+          mi325x: unsupported
+          mi350x: unsupported
       # No mono path for the other checkpoints: an empty mode, so enabling the
       # feature there changes nothing (the recipe schema wants a mode per
       # variant).

tests/models/minimax_m3/amd_specialized_decode/test_runtime.py

@@ -151,8 +151,8 @@ def test_config_refusals_report_every_reason():
     config.cache_config.block_size = 64
     config.aux_output_config.enable_return_routed_experts = True
     config.model_config.hf_text_config.sparse_attention_config["sparse_attention_freq"][
-        2
-    ] = 1
+        10
+    ] = 0
     with pytest.raises(ValueError) as error:
         specialized_decode_config(config)
     for reason in (
@@ -167,6 +167,17 @@ def test_config_refusals_report_every_reason():
         specialized_decode_config(_vllm_config(layers=3))
 
 
+def test_minimax_m31_layout_is_accepted():
+    # MiniMax-M3.1: NVFP4 KV and indexer cache, sparse attention on every layer
+    config = _vllm_config()
+    config.cache_config.cache_dtype = "nvfp4"
+    config.attention_config.resolve_indexer_kv_dtype = lambda _default: "nvfp4"
+    config.model_config.hf_text_config.sparse_attention_config[
+        "sparse_attention_freq"
+    ] = [1] * 60
+    assert specialized_decode_config(config).last_layer == 59
+
+
 @pytest.mark.parametrize("method", ["eagle3", "ngram", "mtp", "draft_model"])
 def test_only_eagle3_speculation_is_accepted(method):
     config = _vllm_config()

vllm/models/minimax_m3/amd/specialized_decode/runtime/spec.py

@@ -40,16 +40,18 @@ def refusal_reasons(self, config: VllmConfig) -> list[str]:
             reasons.append("does not support expert parallelism")
         if pc.enable_eplb:
             reasons.append("does not support EPLB")
-        if cc.cache_dtype not in ("fp8", "fp8_e4m3"):
-            reasons.append("requires an FP8 KV cache")
+        # The fused launch starts after attention and reads no KV cache, so
+        # either cache format the stock attention serves works.
+        if cc.cache_dtype not in ("fp8", "fp8_e4m3", "nvfp4"):
+            reasons.append("requires an FP8 or NVFP4 KV cache")
         if cc.block_size != 128:
             reasons.append("requires KV cache block size 128")
         if cc.kv_offloading_size is not None:
             reasons.append("does not support KV cache offloading")
 
         index_dtype = config.attention_config.resolve_indexer_kv_dtype("bf16")
-        if index_dtype != "fp8":
-            reasons.append('requires indexer_kv_dtype "fp8"')
+        if index_dtype not in ("fp8", "nvfp4"):
+            reasons.append('requires indexer_kv_dtype "fp8" or "nvfp4"')
         if getattr(config, "lora_config", None) is not None:
             reasons.append("does not support LoRA")
         transfer = getattr(config, "kv_transfer_config", None)
@@ -82,11 +84,13 @@ def refusal_reasons(self, config: VllmConfig) -> list[str]:
         sparse = getattr(text_config, "sparse_attention_config", None) or {}
         layers = text_config.num_hidden_layers
         frequencies = sparse.get("sparse_attention_freq", ())[:layers]
-        sparse_layers = [i for i, enabled in enumerate(frequencies) if enabled]
+        sparse_layers = {i for i, enabled in enumerate(frequencies) if enabled}
+        # The leading dense-FFN layers run the stock path, so their attention
+        # type doesn't matter (MiniMax-M3.1 makes every layer sparse).
         if (
             self.last_layer != layers - 1
             or layers <= self.first_layer
-            or sparse_layers != list(range(self.first_layer, layers))
+            or not sparse_layers.issuperset(range(self.first_layer, layers))
         ):
             reasons.append(
                 f"requires sparse attention on layers {self.first_layer} through "

vllm/models/minimax_m3/amd/specialized_decode/runtime/weights.py

@@ -65,6 +65,27 @@ def _tensor(
     return value.detach()
 
 
+def _bf16_projection(linear, shape: tuple[int, int]) -> torch.Tensor:
+    """The projection weight as BF16. An MXFP8 checkpoint (MiniMax-M3.1) gets
+    a dequantized copy; E4M3 times a power of two is exact in BF16, and the
+    stock path keeps serving the MXFP8 weight."""
+    weight = linear.weight
+    scale = getattr(linear, "weight_scale", None)
+    if (
+        weight.dtype == torch.float8_e4m3fn
+        and isinstance(scale, torch.Tensor)
+        and scale.dtype == torch.uint8
+        and tuple(weight.shape) == shape
+        and tuple(scale.shape) == (shape[0], shape[1] // 32)
+    ):
+        from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
+            dequant_mxfp8_to_bf16,
+        )
+
+        return dequant_mxfp8_to_bf16(weight, scale).contiguous()
+    return weight
+
+
 def _bf16_vector(
     value: torch.Tensor, size: int, name: str, require_cuda: bool
 ) -> torch.Tensor:
@@ -213,7 +234,7 @@ def from_layer(cls, layer, *, require_cuda: bool = True) -> MonoLayerWeights:
                 require_cuda,
             ),
             qkv=_tensor(
-                attn.qkv_proj.weight,
+                _bf16_projection(attn.qkv_proj, (QKV_ROWS, HIDDEN)),
                 shape=(QKV_ROWS, HIDDEN),
                 dtype=torch.bfloat16,
                 name=f"layer {layer_id} QKV/index projection",
@@ -238,7 +259,7 @@ def from_layer(cls, layer, *, require_cuda: bool = True) -> MonoLayerWeights:
                 require_cuda,
             ),
             output=_tensor(
-                attn.o_proj.weight,
+                _bf16_projection(attn.o_proj, (HIDDEN, O_K)),
                 shape=(HIDDEN, O_K),
                 dtype=torch.bfloat16,
                 name=f"layer {layer_id} output projection",