Kimi-K3: treat deepep_v2 as an EP all-to-all backend

[whn09]

Motivation

KimiK3MoE._ep_a2a (kimi_k3.py:543) and KimiK3DecoderLayer._sp_moe (:2256) each enumerate the EP all-to-all backends by hand, and both omit deepep_v2. The central helper does include it:

# python/sglang/srt/layers/moe/utils.py:577
def is_deepep_class_backend() -> bool:
    """Return whether A2A combine occurs inside a DeepEP-family dispatcher."""
    b = get_moe_a2a_backend()
    return (
        b.is_deepep()
        or b.is_deepep_v2()      # <-- present here
        or b.is_mooncake()
        or b.is_mori()
        or b.is_pplx()
    )

So the two hand-rolled lists have drifted from the helper rather than expressing a different predicate. _ep_a2a's own comment states exactly the DeepEP-family property:

whatever rows this rank holds — an SP-MoE token shard (attn_tp > 1) or the DP-local batch (DP attention) — with every global token dispatched exactly once. No DP gather and no TP reduce is needed anywhere in the region.

Consequence with --moe-a2a-backend deepep_v2: _ep_a2a stays False, so the MoE region keeps its DP gather and TP reduce even though the v2 dispatcher has already combined. That is precisely the hazard validate_deepep_v2_model_architecture warns about in its own error text — "Other model workflows may require an all-reduce after A2A combine" (arg_groups/moe_hook.py:455).

Modifications

Two lines: or _a2a_backend.is_deepep_v2() added to both lists, next to is_deepep().

Deliberately minimal. I did not replace the lists with is_deepep_class_backend() or is_megamoe() or is_ascend_fuseep(), even though that would remove the drift permanently, because the helper also includes is_pplx() and I have not validated pplx on K3. Turning a missing-entry fix into a behavior change for another backend I cannot vouch for seemed wrong.

Open question for the K3 owners: the same two lists also omit is_pplx(). If pplx belongs there, switching both sites to the helper (plus is_megamoe() / is_ascend_fuseep(), which the helper does not cover) would be the better fix and would stop these lists drifting again. I'm happy to do that in this PR if you confirm pplx is correct for K3.

This does not by itself enable K3 on deepep_v2. The path is still rejected by validate_deepep_v2_model_architecture (K3 is not in validated_architectures, moe_hook.py:426) and by _validate_deepep_v2_quant_method (K3 is MXFP4, fused_moe_triton/layer.py:240). This PR makes the wiring correct for when those are addressed — and it is a prerequisite for them, since enabling K3 on v2 without it runs straight into the double reduce.

Accuracy Tests

Run on 8 × B300 (sm_103), Kimi-K3, lmsysorg/sglang:nightly-dev-cu13-20260901-07c8f729, with this fix plus local relaxations of the two validations above (not proposed here) and #37211.

Reference arm --moe-a2a-backend deepep --deepep-mode auto --moe-runner-backend deep_gemm, test arm the same with deepep_v2. First-token top-5 logprobs are byte-identical: 15/15 tokens and 15/15 logprob values across three fixed prompts.

Since the failure this fixes is a double reduce, a wrong _ep_a2a would corrupt outputs rather than crash — so matching logprobs is the relevant evidence, and it is the reason I am reasonably confident the flag is now set correctly.

Scope of that evidence, stated plainly:

• First-token logprobs exercise the prefill/contiguous path. Masked/decode evidence is weaker: median TPOT 42.05 ms (v2) vs 41.56 ms (v1), long greedy completions agreeing exactly on 2 of 4 prompts.
• Single node, attn_tp = 1. _sp_moe additionally requires attn_tp_group.world_size > 1, so the second hunk was not exercised at runtime — it is included because it is the same predicate at the same drift, and leaving one of the two behind would just reproduce the bug under SP-MoE.
• No standard accuracy benchmark (GSM8K/MMLU) was run.

Models other than K3 are unaffected: this touches only kimi_k3.py, and for every backend other than deepep_v2 both expressions evaluate exactly as before.

Speed Tests and Profiling

No direct perf claim; this changes which code path K3 takes, and removing a redundant DP gather / TP reduce should if anything help. For context, at matched chunk 2048 on the config above: v2 prefill ISL=8K/OSL=1 5903.74 vs v1 5540.26 tok/s (+6.6%). Decode TPOT is at parity, but v2 showed a TTFT mean/median gap (18138 vs 3444 ms, against v1's 3358/3363) that I have not root-caused and that is not attributable to this change.

Checklist

• Format your code according to Format code with pre-commit
• Add unit tests — see note below
• Update documentation — n/a
• Provide accuracy and speed benchmark results
• Follow the SGLang code style guidance

On unit tests: the admissible form here is the "critical-path bookkeeping" category in .claude/rules/unit-test-admission.md — a registry-completeness test asserting that every backend for which is_deepep_class_backend() is true is also handled by K3's two predicates, which is exactly the "someone extended X without updating Y" failure mode that produced this bug. That test only makes sense together with a decision on the pplx question above, so I have left it out pending your answer; say the word and I'll add it.

---

CI States

Latest PR Test (Base): ❌ Run #33584403330
Latest PR Test (Extra): ❌ Run #33584403158
Latest PR Test (AMD ROCm 7.2): ❌ Run #33584403187

[whn09]

@BBuf Hi, can you help to review this PR?