Fused gemm + comm for CP A2A on Blackwell
Maintainers usually reply within 2 days
@cyanguwa is already working on this.
Since Oct 9, 2026.
Assessment
This issue has not been assessed yet.
Description
To make CP A2A more efficient, we can fuse gemms and the communication.
Affected module:
- MHA: QKV Linear + DPA (CP/non-CP) + Proj Linear
Right now: - QKV Linear produces (B,S/cp_rank,H,D), DPA CP gathers S and splits along H to get (B,S,H/cp_rank,D), before the actual attention kernel, ...
After fusion: - QKV Linear produces (B,S,H/cp_rank,D), DPA CP just needs to check if the right shards have arrived, before calling the kernel, ...
We have some Blackwell kernels ready, and will need to look into CP no_load_balancing strategy and integrate those kernels.
- Dominant language
- Python
- Stars
- 3.6k
- Forks
- 851
- Avg merge
- 4d 15h
- Merged PRs (30d)
- 54
Getting set up
- No Dockerfile or Docker Compose file
- Has a pull request template
- Read the contributing 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.
More from NVIDIA/TransformerEngine
-
[Bug] group_quantize_fp8_blockwise: mbarrier invalidated before other threads finish waiting on itOpen
Difficulty 2/5 1-3 hours Newbie friendliness 82/100
NVIDIA/TransformerEngine#3647 ·
Maintainers usually reply within 2 days
-
[PyTorch] fp8_cs_quantize fake implementation returns a vector inverse scale instead of a scalarPossibly taken @sanjana658 claimed this 3 days ago. Open
Difficulty 2/5 1-3 hours Newbie friendliness 82/100
NVIDIA/TransformerEngine#3636 · 2 comments ·
Maintainers usually reply within 2 days
-
[Bug] Backend selection picks FA3 for training with head_dim_qk=192 / v_head_dim=128, but FA3 backward cannot run itMay be free again @modelpath-dev claimed this 37 days ago, and no pull request is open. Openattention
Difficulty 2/5 1-3 hours Newbie friendliness 85/100
NVIDIA/TransformerEngine#3481 · 4 comments ·
Maintainers usually reply within 2 days
-
bug
Difficulty 2/5 1-3 hours Newbie friendliness 68/100
NVIDIA/TransformerEngine#2189 · 7 comments · 5 reactions ·
Maintainers usually reply within 2 days
-
Difficulty 4/5 3-5 days Newbie friendliness 48/100
NVIDIA/TransformerEngine#3668 ·
Maintainers usually reply within 2 days
All issues in NVIDIA/TransformerEngine
Similar issues
-
bug
Difficulty 2/5 1-3 hours Newbie friendliness 76/100
lugassawan/swe-workbench#799 ·
Maintainers usually reply within 1 day
-
Difficulty 2/5 1-3 hours Newbie friendliness 74/100
vllm-project/vllm#61131 ·
Maintainers usually reply within 1 day
-
Non-string form field names raise AttributeError when files are attachedPossibly taken @nightcityblade claimed this today. Open
Difficulty 2/5 1-3 hours Newbie friendliness 82/100
-
Difficulty 2/5 1-3 hours Newbie friendliness 78/100
kubernetes-client/python#2735 ·
Maintainers usually reply within 1 day
-
desktop help wanted test
Difficulty 2/5 1-3 hours Newbie friendliness 78/100
Maintainers usually reply within 1 day