Skip to content

Optimize DeepEP metadata-to-GEMM launch bridge - #68

Draft
A-nnonymous wants to merge 3 commits into
PFCCLab:paddlefrom
A-nnonymous:agent/metadata-up-launch-clean-20260713
Draft

Optimize DeepEP metadata-to-GEMM launch bridge#68
A-nnonymous wants to merge 3 commits into
PFCCLab:paddlefrom
A-nnonymous:agent/metadata-up-launch-clean-20260713

Conversation

@A-nnonymous

@A-nnonymous A-nnonymous commented Jul 13, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • Pack four E8M0 scale bytes into each DeepEP int32 carrier word without numeric conversion.
  • Derive Sonic metadata and allocate gated outputs in one native operation.
  • Use a pointer-safe prepared TVM-FFI launch path that rebuilds all per-call tensor descriptors.
  • Connect router-score gradients directly to the source tensor while preserving the legacy fallback.
  • Fix BF16 down-projection wgrad accumulation into persistent main_grad.

Companion Fleet PR: PaddlePaddle/PaddleFleet#1458

Correctness fix

Fleet converts both down-projection weight and main_grad to Sonic layout [E, H, I]. The BF16 wgrad result has the same logical shape, but the previous path wrote it through an [E, I, H] transposed view. The kernel accepted the mismatched view and silently corrupted the persistent optimizer buffer. This PR writes directly to [E, H, I] and rejects any output or accumulator whose logical shape does not match the GEMM contract.

The fused router edge keeps metadata scores forward-only, scatters dscore to the current source tensor, and does not cache microbatch tensor pointers.

Validation

  • BF16 dw1/dw2 correctness and determinism tests passed.
  • Both beta-accumulate and TMA-add regressions accept [E, H, I] and reject the old transposed shape.
  • Router forward, dscore, dz, and dw2 are bit-exact against the standalone carrier path; metadata receives no gradient.
  • In a single-node 8-GPU 100-step integration run, all samples and learning rates matched the reference. Maximum loss difference was 6.49e-4, versus 1.91e-2 before the fix, under the same nondeterministic settings.
  • Existing metadata, live-context, large-shape, FFI pointer-rebinding, and sanitizer coverage remains in place.

Performance

On the same steady-state integration trace, dispatch-to-gated-GEMM idle latency changed from 208.4 us to 85.5 us at p50 and from 243.3 us to 104.4 us at p95.

Compatibility

Capability, shape, and pointer checks retain the validated fallback paths. No barrier logic or persistent parameter-buffer ownership is changed.

Pack DeepEP scale carriers without materialization, allocate independent gated outputs in the native metadata bridge, and bypass generic gated GEMM launch setup with a pointer-safe TVM FFI path. Add byte-exact, live-context, large-shape, and fallback regression coverage.
@pkuzyc pkuzyc mentioned this pull request Jul 20, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant