NVIDIA / NVIDIA/Megatron-LM

TE CUDA graphs do not support THD packed attention because packed_seq_params is not graph-safe

Open
#5,619 4 comments 0 reactions 1 assignee Claimed by @HaochenYuan View on GitHub
bug community-request module: transformer engine waiting-on-maintainers
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.