NVIDIA / NVIDIA/TransformerEngine
_flash_attn_forward() missing `v_descale` when enabling context parallel
@cyanguwa is already working on this.
Since Dec 4, 2025.
- Dominant language
- Python
- Stars
- 3.5k
- Forks
- 831
- Avg merge
- 3d 11h
- Merged PRs (30d)
- 65
Description
Describe the bug
When enabling context parallel and using FlashAttention3, the flash_attn_forward call in transformer_engine.pytorch.attention.dot_product_attention.context_parallel.py missed the v_descale.
Steps/Code to reproduce bug
Enabling context parallel and using FlashAttn3 at the same time.
Expected behavior
The actual reason should be the get_fa_args() in context_parallel.py mismatching flash_attn_interface._flash_attn_fwd's args in FA3. It requires a seqlens_rotary, while the get_fa_args() does not include this argument. Add an additonal None to the returns should fix it.
Environment details
- PyTorch 2.7.0
- Python 3.12.3
- Transformer Engine 2.4.0
- CUDA 12.8
Device details
- NVIDIA H200
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Assessment
This issue has not been assessed yet.