[MoE][deepep_v2] Fix ep_scatter_from_psum missing positional args to _fwd_kernel_ep_scatter_2

[cynton503]

Fixes #39402：ep_scatter_from_psum calls the Triton kernel _fwd_kernel_ep_scatter_2 without the two positional arguments expert_start and num_experts, so Triton's dynamic_func() raises "missing 2 required positional arguments" on the DeepEP v2 contiguous/extend (prefill) path. Pass expert_start=0 (the psum recv layout already carries local expert ids) and num_experts (already computed locally), mirroring the ep_scatter call.

Motivation

ep_scatter_from_psum calls the Triton kernel _fwd_kernel_ep_scatter_2 without the two positional arguments expert_start and num_experts. The kernel signature declares them as positional params (immediately before topk_num: tl.constexpr):

def _fwd_kernel_ep_scatter_2(
    ...,
    output_index, output_index_stride0, output_index_stride1,
    expert_start,
    num_experts,
    topk_num: tl.constexpr,
    ...
):

The other caller, ep_scatter, passes both. But ep_scatter_from_psum jumps straight from output_index.stride(1) to the topk_num= keyword, so Triton's dynamic_func() fails at argument binding with:

TypeError: dynamic_func() missing 2 required positional arguments: 'expert_start' and 'num_experts'

This path is reached on the DeepEP v2 contiguous/extend (prefill) route via pre_permute_deepep_v2_to_deep_gemm (layers/moe/moe_runner/deep_gemm.py), and is required whenever --moe-a2a-backend deepep_v2 runs an FP8 blockwise-quantized model (the psum path). The masked/decode path goes through expand_to_masked_slab instead,
so it never hits this call — which is why the bug only surfaces on the first extend/prefill forward.

Modifications

In ep_scatter_from_psum, add the two missing positional args to the kernel call, mirroring ep_scatter:

diff
     output_index,
     output_index.stride(0),
     output_index.stride(1),
+        # The psum recv layout already carries local expert ids, so the
+        # global-to-local shift is zero.
+        0,
+        num_experts,
     topk_num=recv_topk.shape[1],

• expert_start=0: the psum input layout is already organized by local expert id, so there is no global-to-local id shift (the kernel computes expert_id = global_expert_id - expert_start). This matches ep_scatter's default expert_start=0.
• num_experts: reuse the local variable already computed at the top of the function (num_experts = psum_num_recv_tokens_per_expert.shape[0]), no new computation needed.

Accuracy Tests

Pure call-site fix — restores the missing kernel args so the kernel runs at all; it does not change the kernel's math. No accuracy regression expected. (Verified manually that health_generate returns HTTP 200 on a DeepSeek-R1 FP8 + deepep_v2 TP16/EP16 run after the fix; before the fix every rank raised the above TypeError on the first extend forward.)

No unit test is added: this is a pure Triton kernel call-site fix, and the only CPU-side observable is the kernel-launch argument binding, which cannot run on CPU (a mock-based assertion would be meaningless).

Speed Tests and Profiling

No speed impact: this only unbricks a previously crashing code path. No benchmark change expected.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the [Run and add unit tests](N/A — pure Triton kernel call-site fix; no CPU-testable logic, proper coverage is a GPU kernel test, left as follow-up.)(https://docs.sglang.io/developer_guide/contribution_guide.html#run-and-add-unit-tests).
• Update documentation according to Write documentations.(N/A — no user-facing API/flag/behavior change.)
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.(N/A — no accuracy/speed impact; only unbricks a previously crashing path.)
• 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.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

---

CI States

Latest PR Test (Base): ❌ Run #34827597020
Latest PR Test (Extra): ❌ Run #34827596722
Latest PR Test (AMD ROCm 10): ❌ Run #34827596916