[Perf] Optimize RMSNorm backward with staged reduction and vec8 I/O

[jhinpan]

Important

This PR is stacked on #854 and includes its head commit 067598ce.

#854 intentionally remains focused on benchmark/configuration coverage and compiled-cache validation. This PR keeps the resulting kernel optimization separate so its behavior, performance, and regressions can be reviewed independently. Once #854 lands, this branch will be rebased onto main.

Motivation

#854 made a substantial RMSNorm backward performance problem visible. Follow-up measurements at the direct-backward-wrapper level, which remove public-autograd host noise, confirmed the same issue against both PyTorch's fused RMSNorm backward and AIter.

The original FlyDSL backward launches one block per row and performs one global FP32 atomic add for every (row, dweight column) pair. At 4096x4096, that is 16,777,216 atomic updates, with 4096 blocks contending for every dweight element. It also loads x, dy, and gamma again in the second pass over each row.

This PR addresses those bottlenecks in a dedicated optimization change because it adds a new execution path, temporary workspace, runtime dispatch, and architecture-sensitive vector I/O. The existing atomic kernel remains the fallback.

Implementation

Hybrid atomic / two-stage backward

For large M, backward is split into two kernels:

1. A bounded persistent grid walks rows in a grid-stride loop, writes dx, and accumulates one FP32 dweight partial per program.
2. A finalizer reduces the logical [num_programs, N] partial workspace and writes dweight directly in the parameter dtype.

This removes the large-M global dweight atomics, the steady-state dweight clear, and the separate FP32-to-parameter-dtype conversion. Plain and fused-add RMSNorm share one builder; fused-add behavior is a compile-time specialization and retains the existing dx == dresidual aliasing contract.

512-thread vec8 fast path

The staged main kernel uses 512 threads. For FP16/BF16 with 4096 <= N <= 8192 and N % 8 == 0, it additionally uses:

• 128-bit vec8 loads and stores;
• one gamma load per program, retained across runtime rows;
• one x and dy load per row, retained across the two row passes;
• FP32 loop-carried dweight accumulation;
• direct vec8 dx stores.

Contiguous inputs with a non-16-byte-aligned storage offset are supported and covered by regression tests. FP32, smaller N, and N % 8 != 0 use the scalar staged path. BF16 vector stores use the gfx95 hardware conversion where available and the existing software round-to-nearest-even conversion elsewhere.

Ablations showed that both caches are material: removing the gamma cache regressed the measured shapes by about 2-7%, while reloading x/dy in the second pass regressed them by about 4-8%.

Hot-cache launch path

Profiling the remaining 512x2048 gap showed that FlyDSL's two device kernels total about 10.25 us, versus about 14.61 us for PyTorch's three kernels. The residual end-to-end gap came from entering a Python device context and constructing a Python Stream object on every cache hit.

A cached FlyDSL callable already accepts a raw cudaStream_t. On the common same-device path, the wrapper now passes the current raw stream directly. If the tensors belong to a non-current device, it retains the device guard and restores the caller's device afterward. Cache misses also keep the guard because compilation and code-object loading are device-context bound.

This preserves non-default-stream and multi-GPU semantics while removing hot-path Python work. A public current_stream(...).cuda_stream fallback is retained for PyTorch builds without the raw-stream getter.

Runtime configuration and ownership

The selector is shared by the plain and fused wrappers:

• M < 512 or N > 8192: existing atomic kernel;
• vec8-compatible FP16/BF16:
  • M < 2048: one program per CU;
  • M >= 2048: 1.5 programs per CU;
• scalar staged fallback:
  • M < 1024: one program per CU;
  • M >= 1024: two programs per CU.

The program count is capped by M. The vec/scalar decision has one semantic owner and is reused by both the builder and selector.

Production cache keys include path, N, dtype, program count, and device. M remains a runtime argument, so compatible row counts reuse one compiled callable. CU count is cached per device. Writable partial storage is allocated per invocation, never stored in a process-global cache, and both kernels launch on the caller's current stream.

Workspace size is num_programs * N * sizeof(float). Representative MI355X cases are:

shape	programs	workspace
512x2048 BF16	256	2 MiB
1024x8192 BF16	256	8 MiB
4096x4096 BF16	384	6 MiB

The largest selected scalar fallback near N=8192 uses about 16 MiB. N > 8192 continues to use the atomic path.

Performance

Environment and measurement protocol:

• AMD Instinct MI355X (gfx950, 256 CUs);
• PyTorch 2.9.1+rocm7.2.0.git7e1940d4;
• BF16 and hot compiled caches;
• 30 warmups, 200 timed iterations, five-run median;
• GPU-event timing;
• forward and rstd construction excluded;
• output/workspace allocation and every launched backward kernel included.

PyTorch uses torch.ops.aten._fused_rms_norm_backward; AIter commit 9127c94a uses aiter.ops.triton.normalization.rmsnorm._rmsnorm_backward. The tables below were collected with a manual direct-wrapper harness; the existing opt-in public-autograd benchmark in tests/kernels/test_rmsnorm.py is the in-tree reproducible end-to-end comparison. The direct-wrapper level isolates the backward implementations from variable Python/autograd-engine overhead.

Plain RMSNorm backward

shape	FlyDSL	PyTorch direct	AIter direct
512x2048	16.05 us	16.83 us	36.48 us
1024x8192	17.16 us	47.97 us	36.34 us
4096x4096	23.13 us	65.74 us	36.99 us

FlyDSL is 1.05x, 2.80x, and 2.84x faster than PyTorch across these shapes, and 1.60-2.27x faster than AIter. Because 512x2048 is host-sensitive, an independent alternating-process A/B on another MI355X was also run twice:

512x2048 path	before launch optimization	after	speedup
plain	19.11 us	15.01 us	1.27x
fused-add	19.87 us	15.75 us	1.26x

For context only, the committed #854 public-autograd benchmark measured about 159 us at 4096x4096, while the optimized direct wrapper here measures about 23-24 us. These use different host paths and are not a like-for-like speedup claim; the direct Torch/AIter rows above are the valid comparison for this PR.

Fused-add RMSNorm backward

AIter does not expose the equivalent full residual-gradient contract. The PyTorch equivalent below composes its direct fused RMSNorm backward with the residual-gradient addition.

shape	FlyDSL fused-add	PyTorch equivalent
512x2048	16.53 us	21.90 us
1024x8192	19.59 us	56.53 us
4096x4096	26.77 us	80.63 us

4096x4096 kernel profile

With the selected 384 programs:

path	persistent main	dweight finalizer	total
plain	18.87 us	4.51 us	23.39 us
fused-add	22.23 us	4.57 us	26.81 us

rocprofv3 reports 32 VGPRs for the main kernel at N=4096 and 52 at N=8192, with zero private scratch/spill in both cases. The N=4096 launch uses 48 SGPRs and 512 B LDS per block. The finalizer uses 16 VGPRs and also has no private scratch.

The explicit dweight partial workspace is global memory and is separate from rocprof's private-scratch metric.

Validation

Local validation against the final 512-thread/vec8 and cached-launch implementation:

• black, ruff check, py_compile, and git diff --check: passed;
• non-benchmark, non-large, single-GPU RMSNorm suite: 17 passed, 4 deselected;
• multi-GPU/device tests: 2 passed;
• raw stream ABI unit suite: 9 passed;
• representative gfx942 and gfx1100 BF16 vec8 kernels: compile-only checks passed.

Coverage includes:

• atomic, scalar staged, and vec8 staged dispatch;
• FP32, FP16, and BF16;
• N=5000 vec/program tails and N=4097 scalar fallback;
• contiguous tensors with non-16-byte-aligned storage offsets;
• runtime row tails and the N=8192 boundary;
• plain and fused-add backward;
• NaN-poisoned workspaces to prove complete overwrite;
• compiled callable reuse across runtime M;
• per-device cache separation and a cache hit on cuda:1 while cuda:0 remains current;
• concurrent non-default streams with per-call workspaces;
• numerical comparison with PyTorch in the committed suite; AIter direct-wrapper outputs were checked in the manual comparison harness.

Representative commands:

HIP_VISIBLE_DEVICES=7 CUDA_VISIBLE_DEVICES=7 \
python3 -m pytest tests/kernels/test_rmsnorm.py -q -s \
  -m 'not benchmark and not large_shape and not multi_gpu'

HIP_VISIBLE_DEVICES=6,7 CUDA_VISIBLE_DEVICES=6,7 \
python3 -m pytest \
  tests/kernels/test_rmsnorm.py::test_rmsnorm_multi_gpu \
  tests/kernels/test_rmsnorm.py::test_fused_add_rmsnorm_device_mismatch \
  -q -s

Scope

This PR deliberately does not add a process-global writable workspace cache, structural autotuning, a next-row software pipeline, or a staged path for N > 8192. A writable workspace cache measured at most about 0.13 us of possible benefit at the production 512x2048 configuration while introducing concurrency ownership. Lowering the vec8 threshold to N=2048 also showed no stable main-kernel improvement at the selected program count and raised main-kernel VGPR use from 28/30 to 58, so it is not included.

The hybrid selector targets the configurations where dweight contention and repeated row traffic dominate while keeping the existing atomic implementation as the fallback.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[coderfeli]

vec and no vec total different code path? only atom and some addr diff?

[jhinpan]

Thanks for pointing this out. You are right that the vector and scalar paths share the same RMSNorm math, and duplicating the complete row loop and both passes was unnecessary.

Addressed in fd6a3b88. The staged partial kernel now has one shared persistent row loop, two-pass RMSNorm math, block reduction, dweight accumulation, barrier, and partial writeback.

The remaining USE_VEC specializations are limited to the parts that actually differ:

• vec8 versus scalar copy atom, layout, indexing, and load/store;
• the vec8 gamma cache and source/dy retention across the two passes;
• vector lane reduction/expansion and the gfx95-aware BF16 vector conversion.

The cache difference is intentional: vec8 retains gamma across runtime rows and source/dy across the two passes, while the scalar fallback reloads them to preserve its FP32/arbitrary-N tail behavior and register profile.

I also checked generated code after the refactor. The BF16 scalar N=4097 ISA is byte-identical to the previous implementation. The BF16 vec8 N=4096 path has the same opcode mix and resource usage (58 VGPRs, zero spill); only equivalent register numbering/scheduling changed. The plain/fused scalar-tail, vec8, and unaligned-storage tests all pass.

[coderfeli]

use run_compile should already have cache?

[jhinpan]

Thanks — you are right that _run_compiled should own the CompiledFunction cache here.

Addressed in 629e0ee8:

• _BWD_CACHE and _FUSED_ADD_BWD_CACHE now retain the generated JitFunction launcher for each (path, N, dtype, num_programs, device) specialization;
• both cold and hot launches go through _run_compiled, so launcher._cf is now the single CompiledFunction cache;
• M remains a runtime argument and is intentionally excluded, so compatible row counts reuse the same launcher and compiled function.

The outer mapping still has a separate role: each builder call creates a new launcher, so without retaining that launcher the per-launcher ._cf cache would be lost. The device remains in the key because the loaded callable is device/context-bound.

I also kept builder creation and the first compile/launch inside the tensor-device guard, since the vec builder has an arch-dependent build-time conversion choice. Cache hits still pass the raw current cudaStream_t; cross-device calls enter the target-device guard and restore the caller device.

The cache test now verifies one launcher and one ._cf are reused across runtime M values for both plain and fused-add paths, and the multi-GPU test covers cold/hot cross-device restoration. The raw-stream unit suite and direct-wrapper A/B measurements also pass with no measurable launch-path regression.