qqtang/sskj
1
3
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[Fix] Support legacy Kimi MoE runner config

Browse Source
This commit is contained in:
Zhiyi Hong 2026-08-18 14:33:42 +08:00
parent 13944079fa
commit d0863501ca
1 changed files with 6 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
experiments/pro6000/kimi3_pro6000_sglang_tp32ep32_moe_backend_prefill/patches/sglang-kimi-image-sm120-flashinfer-mxfp4.patch
View File

@ -95,7 +95,7 @@
if self.use_flashinfer: if self.use_flashinfer:
# Per-expert buffers are local (create_weights uses num_local_experts); # Per-expert buffers are local (create_weights uses num_local_experts);
# the global self.num_experts here breaks EP>1. Mirrors the SM90 path. # the global self.num_experts here breaks EP>1. Mirrors the SM90 path.
@@ -1130,7 +1135,135 @@ @@ -1130,7 +1135,136 @@
layer.w2_weight_bias = Parameter(w2_bias_padded, requires_grad=False) layer.w2_weight_bias = Parameter(w2_bias_padded, requires_grad=False)
torch.cuda.empty_cache() torch.cuda.empty_cache()
@ -192,7 +192,8 @@
+ +
+ activation = layer.moe_runner_config.activation + activation = layer.moe_runner_config.activation
+ alpha = layer.moe_runner_config.gemm1_alpha + alpha = layer.moe_runner_config.gemm1_alpha
+ beta = layer.moe_runner_config.gemm1_beta + # Older Kimi image baselines predate the generic gemm1_beta field.
+ beta = getattr(layer.moe_runner_config, "gemm1_beta", None)
+ limit = layer.moe_runner_config.gemm1_clamp_limit + limit = layer.moe_runner_config.gemm1_clamp_limit
+ if activation == "situ": + if activation == "situ":
+ alpha = 4.0 if alpha is None else alpha + alpha = 4.0 if alpha is None else alpha
@ -231,7 +232,7 @@
def create_moe_runner( def create_moe_runner(
self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig
): ):
@@ -1154,10 +1287,23 @@ @@ -1154,10 +1288,23 @@
or moe_runner_backend.is_deep_gemm() or moe_runner_backend.is_deep_gemm()
): ):
self.runner = MoeRunner(moe_runner_backend, moe_runner_config) self.runner = MoeRunner(moe_runner_backend, moe_runner_config)
@ -258,7 +259,7 @@
# Register the fused func at runner construction so the FusedOpPool # Register the fused func at runner construction so the FusedOpPool
# lookup at `MoeRunner.__init__` finds it. # lookup at `MoeRunner.__init__` finds it.
import sglang.srt.layers.moe.moe_runner.flashinfer_cutlass # noqa: F401 import sglang.srt.layers.moe.moe_runner.flashinfer_cutlass # noqa: F401
@@ -1215,8 +1361,33 @@ @@ -1215,8 +1362,33 @@
quant_info = build_marlin_moe_quant_info(layer) quant_info = build_marlin_moe_quant_info(layer)
return self.runner.run( return self.runner.run(
dispatch_output._replace(hidden_states=x_padded), quant_info dispatch_output._replace(hidden_states=x_padded), quant_info
@ -292,7 +293,7 @@
def apply( def apply(
self, self,
layer: torch.nn.Module, layer: torch.nn.Module,
@@ -1290,6 +1461,8 @@ @@ -1290,6 +1462,8 @@
if self._fi_kernel == "cutlass_sm90": if self._fi_kernel == "cutlass_sm90":
return self._apply_sm90_cutlass(layer, dispatch_output) return self._apply_sm90_cutlass(layer, dispatch_output)