Skip to content

WGMMA+SP Kernel & AGU Serialization Fix - #383

Merged
tinebp merged 1 commit into
masterfrom
feature_wg2k
Jul 30, 2026
Merged

WGMMA+SP Kernel & AGU Serialization Fix#383
tinebp merged 1 commit into
masterfrom
feature_wg2k

Conversation

@NikhilRout

@NikhilRout NikhilRout commented Jul 25, 2026

Copy link
Copy Markdown
Collaborator

1. Sparse WGMMA kernel

  • old kernel redundantly loaded every warp’s compressed A and metadata through SMEM causing a lot of GMEM traffic
  • now each warp only loads its own compressed A tile instead of cooperatively serializing loads across warps
  • sp metadata is now directly loaded from GMEM -> TCU SRAM instead of old GMEM -> SMEM -> TCU SRAM
  • now uses tilek/2 (A) tilek (B) strided descriptor layouts matching the TCU buffer

2. TCU metadata AGU serialization bug

  • old AGU serialized all TCU blocks execution behind all warps' metadata memory round-trip latency. this prevented any memory-latency overlap
  • changed to one in-flight TCU_LD transaction per TCU block with round-robin issue/commit and tagged response routing (just pipelining metadata load/per TCU)

Perf improvements

1.5x speedup on wg+sp+dxa
1.9x speedup on wg+sp

@NikhilRout
NikhilRout marked this pull request as draft July 25, 2026 13:34
@NikhilRout
NikhilRout marked this pull request as ready for review July 28, 2026 13:54
@NikhilRout NikhilRout changed the title WGMMA+SP Kernel + AGU Serialization Fix WGMMA+SP Kernel AGU Serialization Fix Jul 28, 2026
@NikhilRout NikhilRout changed the title WGMMA+SP Kernel AGU Serialization Fix WGMMA+SP Kernel & AGU Serialization Fix Jul 28, 2026
@NikhilRout
NikhilRout requested a review from tinebp July 28, 2026 17:37
@tinebp
tinebp merged commit d76b7f2 into master Jul 30, 2026
211 of 212 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants