Repository navigation
Conversation
On gfx1250 the RMSNorm kernels from #848 lose time in four places: - BF16 outputs use integer rounding, although gfx1250 has v_cvt_pk_bf16_f32. The integer rounding also turns NaN payloads >= 0x7FFF8000 into +-0. - The compiler cannot move a load above a store that may alias it. The plain and small-N kernels load gamma in pass 2, after the first output store, and the fused kernel interleaves its x and residual loads with the residual_out stores. Each such load waits one full memory latency. - At N = 32768 the bounds checks of full tiles raise the plain kernel to 258 VGPRs, which needs s_set_vgpr_msb bank switches. - For 16-bit rows with N <= 2048, every M <= 8192 takes one workgroup per row. At N = 512 only 64 of its 256 threads have work. On gfx1250 only: - has_hw_cvt_pk_bf16_f32() selects the hardware conversion in every RMSNorm builder (forward, small N, fused add, fused-add quant, backward). - The plain and small-N kernels load gamma in pass 1. The fused kernel issues every x, residual and gamma load before the first residual_out store. preload_weight_vec() keeps 16-bit gamma unconverted in registers. - The plain and fused kernels skip the bounds checks of full tiles. - The small-N module takes the multi-row kernel at every M when each lane handles at most three vectors of a row (N <= 768), and for M > N above that. The LLVM IR for gfx942 and gfx950 does not change. test_rmsnorm_bf16_output_keeps_nan checks that NaN payloads in FP32 weights stay NaN in BF16 outputs on targets with the hardware conversion. Signed-off-by: Jin Pan <jin.pan@amd.com>
986b766 to
0f0ddbc
Compare
jhinpan
changed the title
gfx1250: RMSNorm 4096×1536 and 4096×2048 (BF16) are slower than mainI benchmarked this PR against its base (main @ 1bd7ebd) on gfx1250 (MI450-class, 256 CUs). Most rows of the table reproduce, for example RMSNorm 1×32768 at 3.4–3.9x and fused add 4096×32768 at 2.1x. The two small-N rows listed at about 2x run slower than main here:
Cause: the new small-N kernel choice. On main, M = 4096 ≤
Why the multi-row kernel is slow here. In the final ISA it serializes its x loads:
The compiler is the pinned LLVM (e2a39f504), so the ISA should match yours. The difference is probably memory latency. Here, latency-bound rows run 1.4–2.1x slower than in your table (RMSNorm 1×32768 on main: 14.3–15.8 µs vs 8.61 µs), while bandwidth-bound rows match (32768×8192 on main: 81.8–84.4 µs vs 86.35 µs). Serialized loads cost more as latency grows, so the M > N rule that wins on your machine loses on this one. Possible fixes, not tested: keep one workgroup per row for N > 768 at M ≤ 8192 on gfx1250, as main does, or make the multi-row kernel issue all its x loads before the first wait. It may also be worth checking the rule on a second gfx1250 machine. Method and a minimal check
Run in each checkout: import torch
from torch.profiler import ProfilerActivity, profile
from kernels.norm.rmsnorm_kernel import rmsnorm_fwd
for M, N in [(4096, 1536), (4096, 2048)]:
x = torch.randn(M, N, device="cuda", dtype=torch.bfloat16)
w = torch.rand(N, device="cuda", dtype=torch.bfloat16)
for _ in range(20):
rmsnorm_fwd(x, w)
torch.cuda.synchronize()
with profile(activities=[ProfilerActivity.CUDA]) as prof:
for _ in range(100):
rmsnorm_fwd(x, w)
torch.cuda.synchronize()
kernels = [e for e in prof.profiler.kineto_results.events() if str(e.device_type()) == "DeviceType.CUDA"]
us = sorted(e.duration_ns() / 1e3 for e in kernels)
print(f"{M}x{N}: {kernels[0].name()} median {us[len(us) // 2]:.2f} us")On this machine it prints |
On gfx1250 the scheduler sank the gamma loads, and the x load of a partial last tile, below the first wait of the plain kernel, so each row waited two or three memory latencies. Issue every x and gamma load before a scheduling barrier. The plain and fused kernels now reduce the partial sums after one barrier, with every wave reading them; the fused kernel no longer reduces a dummy zero; and the mean uses the reciprocal of N instead of a correctly rounded division. The LLVM IR for gfx942 and gfx950 does not change. Signed-off-by: Jin Pan <jin.pan@amd.com>
…N kernel From four 128-bit vectors per lane on (N > 768 for 16-bit rows), the multi-row small-N kernel converted each x vector to FP32 before the next load issued, so a lane waited one memory latency per vector of its row. At N = 1536 and 2048 one load issued before the first wait, instead of all 12 and 16 x and gamma loads. Issue the loads first behind a scheduling barrier, as the one-row kernel does. With three or fewer vectors the loads already issue together and the barrier only costs 3-5%, so those rows keep the old code. gfx1250, against 824a384: 4096 x 1536 1.24x, 4096 x 2048 1.24x, M = N + 1 1.41-1.46x, M = 8193 to 65536 1.15-1.60x; the 69 rows with N >= 1000 1.106x geomean (min 0.975). gfx942 and gfx950 IR is unchanged. Signed-off-by: Jin Pan <jin.pan@amd.com>
|
Your ISA reading was right. bb146e1 makes the multi-row small-N kernel issue every x and gamma load before the first wait when a lane holds more than three vectors (N > 768 for 16-bit rows). Before the first On our gfx1250, with 824a384 and bb146e1 in one process plus an A/A copy (hot device time, median of 25 interleaved samples):
M = 8192 to 65536 gains 1.15x to 1.60x at N = 1536 and 2048, so N = 1536 is now ahead of main there too (65536 × 1536 on main: 34.08 µs). Our machine has lower memory latency than yours, so could you rerun your two rows and the M = N / N + 1 pairs on it? If 4096 × 1536 or 4096 × 2048 still lose to main there, the fallback is your first suggestion: one workgroup per row for N > 768 at M ≤ 8192. |
This branch has not been deployed
2 participants
Summary
#848 vectorized RMSNorm for partial tiles, small N, and
N % 8 != 0, and tuned it for gfx942 and gfx950. On gfx1250 its kernels still lose time in five places. This PR fixes them on gfx1250 only. The LLVM IR for gfx942 and gfx950 does not change. On gfx1250 the geomean speedup over main is 1.43x for RMSNorm (159 rows) and 1.57x for fused-add RMSNorm (40 rows).Motivation
On gfx1250, main @ 1bd7ebd:
v_cvt_pk_bf16_f32. The integer rounding also turns NaN payloads >= 0x7FFF8000 into ±0.residual_outstores. Loads written in pass 1 are not enough: to save VGPRs, the scheduler sinks the gamma loads, and the x load of a partial last tile, below the first wait. Example: RMSNorm 1 x 32768 takes 8.29 µs.s_set_vgpr_msbbank switches.FWD_MULTI_ROW_THRESHOLD) takes one 256-thread workgroup per row. At N = 512 only 64 of the 256 threads have work. Example: RMSNorm 4096 x 512 takes 4.81 µs, but 16384 x 512, which takes the multi-row kernel, takes 3.58 µs.Changes
All changes apply on gfx1250 only.
rmsnorm_common.has_hw_cvt_pk_bf16_f32()selectsv_cvt_pk_bf16_f32in every RMSNorm builder: forward, small N, fused add, fused-add quant, and backward.rocdl.sched_barrier(0)). The fused kernel issues every x, residual, and gamma load before the firstresidual_outstore.rmsnorm_common.preload_weight_vec()keeps 16-bit gamma unconverted in registers. FP32 gamma takes two 128-bit loads, as before.sum_sq * (1 / N).test_rmsnorm_bf16_output_keeps_nanputs NaN payloads (0xFFFFFFFF, 0x7FFFFFFF, 0x7FC00000) into FP32 weights. It checks the BF16 output of the plain, small-N (one-row and multi-row), partial-tile, and fused-add paths. It runs where the hardware conversion exists. On gfx1250 its 8 cases fail on main and pass with this PR.Performance (gfx1250, MI450-class, 256 CUs)
BF16 input, weight, and output. Main @ 1bd7ebd and this PR run in one process. Hot device time, median of 25 interleaved samples.
N % 8 != 0, and N from 8 to 32768.Small-N kernel choice. Hot µs, one workgroup per row / multi-row kernel, BF16, both kernels of this PR forced:
Over 131 sweep rows (N = 8 to 2048, M = 1 to 16384), the rule picks a kernel at most 2.4% slower than the faster one (N = 2048, M = 2048). In this sweep the multi-row kernel also skipped its full-tile checks; the PR keeps them there, which changes it by 2 VGPRs. The fixes for items 2 and 5 came after this sweep. They do not change the multi-row kernel, and they make the RMSNorm rows with N <= 2048 and M <= 256 3.9% faster in geomean. bb146e1 then makes the multi-row kernel faster for N > 768 (4096 x 1536: 3.33 to 2.69 µs; 4096 x 2048: 4.10 to 3.32 µs). The rule was not re-tuned after either change.
Comparison with AITER
AITER a6d2b564f (the installed package) on the same gfx1250 part, with the same inputs. Outputs are checked against an FP32 reference. Hot device time, median of 25 interleaved samples. AITER uses its C++ HIP kernel for 16-bit rows with N <= 8192, and opus otherwise (
rmsnorm.py). AITER time / this PR time, geomean (range); below 1.0x, AITER is faster:Testing
Dependencies
Breaking Changes
None.
Related PRs
N % 8 != 0; [Perf] Optimize rmsnorm/layernorm #848 covers them.v_cvt_pk_bf16_f32(gfx942), the 128-bit BF16 paths of [Perf] Optimize rmsnorm/layernorm #848 use the integer rounding, which turns FP32-weight NaN payloads >= 0x7FFF8000 into ±0. LayerNorm keeps its own gfx95-only check for the conversion.