[minimax-m3] Split 1/4: sparse attention ops + JIT kernels + config foundation (#28712)

benchmark/kernels/fused_moe_triton/common_utils.py

@@ -1,4 +1,5 @@
 import json
+from types import SimpleNamespace
 from typing import Dict, List, TypedDict
 
 import torch
@@ -7,7 +8,7 @@ from sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe import get_config_d
 from sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe_triton_config import (
     get_config_file_name,
 )
-from sglang.srt.utils import is_hip
+from sglang.srt.utils import is_gfx95_supported, is_hip
 from sglang.srt.utils.hf_transformers_utils import get_config
 
 
@@ -59,7 +60,13 @@ def get_model_config(
         assert len(block_shape) == 2
     # Replace config with text_config for encoder-decoder models after getting block_shape and architecture
     if hasattr(config, "text_config"):
-        config = config.get_text_config()
+        text_config = config.get_text_config()
+        # Some models (e.g. MiniMax-M3) carry text_config as a plain dict; wrap
+        # it so downstream attribute access works uniformly.
+        if isinstance(text_config, dict):
+            config = SimpleNamespace(**text_config)
+        else:
+            config = text_config
 
     hidden_size = config.hidden_size
     if architecture == "DbrxForCausalLM":
@@ -148,6 +155,17 @@ def get_model_config(
         E = config.num_experts // ep_size
         topk = config.num_experts_per_tok
         intermediate_size = config.moe_intermediate_size
+    elif architecture == "MiniMaxM3SparseForConditionalGeneration":
+        # Serving fuses the shared expert into the routed-expert tensor by
+        # default (E = num_local_experts + 1), so tune for that shape unless
+        # fusion is explicitly disabled.
+        E = config.num_local_experts // ep_size + (
+            0 if disable_shared_experts_fusion else 1
+        )
+        topk = config.num_experts_per_tok + (
+            0 if disable_shared_experts_fusion or topk_ids_dir is None else 1
+        )
+        intermediate_size = config.intermediate_size
     else:
         # Default: Mixtral
         E = config.num_local_experts // ep_size
@@ -158,12 +176,25 @@ def get_model_config(
         intermediate_size, tp_size, ep_size
     )
 
+    # gfx942 (MI300X) remaps MXFP8 [1,32]->[128,128] at load; tune must match.
+    # sm100/gfx95 run native [1,32] and must not remap (mirrors fp8.py).
+    if (
+        architecture == "MiniMaxM3SparseForConditionalGeneration"
+        and is_hip()
+        and not is_gfx95_supported()
+        and block_shape == [1, 32]
+    ):
+        block_shape = [128, 128]
+
+    # text_config may not carry torch_dtype; fall back to bf16.
+    torch_dtype = getattr(config, "torch_dtype", None) or torch.bfloat16
+
     return {
         "num_experts": E,
         "topk": topk,
         "hidden_size": hidden_size,
         "shard_intermediate_size": shard_intermediate_size,
-        "dtype": config.torch_dtype,
+        "dtype": torch_dtype,
         "block_shape": block_shape,
         "architecture": architecture,
     }

benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton.py

@@ -248,8 +248,10 @@ class BenchmarkWorker:
         torch.get_device_module().manual_seed_all(0)
         self.seed = seed
         # Get the device ID to allocate tensors and kernels
-        # on the respective GPU.
-        self.device_id = int(ray.get_gpu_ids()[0])
+        # on the respective GPU. Ray isolates each worker to a single visible
+        # GPU via CUDA_VISIBLE_DEVICES, so the local ordinal is always 0. On
+        # ROCm using the global ray gpu id here raises "invalid device ordinal".
+        self.device_id = 0 if is_hip() else int(ray.get_gpu_ids()[0])
         set_global_server_args_for_scheduler(server_args)
 
     def benchmark(