MiniMax-M3: allow shared-experts fusion on ROCm gfx942 and newer (#36576)

Co-authored-by: Kevin Mi <kevin.mi@radixark.ai>
This commit is contained in:
Chunan Zeng
2026-09-16 19:55:47 -07:00
committed by GitHub
co-authored by Kevin Mi
parent 329ffc89b9
commit 241a5b9823
3 changed files with 21 additions and 7 deletions
+2 -1
View File
@@ -2404,9 +2404,10 @@ def select_experts(
# slots on the marker) and places that marker at id num_experts, which the
# DeepEP remap shifts one past the end of the expert space -- 384 -> 392 for
# 384 routed experts on EP8, where the valid ids are 0..391.
# aiter appends the shared expert in _post_process_topk_ids, so the gate must not.
num_fused_shared_experts_for_gate = (
0
if has_per_rank_fused_shared_slots(num_fused_shared_experts)
if (has_per_rank_fused_shared_slots(num_fused_shared_experts) or _use_aiter)
else num_fused_shared_experts
)
if dynamic_expert_bias is not None:
+4 -2
View File
@@ -1612,10 +1612,12 @@ class MiniMaxM3SparseForCausalLM(nn.Module):
"Shared and routed experts may use different quantization formats "
"in ModelOpt mixed-precision checkpoints."
)
if not _is_cuda:
return "Shared experts fusion currently requires CUDA devices."
if not (_is_cuda or _is_hip):
return "Shared experts fusion currently requires CUDA or ROCm devices."
if _is_cuda and (_device_sm is not None) and (_device_sm < 80):
return "Shared experts fusion requires SM80 or newer GPUs."
if _is_hip and not _is_gfx95_supported:
return "Shared experts fusion on ROCm is validated on gfx950 only."
if get_parallel().moe_ep_size > 1:
return "Shared experts fusion is not supported together with expert parallelism yet."
if get_moe_a2a_backend().is_deepep():
+15 -4
View File
@@ -47,13 +47,22 @@ from sglang.srt.models.minimax_vl_common import (
)
from sglang.srt.models.utils import WeightsMapper
from sglang.srt.runtime_context import get_mm, get_parallel
from sglang.srt.utils import add_prefix, get_device_sm, is_cuda, log_info_on_rank0
from sglang.srt.utils import (
add_prefix,
get_device_sm,
is_cuda,
is_gfx95_supported,
is_hip,
log_info_on_rank0,
)
from sglang.srt.utils.hf_transformers_utils import get_rope_config
logger = logging.getLogger(__name__)
_is_cuda = is_cuda()
_is_hip = is_hip()
_is_gfx95_supported = is_gfx95_supported()
_device_sm = get_device_sm()
@@ -152,10 +161,12 @@ class MiniMaxM3SparseForConditionalGeneration(nn.Module):
"Shared and routed experts may use different quantization formats "
"in ModelOpt mixed-precision checkpoints."
)
if not _is_cuda:
return "Shared experts fusion currently requires CUDA devices."
if (_device_sm is not None) and (_device_sm < 80):
if not (_is_cuda or _is_hip):
return "Shared experts fusion currently requires CUDA or ROCm devices."
if _is_cuda and (_device_sm is not None) and (_device_sm < 80):
return "Shared experts fusion requires SM80 or newer GPUs."
if _is_hip and not _is_gfx95_supported:
return "Shared experts fusion on ROCm is validated on gfx950 only."
if get_parallel().moe_ep_size > 1:
return (
"Shared experts fusion is not supported together with expert "