[Performance] Reduce SonicMoE metadata launch overhead

[A-nnonymous]

PR Category

Performance Optimization

PR Types

Performance

Description

This PR removes avoidable host and GPU launch overhead between DeepEP dispatch and SonicMoE gated GEMM.

Root causes addressed:

• Four opaque E8M0 scale bytes are packed into each DeepEP int32 carrier word without numeric conversion.
• Metadata conversion and gated-output allocation are combined when the companion Sonic capability is available.
• Router-score gradients connect directly to the current source tensor instead of using a per-microbatch carrier.
• The caller reads stop_gradient before entering the Sonic PyLayer. Stopped scores remain forward-only and do not receive an invalid backward edge; differentiable scores retain the fused scatter path.

Compatibility is capability-gated. Legacy quantization, metadata, and router-gradient paths remain available with older Sonic versions. No barrier logic, optimizer bypass, layout conversion, or persistent parameter/gradient buffer ownership is changed.

Companion PR: PFCCLab/supersonic-moe#68

The Sonic submodule pointer remains at the current upstream revision. It can advance after the companion PR lands.

Validation

• Focused tests passed for the fused differentiable edge, stopped-score semantics, and legacy carrier fallback.
• Repository pre-commit hooks passed, including ruff, formatting, private-key, conflict, and AST checks.
• Companion tests cover bit-exact router forward/backward, metadata ownership, FFI pointer rebinding, large shapes, and BF16 persistent main_grad accumulation.
• A single-node 8-GPU 100-step integration run matched all reference samples and learning rates. Maximum loss difference was 6.49e-4, versus 1.91e-2 before the companion wgrad fix, under the same nondeterministic settings.

[risemeup1111]

已完成初轮代码审查。当前发现一个需要提交新 commit 修复的阻塞问题，另有一个关于新 packed-scale combine 分支可达性的中优先级问题，详情已放在行级评论里。

CI 里仍有多项测试在运行中；请等相关检查结束后再合入。

^{Powered by Nyanpasu with gpt-5.5 xhigh, please check the suggestions carefully.}

[risemeup1111]

优先级：P1

问题：RR DeepEP combine 分支没有把 fp8_dispatch / combine_grad_handle 继续传给 DeepEPCombineAsyncFunctor。

证据与影响： 这里新增的 _record_fp8_combine_grad(...) 只有在 ctx.fp8_dispatch=True 时才会把 DeepEP combine backward 的 FP8 data/scale 写回 handle。但 fused_combine(..., use_rr_deepep_combine=True) 当前调用 _rr_fusedcombined(..., fn=...) 时没有传 fp8_dispatch 和 combine_grad_handle，而 DeepEPCombineAsyncRefinedRecompute._second_fwd() 又用默认参数调用 DeepEPCombineAsyncFunctor.apply(..., fn=fn)。因此 fp8_dispatch_bwd + use_rr_deepep_combine 的实际路径会跳过这段写回逻辑，_DownProjection.backward 收到的 fp8_combine_grad_handle 无法拿到 DeepEP combine backward 产生的 FP8 梯度/scale，至少会丢掉该配置下的 FP8 backward 通信优化，并可能和 Sonic backward 对 handle 的预期不一致。

期望处理： 请把这两个参数从 fused_combine 的 RR 分支一路传到 DeepEPCombineAsyncFunctor.apply，并补一个覆盖 use_rr_deepep_combine=True, fp8_dispatch=True, combine_grad_handle=... 的测试。修复形态可以类似：

# fused_combine RR branch
combined_x, *fn_out = _rr_fusedcombined(
    x,
    group,
    states,
    *(combine_overlap_handle["fn_args"]),
    fn=combine_overlap_handle["fn"],
    fp8_dispatch=fp8_dispatch,
    combine_grad_handle=combine_grad_handle,
)

class DeepEPCombineAsyncRefinedRecompute:
    def forward(
        self,
        x,
        group,
        states,
        *fn_args,
        fn,
        fp8_dispatch=False,
        combine_grad_handle=None,
    ):
        ...
        output = self._second_fwd(
            hold_tensors,
            x,
            group,
            states,
            fn,
            *fn_args,
            fp8_dispatch=fp8_dispatch,
            combine_grad_handle=combine_grad_handle,
        )

    def _second_fwd(
        self,
        hold_tensors,
        x,
        group,
        states,
        fn,
        *fn_args,
        fp8_dispatch=False,
        combine_grad_handle=None,
    ):
        return DeepEPCombineAsyncFunctor.apply(
            hold_tensors,
            x,
            group,
            states,
            *fn_args,
            fn=fn,
            fp8_dispatch=fp8_dispatch,
            combine_grad_handle=combine_grad_handle,
        )

处理要求：请针对该评论修复并提交新的 commit。

[risemeup1111]

优先级：P2

问题：这个 Sonic packed-scale combine 分支依赖 combine_grad_handle["using_sonic_moe"]，但生产路径创建的是普通 {}，目前没有看到任何调用链会写入这个字段。

证据与影响： MoELayer.fusion_moe_forward() 在 self.fp8_dispatch_bwd 时只创建 fp8_combine_grad_handle = {}，随后传给 run_sonic_moe() 和 _DeepEPManager.combine()；全仓搜索也只有本次新增单测手工构造了 {"using_sonic_moe": True}。这会让 _quantize_combine_grad_for_deepep() 在真实 Sonic+DeepEP combine backward 中始终走 scale_dtype=paddle.int32 的旧路径，新增的 pack_scale_words / _unpack_sonic_fp8_scale_from_deepep() 分支基本不可达，PR 里关于 combine-grad scale 字节打包的优化和覆盖容易产生误判。

期望处理： 请在生产调用链里明确标记该 handle，或改成由调用方直接传入 using_sonic_moe 布尔值，避免依赖测试专用状态。例如可以在创建 handle 时带上标记，并补一个不 mock handle 内容的端到端单测：

fp8_combine_grad_handle = (
    {"using_sonic_moe": self.using_sonic_moe}
    if self.fp8_dispatch_bwd
    else None
)

如果你的设计是只优化 dispatch scale 而不优化 combine-grad scale，也请删掉这部分不可达分支或调整测试断言，避免后续维护者以为真实路径已覆盖。

处理要求：请针对该评论进行回复（同意并已修改请回复 Done，不同意请说明理由）。

[codecov-commenter]

Codecov Report

❌ Patch coverage is 62.00000% with 38 lines in your changes missing coverage. Please review.
⚠️ Please upload report for BASE (develop@1892981). Learn more about missing BASE report.

Files with missing lines	Patch %	Lines
.../paddlefleet/transformer/moe/fusion_layer_utils.py	25.64%	29 Missing ⚠️
src/paddlefleet/transformer/moe/fused_a2a.py	85.24%	7 Missing and 2 partials ⚠️

❌ Your patch status has failed because the patch coverage (62.00%) is below the target coverage (90.00%). You can increase the patch coverage or adjust the target coverage.

Additional details and impacted files

@@            Coverage Diff             @@
##             develop    #1458   +/-   ##
==========================================
  Coverage           ?   62.00%           
==========================================
  Files              ?        2           
  Lines              ?      100           
  Branches           ?       21           
==========================================
  Hits               ?       62           
  Misses             ?       36           
  Partials           ?        2           

Flag	Coverage Δ
coverage_combine	62.00% <62.00%> (?)

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines	Coverage Δ
src/paddlefleet/transformer/moe/fused_a2a.py	85.24% <85.24%> (ø)
.../paddlefleet/transformer/moe/fusion_layer_utils.py	25.64% <25.64%> (ø)

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[Paddle-CI-Bot]

PaddleFleet Log Analysis

Run #29274536588 · Attempt 1

日志分析报告

流水线名称	问题标签	修复建议	日志片段
Integration test (H20, multi-card) — job 86884643801	NCCL 通信错误 / MoE fused_dispatch 崩溃	检查本 PR 对 fused_a2a.py 的修改，尤其是 barrier_ep 路径引入的 EP group 传参变化；确认 fused_dispatch_forward_func 在多卡 EP 场景下 ep_group 非空	报错代码
Coverage Upload And Check — job 86900089010	增量覆盖率不达标（exit 9）	为 fusion_layer_utils.py 新增行（3276–3464）和 fused_a2a.py（171,651–787）补充单测，将增量覆盖率从 61% 提升到 ≥90%	报错代码

失败的测试 case:

Integration test (H20, multi-card) — Qwen LoRA 多卡训练（qwen_lora）：
  OSError: NCCL error 'unhandled cuda error' at process_group_nccl.cc:912
  调用链：moe_layer.dispatch → token_dispatcher.token_dispatch → fused_a2a.fused_dispatch → barrier_ep → paddle.distributed.barrier
  exit_code=241，check_exit_code=1，Training completed 未出现在日志中

Coverage Upload And Check — 增量覆盖率 61%（阈值 90%）：
  src/paddlefleet/transformer/moe/fusion_layer_utils.py  23.3%  (Missing lines 52-53,56-57,3276,3290,3292-3293,3296,3303-3304,3318,3340,3365-3367,3370,3374-3376,3379-3382,3399-3401,3404,3406,3436,3456-3457,3464)
  src/paddlefleet/transformer/moe/fused_a2a.py           88.5%  (Missing lines 171,651,666,714,730,771,787)
  Total: 104 changed lines, 40 missing, Coverage: 61%

根本原因分析:

本 PR（agent/sonic-metadata-launch-bridge-20260713）改动了 fused_a2a.py 和 fusion_layer_utils.py：fused_a2a.py 中 barrier_ep 调用的 EP group 在 Qwen LoRA 多卡 EP 场景下触发了 NCCL unhandled CUDA error（进程以 signal 241 退出）；同时新增的 fusion_layer_utils.py 逻辑（Sonic MoE 相关，~3276–3464 行）未被任何测试覆盖，导致增量覆盖率仅 61%。

修复建议:

1. NCCL 崩溃：在 fused_a2a.py 的 barrier_ep 前加护卫，确保 ep_group 不为空或 group.nranks > 1 时才执行 barrier；在 Qwen LoRA 多卡配置下本地复现（EP size > 1），用 NCCL_DEBUG=INFO 抓完整 NCCL 错误再定位根本 CUDA error 类型。
2. 覆盖率：针对 fusion_layer_utils.py 新增的 Sonic MoE 分支（_resolve_sonic_config_bool、run_sonic_moe 相关路径 3276–3464 行）补充单元测试，参考已有的 test_run_sonic_moe_* 系列 test case 扩展覆盖缺失分支；fused_a2a.py 缺失的 7 行（171,651,666,714,730,771,787）同样需补 mock 测试。

---

🔍 准确性记录：请点击评论底部 😊 图标，选择 👍（准确）或 👎（有误），将自动记录到 CI 监控系统

_{🔄 每次 Re-run 后自动更新}

[PaddlePaddle-bot]

🤖 Paddle-CI-Agent | pr_review | 2026-07-14 01:53:40

📋 Review 摘要

PR 概述：优化 DeepEP dispatch 到 SonicMoE gated GEMM 之间的 FP8 scale carrier、metadata 和 router-score 梯度桥接路径。

变更范围：src/paddlefleet/transformer/moe/fused_a2a.py、src/paddlefleet/transformer/moe/fusion_layer_utils.py 以及对应单卡桥接测试。

影响面 Tag：MoE FP8

问题

未发现新的阻塞性问题。PR 规范问题在下面章节报,不要在这里重复。

历史 Findings 修复情况

Finding	问题	状态
F1	fp8_combine_grad_handle 未携带 using_sonic_moe，导致 combine backward 的 Sonic pack/unpack 分支无法命中	⚠️ 仍存在

📝 PR 规范检查

符合规范。

总体评价

本轮重点复核了 MoE Sonic FP8 dispatch/combine 对称性、metadata bridge 的 router-score 梯度路径，以及新增测试对新旧 Sonic capability fallback 的覆盖。最新提交已补上 stopped router score 不应挂梯度边的判断，但历史 F1 指向的真实 handle 标记仍未在当前 diff 中修复，本轮不重复发 inline comment。