Problem
The SM100/SM103 blk64 BSA backward path always accumulates dK and dV through FP32 global workspaces and atomic adds, then launches separate conversion work, even when the bucketed reverse CSR has exactly one Q group. In that case each (batch, head, KV64) tile has one unique CTA owner, so cross-group dK/dV accumulation is unnecessary.
Proposed optimization
Mirror the existing SM100 blk128 policy in the blk64 path:
- specialize at compile time when
num_q_groups == 1;
- keep dK/dV in FP32 TMEM across the complete reverse row;
- scale dK, convert dK/dV to BF16, and write caller outputs directly;
- retain the current FP32 atomic workspace and postprocess path for multiple Q groups.
This does not require a public API change. Empty KV rows, partial block handling, BHSD/BSHD layouts, and the generic multi-group path must retain current behavior.
Prior evidence
An independent SM103 Q64/KV64 implementation measured the isolated direct-epilogue change at approximately 4.5% lower end-to-end latency, with two 128 MiB FP32 workspaces removed at B=1,H=32,S=8192,D=128 and about 31.6% fewer measured DRAM writes. These numbers are directional; the change will be benchmarked directly in this repository before making an upstream performance claim.
I plan to submit a focused PR with single- versus multi-group correctness coverage and short/sustained B30Z measurements.
Problem
The SM100/SM103 blk64 BSA backward path always accumulates dK and dV through FP32 global workspaces and atomic adds, then launches separate conversion work, even when the bucketed reverse CSR has exactly one Q group. In that case each
(batch, head, KV64)tile has one unique CTA owner, so cross-group dK/dV accumulation is unnecessary.Proposed optimization
Mirror the existing SM100 blk128 policy in the blk64 path:
num_q_groups == 1;This does not require a public API change. Empty KV rows, partial block handling, BHSD/BSHD layouts, and the generic multi-group path must retain current behavior.
Prior evidence
An independent SM103 Q64/KV64 implementation measured the isolated direct-epilogue change at approximately 4.5% lower end-to-end latency, with two 128 MiB FP32 workspaces removed at
B=1,H=32,S=8192,D=128and about 31.6% fewer measured DRAM writes. These numbers are directional; the change will be benchmarked directly in this repository before making an upstream performance claim.I plan to submit a focused PR with single- versus multi-group correctness coverage and short/sustained B30Z measurements.