Skip to content

[MoonEP] BF16 PoC integration for Kimi-K3 #33249

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Fridge003 merged 7 commits into from
Aug 20, 2026
Merged

[MoonEP] BF16 PoC integration for Kimi-K3 #33249

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
bb5c7f8
Recognize MoonEP MoE A2A backend
wirybeaver Aug 2, 2026
d131f24
Add MoonEP dispatcher data contract
wirybeaver Aug 2, 2026
5765a50
Add MoonEP buffer facade
wirybeaver Aug 2, 2026
7fc1832
Add MoonEP BF16 weight layout
wirybeaver Aug 2, 2026
87f3e47
Add MoonEP BF16 expert runner
wirybeaver Aug 2, 2026
88e185d
Wire MoonEP BF16 runtime path
wirybeaver Aug 2, 2026
ffdf6f9
Add MoonEP BF16 validation script
wirybeaver Aug 2, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions python/sglang/srt/arg_groups/overrides.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2504,6 +2504,7 @@ def _cutlass_moe_env_override(view: Any) -> dict:
{
"megamoe",
"deepep",
"moonep",
"mooncake",
"nixl",
"ascend_fuseep",
Expand Down
8 changes: 8 additions & 0 deletions python/sglang/srt/batch_overlap/two_batch_overlap.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@
)
from sglang.srt.layers.moe.token_dispatcher import (
DeepEPDispatcher,
MoonEPDispatcher,
MooncakeEPDispatcher,
MoriEPDispatcher,
NixlEPDispatcher,
Expand Down Expand Up @@ -1077,6 +1078,10 @@ def __init__(self, **kwargs):
self._inners = [
DeepEPDispatcher(**kwargs) for _ in range(num_inner_dispatchers)
]
elif get_moe_a2a_backend().is_moonep():
self._inners = [
MoonEPDispatcher(**kwargs) for _ in range(num_inner_dispatchers)
]
elif get_moe_a2a_backend().is_mooncake():
self._inners = [
MooncakeEPDispatcher(**kwargs) for _ in range(num_inner_dispatchers)
Expand Down Expand Up @@ -1120,6 +1125,9 @@ def combine_a(self, **kwargs):
def combine_b(self, **kwargs):
return self._execute("combine_b", **kwargs)

def prefetch_weight(self, **kwargs):
return self._execute("prefetch_weight", **kwargs)

def register_deepep_dispatch_hook(self, hook):
handle_list = []
for inner in self._inners:
Expand Down
5 changes: 5 additions & 0 deletions python/sglang/srt/environ.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -976,6 +976,11 @@ class Envs:
SGLANG_DISABLE_STATIC_WATERFILL = EnvBool(False)
SGLANG_NIXL_EP_BF16_DISPATCH = EnvBool(False)
SGLANG_NIXL_EP_NUM_MAX_DISPATCH_TOKENS_PER_RANK = EnvInt(128)
SGLANG_MOONEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK = EnvInt(128)
# -1 uses MoonEP's training-safe default B = E / EP.
SGLANG_MOONEP_NUM_PREFETCH_SLOTS = EnvInt(-1)
SGLANG_MOONEP_TOKEN_PADDING = EnvInt(128)
SGLANG_MOONEP_NUM_SMS = EnvInt(32)
SGLANG_PPLX_NUM_MAX_DISPATCH_TOKENS_PER_RANK = EnvInt(128)
SGLANG_ENABLE_MOE_DEFERRED_FINALIZE = EnvBool(True)
# DeepSeek/GLM MoE (deepseek_v2.py): quantize the (dp-gathered) MoE input
Expand Down
25 changes: 25 additions & 0 deletions python/sglang/srt/layers/moe/ep_moe/layer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,10 @@
DeepEPLLCombineInput,
DeepEPNormalCombineInput,
)
from sglang.srt.layers.moe.token_dispatcher.moonep import (
get_moonep_expert_weight_layout,
run_moonep_bf16_expert,
)
from sglang.srt.layers.moe.topk import (
StandardTopKOutput,
TopKOutput,
Expand Down Expand Up @@ -299,6 +303,26 @@ def run_moe_core(
output = self.forward_cutlass_w4afp8_masked(dispatch_output)
else:
assert False, "forward_deepgemm_masked is deprecated"
elif DispatchOutputChecker.format_is_moonep(dispatch_output):
weight_layout = get_moonep_expert_weight_layout(
self,
num_prefetch_slots=(
int(dispatch_output.expert_ids.numel()) - self.num_experts
),
)
self.dispatcher.prefetch_weight(
plan=dispatch_output.plan,
weight_layout=weight_layout,
)
return run_moonep_bf16_expert(
dispatch_output,
weight_layout,
activation=self.moe_runner_config.activation,
)
else:
raise NotImplementedError(
f"Unsupported DeepEPMoE dispatch format: {dispatch_output.format}"
)

combine_input_wrapper = (
DeepEPNormalCombineInput
Expand Down Expand Up @@ -354,6 +378,7 @@ def get_moe_impl_class(quant_config: Optional[QuantizationConfig]):
if (
get_moe_a2a_backend().is_mori()
or get_moe_a2a_backend().is_deepep()
or get_moe_a2a_backend().is_moonep()
or get_moe_a2a_backend().is_mooncake()
or get_moe_a2a_backend().is_nixl()
or get_moe_a2a_backend().is_pplx()
Expand Down
25 changes: 24 additions & 1 deletion python/sglang/srt/layers/moe/fused_moe_triton/layer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -149,6 +149,7 @@ def create_moe_dispatcher(moe_runner_config: MoeRunnerConfig) -> BaseDispatcher:
return StandardDispatcher(moe_runner_config)
elif (
a2a_backend.is_deepep()
or a2a_backend.is_moonep()

@charliechenye charliechenye Aug 12, 2026 •
edited
Loading

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

One concern with routing MoonEP through MaybeTboDeepEPDispatcher: when TBO is enabled, the wrapper creates two MoonEPDispatcher instances, but both appear to resolve through the same process wide MoonEPBuffer cache.

The buffer key is based on the static MoonEP dimensions + id(group), without dispatcher/in-flight identity, so two otherwise identical inner dispatchers will reuse the same underlying moonep.Buffer. That seems unsafe for overlapping sub-batches, especially since MoonEP's dispatch_a/b and combine_a/b are intentionally not implemented yet.

For this PoC, could we keep MoonEP on the single synchronous path and explicitly reject TBO/SBO at startup? Then overlap support can be added once there is an explicit ownership model for multiple in-flight MoonEP buffers/plans.

or a2a_backend.is_mooncake()
or a2a_backend.is_mori()
or a2a_backend.is_nixl()
Expand Down Expand Up @@ -396,9 +397,24 @@ def __init__(
f"quant_method={type(self.quant_method).__name__})."
)

moonep_global_weight_storage = get_moe_a2a_backend().is_moonep()
if moonep_global_weight_storage:
if quant_config is not None:
raise NotImplementedError(
"MoonEP PoC supports unquantized BF16 MoE weights only."
)
if num_fused_shared_experts != 0:
raise NotImplementedError(
"MoonEP PoC does not support fused shared experts yet."
)

self.quant_method.create_weights(
layer=self,
num_experts=self.num_local_experts,
num_experts=(
self.num_experts
if moonep_global_weight_storage
else self.num_local_experts
),
hidden_size=hidden_size,
intermediate_size_per_partition=self.intermediate_size_per_partition,
params_dtype=params_dtype,
Expand All @@ -410,6 +426,13 @@ def __init__(
with_bias=with_bias,
moe_intermediate_size=intermediate_size,
)
if moonep_global_weight_storage:
self.w13_weight._sglang_require_global_experts = True
self.w2_weight._sglang_require_global_experts = True
if getattr(self, "w13_weight_bias", None) is not None:

@charliechenye charliechenye Aug 12, 2026 •
edited
Loading

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should we explicitly reject expert bias for the MoonEP PoC?

We mark the bias tensors as requiring global experts when they exist, but MoonEPExpertWeightLayout only carries gate/up/down weights and run_moonep_bf16_expert() uses the biasless F.linear(...) path. So a bias containing MoE could appear to construct/load successfully but then silently ignore those biases during expert compute.

I don't think bias support needs to be added in this PoC: a fail fast NotImplementedError when with_bias=True would make the supported contract much clearer.

self.w13_weight_bias._sglang_require_global_experts = True
if getattr(self, "w2_weight_bias", None) is not None:
self.w2_weight_bias._sglang_require_global_experts = True

self.quant_method.create_moe_runner(self, self.moe_runner_config)
self.dispatcher = create_moe_dispatcher(self.moe_runner_config)
Expand Down
20 changes: 20 additions & 0 deletions python/sglang/srt/layers/moe/token_dispatcher/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,17 @@
MooncakeDispatchOutput,
MooncakeEPDispatcher,
)
from sglang.srt.layers.moe.token_dispatcher.moonep import (
MoonEPBuffer,
MoonEPBufferKey,
MoonEPCombineInput,
MoonEPDispatcher,
MoonEPDispatchOutput,
MoonEPExpertWeightLayout,
get_moonep_expert_weight_layout,
get_moonep_num_prefetch_slots,
run_moonep_bf16_expert,
)
from sglang.srt.layers.moe.token_dispatcher.moriep import (
MoriEPDispatcher,
MoriEPLLCombineInput,
Expand Down Expand Up @@ -67,6 +78,15 @@
"MooncakeCombineInput",
"MooncakeDispatchOutput",
"MooncakeEPDispatcher",
"MoonEPCombineInput",
"MoonEPBuffer",
"MoonEPBufferKey",
"MoonEPDispatcher",
"MoonEPDispatchOutput",
"MoonEPExpertWeightLayout",
"get_moonep_expert_weight_layout",
"get_moonep_num_prefetch_slots",
"run_moonep_bf16_expert",
"MoriEPNormalDispatchOutput",
"MoriEPNormalCombineInput",
"MoriEPLLDispatchOutput",
Expand Down
19 changes: 19 additions & 0 deletions python/sglang/srt/layers/moe/token_dispatcher/base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,8 @@
DeepEPNormalDispatchOutput,
FlashinferCombineInput,
FlashinferDispatchOutput,
MoonEPCombineInput,
MoonEPDispatchOutput,
StandardCombineInput,
StandardDispatchOutput,
)
Expand Down Expand Up @@ -165,12 +167,19 @@ def format_is_flashinfer(
) -> TypeGuard[FlashinferDispatchOutput]:
return dispatch_output.format.is_flashinfer()

@staticmethod
def format_is_moonep(
dispatch_output: DispatchOutput,
) -> TypeGuard[MoonEPDispatchOutput]:
return dispatch_output.format.is_moonep()


class DispatchOutputFormat(Enum):

STANDARD = "standard"
DEEPEP_NORMAL = "deepep_normal"
DEEPEP_LL = "deepep_ll"
MOONEP = "moonep"
FLASHINFER = "flashinfer"
ASCEND_TP = "ascend_tp"

Expand All @@ -195,6 +204,9 @@ def is_deepep(self) -> bool:
def is_flashinfer(self) -> bool:
return self == DispatchOutputFormat.FLASHINFER

def is_moonep(self) -> bool:
return self == DispatchOutputFormat.MOONEP


@runtime_checkable
class DispatchOutput(Protocol):
Expand Down Expand Up @@ -249,11 +261,18 @@ def format_is_flashinfer(
) -> TypeGuard[FlashinferCombineInput]:
return combine_input.format == CombineInputFormat.FLASHINFER

@staticmethod
def format_is_moonep(
combine_input: CombineInput,
) -> TypeGuard[MoonEPCombineInput]:
return combine_input.format == CombineInputFormat.MOONEP


class CombineInputFormat(Enum):
STANDARD = "standard"
DEEPEP_NORMAL = "deepep_normal"
DEEPEP_LL = "deepep_ll"
MOONEP = "moonep"
FLASHINFER = "flashinfer"
ASCEND_TP = "ascend_tp"

Expand Down
Loading
Loading