NVIDIA / NVIDIA/TransformerEngine
Why use two streams for context parallel
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 3.5k
- Forks
- 831
- Avg merge
- 3d 11h
- Merged PRs (30d)
- 65
Description
Hi,
I see in https://github.com/NVIDIA/TransformerEngine/blob/29e8bfc99d803770ad82ae9351db63673bc34f69/transformer_engine/pytorch/attention.py#L624
that you used two cuda streams to resolve "wave quantization" in flash attention. Could you clarify what "wave quantization" means? I think flash attention just uses fp16/bf16
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.
Research direction
Start with transformer_engine/pytorch/attention.py around line 624 and trace the context-parallel code that creates the two CUDA streams. Read the linked implementation and any nearby comments to determine what “wave quantization” refers to and why the streams are used. Done means the terminology and stream rationale are clearly documented for users.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning, performance
- Issue type
- Documentation
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Needs clarification
- Newbie friendliness
- 35/100