[PyT] [Common] add support for enabling cuda graph under thd format in megatron.

[HaochenYuan]

Description

This PR supports enabling cuda graph under thd format in megatron.
Related Megatron PR 4359.

Two root causes block THD under CUDA Graph capture:

1. GPU→CPU sync forbidden during capture — torch.equal() on CUDA tensors and
   cu_seqlens[-1] as a Python int both trigger implicit device sync.
2. Garbage at padded positions — fused attention and CP ring kernels leave
   padded slots undefined; in MLA + CP these survive into gradient accumulation.

Patches

#	File	Change
1	attention/dot_product_attention/dot_product_attention.py	Drop torch.equal() check; force pad_between_seqs=True for THD (always safe).
2	attention/dot_product_attention/context_parallel.py (AttnFuncWithCPAndKVP2P.forward)	Zero-fill out/out_ret past cu_seqlens_q_padded[-1]//cp_size after CP assembly. arange-based mask, no sync.
3	attention/dot_product_attention/context_parallel.py (AttnFuncWithCPAndKVP2P.backward, existing zero-fill)	Under is_current_stream_capturing(), use dq.shape[0] instead of cu_seqlens[-1] for the slice end.
4	attention/dot_product_attention/context_parallel.py (AttnFuncWithCPAndKVP2P.backward, before return)	Zero-fill dq/dk/dv past cu_seqlens_q_padded[-1]//cp_size via arange mask.
5	attention/dot_product_attention/backends.py (FusedAttnFunc.backward)	empty_like → zeros_like for dq/dk/dv so MLA+CP path doesn't leak garbage.
6	cpp_extensions/fused_attn.py (fused_attn_fwd/fused_attn_bwd)	Zero-fill output tensors past cu_seqlens_q[-1] for t3hd/th3d/thd_* layouts. arange mask.

All zero-fills use torch.arange(N, device=...) >= bound so the bound stays a CUDA
tensor — graph-safe.

Validation

Both Moonlight-16B (MLA + MoE, CP=2) and Qwen3-8B (GQA, CP=2) produce noGraph vs
cudaGraph bitwise identical results in --deterministic-mode against this branch. See details in test.

Fixes # (issue)

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

Please list the changes introduced in this PR:

• Change A
• Change B

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[greptile-apps]

Greptile Summary

This PR enables CUDA Graph capture under THD format in Megatron by eliminating two root causes: GPU→CPU syncs triggered by torch.equal() and cu_seqlens[-1] used as a Python int, and undefined garbage at padded positions surviving into gradient accumulation. It also adds a device-tensor path for the MoE aux-loss forward so total_num_tokens stays dynamic across graph replays.

• Attention (THD + CP): torch.equal() replaced by identity-based pad_between_seqs auto-detect; backward dq/dk/dv zero-fills converted to sync-free arange-based masks in both FusedAttnFunc and AttnFuncWithCPAndKVP2P.
• Fused router aux loss: New fused_moe_aux_loss_fwd_graph_safe C++ function + CUDA kernel reads total_num_tokens from a 0-dim device tensor, computing the coefficient on-device; dispatched from Python when total_num_tokens is a torch.Tensor.
• Tests: CUDA-graph replay test for the aux-loss (three distinct token counts, each compared against a PyTorch reference) and explicit pad_between_seqs plumbing in the CP attention test runner.

Confidence Score: 5/5

The changes are well-scoped and the core graph-safety fixes are correct; the only nits are efficiency issues in the non-fused attention backward where zero-fills run redundantly.

All graph-safety invariants are correctly maintained: arange-based masks stay on-device, the device-tensor coefficient path never reads the scalar on the host, and the existing eager-mode fill_ paths are preserved. The new second zero-fill block in the CP backward is redundant for the non-fused path and skips the is_graph_capturing() branch that would give cheaper fill_ in eager mode, but neither issue affects correctness.

transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py — the new unconditional arange-mask zero-fill block after nvtx_range_pop runs twice in the non-fused attention THD path and lacks the is_graph_capturing() branch used by the existing block above it.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py	Adds graph-safe backward zero-fills for THD format in AttnFuncWithCPAndKVP2P; the new unconditional arange-mask zero-fill block is redundant in the non-fused attention path where the existing block already zeroes dq/dk/dv, and unlike the existing block it does not branch on is_graph_capturing() for eager efficiency.
transformer_engine/pytorch/attention/dot_product_attention/backends.py	Adds sync-free arange-mask zero-fill for dQ/dK/dV at padded THD positions in FusedAttnFunc.backward; guarded correctly on cu_seqlens_{q,kv}_padded being non-None and tensor shape > 0.
transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py	Replaces torch.equal() (causes GPU→CPU sync) with identity-based pad_between_seqs auto-detect; the trade-off of being slightly more conservative (may yield True when objects differ but values match) is intentional and well-commented.
transformer_engine/common/fused_router/fused_moe_aux_loss.cu	New graph-safe forward kernel reads total_num_tokens from a device pointer so the coefficient stays dynamic across replays; only Coeff_buf[1] is pre-zeroed (Coeff_buf[0] is unconditionally overwritten by block-0 lane-0), which is correct.
transformer_engine/pytorch/router.py	Dispatches to fused_moe_aux_loss_fwd_graph_safe when total_num_tokens is a Tensor, falling back to the original int path otherwise; the int() cast in the fallback correctly handles numpy int and other int-like types.
tests/pytorch/test_fused_router.py	Adds CUDA-graph capture test that replays with three distinct token counts and compares against a PyTorch reference; warmup is correctly done on a side stream to satisfy graph capture requirements.
tests/pytorch/attention/run_attention_with_cp.py	Explicitly passes pad_between_seqs to suppress the conservative auto-detect for FlashAttention paths where the test constructs no real inter-sequence padding.
transformer_engine/pytorch/csrc/extensions/router.cpp	Adds fused_moe_aux_loss_fwd_graph_safe C++ wrapper; validates tensor shape/dtype on the host without reading the value, then launches the new CUDA kernel.
transformer_engine/common/include/transformer_engine/fused_router.h	Adds nvte_fused_moe_aux_loss_forward_graph_safe declaration and improves documentation for the existing host-int path, distinguishing the two API variants clearly.
transformer_engine/pytorch/csrc/extensions.h	Adds declaration for fused_moe_aux_loss_fwd_graph_safe; straightforward header update.
transformer_engine/pytorch/csrc/extensions/pybind.cpp	Registers fused_moe_aux_loss_fwd_graph_safe in the Python module; docstrings clearly distinguish the host-int and device-tensor variants.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[Backward: AttnFuncWithCPAndKVP2P] --> B{qkv_format == 'thd' AND\nnot use_fused_attention AND\nboth padded seqlens non-None?}
    B -- Yes --> C{is_graph_capturing?}
    B -- No --> E
    C -- Yes --> D1[arange mask zero-fill dq/dk/dv]
    C -- No --> D2[fill_ zero-fill dq/dk/dv]
    D1 --> E[FP8 quantize if needed]
    D2 --> E
    E --> F[A2A communicate if cp_size_a2a > 1]
    F --> G{qkv_format == 'thd'?}
    G -- Yes --> H[arange mask zero-fill always\ndq / dk / dv independently guarded]
    G -- No --> I[return dq/dk/dv]
    H --> I
    style H fill:#ffe0b2
    style B fill:#e3f2fd

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
flowchart TD
    A[Backward: AttnFuncWithCPAndKVP2P] --> B{qkv_format == 'thd' AND\nnot use_fused_attention AND\nboth padded seqlens non-None?}
    B -- Yes --> C{is_graph_capturing?}
    B -- No --> E
    C -- Yes --> D1[arange mask zero-fill dq/dk/dv]
    C -- No --> D2[fill_ zero-fill dq/dk/dv]
    D1 --> E[FP8 quantize if needed]
    D2 --> E
    E --> F[A2A communicate if cp_size_a2a > 1]
    F --> G{qkv_format == 'thd'?}
    G -- Yes --> H[arange mask zero-fill always\ndq / dk / dv independently guarded]
    G -- No --> I[return dq/dk/dv]
    H --> I
    style H fill:#ffe0b2
    style B fill:#e3f2fd

Loading

_{Reviews (11): Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." | Re-trigger Greptile}

[cyanguwa]

Do you only have pad tokens at the end of the batch? Would the pads between sequences in the batch also need to be zeroed out?

[HaochenYuan]

Yes, Megatron PR #4359 only adds tail padding to the packed THD buffer. These zero-fills only clear positions >= cu_seqlens_*_padded[-1]; inter-sequence padding stays with the existing TE handling.

[cyanguwa]

Tail padding is probably a very special case of pad_between_seqs=True. If we want to claim that THD + CUDA graph is working, could we please zero out all the inter-sequence padding as well (here and in a few other places)? Thanks!

[HaochenYuan]

In current megatron thd e2e test, due to accuracy problem, we let cu_seqlens_* to be equal to cu_seqlens_*_padded, and address the inter-sequence padding by loss mask. So here we want to only focus on tail padding.

[cyanguwa]

Both Moonlight-16B (MLA + MoE, CP=2) and Qwen3-8B (GQA, CP=2) produce noGraph vs
cudaGraph bitwise identical results in --deterministic-mode against this branch. See details in test.

Do you know if these tests were running with FlashAttention backend or FusedAttention backend?

[HaochenYuan]

Both Moonlight-16B (MLA + MoE, CP=2) and Qwen3-8B (GQA, CP=2) produce noGraph vs
cudaGraph bitwise identical results in --deterministic-mode against this branch. See details in test.

Do you know if these tests were running with FlashAttention backend or FusedAttention backend?

The Megatron E2E bitwise tests run with FusedAttention, not FlashAttention. In the Qwen3 THD config, padded cu_seqlens make TE infer pad_between_seqs=True, and FlashAttention is not available for THD with inter-sequence padding, so auto-selection picks FusedAttention. Moonlight is also FusedAttention-only because head_dim_qk != head_dim_v.

The unit tests cover the backends separately: test_cp_with_flash_attention exercises the FA path, and the CP FusedAttention tests exercise the Fused path. The THD FA unit test explicitly passes pad_between_seqs=False because that test has no inter-sequence padding.

[timmoon10]

/te-ci pytorch L1

[timmoon10]

LGTM, pending CI

[HaochenYuan]

The remaining CI failure looks like an infra/runner issue？: GitHub reports "The hosted runner lost communication with the server" during the Build step.

[cyanguwa]

The remaining CI failure looks like an infra/runner issue？: GitHub reports "The hosted runner lost communication with the server" during the Build step.

The failed tests are due to a different issue which we're fixing, so I think you can ignore those. Thanks!

[HaochenYuan]

Add one fuesd_moe_aux_loss kernel in this PR, the only change is allow the total_num_tokens to be a tensor type, so when we enable router fusion in megatron with cuda graph in thd format, the total_num_tokens can vary across different micro-batch. The previous int type couldn't allow value change when replay graph.

[timmoon10]

LGTM, pending CI.

[timmoon10]

/te-ci L1

[timmoon10]

/te-ci L1