[Fix] Support legacy Kimi MoE runner config
This commit is contained in:
parent
13944079fa
commit
d0863501ca
@ -95,7 +95,7 @@
|
||||
if self.use_flashinfer:
|
||||
# Per-expert buffers are local (create_weights uses num_local_experts);
|
||||
# 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)
|
||||
|
||||
torch.cuda.empty_cache()
|
||||
@ -192,7 +192,8 @@
|
||||
+
|
||||
+ activation = layer.moe_runner_config.activation
|
||||
+ 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
|
||||
+ if activation == "situ":
|
||||
+ alpha = 4.0 if alpha is None else alpha
|
||||
@ -231,7 +232,7 @@
|
||||
def create_moe_runner(
|
||||
self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig
|
||||
):
|
||||
@@ -1154,10 +1287,23 @@
|
||||
@@ -1154,10 +1288,23 @@
|
||||
or moe_runner_backend.is_deep_gemm()
|
||||
):
|
||||
self.runner = MoeRunner(moe_runner_backend, moe_runner_config)
|
||||
@ -258,7 +259,7 @@
|
||||
# Register the fused func at runner construction so the FusedOpPool
|
||||
# lookup at `MoeRunner.__init__` finds it.
|
||||
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)
|
||||
return self.runner.run(
|
||||
dispatch_output._replace(hidden_states=x_padded), quant_info
|
||||
@ -292,7 +293,7 @@
|
||||
def apply(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
@@ -1290,6 +1461,8 @@
|
||||
@@ -1290,6 +1462,8 @@
|
||||
|
||||
if self._fi_kernel == "cutlass_sm90":
|
||||
return self._apply_sm90_cutlass(layer, dispatch_output)
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user