[Dev] Reuse caller-owned HybridEP dispatch and expert output buffers

[lhb8125]

Summary

This draft bounds two HybridEP cross-stream allocation patterns with caller-owned persistent buffer rings:

• dispatch output: HybridEP writes permuted tokens into a reusable buffer; the expert consumer records the final-read event;
• expert FC2 output: Transformer Engine fused grouped MLP writes FC2 directly into a reusable buffer; HybridEP combine records the final-read event;
• provide a shared schedule-aware buffer pool that keeps storage persistent, returns a fresh TensorImpl per acquisition, and uses external CUDA events for eager warmup → full-iteration CUDA Graph capture/replay;
• keep both optimizations opt-in, validate static shapes and required backends, and fail explicitly on early reuse or double release.

The final-consumer release points are intentional. A generic ScheduleNode.free_input ownership probe must not mutate ring state because an old alias can be encountered after the same storage has already been acquired by a later operation.

Configuration

# Dispatch-output reuse (independent)
moe_hybridep_reuse_dispatch_output_buffers: true
moe_hybridep_num_dispatch_output_buffers: 4

# Expert FC2-output reuse (independent; zero disables)
moe_hybridep_num_expert_output_buffers: 4

Both paths require a static HybridEP output bound from moe_expert_rank_capacity_factor or moe_pad_expert_input_to_capacity.

Dispatch-output reuse additionally requires HybridEP flex dispatch, EP communication overlap, FP8/FP4 expert GEMMs, and a DeepEP dispatch_with_permute(..., output_token=...) API.

Expert-output reuse requires the Transformer Engine op fuser and grouped MLP. It depends on Transformer Engine PR #3161 (GroupedLinear caller-provided output), plus a cuDNN frontend wrapper that forwards the caller output pointer. Megatron fails early if the TE API is unavailable.

Motivation

A 64-GPU DeepSeek-V3 MBS2 allocator snapshot showed 9 HybridEP dispatch outputs × 1,078 MiB = 9.47 GiB in the caching allocator pending-free queue at the max-active point. Those buffers are logically dead after grouped-MLP consumption, but the generic allocator cannot infer the higher-level producer/consumer ordering across streams.

FC2/grouped-GEMM outputs have the analogous lifetime:

grouped FC2 on compute stream writes expert_output
    -> HybridEP combine on communication stream performs the final read
    -> caching allocator waits for cross-stream retirement

The application-owned rings make these final-consumer dependencies explicit without a CPU or stream-wide synchronization and bound the amount of persistent output storage.

Validation

• Megatron autoformat against current dev: Black, isort, pylint (10.00/10), and ruff passed.
• Copyright check passed for all changed Python files.
• Unit coverage includes ring rotation, ownership probes, early/double release, eager-warmup-to-CUDA-Graph transition, FC2 output routing, and combine final-consumer release.
• Existing dispatch-output validation:
  • 8-GPU full-iteration CUDA Graph control/feature runs completed 10/10;
  • four 64-GPU, 30-step slot-sweep runs completed 30/30;
  • 4 slots matched control throughput and reduced max reserved from 169,326 to 159,026 MiB on the measured MBS2 proxy.
• Expert FC2-output validation:
  • 64-GPU GB200 proxy completed full-CG replay with combined-1F1B and with the regular schedule, including dense-inner symmetric UBR;
  • 256-GPU DeepSeek-V3, MBS1, combined-1F1B/full-CG/UBR completed 30/30 (SLURM 2796654): mean/median 1112.34/1115.00 TFLOP/s/GPU, peak allocated/reserved/device-used 154701/175300/189420 MiB;
  • the comparable no-pool baseline was 1124.86/1128.10 TFLOP/s/GPU and 152545/175294/187430 MiB. This establishes compatibility and stability, but does not show a throughput or reserved-memory benefit for MBS1;
  • exploratory full-model regular-schedule one-step controls completed with 0 and 10 expert-output slots, while four slots spun before the first iteration at EP64. Because HybridEP does not expose a safe consumer-command depth bound, production configuration now rejects expert-output reuse without combined-1F1B overlap; the pending 30-step regular runs are characterization only.

Known limitations

• Ring size is schedule- and pipeline-depth-dependent. Reuse before the prior consumer release fails explicitly.
• The current measured MBS1 expert FC2 output is about 539 MiB per slot; persistent allocated memory therefore rises linearly with the configured slot count.
• Expert-output reuse intentionally requires combined-1F1B overlap. DeepEP's dispatch kernel stage-count environment variables are not a general cross-API lifetime bound, so regular-schedule reuse fails fast until the backend exposes one.
• Throughput measurements used different node sets. Small deltas are treated as parity/regression signals, not speedup claims.
• Dispatch-output reuse is rejected with fine-grained fused_group_mlp activation offload because that path releases the persistent dispatch input storage.
• This remains a draft until the final-code full-model replay and clean-dev GPU CI are complete.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.