[Fix] deepep_v2: pass expert_start/num_experts in ep_scatter_from_psum

[whn09]

Motivation

ep_scatter_from_psum launches _fwd_kernel_ep_scatter_2 without the expert_start
and num_experts parameters the kernel has required since 5f216fc, so every rank
dies on the first forward pass when --moe-a2a-backend deepep_v2 is used:

sglang/srt/layers/moe/moe_runner/deep_gemm.py:1677 in pre_permute_deepep_v2_to_deep_gemm
    ep_scatter_from_psum(
sglang/kernels/ops/moe/ep_moe_kernels.py:1297 in ep_scatter_from_psum
    _fwd_kernel_ep_scatter_2[(grid,)](
TypeError: dynamic_func() missing 2 required positional arguments: 'expert_start' and 'num_experts'

Still present on main @ 29578d5.

5f216fc ("qwen 3.8 rebase", #35758) generalised _fwd_kernel_ep_scatter_2 from local
to global expert ids:

-    expert_id = tl.load(recv_topk + ...)
-    if expert_id >= 0:
+    global_expert_id = tl.load(recv_topk + ...)
+    expert_id = global_expert_id - expert_start
+    valid = (expert_id >= 0) & (expert_id < num_experts)
+    tl.store(output_index_ptr, -1)
+    if valid:

It added the two parameters and updated the ep_scatter call site. The other call
site, ep_scatter_from_psum, had landed two days earlier in a3ae667 (#35634, DeepEPv2
ElasticBuffer MoE A2A backend) and was not updated. ep_scatter_from_psum is reached
only from pre_permute_deepep_v2_to_deep_gemm, i.e. only with
--moe-a2a-backend deepep_v2, so the omission is invisible to every other backend and
fatal for that one. Nothing in CI exercises the path.

Reproduces on any deepep_v2 server that reaches the non-masked permute path. Ours:
GLM-4.6-style MoE (GLM-5.1-FP8) on 4x p5en.48xlarge (H200, EFA), prefill TP=16 EP=16,
--moe-a2a-backend deepep_v2 --deepep-v2-mode hybrid. The crash lands after weight
load and ElasticBuffer init, so everything up to it looks healthy. The masked path
(expand_to_masked_slab) returns early and does not reach this launch.

Modifications

1. ep_scatter_from_psum passes expert_start and num_experts to
   _fwd_kernel_ep_scatter_2, and takes expert_start: int = 0 — the same defaulted
   parameter ep_scatter already has.

   expert_start=0 restores the pre-5f216fc3 behaviour rather than assuming it. The
   old kernel indexed expert_start_loc + expert_id directly, and this function is
   called with an expert_start_loc of length num_local_experts: deep_gemm.py
   builds it as torch.empty_like(psum_num_recv_tokens_per_expert) and reads that same
   shape as num_local_experts a few lines earlier. Global ids would therefore have
   indexed out of bounds on every rank but rank 0 from the day [Feature] Add DeepEPv2 (ElasticBuffer) MoE A2A backend  #35634 merged — the
   DeepEPv2 dispatch delivers already-local ids with -1 padding, so local_id - 0
   reproduces the old behaviour exactly and -1 still fails the >= 0 test. Exposing
   the parameter rather than hard-coding 0 gives a future caller holding global ids
   somewhere to say so; happy to pass a literal 0 instead if you would rather not
   widen the signature.

   num_experts is psum_num_recv_tokens_per_expert.shape[0], which the function
   already computes for the _fwd_kernel_ep_scatter_psum_init grid. The psum is an
   inclusive prefix sum of length num_local_experts
   (_fwd_kernel_ep_scatter_psum_init reads psum[e-1]..psum[e] as expert e's
   span), so the new expert_id < num_experts bound is exact — no off-by-one.

   This also satisfies the requirement that output_index be -1-initialised, which
   the same commit introduced when it switched post_reorder_deepgemm_triton_kernel
   from gating on expert_id >= 0 to dst_idx >= 0 (output_index = torch.empty_like(topk_ids) is otherwise uninitialised).

2. New test/registered/kernels/ops/moe/test_ep_scatter.py. Neither ep_scatter nor
   ep_scatter_from_psum had any coverage, which is exactly why a signature change
   could update one call site and leave the other broken.

   Destination slots are chosen by atomic_add, so the order within one expert is not
   deterministic. The test asserts the invariants that pin the result down regardless:
   the -1 sentinel lands exactly on the padding lanes; every valid lane lands inside
   its own expert's slab; the slabs form a bijection onto [0, total); m_indices
   agrees with the routing and its BLOCK_E tail stays -1; expert_start_loc ends
   at the prefix sum; and the rows and FP8 scales themselves arrive. Parametrised over
   num_tokens in {1, 8, 256, 4096}, padding fraction, and bf16 / float8_e4m3fn,
   plus cases for a nonzero expert_start and for expert ids this rank does not own.

   It leaves expert_start at its default in the local-id case, so it exercises the
   call shape the DeepEPv2 permute path actually uses — which is what makes it a
   regression test for this bug rather than only for the signature.

Accuracy Tests

The new unit test passes (20/20) and fails on an unpatched tree with exactly the
TypeError above, so the before/after is observed rather than asserted (H200, one GPU):

A: main + this PR      ->  20 passed in 8.57s
B: main, unpatched     ->  TypeError: dynamic_func() missing 2 required positional
                           arguments: 'expert_start' and 'num_experts'

Also checked structurally: both launch sites now supply all 22 parameters preceding the
first tl.constexpr, and the two added arguments land on the kernel's expert_start /
num_experts formals rather than merely making the count add up.

There is no end-to-end accuracy comparison to offer, because the path cannot execute a
single forward pass without this change — there is no pre-fix baseline to diff against.
Correctness rests on the argument above that this restores pre-5f216fc3 semantics
exactly, plus the invariant checks in the test. Glad to add a gsm8k run on
--moe-a2a-backend deepep_v2 if you would like one attached.

Speed Tests and Profiling

None: the change passes two additional scalar arguments to an existing kernel launch
and adds no work. Not applicable.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations. — not applicable, no user-facing behaviour change
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed. — see above; speed is not applicable and end-to-end accuracy has no pre-fix baseline
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

cc @Qiaolin-Yu (5f216fc, #35758) @MengYu10151 (a3ae667, #35634) — and whoever on the
/python/sglang/kernels CODEOWNERS list is on rotation. I cannot trigger CI myself.

---

CI States

Latest PR Test (Base): ❌ Run #33366659160
Latest PR Test (Extra): ❌ Run #33366659073
Latest PR Test (AMD ROCm 7.2): ❌ Run #33366659163