[wirybeaver]

It is opened as a draft for partial task of #32607's Integrate MoonEP, not as a final production-ready merge request.

What is included

• Distinct moonep MoE A2A backend recognition.
• MoonEP dispatch/combine data contracts and process-wide buffer facade.
• BF16 contiguous expert layout and correctness-first expert segment runner.
• Eager runtime path: dispatch -> weight prefetch -> BF16 compute -> combine.
• Distributed validation script: scripts/moonep/validate_moonep_bf16_poc.py.

Validation

• 8x H100 SXM (NV18) passed BF16 validation on 8 ranks with global_ok=true, max_abs_err=0.0, relative_err=0.0.
• docker-ncu-benchmark preflight was run. Suitable Vast.ai H100 NVSwitch hosts were container instances without Docker/NVIDIA Container Toolkit; host NCU counters were restricted (ERR_NVGPUCTRPERM), so no NCU performance claims are made.

Known limitations

• BF16/unquantized only.
• Eager-only; CUDA graph disabled.
• Correctness-first Python segment loop, not final grouped GEMM kernel.
• PoC global expert-row storage is memory-heavy; Kimi-K3 production needs sharded symmetric-memory expert ownership.
• Requires NVSwitch/multicast-capable hosts (8x H100/H200/B200 class).

Notes

Local agent workflow files, scratch artifacts, and spec notes are intentionally omitted from this upstream PR diff.

Fork PoC PR: wirybeaver#17

---

CI States

Latest PR Test (Base): ❌ Run #31810458189
Latest PR Test (Extra): ❌ Run #31810457742

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[charliechenye]

Thanks for putting this together. I spent some time walking through the SGLang side and comparing the integration boundary against MoonEP's dispatch / plan / prefetch / combine model.

Overall I like the direction. In particular, making MoonEP a distinct dispatch/combine format rather than trying to make it look wire compatible with DeepEP feels like the right abstraction. I also like keeping weight prefetch explicit between dispatch and expert compute; that seems like a useful boundary to preserve when the correctness first runner is eventually replaced by the production kernel.

Since this is explicitly a draft PoC, I'm intentionally not asking for grouped GEMM, CUDA graphs, quantized weights, or production memory efficiency here. The inline comments are mostly about making sure the current narrow contract fails safely instead of accidentally advertising capabilities it doesn't have yet.

One broader design thought for the next iteration: I'd be careful about using is_deepep_class_backend() as a growing capability bucket. Sharing the DeepEPMoE execution skeleton makes sense, but being on that execution branch doesn't necessarily imply MoonEP shares every DeepEP runtime capability. As this grows, capability-oriented checks such as supports_overlap, uses_per_rank_shared_slots, or supports_cuda_graph may keep the backend boundaries easier to reason about.

I also think the replicated global weight layout is reasonable for proving the first end-to-end path, as long as we keep treating it as disposable PoC infrastructure rather than the long term storage contract. The production path should ideally make MoonEP's actual expert ownership / symmetric memory layout canonical instead of allocating global SGLang weights and then building another MoonEP representation.

Finally, it would be useful to record the exact MoonEP revision used for the 8xH100 validation. I would currently interpret that result as good evidence for the synthetic BF16 MoonEP dispatch -> prefetch -> compute -> combine path, but not yet full Kimi K3 integration validation because the script bypasses the real FusedMoE weight loading/layout construction path.

Overall, promising direction. My main suggestion for this stage is to keep the supported surface narrow and explicit. If unsupported combinations fail early, the underlying storage and kernels can evolve later without having to unwind accidental compatibility assumptions.

[charliechenye]

Could we avoid falling back to rank 0 here?

plan.experts_to_copy is rank local, and _expert_ids_from_plan() uses this value to select the current rank's row. If rank lookup fails on a nonzero rank, we would keep running but associate dispatched segments with rank 0's prefetched expert mapping. That looks like a silent correctness failure rather than a safe fallback.

I think it would be safer to let dist.get_rank(group=...) fail here and mock/inject rank explicitly in unit tests if needed.

[charliechenye]

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.

[charliechenye]

Could we avoid silently switching to stubbed SGLang modules in this validation path?

My concern is about what a passing result proves: if the real SGLang import breaks, this catches the failure, installs lightweight stand ins for parts of SGLang, and can still complete the MoonEP validation.

I think both tests are useful, but they prove different things. Could we split this into a lightweight MoonEP adapter/component validation that may use stubs, and a real SGLang integration validation where an import failure is a test failure?

Ideally the latter would also exercise the actual FusedMoE weight creation/loading -> MoonEP weight-layout path, since this script currently constructs MoonEPExpertWeightLayout directly.

[charliechenye]

Small planning/perf question on the padding strategy: all padded routes currently get expert id 0.

Their route weights are zero, so the final numerical contribution should be neutral. But _tokens_per_expert() is computed from these padded IDs, which means the padding still contributes to MoonEP's planning statistics.

For a decode-like case such as 1 real token padded to S=128, expert 0 could therefore look artificially hot and affect migration/prefetch decisions. Could we avoid concentrating dummy routes on one expert, or otherwise ensure padding does not skew the planner? At minimum, a num_tokens << capacity regression test would be useful.

[charliechenye]

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.

[lyppg]

Hi @wirybeaver, thanks for putting this up, the backend plumbing here is exactly the right shape and it made the path forward much clearer for us.

Three things, all of which relax constraints this PoC had to work around:

1. MoonEP is no longer bf16-only. We submitted PRs to support MXFP4 expert weight prefetch in upstream (packed e2m1 + ue8m0 block scales) plus multi-node NVLink fabric support. That matters for Kimi-K3 specifically — at EP16 its 896 experts are ~91 GB/rank in MXFP4 but ~344 GB/rank in bf16, so bf16 does not fit.

2. run_moonep_bf16_expert could go through the DeepGEMM m-grouped contiguous runner instead of the Python segment loop.

3. Hardware. This PR looks like it was validated on H100. We have GB200 (GB200 NVL, multi-node NVLink) available and would be glad to run K3 on it for you. Happy to just report numbers on your current branch as a first step, independent of anything else here.

One smaller note, following @charliechenye's "keep the supported surface narrow and explicit": that may also argue against folding moonep into is_deepep_class_backend(). MoonEP is a third dispatch format rather than a member of the deepep_normal/deepep_ll family. Happy to be wrong if some call site required it.

Are you still actively working on this? We'd like to help push it toward a production path (MXFP4 weights, the DeepGEMM runner, and real symmetric-memory weight residency instead of per-rank full copies). Happy to contribute commits to your branch, or to open a follow-up that builds on this one — whichever you prefer.

[wirybeaver]

@lyppg Yeah, I am looking for this contribution and ask @charliechenye who is my Uber colleague for co-authoring. Let me touch with @Fridge003 for coordination. The stacking PR approach sounds good to me for review-ability

Though I don't know which git development mode we will adopt, I already give you folks the collaborator access to my forked repo for your convenience: https://github.com/wirybeaver/sglang

[lyppg]

@lyppg Yeah, I am looking for this contribution and ask @charliechenye who is my Uber colleague for co-authoring. Let me touch with @Fridge003 for coordination. The stacking PR approach sounds good to me for review-ability

Though I don't know which git development mode we will adopt, I already give you folks the collaborator access to my forked repo for your convenience: https://github.com/wirybeaver/sglang

Yes I saw it, thanks. It seems you also use slack, just sent you a message, let's discuss it there