minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[AMD] Enable Mori-EP on kimi-k3 (#35630)

Browse Source
This commit is contained in:
Xinyi Song
2026-08-25 01:48:09 -07:00
committed by GitHub
parent d067622820
commit 2e3934f4cb
2 changed files with 85 additions and 64 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/quantization/mxfp4.py
+77 -58
View File
@@ -1412,7 +1412,10 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
dispatch_output: StandardDispatchOutput,
) -> CombineInput:
from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput
from sglang.srt.layers.moe.token_dispatcher import (
DispatchOutputChecker,
StandardCombineInput,
)
from sglang.srt.layers.moe.topk import TopKOutputChecker
if self.use_deep_gemm:
@@ -1433,6 +1436,12 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
)
return self.runner.run(dispatch_output, quant_info)
# Same constraint as the deep_gemm branch above: the AITER runner also
# serves the DeepEP formats (deepep_normal / deepep_ll), which carry
# topk_ids/topk_weights directly and have no `.topk_output` to unpack.
if _use_aiter and DispatchOutputChecker.format_is_deepep(dispatch_output):
return self._apply_aiter(layer, dispatch_output)
x = dispatch_output.hidden_states
topk_output = dispatch_output.topk_output
if _is_cpu:
@@ -1710,63 +1719,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
)[0]
return StandardCombineInput(hidden_states=trtllm_gen_output)
if _use_aiter:
from sglang.srt.layers.moe.moe_runner.aiter import (
AiterMoeQuantInfo,
AiterQuantType,
)
if hasattr(torch, "float4_e2m1fn_x2"):
w13_weight = layer.w13_weight.view(torch.float4_e2m1fn_x2)
w2_weight = layer.w2_weight.view(torch.float4_e2m1fn_x2)
else:
w13_weight = layer.w13_weight
w2_weight = layer.w2_weight
# `.view()` creates a fresh tensor that drops the `is_shuffled`
# marker we set in process_weights_after_loading. Re-tag it so the
# downstream aiter.fused_moe selects preshuffle_on kernels.
if getattr(layer.w13_weight, "is_shuffled", False):
w13_weight.is_shuffled = True
w2_weight.is_shuffled = True
# Skip the explicit pad if x already arrives at the padded
# hidden_size (the upstream RMSNorm fused the pad into its
# output — see RMSNorm.x_pad_to_multiple). Saves a separate
# zero-pad kernel launch per layer.
if x.shape[-1] == self.hidden_size:
x_padded = x
else:
x_padded = torch.nn.functional.pad(
x, (0, self.hidden_pad), mode="constant", value=0.0
)
quant_info = AiterMoeQuantInfo(
w13_weight=w13_weight,
w2_weight=w2_weight,
quant_type=AiterQuantType.PER_1X32,
w13_scale=layer.w13_weight_scale,
w2_scale=layer.w2_weight_scale,
b13=layer.w13_weight_bias if self.with_bias else None,
b2=layer.w2_weight_bias if self.with_bias else None,
expert_mask=layer.dispatcher.expert_mask_gpu,
doweight_stage1=self.moe_runner_config.apply_router_weight_on_input,
hidden_pad=self.hidden_pad,
intermediate_pad=self.intermediate_pad,
# Applies swiglu clamp for GPT-OSS-style activations. K3 SiTU
# uses gemm1_clamp_limit as linear_beta, which is forwarded by
# the AITER runner and must not be treated as swiglu_limit.
swiglu_limit=(
0.0
if self.moe_runner_config.activation == "situ"
else (
self.moe_runner_config.gemm1_clamp_limit
or self.moe_runner_config.swiglu_limit
or 0.0
)
),
)
return self.runner.run(
dispatch_output._replace(hidden_states=x_padded), quant_info
)
return self._apply_aiter(layer, dispatch_output)
backend = self.runner.runner_backend
if backend.is_triton_kernels():
@@ -1802,6 +1755,72 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
)
return self.runner.run(dispatch_output, quant_info)
def _apply_aiter(self, layer, dispatch_output) -> CombineInput:
"""MXFP4 MoE via the AITER runner.
Reads only ``hidden_states`` off the dispatch output, so it serves the
standard and the DeepEP formats alike; the routing tensors are resolved
by the runner's registered permute hooks.
"""
from sglang.srt.layers.moe.moe_runner.aiter import (
AiterMoeQuantInfo,
AiterQuantType,
)
x = dispatch_output.hidden_states
if hasattr(torch, "float4_e2m1fn_x2"):
w13_weight = layer.w13_weight.view(torch.float4_e2m1fn_x2)
w2_weight = layer.w2_weight.view(torch.float4_e2m1fn_x2)
else:
w13_weight = layer.w13_weight
w2_weight = layer.w2_weight
# `.view()` creates a fresh tensor that drops the `is_shuffled`
# marker we set in process_weights_after_loading. Re-tag it so the
# downstream aiter.fused_moe selects preshuffle_on kernels.
if getattr(layer.w13_weight, "is_shuffled", False):
w13_weight.is_shuffled = True
w2_weight.is_shuffled = True
# Skip the explicit pad if x already arrives at the padded
# hidden_size (the upstream RMSNorm fused the pad into its
# output — see RMSNorm.x_pad_to_multiple). Saves a separate
# zero-pad kernel launch per layer.
if x.shape[-1] == self.hidden_size:
x_padded = x
else:
x_padded = torch.nn.functional.pad(
x, (0, self.hidden_pad), mode="constant", value=0.0
)
quant_info = AiterMoeQuantInfo(
w13_weight=w13_weight,
w2_weight=w2_weight,
quant_type=AiterQuantType.PER_1X32,
w13_scale=layer.w13_weight_scale,
w2_scale=layer.w2_weight_scale,
b13=layer.w13_weight_bias if self.with_bias else None,
b2=layer.w2_weight_bias if self.with_bias else None,
expert_mask=layer.dispatcher.expert_mask_gpu,
doweight_stage1=self.moe_runner_config.apply_router_weight_on_input,
hidden_pad=self.hidden_pad,
intermediate_pad=self.intermediate_pad,
# Applies swiglu clamp for GPT-OSS-style activations. K3 SiTU
# uses gemm1_clamp_limit as linear_beta, which is forwarded by
# the AITER runner and must not be treated as swiglu_limit.
swiglu_limit=(
0.0
if self.moe_runner_config.activation == "situ"
else (
self.moe_runner_config.gemm1_clamp_limit
or self.moe_runner_config.swiglu_limit
or 0.0
)
),
)
return self.runner.run(
dispatch_output._replace(hidden_states=x_padded), quant_info
)
class Mxfp4DynamicQuantMoEMethod(FusedMoEMethodBase):
def create_weights(
python/sglang/srt/models/kimi_k3.py
+8 -6
View File
@@ -523,16 +523,17 @@ class KimiK3MoE(nn.Module):
"got a checkpoint with different constants"
)
# EP a2a backends (megamoe / DeepEP) move each row to its experts
# directly, so the MoE region can consume whatever rows this rank
# holds — an SP-MoE token shard (attn_tp > 1) or the DP-local batch
# (DP attention) — with every global token dispatched exactly once.
# No DP gather and no TP reduce is needed anywhere in the region.
# EP a2a backends (megamoe / DeepEP / MoRI) move each row to its
# experts directly, so the MoE region can consume whatever rows this
# rank holds — an SP-MoE token shard (attn_tp > 1) or the DP-local
# batch (DP attention) — with every global token dispatched exactly
# once. No DP gather and no TP reduce is needed anywhere in the region.
_a2a_backend = get_moe_a2a_backend()
self._ep_a2a = (
_a2a_backend.is_megamoe()
or _a2a_backend.is_deepep()
or _a2a_backend.is_ascend_fuseep()
or _a2a_backend.is_mori()
)
# Defer the trtllm-gen finalize (top-k weighted unpermute) out of the
@@ -2081,7 +2082,7 @@ class KimiK3DecoderLayer(nn.Module):
and layer_idx >= config.first_k_dense_replace
and layer_idx % config.moe_layer_freq == 0
)
# SP-MoE (EP a2a backend — megamoe or DeepEP): o_proj defers its
# SP-MoE (EP a2a backend — megamoe, DeepEP or MoRI): o_proj defers its
# attention-TP reduction; this layer completes it as a reduce-scatter
# so the whole MoE region (agg2, norms, gate, latent projs, tp1
# shared experts, EP a2a dispatch) runs on 1/attn_tp of the rows,
@@ -2104,6 +2105,7 @@ class KimiK3DecoderLayer(nn.Module):
_a2a_backend.is_megamoe()
or _a2a_backend.is_deepep()
or _a2a_backend.is_ascend_fuseep()
or _a2a_backend.is_mori()
)
and self._is_moe_layer
and get_parallel().attn_tp_group.world_size > 1