TE CUDA graphs do not support THD packed attention because packed_seq_params is not graph-safe
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 3h
- Merged PRs (30d)
- 272
Description
**Describe the bug**
`cuda_graph_impl="transformer_engine"` currently cannot be used with THD packed attention because `packed_seq_params` is rejected by the TE CUDA graph path.
Tag @NVIDIA/mcore-oncall.
In the current code path, THD attention requires `PackedSeqParams` to flow through:
`TransformerBlock.forward -> TransformerLayer.forward -> Attention.forward -> TEDotProductAttention.forward`
However, `TransformerLayer._te_cuda_graph_replay()` asserts that `packed_seq_params is None`, and `GraphableMegatronModule._te_cuda_graph_replay()` only accepts Tensor/None inputs. `TECudaGraphHelper._get_sample_arguments()` also builds sample kwarg signatures assuming Tensor values (`shape`, `dtype`, `layout`).
This makes THD packed attention and TE CUDA graphs mutually exclusive, even though the actual dynamic fields needed by TE attention are mostly Tensor fields such as `cu_seqlens_q`, `cu_seqlens_kv`, and padded variants.
**Steps/Code to reproduce bug**
1. Use a Transformer Engine attention layer with `qkv_format="thd"`.
2. Pass `packed_seq_params=PackedSeqParams(qkv_format="thd", cu_seqlens_q=..., cu_seqlens_kv=..., cu_seqlens_q_padded=..., cu_seqlens_kv_padded=..., max_seqlen_q=..., max_seqlen_kv=...)`.
3. Enable `cuda_graph_impl="transformer_engine"` and capture/replay TE CUDA graphs.
4. The TE graph path rejects `packed_seq_params`.
Relevant files:
- `megatron/core/packed_seq_params.py`
- `megatron/core/transformer/module.py`
- `megatron/core/transformer/transformer_layer.py`
- `megatron/core/transformer/cuda_graphs.py`
- `megatron/core/extensions/transformer_engine.py`
**Expected behavior**
TE CUDA graphs should support THD packed attention when tensor shapes are static.
A possible design is:
- Flatten Tensor fields from `PackedSeqParams` into graph-safe Tensor kwargs.
- Treat non-Tensor metadata (`qkv_format`, `max_seqlen_*`, `cp_group`, `local_cp_size`) as static graph metadata.
- Reconstruct `PackedSeqParams` inside the capture/replay wrapper before calling the normal attention path.
- Require recapture or graph cache separation when non-Tensor metadata changes.
This would allow sequence packing / THD / CP use cases to benefit from TE CUDA graphs.
**Additional context**
Contributor guide
Assessment
This issue has not been assessed yet.