[WANDY666]

普通 decode 原来先算 shared expert，再进入 routed experts；routed 输出的 combine 通信期间没有利用这段 shared 计算。本次在支持的 DeepGEMM/DeepEP 路径中先发起 combine send，再计算 shared expert，最后调用 receive hook 等待 routed 结果并相加。不需要开启 Microbatch。

沿用 main 现有 fused_experts 调用链，只增加可选的 shared-expert 回调。路由捕获仍由 weight owner 处理；prefill、原地修改输入的 backend 和 SM100 FP4 Mega MoE 保持原路径，不把回调传给不支持的实现。

提升效果来自之前 H100 1P1D、TP8/DP8、MTP3、16K 输入／4096 输出的服务对照，以下不是本次重跑新版 main 的结果：

并发	平均 TPOT	整轮 output TPS
128	8.1817→7.6108 ms（下降 6.98%）	13,900.8→15,342.5（增加 10.37%）
512	13.0199→12.4207 ms（下降 4.60%）	30,066.7→32,115.6（增加 6.81%）

同轮 profile 中，通信／collective 区间约从 7.55 降到 6.63 ms/周期。历史测量有负载与批次波动，不能把这些比例当成所有配置的固定收益。原始记录见 分支性能记录。

本次补充：

• FusedMoeWeight 对外暴露 backend 能力，模型不再穿透访问实现对象。
• 明确 callback 的只读输入、当前 compute stream、恰好一次调用和已包含 shared 的返回值；Graph replay 只执行捕获的 GPU 工作。
• _ffn() 合并为一次 experts 调用和一次 reduce，保留 shared 先于原地 backend 的回退顺序。
• 新增 LIGHTLLM_DSV4_SHARED_EXPERT_OVERLAP，默认 1；设为 0 回到先 shared 后 routed，能力判断不受开关影响。开关在层初始化时读取，A/B 需重新创建层/服务与 Graph。
• 增加 test/benchmark/kernel/benchmark_shared_expert_overlap.py，使用固定路由、合成 FP8 权重验证实际模型 FFN → weight → DeepGEMM → DeepEP 调用链，并交错测量开启/关闭的 eager 与 Graph。单元回归已提交，详见下文。

本次在 8×H200、PyTorch 2.11.0+cu130 上实跑。DeepEP 为 Dockerfile 固定的 60d44037，DeepGEMM 为 891d57b4。采用 DS4 Flash 的 hidden=4096、intermediate=2048、256 experts、top-6；UE8M0 scale，dispatch 容量为服务默认的 256。四组在所有 rank 上的 eager、Graph replay、更新输入与路由后的 replay 均逐元素一致，输入未被改写；capture 时 callback 执行一次，replay 不再调用 Python。

当前 head 的单层实测如下（μs，关闭→开启；每组 4 轮交错 A/B，每轮 50 次，取各轮最慢 rank 的中位数）：

每 rank 请求数	每请求展开行数	eager FFN	Graph FFN
16	1	806.87→799.44	278.76→280.34
16	4	389.18→391.37	294.95→295.93
64	1	381.06→389.68	295.52→296.22
64	4	409.14→406.16	366.34→367.94

这次微测没有复现历史整机的明显加速，Graph 慢约 0.2%～0.6%；不能宣称本次 head 已获得历史表中的服务收益。固定路由不含路由计算，4 行仅模拟 MTP 展开形状，不包含完整 speculative 流程；单机 H200、合成权重和单层测试也不能替代 H100 整模型精度及端到端吞吐对照。保留开关以便在目标部署按实际负载复测。

复跑（Hopper、已安装仓库固定依赖、具备 RDMA 访问）：

PYTHONPATH=. LIGHTLLM_TRITON_AUTOTUNE_LEVEL=3 NVSHMEM_QP_DEPTH=1024 \
python test/benchmark/kernel/benchmark_shared_expert_overlap.py \
    --world-size 8 --requests 16 64 --widths 1 4 \
    --iterations 50 --warmup 5 \
    --output /mnt/vqa/wanzihao/shared-expert-overlap/result.json

此前另有 20 项受控回归通过，覆盖回调契约、能力选择、开关、prefill/原地 backend 回退、路由捕获及 reduce 分支；Black、Flake8 通过。原始验证文件保存在 /mnt/vqa/wanzihao/bench-logs/20261030/h100-ttft-night-2305/seven-main-prs-20261102/validation/overlap-review-20261103/。

已提交的单元测试

unit_tests/models/deepseek_v4/test_shared_expert_overlap.py：24 项通过。直接导入实际模型类，覆盖 SEND→shared→RECV 顺序、结果相加、输入只读、weight owner 的能力与路由捕获、prefill/原地 backend 回退、reduce 次数，以及默认开关和初始化时读取环境变量。

这些是受控单元回归；真实 DeepEP/Graph 数值与性能验证仍使用上面的独立 benchmark，二者不混为一项测试。

python -m pytest -q unit_tests/models/deepseek_v4/test_shared_expert_overlap.py

Black、Flake8 通过。本次只补充测试，未改变推理实现。

Summary by CodeRabbit

• New Features
  • DeepSeek V4 decoding can overlap shared-expert computation with routed-expert processing in supported configurations. The feature is enabled by default and can be controlled with the LIGHTLLM_DSV4_SHARED_EXPERT_OVERLAP setting.
  • When overlap is unavailable or disabled, shared-expert computation follows the existing sequential path. Prefill behavior is unchanged.
  • With overlap enabled, the shared-expert result is combined with the routed-expert result; unsupported configurations continue to use the sequential path.

[coderabbitai]

No actionable comments were generated in the recent review.

Walkthrough

DeepSeek V4 can overlap shared-expert computation with routed-expert combine during supported decode. The change adds a callback and backend capability interface, selects the callback path through an environment setting, and adds tests and a distributed benchmark.

Changes: Decode Shared-Expert Overlap

Layer / File(s)	Summary
Callback API and backend capability lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py, lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/*	The MoE API exposes backend overlap support and accepts an optional shared-expert callback. DeepGEMM sets its capability based on the SM100 Mega MoE check. The Triton wrapper forwards callback arguments to _fused_experts; the shown _fused_experts declaration does not accept arbitrary keyword arguments.
Callback execution during combine lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py	The grouped fused-MoE path restricts callback use to decode and rejects it on the SM100 Mega MoE path. During low-latency combine, it requests a receive hook when a callback is supplied, runs the callback, invokes the hook, and adds the callback result to the combined expert output.
DeepSeek V4 decode integration lightllm/models/deepseek_v4/layer_infer/transformer_layer_infer.py	An environment setting enables overlap by default. During decode, DeepSeek V4 passes the shared FFN as a callback when the setting and backend capability allow it. Otherwise, it computes the shared result before routed experts and adds it afterward.
Overlap tests and benchmark unit_tests/models/deepseek_v4/test_shared_expert_overlap.py, test/benchmark/kernel/benchmark_shared_expert_overlap.py	Unit tests cover callback behavior, backend capability, fallback ordering, and environment parsing. The distributed benchmark checks eager output equality, input preservation, and CUDA graph replay behavior for overlap-enabled and disabled modes. It measures eager and graph timings and writes results and metadata.

Estimated code review effort: 3 (Moderate) | ~25 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Transformer as DeepseekV4TransformerLayerInfer
  participant Weights as FusedMoeWeight
  participant Triton as FuseMoeTriton
  participant GroupedMoE as grouped_fused_moe_ep
  participant SharedFFN as Shared FFN callback
  Transformer->>Weights: Pass callback during supported decode
  Weights->>Triton: Forward callback to fused_experts_with_topk
  Triton->>GroupedMoE: Forward callback to fused_experts
  GroupedMoE->>SharedFFN: Run callback during low-latency combine
  GroupedMoE-->>Transformer: Return combined expert and shared output

Loading

Suggested reviewers: shihaobai, hiworldwzj

Merge Risk: Minimal · up to 1d75d

No merge-blocking issue was established in the overlap or fallback paths.

Pre-merge checks: 4 · 1

Failed checks: 1 warning

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 7.14% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentat…	Write docstrings for the functions missing them to satisfy the coverage threshold.

Passed checks: 4 passed

Check name	Status	Explanation
Description Check		Check skipped - CodeRabbit’s high-level summary is enabled.
Title check		The title "perf(dsv4): overlap shared expert computation with DeepEP combine" matches the main change. The PR adds an optional shared-expert callback to the fused MoE call chain. It starts the routed-…
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.

Full details: Title check

Explanation

The title "perf(dsv4): overlap shared expert computation with DeepEP combine" matches the main change. The PR adds an optional shared-expert callback to the fused MoE call chain. It starts the routed-expert combine send before the shared expert runs on supported DeepGEMM/DeepEP decode paths. The title is short, specific, and names the DeepSeek V4 scope. It does not mention the env switch or the benchmark, which is expected for a title.

Full details: Docstring Coverage

Explanation

Docstring coverage is 7.14% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentation-only edits and deletions within surviving functions are included; unchanged diff context and fully deleted functions are excluded. Analyzed 14 functions across 8 files.

Resolution

Write docstrings for the functions missing them to satisfy the coverage threshold.

Finishing Touches

Generate docstrings

• Commit to this branch
• Create a new PR

Generate unit tests (beta)

• Commit to this branch
• Create a new PR

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

Share

• X
• Mastodon
• Reddit
• LinkedIn

CodeRabbit Verify

• Run CodeRabbit Verify to test the changes in this pull request.

Recent review info

Run configuration

• Source mode: Auto
• Repository: Unavailable
• Branch: Unavailable
• File: No configuration file from this repository applied
• Commit: Unavailable
• Layers: defaults
• Review profile: CHILL
• Plan: Advanced
• Run ID: 32704445-2927-4f6b-8cc9-2a0afe7a2620

Commits

Reviewing files that changed from the base of the PR and between 5718695 and 1d75d2b.

Files selected for processing: 1

• unit_tests/models/deepseek_v4/test_shared_expert_overlap.py

Included review availability: This review used your included allowance. Your plan provides up to 8 included reviews per hour; 5 remain after this review.