Repository navigation
perf(nf12): specialize the RDNA4 fused small-M kernel on the token count - #3645
Merged
Merged
Conversation
Mirrors dphnAI/nf12#3. The m <= 15 fused kernel took m as a runtime argument, so acc[16][2] was indexed dynamically -- which keeps it out of VGPRs -- and the warp-partial plane was always sized for 16 tokens. Making the token count a template parameter fixes both and unrolls the token loop. At m=1 that raises occupancy enough to hurt: eight CTAs per CU and 256 concurrent packed streams thrash the 64 MiB LLC. Over-allocating the partial plane to 32 KiB throttles it to two CTAs per CU. Same-process A/B, old -> new: 4096x4096 m=1 28.6 -> 20.5 us (1.40x), m=8 32.6 -> 24.2 (1.35x); 2048x8192 m=1 19.4 -> 18.5 (1.05x), m=8 32.0 -> 23.9 (1.34x); 896x4864 m=8 15.4 -> 11.6 (1.33x).
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Mirrors the small-batch work from dphnAI/nf12#3.
What changed
m <= 15runs the fused kernel, which decodes each 16×64 super-atom once anddots it against every token in flight. It took
mas a runtime argument, soacc[16][2]was indexed dynamically — which keeps the array out of VGPRs — andthe warp-partial plane was always sized for 16 tokens even when one was in
flight. The token count is now a template parameter: both are compile-time
sized and the token loop fully unrolls.
At
m=1that raises occupancy enough to hurt: the kernel then runs eight CTAsper CU and 256 concurrent packed streams, which thrashes the 64 MiB LLC.
Over-allocating the partial plane to 32 KiB throttles it to two CTAs per CU.
Same-process A/B of the old and new fused kernel (µs, lower is better)
Result against dense BF16 (warm, R9700 gfx1201)
Every shape is now above parity at
m=1(1.20–1.71×);m=8is 0.94–1.36×.m=16still routes to the matrix-core kernel, which is 1.2–2.3× faster than thefused one there on every measured shape.
scripts/validate_nf12_ad.py --backend hipon this branch:4096×4096m=11.32×, m=8 0.96×, m=16 0.90×;
11008×4096m=1 3.75×.Accuracy
Unchanged. The kernel still decodes to exact BF16 and accumulates in FP32; the
only change is codegen shape. The validator reports 0 violations of
|err| <= 2e-3 + 2e-3*|ref|,rms_rel ≈ 1.66e-3(the BF16 rounding floor).Tests