Support PackedSeqParams in TE CUDA graphs

[SakshamKapoor2911]

• I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Adds a TE CUDA graph adapter for THD PackedSeqParams by flattening graph-dynamic tensor fields into keyword graph inputs, storing non-tensor metadata on the captured TransformerLayer, and rebuilding PackedSeqParams inside the layer capture callable.

Issue tracking

Linked issue: Related to #5619

Contribution process

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests (local GPU functional smoke; no always-on in-tree functional test because this path requires a TE/CUDA runtime)
• I have added proper typing to my code Typing guidelines
• I have added relevant documentation (internal helper docstrings and PR contract; no user-facing docs were applicable)
• I have run the autoformatter.sh on my PR

Details

This PR implements the dynamic/static contract requested in #5619:

• Dynamic TE graph inputs: cu_seqlens_q, cu_seqlens_kv, cu_seqlens_q_padded, and cu_seqlens_kv_padded.
• Static metadata: qkv_format, max_seqlen_q, max_seqlen_kv, local_cp_size, and cp_group.
• Mamba-only fields total_tokens and seq_idx stay outside the TE attention CUDA graph boundary, matching TEDotProductAttention's existing filtering.
• Replay validates that static metadata and flattened tensor field presence match capture; changed metadata or field presence requires recapture.
• TECudaGraphHelper accepts an optional sample_packed_seq_params and injects its flattened tensor fields into sample kwargs.
• RL training/logprob paths now share a small helper to create the fixed-shape THD PackedSeqParams sample used for CUDA graph signature consistency.

Validation

• python3 -m py_compile megatron/core/packed_seq_params.py megatron/core/transformer/transformer_layer.py megatron/core/transformer/cuda_graphs.py megatron/rl/rl_utils.py megatron/rl/sequence_packing_utils.py megatron/training/training.py train_rl.py tests/unit_tests/transformer/test_packed_seq_params_cuda_graph.py tests/unit_tests/rl/test_rl_utils.py
• git diff --check over all touched files plus the new unit test file
• tools/autoformat.sh and CHECK_ONLY=true tools/autoformat.sh with repo-pinned formatter versions in a disposable CPU-only venv
• Focused pytest with a TensorBoard import shim in the lambda environment: 18 passed, 19 warnings
• Local GPU functional smoke in /tmp/te5619-venv on RTX 3090: Megatron TECudaGraphHelper created TE graphs, replayed a GPT layer with THD PackedSeqParams, and completed backward with fused RoPE (MEGATRON_TE_PACKED_SEQ_CUDAGRAPH_OK_FUSED_ROPE).
  • tests/unit_tests/transformer/test_packed_seq_params_cuda_graph.py
  • the 3 new TestRLUtils::test_get_rl_packed_seq_params_for_cuda_graph_* tests

Additional local evidence from a disposable TE environment:

• Transformer Engine 1.13.0 imported with flash-attn 2.6.3 in /tmp/te5619-venv.
• Eager TE THD DotProductAttention passed for bf16 and fp16.
• A raw torch.cuda.CUDAGraph THD attention wrapper replayed successfully after changing cu_seqlens, with max_abs_diff_vs_eager = 0.0.
• The Megatron helper smoke used validation-only compatibility shims for the older TE 1.13 stack: an identity make_weak_ref shim and dropping the newer unsupported retain_graph_in_backward kwarg. These should not be needed in NVIDIA's pinned/current TE runtime.

Remaining draft caveat: this is still draft until NVIDIA external-contributor validation/CI runs. In this disposable lambda stack, unfused THD RoPE is not graph-safe because it calls .tolist()/.item() on CUDA cu_seqlens during capture; the Megatron helper smoke therefore validates the graph path with fused RoPE enabled. Maintainers should still validate the full path in the supported pinned TE runtime before review-ready state.

[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.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.