Skip to content

Tune BF16 wgrad and fuse DeepEP metadata scale packing - #64

Draft
A-nnonymous wants to merge 3 commits into
PFCCLab:paddlefrom
pkuzyc:agent/bf16-wgrad-2cta-default
Draft

Tune BF16 wgrad and fuse DeepEP metadata scale packing#64
A-nnonymous wants to merge 3 commits into
PFCCLab:paddlefrom
pkuzyc:agent/bf16-wgrad-2cta-default

Conversation

@A-nnonymous

Copy link
Copy Markdown
Collaborator

Summary

  • tunes the Blackwell BF16 wgrad default path and adds coverage for the varlen-K BF16 wgrad path
  • adds fused DeepEP topk metadata + raw-scale ISA packing support with fallback to the existing metadata path
  • plumbs prepacked activation scales into the Sonic up-projection path so the standalone raw-scale gather can be skipped when the fused path is active
  • extends the metadata NCU harness with production A35B shape and pack-mode switches
  • adds a PR-facing metadata/prelaunch bubble breakdown and experiment/test matrix in artifacts/deepep_sonic_metadata_gap_breakdown_20260709.md

Why

The latest A35B timeline shows Sonic compute is no longer the primary blocker after BF16 wgrad tuning. The remaining DeepEP -> Sonic up-projection path has a stable host-side bubble: the old sequence still shows _gather_raw_scales_1x32_to_isa_kernel, followed by about 485 us p50 before the Sonic up GEMM launch. Same-launch-thread NVTX attributes stable portions of that bubble to view_dtype, empty, transpose, allocator work, and generic CUTE/QuACK wrapper setup.

The fused metadata+scale path is the first step: it removes the separate host transition into raw-scale packing and gives the next optimizer a clear test/perf matrix for eliminating the remaining wrapper/carrier costs without changing numerics.

Validation

  • CUDA_VISIBLE_DEVICES=0 python -m pytest tests/ops/test_deepep_topk_metadata.py::TestCudaScalePacking -q --tb=short -> 9 passed
  • CUDA_VISIBLE_DEVICES=0 python -m pytest tests/ops/test_deepep_topk_metadata.py -q --tb=short -> 74 passed
  • CUDA_VISIBLE_DEVICES=0 python tests/ops/ncu_deepep_topk.py --config prod_a35b --with-scales --repeat 10 -> avg_gpu_ms=0.356854, packed_scales_shape=(1, 208044032)

Notes For Reviewers

  • This PR does not disable or modify the EP barrier.
  • The current A35B path is DeepEP, not MoE allgather.
  • The new analysis document distinguishes measured facts from proposed follow-up experiments and lists the required bit-exact tests for each follow-up.
  • Untracked local NCU .ncu-rep artifacts were intentionally not committed.

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.

2 participants