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:
@@ -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:
|
||||
|
||||
@@ -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():
|
||||
|
||||
@@ -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 "
|
||||
|
||||
Reference in New Issue
Block a user