Skip to content

perf(nf12): specialize the RDNA4 fused small-M kernel on the token count - #3645

Merged
Qubitium merged 1 commit into
ModelCloud:mainfrom
Qubitium:feat/nf12-fused-tmpl
Nov 8, 2026
Merged

Qubitium merged 1 commit into
ModelCloud:mainfrom
Qubitium:feat/nf12-fused-tmpl

Conversation

@Qubitium

@Qubitium Qubitium commented Nov 8, 2026

Copy link
Copy Markdown
Collaborator

Mirrors the small-batch work from dphnAI/nf12#3.

What changed

m <= 15 runs the fused kernel, which decodes each 16×64 super-atom once and
dots it against every token in flight. It took m as a runtime argument, so
acc[16][2] was indexed dynamically — which keeps the array out of VGPRs — and
the 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=1 that raises occupancy enough to hurt: the kernel then runs eight CTAs
per 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)

k × n m before after gain
4096×4096 1 28.57 20.45 1.40×
4096×4096 8 32.58 24.15 1.35×
2048×8192 1 19.44 18.49 1.05×
2048×8192 8 31.95 23.87 1.34×
896×4864 1 8.36 8.08 1.03×
896×4864 8 15.40 11.57 1.33×
4096×14336 8 161.67 149.18 1.08×
25600×5120 8 364.93 323.03 1.13×

Result against dense BF16 (warm, R9700 gfx1201)

k × n m=1 before → after m=8 before → after
4096×4096 0.84 → 1.20 0.67 → 0.94
2048×8192 1.33 → 1.38 0.80 → 0.97
896×4864 1.63 → 1.69 0.83 → 0.97
4096×28672 1.33 → 1.34 1.22 → 1.30
25600×5120 1.33 → 1.33 1.15 → 1.31
8192×8192 1.64 → 1.71 1.09 → 1.27

Every shape is now above parity at m=1 (1.20–1.71×); m=8 is 0.94–1.36×.
m=16 still routes to the matrix-core kernel, which is 1.2–2.3× faster than the
fused one there on every measured shape.

scripts/validate_nf12_ad.py --backend hip on this branch: 4096×4096 m=1
1.32×, m=8 0.96×, m=16 0.90×; 11008×4096 m=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

pytest tests/test_nf12_ad_codec.py tests/test_backend_naming.py \
       tests/test_quant_bits_validation.py tests/test_rocm_llc.py
# 235 passed, 1 skipped

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).
@Qubitium
Qubitium merged commit ebf6d04 into ModelCloud:main Nov 8, 2026
1 check 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.

1 participant