[MoonEP] MXFP4 experts for Kimi K3 on the DeepGEMM runner in symmetric memory

[lyppg]

Stacked on #33249, this PR adds initial support for MoonEP a2a backend and DeepGemm runner for MXFP4 experts K3.
Please review the following commits:

• Integrate with MoonEP and MXFP4 experts on the DeepGEMM runner

Current implementation depends on MoonEP update prefetch api and support rotating local to first.

Validation

Validated end to end on Kimi-K3, 4x GB200, TP16/EP16, MXFP4 weights.
gsm8k test:

python -m sglang.test.run_eval --eval-name gsm8k --port 30000 --num-examples 500 --num-shots 5  --max-tokens 2048
Total latency: 643.858 s
Score: 0.986
Output throughput: 157.187 token/s
[METRIC] gsm8k_score=0.986 labels={"model": "Kimi-K3", "eval": "gsm8k"}
[METRIC] gsm8k_latency=643.8576721129939 labels={"model": "Kimi-K3", "eval": "gsm8k"}

---

CI States

Latest PR Test (Base): ❌ Run #33569351386
Latest PR Test (Extra): ❌ Run #33569351091
Latest PR Test (AMD ROCm 7.2): ❌ Run #33569351178

[ch-wan]

[bug] MoonEPWeightPool calls create_nvl_dist_tensor(..., local_first=True) and later local_first_chunk_index. Published MoonEP (MoonshotAI/MoonEP master buffer.py) defines create_nvl_dist_tensor(chunk_shape, dtype, local_rank, world_size, group=None) with no local_first, and has no local_first_chunk_index. This TypeErrors as soon as MXFP4 create_weights allocates the pool. The comment about DeepGEMM tvm_ffi resolving device from the unrotated rank-0 base is exactly why the public mapper would then fail the device check on every non-zero rank.

Suggestion: Either vendor/pin the MoonEP commit that actually provides local_first / local_first_chunk_index, or implement the rotation in sglang on top of the public mapper and call only public symbols. Do not land this against the same MoonEP the parent PR documents.

[ch-wan]

[bug] Quantized pre-permute calls MoonEPBuffer.get_existing_buffer().prefetch_weight(plan=..., weight_pairs=..., scale_pairs=..., experts_to_copy=...). get_existing_buffer() returns a raw moonep.Buffer. Published Buffer.prefetch_weight is keyword-only full_gate_weight / full_up_weight / full_down_weight (and matching scales), asserts shape[0] == E+B, and always uses plan.experts_to_copy as global expert ids in [0, E). Unexpected kwargs raise TypeError; even rewriting to full_*_weight=pool.ranges[...] still fails the E+B assert because the pool's dim0 is ep_size * chunk_rows. Parent #33249 already uses the public full_*_weight convention for BF16, so the two paths cannot share one Buffer class.

Suggestion: Add a sglang-side wrapper that drives published launch_prefetch(remote_expert, prefetch_buffers, experts_to_copy, num_sms) on each pool range (ndim-3 source, dest = slot_view, ids = remapped pool rows), or require a documented MoonEP API that accepts pair lists. Keep BF16 on full_*_weight behind that wrapper so both paths stay on one contract.

[ch-wan]

[suggestion] expert_rows hardcodes chunk = (owner - ep_rank) % ep_size while chunk_start delegates to local_first_chunk_index. If those ever diverge (or the unpublished helper is not (owner - rank) % size), m_indices and prefetch sources point at the wrong VMM chunk while loader views still use chunk_start. This is the exact class of silent correctness bug the rotation exists to prevent.

Suggestion: Tensorize one helper (the same function chunk_start uses) and compute rows = chunk_index * chunk_rows + layer_offset + expert_ids % epn from it.

[lyppg]

added checking in MoonEPWeightPool init, should be able to avoid the divergence.

[ch-wan]

[suggestion] _num_moe_layers uses getattr(config.hf_text_config, "moe_layer_freq", 1) (defensive getattr, forbidden by repo no-getattr-defensive) while first_k_dense_replace is already a ModelConfig field. The i % freq == 0 test also disagrees with models that use (i - first_dense) % freq == 0. It happens to match K3 (first_k_dense_replace=1, 92 MoE layers, freq default 1). A config with a list-valued moe_layer_freq or MTP layers that also construct MoE will either TypeError or undersize the VMM mapping until layer_offset raises mid-build.

Suggestion: Read config.first_k_dense_replace / a real moe_layer_freq off ModelConfig (default on the config object, not getattr), count only layers this rank will actually create_weights for (PP), and include MTP MoE if those layers allocate from the same pool.

[ch-wan]

[suggestion] _moonep_m_indices / expert_rows / group_rows are derived-property mapping (segment ends, empty groups, local-first row ids, prefetch-slot overwrite). The only new tests in this PR are prefetch-slot defaults in test_moonep_buffer.py. A rewrite that uses right=False, forgets the group < num_groups pad, or maps slot groups through expert_rows instead of local slots would still look equivalent and only fail on K3.

Suggestion: Add CPU unit tests for (1) empty groups + tail padding -> -1, (2) live group id is expert_ids[g] not g, (3) group_rows leaves home groups on owner rows and remaps only the tail to slot_base + i.

[lyppg]

Unit tests added: test/registered/unit/layers/moe/test_moonep_weights.py