From fd031f4d9192834bbd4b4c58f6173137ff848055 Mon Sep 17 00:00:00 2001 From: Jin Pan <47354855+jhinpan@users.noreply.github.com> Date: Wed, 15 Jul 2026 00:46:58 +0000 Subject: [PATCH 1/5] [Perf] Add staged dweight reduction for RMSNorm backward --- kernels/norm/rmsnorm_bwd_kernel.py | 325 +++++++++++++++++++++++++++++ kernels/norm/rmsnorm_kernel.py | 113 ++++++++-- tests/kernels/test_rmsnorm.py | 236 ++++++++++++++++++--- 3 files changed, 631 insertions(+), 43 deletions(-) diff --git a/kernels/norm/rmsnorm_bwd_kernel.py b/kernels/norm/rmsnorm_bwd_kernel.py index b630da32e..cbc914d39 100644 --- a/kernels/norm/rmsnorm_bwd_kernel.py +++ b/kernels/norm/rmsnorm_bwd_kernel.py @@ -23,6 +23,14 @@ store_scalar, ) +# The staged path is selected by the Python wrapper only when M is large +# enough to keep one persistent program per CU busy. NUM_PROGRAMS is a build +# parameter (a small multiple of the device CU count); M remains a runtime +# argument so a compiled callable can be reused across row counts. +DWEIGHT_REDUCE_COLS = 64 +DWEIGHT_REDUCE_ROW_LANES = 4 +DWEIGHT_REDUCE_THREADS = DWEIGHT_REDUCE_COLS * DWEIGHT_REDUCE_ROW_LANES + def build_rmsnorm_bwd_module(N: int, dtype_str: str): """Fused RMSNorm backward: grid=(M,), one block per row. @@ -331,3 +339,320 @@ def launch_fused_add_rmsnorm_bwd( launcher.launch(grid=(m_in, 1, 1), block=(BLOCK_THREADS, 1, 1), stream=stream) return launch_fused_add_rmsnorm_bwd + + +def _build_rmsnorm_bwd_two_stage_module( + N: int, + dtype_str: str, + num_programs: int, + *, + fused_add: bool, +): + """Build the large-M persistent backward + dweight finalizer. + + The main kernel launches ``num_programs`` blocks. Each block walks rows in + a grid-stride loop, writes every assigned ``dx`` row, and accumulates one + FP32 dweight partial in registers. It then writes exactly one ``N``-wide + partial row. A second kernel reduces the ``num_programs`` partial rows and + stores dweight directly in the parameter dtype. + + This replaces M competing atomic adds per dweight element with a bounded + ``num_programs x N`` workspace and a deterministic final reduction. The + caller owns the workspace for the duration of one invocation; the builder + never caches writable storage. + """ + if num_programs <= 0: + raise ValueError(f"num_programs must be positive, got {num_programs}") + + RED_SLOTS = max(1, (BLOCK_THREADS + WARP_SIZE - 1) // WARP_SIZE) + NUM_COL_TILES = (N + BLOCK_THREADS - 1) // BLOCK_THREADS + elem_bits = 32 if dtype_str == "f32" else 16 + SharedStorage = make_single_reduction_storage(RED_SLOTS) + DWeightReduceStorage = make_single_reduction_storage(DWEIGHT_REDUCE_THREADS) + + @flyc.kernel + def rmsnorm_bwd_partial_kernel( + Source: fx.Tensor, + Gamma: fx.Tensor, + DY: fx.Tensor, + DResidualOut: fx.Tensor, + Rstd: fx.Tensor, + DX: fx.Tensor, + DWeightPartial: fx.Tensor, + MIn: fx.Int32, + ): + bid = fx.block_idx.x + tid = fx.thread_idx.x + + elem_dtype = dtype_to_elem_type(dtype_str) + fm_fast = arith.FastMathFlags.fast + n_float = float(N) + c_zero_f = fx.Float32(0.0) + + lds = fx.SharedAllocator().allocate(SharedStorage).peek() + s_red = lds.s_red.view(fx.make_layout(RED_SLOTS, 1)) + + def wave_reduce_add(x): + w = x + for _sh_exp in range_constexpr(int(math.log2(WARP_SIZE))): + off = WARP_SIZE // (2 << _sh_exp) + peer = w.shuffle_xor(off, WARP_SIZE) + w = w.addf(peer, fastmath=fm_fast) + return w + + def block_reduce_add(val): + if const_expr(RED_SLOTS == 1): + return wave_reduce_add(val) + lane = tid % WARP_SIZE + wave = tid // WARP_SIZE + w = wave_reduce_add(val) + if lane == 0: + fx.memref_store(w, s_red, wave) + gpu.barrier() + if wave == 0: + in_range = lane < RED_SLOTS + lane_safe = in_range.select(lane, 0) + v = fx.memref_load(s_red, lane_safe) + ww = in_range.select(v, c_zero_f) + ww = wave_reduce_add(ww) + if lane == 0: + fx.memref_store(ww, s_red, 0) + gpu.barrier() + return fx.memref_load(s_red, 0) + + Source_buf = fx.rocdl.make_buffer_tensor(Source) + Gamma_buf = fx.rocdl.make_buffer_tensor(Gamma) + DY_buf = fx.rocdl.make_buffer_tensor(DY) + DResidualOut_buf = fx.rocdl.make_buffer_tensor(DResidualOut) + Rstd_buf = fx.rocdl.make_buffer_tensor(Rstd) + DX_buf = fx.rocdl.make_buffer_tensor(DX) + DWeightPartial_buf = fx.rocdl.make_buffer_tensor(DWeightPartial) + + gamma_div = fx.logical_divide(Gamma_buf, fx.make_layout(1, 1)) + rstd_div = fx.logical_divide(Rstd_buf, fx.make_layout(1, 1)) + partial_div = fx.logical_divide(DWeightPartial_buf, fx.make_layout(1, 1)) + + copy_atom_s = fx.make_copy_atom( + fx.rocdl.BufferCopy16b() if elem_bits <= 16 else fx.rocdl.BufferCopy32b(), + elem_bits, + ) + copy_atom_f32 = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), 32) + + # One scalar per compile-time column tile. Carrying a single Vector + # through the runtime row loop keeps the accumulator in VGPRs and gives + # scf.for one well-defined loop-carried value. + dweight_partial = fx.Vector.filled(NUM_COL_TILES, 0.0, fx.Float32) + + for row in range(fx.Int32(bid), MIn, num_programs): + row_source = fx.slice(Source_buf, (row, None)) + row_dy = fx.slice(DY_buf, (row, None)) + row_dx = fx.slice(DX_buf, (row, None)) + if const_expr(fused_add): + row_dres_out = fx.slice(DResidualOut_buf, (row, None)) + + source_div = fx.logical_divide(row_source, fx.make_layout(1, 1)) + dy_div = fx.logical_divide(row_dy, fx.make_layout(1, 1)) + dx_div = fx.logical_divide(row_dx, fx.make_layout(1, 1)) + if const_expr(fused_add): + dres_out_div = fx.logical_divide(row_dres_out, fx.make_layout(1, 1)) + + rstd = load_scalar(copy_atom_f32, fx.Float32, rstd_div, row) + + thread_acc = c_zero_f + for base in range_constexpr(0, N, BLOCK_THREADS): + idx = tid + base + is_valid = idx < N + idx_safe = is_valid.select(idx, 0) + source_e = load_scalar(copy_atom_s, elem_dtype, source_div, idx_safe) + dy_e = load_scalar(copy_atom_s, elem_dtype, dy_div, idx_safe) + g_e = load_scalar(copy_atom_s, elem_dtype, gamma_div, idx_safe) + source = source_e if dtype_str == "f32" else source_e.to(fx.Float32) + dy = dy_e if dtype_str == "f32" else dy_e.to(fx.Float32) + g = g_e if dtype_str == "f32" else g_e.to(fx.Float32) + source_hat = source * rstd + wdy = dy * g + thread_acc = thread_acc + is_valid.select(source_hat * wdy, c_zero_f) + + sum_prod = block_reduce_add(thread_acc) + c1 = sum_prod / n_float + + row_dweight = [] + for base in range_constexpr(0, N, BLOCK_THREADS): + idx = tid + base + is_valid = idx < N + idx_safe = is_valid.select(idx, 0) + source_e = load_scalar(copy_atom_s, elem_dtype, source_div, idx_safe) + dy_e = load_scalar(copy_atom_s, elem_dtype, dy_div, idx_safe) + g_e = load_scalar(copy_atom_s, elem_dtype, gamma_div, idx_safe) + source = source_e if dtype_str == "f32" else source_e.to(fx.Float32) + dy = dy_e if dtype_str == "f32" else dy_e.to(fx.Float32) + g = g_e if dtype_str == "f32" else g_e.to(fx.Float32) + source_hat = source * rstd + wdy = dy * g + dx = (wdy - source_hat * c1) * rstd + if const_expr(fused_add): + dres_e = load_scalar(copy_atom_s, elem_dtype, dres_out_div, idx_safe) + dres = dres_e if dtype_str == "f32" else dres_e.to(fx.Float32) + dx = dx + dres + + if idx < N: + dx_e = dx if dtype_str == "f32" else dx.to(elem_dtype) + store_scalar(copy_atom_s, elem_dtype, elem_dtype, dx_div, idx, dx_e) + + row_dweight.append(is_valid.select(dy * source_hat, c_zero_f)) + + dweight_partial = dweight_partial + fx.Vector.from_elements(row_dweight, fx.Float32) + + # c1 lives in shared storage. All waves must finish consuming it + # before a faster wave enters the next row and overwrites s_red. + gpu.barrier() + + # Every program overwrites its entire partial row, including programs + # with no assigned rows (their accumulator remains zero). + for tile_i in range_constexpr(NUM_COL_TILES): + idx = tid + tile_i * BLOCK_THREADS + if idx < N: + partial_idx = bid * N + idx + store_scalar( + copy_atom_f32, + fx.Float32, + fx.Float32, + partial_div, + partial_idx, + dweight_partial[tile_i], + ) + + @flyc.kernel + def rmsnorm_bwd_dweight_reduce_kernel( + DWeightPartial: fx.Tensor, + DWeight: fx.Tensor, + ): + bid = fx.block_idx.x + tid = fx.thread_idx.x + col_lane = tid % DWEIGHT_REDUCE_COLS + partial_lane = tid // DWEIGHT_REDUCE_COLS + col = bid * DWEIGHT_REDUCE_COLS + col_lane + is_valid = col < N + col_safe = is_valid.select(col, 0) + + elem_dtype = dtype_to_elem_type(dtype_str) + DWeightPartial_buf = fx.rocdl.make_buffer_tensor(DWeightPartial) + DWeight_buf = fx.rocdl.make_buffer_tensor(DWeight) + partial_div = fx.logical_divide(DWeightPartial_buf, fx.make_layout(1, 1)) + dweight_div = fx.logical_divide(DWeight_buf, fx.make_layout(1, 1)) + copy_atom_f32 = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), 32) + copy_atom_s = fx.make_copy_atom( + fx.rocdl.BufferCopy16b() if elem_bits <= 16 else fx.rocdl.BufferCopy32b(), + elem_bits, + ) + + lds = fx.SharedAllocator().allocate(DWeightReduceStorage).peek() + s_partial = lds.s_red.view(fx.make_layout(DWEIGHT_REDUCE_THREADS, 1)) + + acc = fx.Float32(0.0) + for partial_base in range(0, num_programs, DWEIGHT_REDUCE_ROW_LANES): + partial_row = partial_base + partial_lane + partial_valid = partial_row < num_programs + partial_row_safe = partial_valid.select(partial_row, 0) + partial_idx = partial_row_safe * N + col_safe + value = load_scalar(copy_atom_f32, fx.Float32, partial_div, partial_idx) + acc = acc + partial_valid.select(value, fx.Float32(0.0)) + fx.memref_store(acc, s_partial, tid) + gpu.barrier() + + if partial_lane == 0: + if col < N: + total = fx.Float32(0.0) + for lane in range_constexpr(DWEIGHT_REDUCE_ROW_LANES): + total = total + fx.memref_load(s_partial, lane * DWEIGHT_REDUCE_COLS + col_lane) + out = total if dtype_str == "f32" else total.to(elem_dtype) + store_scalar(copy_atom_s, elem_dtype, elem_dtype, dweight_div, col, out) + + reduce_grid = (N + DWEIGHT_REDUCE_COLS - 1) // DWEIGHT_REDUCE_COLS + + if fused_add: + + @flyc.jit + def launch_fused_add_rmsnorm_bwd_two_stage( + Added: fx.Tensor, + Gamma: fx.Tensor, + DY: fx.Tensor, + DResidualOut: fx.Tensor, + Rstd: fx.Tensor, + DX: fx.Tensor, + DWeight: fx.Tensor, + DWeightPartial: fx.Tensor, + m_in: fx.Int32, + stream: fx.Stream = fx.Stream(None), + ): + partial_launcher = rmsnorm_bwd_partial_kernel( + Added, + Gamma, + DY, + DResidualOut, + Rstd, + DX, + DWeightPartial, + m_in, + ) + partial_launcher.launch( + grid=(num_programs, 1, 1), + block=(BLOCK_THREADS, 1, 1), + stream=stream, + ) + reduce_launcher = rmsnorm_bwd_dweight_reduce_kernel(DWeightPartial, DWeight) + reduce_launcher.launch( + grid=(reduce_grid, 1, 1), + block=(DWEIGHT_REDUCE_THREADS, 1, 1), + stream=stream, + ) + + return launch_fused_add_rmsnorm_bwd_two_stage + + @flyc.jit + def launch_rmsnorm_bwd_two_stage( + Input: fx.Tensor, + Gamma: fx.Tensor, + DY: fx.Tensor, + Rstd: fx.Tensor, + DX: fx.Tensor, + DWeight: fx.Tensor, + DWeightPartial: fx.Tensor, + m_in: fx.Int32, + stream: fx.Stream = fx.Stream(None), + ): + # DResidualOut is compile-time dead in the plain variant; DY fills the + # otherwise-unused tensor slot so both variants share one main kernel. + partial_launcher = rmsnorm_bwd_partial_kernel( + Input, + Gamma, + DY, + DY, + Rstd, + DX, + DWeightPartial, + m_in, + ) + partial_launcher.launch( + grid=(num_programs, 1, 1), + block=(BLOCK_THREADS, 1, 1), + stream=stream, + ) + reduce_launcher = rmsnorm_bwd_dweight_reduce_kernel(DWeightPartial, DWeight) + reduce_launcher.launch( + grid=(reduce_grid, 1, 1), + block=(DWEIGHT_REDUCE_THREADS, 1, 1), + stream=stream, + ) + + return launch_rmsnorm_bwd_two_stage + + +def build_rmsnorm_bwd_two_stage_module(N: int, dtype_str: str, num_programs: int): + """Large-M plain RMSNorm backward without global dweight atomics.""" + return _build_rmsnorm_bwd_two_stage_module(N, dtype_str, num_programs, fused_add=False) + + +def build_fused_add_rmsnorm_bwd_two_stage_module(N: int, dtype_str: str, num_programs: int): + """Large-M fused-add RMSNorm backward without global dweight atomics.""" + return _build_rmsnorm_bwd_two_stage_module(N, dtype_str, num_programs, fused_add=True) diff --git a/kernels/norm/rmsnorm_kernel.py b/kernels/norm/rmsnorm_kernel.py index 424f06468..8fe932a95 100644 --- a/kernels/norm/rmsnorm_kernel.py +++ b/kernels/norm/rmsnorm_kernel.py @@ -24,7 +24,9 @@ # here so existing importers (tests, callers) keep working unchanged. from kernels.norm.rmsnorm_bwd_kernel import ( # noqa: E402,F401 build_fused_add_rmsnorm_bwd_module, + build_fused_add_rmsnorm_bwd_two_stage_module, build_rmsnorm_bwd_module, + build_rmsnorm_bwd_two_stage_module, ) from kernels.norm.rmsnorm_common import ( BLOCK_THREADS, @@ -1458,6 +1460,33 @@ def build_fused_add_rmsnorm_smoothquant_module( _FWD_CACHE: dict = {} _BWD_CACHE: dict = {} + # The staged path needs enough rows to amortize its second launch. Wide + # and FP32 rows carry more live accumulators, so cap their program count to + # avoid excess register pressure and scratch. Keep this selector as the + # single source of truth shared by plain and fused-add wrappers. + _BWD_TWO_STAGE_MIN_ROWS = 512 + _BWD_TWO_STAGE_MAX_N = 8192 + _BWD_CU_COUNT_CACHE: dict = {} + + def _get_rmsnorm_bwd_cu_count(device): + num_cus = _BWD_CU_COUNT_CACHE.get(device) + if num_cus is None: + num_cus = torch.cuda.get_device_properties(device).multi_processor_count + _BWD_CU_COUNT_CACHE[device] = num_cus + return num_cus + + def _select_rmsnorm_bwd_config(M, N, dtype_str, device): + if M >= _BWD_TWO_STAGE_MIN_ROWS and N <= _BWD_TWO_STAGE_MAX_N: + if M < 1024: + programs_per_cu = 1 + elif M < 2048 or dtype_str == "f32" or N <= 2048 or N > 4096: + programs_per_cu = 2 + else: + programs_per_cu = 4 + num_programs = min(M, _get_rmsnorm_bwd_cu_count(device) * programs_per_cu) + return ("two_stage", num_programs) + return ("atomic", None) + def _get_fwd_compiled(x, weight, out, rstd, M, N, dtype_str, store_rstd, eps, stream): key = (N, dtype_str, store_rstd, float(eps), x.device) entry = _FWD_CACHE.get(key) @@ -1510,20 +1539,44 @@ def rmsnorm_bwd(x, weight, dout, rstd, eps=EPS): M, N = x.shape dtype_str = _torch_dtype_to_str(x.dtype) dx = torch.empty_like(x) - dweight = torch.zeros((N,), device=x.device, dtype=torch.float32) - key = (N, dtype_str, x.device) # Bind compile + launch to the tensors' device (multi-GPU correctness). with torch.cuda.device(x.device): stream = torch.cuda.current_stream() + path, num_programs = _select_rmsnorm_bwd_config(M, N, dtype_str, x.device) + if path == "two_stage": + dweight = torch.empty_like(weight) + partial = torch.empty((num_programs * N,), device=x.device, dtype=torch.float32) + key = (path, N, dtype_str, num_programs, x.device) + compiled = _BWD_CACHE.get(key) + if compiled is None: + launch_fn = build_rmsnorm_bwd_two_stage_module(N, dtype_str, num_programs) + compiled = flyc.compile( + launch_fn, + x, + weight, + dout, + rstd, + dx, + dweight, + partial, + M, + stream, + ) + _BWD_CACHE[key] = compiled + else: + compiled(x, weight, dout, rstd, dx, dweight, partial, M, stream) + return dx, dweight + + dweight = torch.empty((N,), device=x.device, dtype=torch.float32) + key = (path, N, dtype_str, x.device) compiled = _BWD_CACHE.get(key) + dweight.zero_() if compiled is None: launch_fn = build_rmsnorm_bwd_module(N, dtype_str) - # flyc.compile executes the kernel once during tracing, which would - # accumulate into DWeight; zero it AFTER compiling. compiled = flyc.compile(launch_fn, x, weight, dout, rstd, dx, dweight, M, stream) _BWD_CACHE[key] = compiled - dweight.zero_() - compiled(x, weight, dout, rstd, dx, dweight, M, stream) + else: + compiled(x, weight, dout, rstd, dx, dweight, M, stream) return dx, dweight.to(weight.dtype) class RMSNormFunction(torch.autograd.Function): @@ -1637,19 +1690,55 @@ def fused_add_rmsnorm_bwd(added, weight, dout, rstd, dresidual_out=None, eps=EPS if dresidual_out is None: dresidual_out = torch.zeros_like(added) dx = torch.empty_like(added) - dweight = torch.zeros((N,), device=added.device, dtype=torch.float32) - key = (N, dtype_str, added.device) with torch.cuda.device(added.device): stream = torch.cuda.current_stream() + path, num_programs = _select_rmsnorm_bwd_config(M, N, dtype_str, added.device) + if path == "two_stage": + dweight = torch.empty_like(weight) + partial = torch.empty((num_programs * N,), device=added.device, dtype=torch.float32) + key = (path, N, dtype_str, num_programs, added.device) + compiled = _FUSED_ADD_BWD_CACHE.get(key) + if compiled is None: + launch_fn = build_fused_add_rmsnorm_bwd_two_stage_module(N, dtype_str, num_programs) + compiled = flyc.compile( + launch_fn, + added, + weight, + dout, + dresidual_out, + rstd, + dx, + dweight, + partial, + M, + stream, + ) + _FUSED_ADD_BWD_CACHE[key] = compiled + else: + compiled( + added, + weight, + dout, + dresidual_out, + rstd, + dx, + dweight, + partial, + M, + stream, + ) + return dx, dx, dweight + + dweight = torch.empty((N,), device=added.device, dtype=torch.float32) + key = (path, N, dtype_str, added.device) compiled = _FUSED_ADD_BWD_CACHE.get(key) + dweight.zero_() if compiled is None: launch_fn = build_fused_add_rmsnorm_bwd_module(N, dtype_str) - # flyc.compile executes the kernel once during tracing, which would - # accumulate into DWeight; zero it AFTER compiling. compiled = flyc.compile(launch_fn, added, weight, dout, dresidual_out, rstd, dx, dweight, M, stream) _FUSED_ADD_BWD_CACHE[key] = compiled - dweight.zero_() - compiled(added, weight, dout, dresidual_out, rstd, dx, dweight, M, stream) + else: + compiled(added, weight, dout, dresidual_out, rstd, dx, dweight, M, stream) # dx == dresidual by construction; return dx as both (aliased). return dx, dx, dweight.to(weight.dtype) diff --git a/tests/kernels/test_rmsnorm.py b/tests/kernels/test_rmsnorm.py index 4d9b424c0..23291a23a 100644 --- a/tests/kernels/test_rmsnorm.py +++ b/tests/kernels/test_rmsnorm.py @@ -29,13 +29,16 @@ # Imported after the torch guard: rmsnorm() is only defined when torch is present, # so importing it earlier makes a torch-less collection fail (ImportError) instead of skip. import flydsl.compiler as flyc # noqa: E402 +import kernels.norm.rmsnorm_kernel as rmsnorm_kernel_impl # noqa: E402 from kernels.common.tensor_shim import _run_compiled # noqa: E402 from kernels.norm.rmsnorm_kernel import ( # noqa: E402 build_fused_add_rmsnorm_bwd_module, + build_fused_add_rmsnorm_bwd_two_stage_module, build_fused_add_rmsnorm_dynamicquant_module, build_fused_add_rmsnorm_module, build_fused_add_rmsnorm_smoothquant_module, build_rmsnorm_bwd_module, + build_rmsnorm_bwd_two_stage_module, build_rmsnorm_dynamicquant_module, build_rmsnorm_module, build_rmsnorm_smoothquant_module, @@ -67,10 +70,12 @@ _RMSNORM_BACKWARD_CONFIGS = ( (64, 256, "f32"), # small-N path, f32 (16, 512, "bf16"), # small-N path, bf16 + (512, 2048, "f16"), # staged finalizer cast/store (4096, 4096, "bf16"), # model-relevant large shape (64, 2000, "f32"), # small-N path, unaligned (128, 4096, "f16"), # f16, large hidden size (64, 3000, "f32"), # generic scalar path, unaligned N > 2048 + (1537, 3001, "f32"), # two-stage row/program tail + unaligned hidden size ) _RMSNORM_AUTOGRAD_CONFIGS = ( @@ -78,6 +83,7 @@ (128, 4096, "bf16"), # large hidden size (128, 4096, "f16"), # large hidden size, f16 (128, 3000, "f32"), # generic scalar path, unaligned N > 2048 + (1024, 8192, "bf16"), # staged path at the supported hidden-size limit ) # Opt-in backward-only performance sweep. Repeated N=4096 entries exercise @@ -644,7 +650,6 @@ def run_bwd_test(M: int, N: int, dtype: str = "f32"): torch_dtype = _torch_dtype(dtype) try: fwd_fn = build_rmsnorm_module(N, dtype, store_rstd=True) - bwd_fn = build_rmsnorm_bwd_module(N, dtype) except Exception as e: print(f"[FAIL] Compile failed for bwd (M={M}, N={N}, dtype={dtype}): {type(e).__name__}: {e}") return False @@ -654,6 +659,16 @@ def run_bwd_test(M: int, N: int, dtype: str = "f32"): weight = torch.rand((N,), device="cuda", dtype=DTYPE_FP32).to(torch_dtype).contiguous() dy = torch.randn((M, N), device="cuda", dtype=DTYPE_FP32).to(torch_dtype).contiguous() + path, num_programs = rmsnorm_kernel_impl._select_rmsnorm_bwd_config(M, N, dtype, x.device) + try: + if path == "two_stage": + bwd_fn = build_rmsnorm_bwd_two_stage_module(N, dtype, num_programs) + else: + bwd_fn = build_rmsnorm_bwd_module(N, dtype) + except Exception as e: + print(f"[FAIL] Build failed for {path} bwd (M={M}, N={N}, dtype={dtype}): {type(e).__name__}: {e}") + return False + dx_ref, dw_ref, rstd_ref = _reference_rmsnorm_bwd(x, weight, dy) stream = torch.cuda.current_stream() @@ -667,9 +682,16 @@ def run_bwd_test(M: int, N: int, dtype: str = "f32"): # --- backward: dx + dweight --- dx = torch.empty((M, N), device="cuda", dtype=torch_dtype) - dweight = torch.empty((N,), device="cuda", dtype=DTYPE_FP32) - dweight.zero_() - _run_compiled(bwd_fn, x, weight, dy, rstd, dx, dweight, M, stream) + if path == "two_stage": + dweight = torch.empty((N,), device="cuda", dtype=torch_dtype) + partial = torch.full((num_programs * N,), float("nan"), device="cuda", dtype=DTYPE_FP32) + _run_compiled(bwd_fn, x, weight, dy, rstd, dx, dweight, partial, M, stream) + partial_ok = torch.isfinite(partial).all().item() + else: + dweight = torch.empty((N,), device="cuda", dtype=DTYPE_FP32) + dweight.zero_() + _run_compiled(bwd_fn, x, weight, dy, rstd, dx, dweight, M, stream) + partial_ok = True torch.cuda.synchronize() dx_err = (dx.to(DTYPE_FP32) - dx_ref).abs().max().item() @@ -681,23 +703,26 @@ def run_bwd_test(M: int, N: int, dtype: str = "f32"): dw_rtol = {"f32": 1e-4, "f16": 3e-2, "bf16": 1e-1}[dtype] dw_atol = {"f32": 1e-2, "f16": 1e-1, "bf16": 5e-1}[dtype] + print(f" path = {path}") print(f" rstd max abs err = {rstd_err:.3e} (atol={rstd_atol})") print(f" dx max abs err = {dx_err:.3e} (atol={dx_atol})") print(f" dweight |max| = {dw_mag:.3e}") dw_ok = True try: - torch.testing.assert_close(dweight, dw_ref, rtol=dw_rtol, atol=dw_atol) + torch.testing.assert_close(dweight.to(DTYPE_FP32), dw_ref, rtol=dw_rtol, atol=dw_atol) except AssertionError as e: dw_ok = False - dw_err = (dweight - dw_ref).abs().max().item() + dw_err = (dweight.to(DTYPE_FP32) - dw_ref).abs().max().item() print(f" dweight max abs err = {dw_err:.3e} (rtol={dw_rtol}, atol={dw_atol})") print(f" [dweight mismatch] {e}") else: - dw_err = (dweight - dw_ref).abs().max().item() + dw_err = (dweight.to(DTYPE_FP32) - dw_ref).abs().max().item() print(f" dweight max abs err = {dw_err:.3e} (rtol={dw_rtol}, atol={dw_atol})") - ok = rstd_err < rstd_atol and dx_err < dx_atol and dw_ok + ok = rstd_err < rstd_atol and dx_err < dx_atol and dw_ok and partial_ok + if path == "two_stage": + print(f" partial fully written= {partial_ok}") print(f" -> {'PASSED' if ok else 'FAILED'}") return ok @@ -737,7 +762,6 @@ def run_fused_add_bwd_test(M: int, N: int, dtype: str = "f32"): torch_dtype = _torch_dtype(dtype) try: fwd_fn = build_fused_add_rmsnorm_module(N, dtype, store_rstd=True) - bwd_fn = build_fused_add_rmsnorm_bwd_module(N, dtype) except Exception as e: print(f"[FAIL] Compile failed for fused_add bwd (M={M}, N={N}, dtype={dtype}): {type(e).__name__}: {e}") return False @@ -748,6 +772,18 @@ def run_fused_add_bwd_test(M: int, N: int, dtype: str = "f32"): weight = torch.rand((N,), device="cuda", dtype=DTYPE_FP32).to(torch_dtype).contiguous() dy = torch.randn((M, N), device="cuda", dtype=DTYPE_FP32).to(torch_dtype).contiguous() + path, num_programs = rmsnorm_kernel_impl._select_rmsnorm_bwd_config(M, N, dtype, x.device) + try: + if path == "two_stage": + bwd_fn = build_fused_add_rmsnorm_bwd_two_stage_module(N, dtype, num_programs) + else: + bwd_fn = build_fused_add_rmsnorm_bwd_module(N, dtype) + except Exception as e: + print( + f"[FAIL] Build failed for fused_add {path} bwd " f"(M={M}, N={N}, dtype={dtype}): {type(e).__name__}: {e}" + ) + return False + stream = torch.cuda.current_stream() # --- forward with store_rstd: residual_out + rstd from the kernel --- @@ -776,10 +812,29 @@ def run_fused_add_bwd_test(M: int, N: int, dtype: str = "f32"): # --- backward: dx (== dresidual by construction), dweight --- # The kernel writes dx only; dresidual is dx aliased in the wrapper. dx = torch.empty((M, N), device="cuda", dtype=torch_dtype) - dweight = torch.empty((N,), device="cuda", dtype=DTYPE_FP32) dres_out_arg = dres_out if dres_out is not None else torch.zeros((M, N), device="cuda", dtype=torch_dtype) - dweight.zero_() - _run_compiled(bwd_fn, residual_out, weight, dy, dres_out_arg, rstd, dx, dweight, M, stream) + if path == "two_stage": + dweight = torch.empty((N,), device="cuda", dtype=torch_dtype) + partial = torch.full((num_programs * N,), float("nan"), device="cuda", dtype=DTYPE_FP32) + _run_compiled( + bwd_fn, + residual_out, + weight, + dy, + dres_out_arg, + rstd, + dx, + dweight, + partial, + M, + stream, + ) + partial_ok = torch.isfinite(partial).all().item() + else: + dweight = torch.empty((N,), device="cuda", dtype=DTYPE_FP32) + dweight.zero_() + _run_compiled(bwd_fn, residual_out, weight, dy, dres_out_arg, rstd, dx, dweight, M, stream) + partial_ok = True if cached_bwd is None: cached_bwd = bwd_fn._cf else: @@ -791,23 +846,26 @@ def run_fused_add_bwd_test(M: int, N: int, dtype: str = "f32"): dres_err = (dx.to(DTYPE_FP32) - dres_ref).abs().max().item() print(f" [{case_name}]") + print(f" path = {path}") print(f" rstd max abs err = {rstd_err:.3e} (atol={rstd_atol})") print(f" dx max abs err = {dx_err:.3e} (atol={dx_atol})") print(f" dresidual max abs err= {dres_err:.3e} (atol={dx_atol})") dw_ok = True try: - torch.testing.assert_close(dweight, dw_ref, rtol=dw_rtol, atol=dw_atol) + torch.testing.assert_close(dweight.to(DTYPE_FP32), dw_ref, rtol=dw_rtol, atol=dw_atol) except AssertionError as e: dw_ok = False - dw_err = (dweight - dw_ref).abs().max().item() + dw_err = (dweight.to(DTYPE_FP32) - dw_ref).abs().max().item() print(f" dweight max abs err = {dw_err:.3e} (rtol={dw_rtol}, atol={dw_atol})") print(f" [dweight mismatch] {e}") else: - dw_err = (dweight - dw_ref).abs().max().item() + dw_err = (dweight.to(DTYPE_FP32) - dw_ref).abs().max().item() print(f" dweight max abs err = {dw_err:.3e} (rtol={dw_rtol}, atol={dw_atol})") - case_ok = rstd_err < rstd_atol and dx_err < dx_atol and dres_err < dx_atol and dw_ok + case_ok = rstd_err < rstd_atol and dx_err < dx_atol and dres_err < dx_atol and dw_ok and partial_ok + if path == "two_stage": + print(f" partial fully written = {partial_ok}") print(f" -> {'PASSED' if case_ok else 'FAILED'}") all_ok = all_ok and case_ok @@ -939,9 +997,7 @@ def run_torch(): def _print_rmsnorm_torch_perf_table(rows): print("\n" + "=" * 100) print("RMSNorm public-autograd backward perf (gpu us): FlyDSL vs PyTorch eager") - print( - "Forward graph construction is excluded; backward allocations, dweight zeroing, and dtype casts are included." - ) + print("Forward graph construction is excluded; backward allocations and all launched kernels are included.") print( f"Device: {torch.cuda.get_device_name()} | PyTorch: {torch.__version__} | " f"warmup={WARMUP_ITERS}, iters={BENCH_ITERS}, hot cache" @@ -1340,6 +1396,109 @@ def test_rmsnorm_eps_honored(): print(" -> PASSED") +def test_rmsnorm_bwd_dispatch_boundary(): + """The hybrid selector is the single authority for atomic vs two-stage.""" + device = torch.device("cuda", torch.cuda.current_device()) + + cases = ( + (511, 4096, "atomic"), + (512, 4096, "two_stage"), + (4096, 8192, "two_stage"), + (4096, 8193, "atomic"), # staged-kernel register-pressure guard + ) + for M, N, expected_path in cases: + path, num_programs = rmsnorm_kernel_impl._select_rmsnorm_bwd_config(M, N, "bf16", device) + assert path == expected_path, f"unexpected path for M={M}, N={N}: {path}" + if path == "two_stage": + assert num_programs > 0 + else: + assert num_programs is None + + +@pytest.mark.parametrize("fused_add", (False, True), ids=("plain", "fused_add")) +def test_rmsnorm_bwd_two_stage_cache_reuse_across_m(monkeypatch, fused_add): + """One staged compiled callable must serve multiple runtime row counts.""" + device = torch.device("cuda", torch.cuda.current_device()) + dtype = DTYPE_BF16 + N = 512 + first_m, second_m = 1024, 1025 + path, _ = rmsnorm_kernel_impl._select_rmsnorm_bwd_config(first_m, N, "bf16", device) + assert path == "two_stage" + + builder_name = "build_fused_add_rmsnorm_bwd_two_stage_module" if fused_add else "build_rmsnorm_bwd_two_stage_module" + cache = rmsnorm_kernel_impl._FUSED_ADD_BWD_CACHE if fused_add else rmsnorm_kernel_impl._BWD_CACHE + original_builder = getattr(rmsnorm_kernel_impl, builder_name) + build_count = 0 + + def counted_builder(*args, **kwargs): + nonlocal build_count + build_count += 1 + return original_builder(*args, **kwargs) + + monkeypatch.setattr(rmsnorm_kernel_impl, builder_name, counted_builder) + saved_cache = cache.copy() + cache.clear() + + def run(M, seed): + torch.manual_seed(seed) + x = torch.randn((M, N), device=device, dtype=dtype) + weight = torch.rand((N,), device=device, dtype=dtype) + dout = torch.randn((M, N), device=device, dtype=dtype) + rstd = torch.rsqrt(x.to(DTYPE_FP32).square().mean(1) + EPS) + if fused_add: + dresidual_out = torch.randn_like(x) + dx, dresidual, dweight = rmsnorm_kernel_impl.fused_add_rmsnorm_bwd(x, weight, dout, rstd, dresidual_out) + assert dx.data_ptr() == dresidual.data_ptr() + else: + dresidual_out = None + dx, dweight = rmsnorm_kernel_impl.rmsnorm_bwd(x, weight, dout, rstd) + torch.cuda.synchronize(device) + dx_ref, dweight_ref, _ = _reference_rmsnorm_bwd(x, weight, dout) + if dresidual_out is not None: + dx_ref = dx_ref + dresidual_out.to(DTYPE_FP32) + torch.testing.assert_close(dx.to(DTYPE_FP32), dx_ref, rtol=1e-1, atol=2e-1) + torch.testing.assert_close(dweight.to(DTYPE_FP32), dweight_ref, rtol=1e-1, atol=1.0) + + try: + run(first_m, 11) + run(second_m, 12) + assert build_count == 1 + finally: + cache.clear() + cache.update(saved_cache) + + +def test_rmsnorm_bwd_two_stage_multistream_workspace(): + """Per-call partials must remain independent when two streams overlap.""" + device = torch.device("cuda", torch.cuda.current_device()) + M, N = 1024, 4096 + path, _ = rmsnorm_kernel_impl._select_rmsnorm_bwd_config(M, N, "bf16", device) + assert path == "two_stage" + streams = (torch.cuda.Stream(device=device), torch.cuda.Stream(device=device)) + cases = [] + + for seed in (21, 22): + torch.manual_seed(seed) + x = torch.randn((M, N), device=device, dtype=DTYPE_BF16) + weight = torch.rand((N,), device=device, dtype=DTYPE_BF16) + dout = torch.randn((M, N), device=device, dtype=DTYPE_BF16) + rstd = torch.rsqrt(x.to(DTYPE_FP32).square().mean(1) + EPS) + cases.append((x, weight, dout, rstd)) + + results = [] + producer = torch.cuda.current_stream(device) + for stream, (x, weight, dout, rstd) in zip(streams, cases, strict=True): + stream.wait_stream(producer) + with torch.cuda.stream(stream): + results.append(rmsnorm_kernel_impl.rmsnorm_bwd(x, weight, dout, rstd)) + torch.cuda.synchronize(device) + + for (x, weight, dout, _), (dx, dweight) in zip(cases, results, strict=True): + dx_ref, dw_ref, _ = _reference_rmsnorm_bwd(x, weight, dout) + torch.testing.assert_close(dx.to(DTYPE_FP32), dx_ref, rtol=1e-1, atol=2e-1) + torch.testing.assert_close(dweight.to(DTYPE_FP32), dw_ref, rtol=1e-1, atol=1.0) + + @pytest.mark.multi_gpu def test_rmsnorm_multi_gpu(): """Compiled-fn cache must not reuse a device-0 kernel on device 1 (would fault).""" @@ -1350,18 +1509,33 @@ def test_rmsnorm_multi_gpu(): pytest.skip("needs >=2 GPUs") torch.manual_seed(0) - N = 256 - for dev in ("cuda:0", "cuda:1"): - x = torch.randn((16, N), device=dev, dtype=DTYPE_FP32, requires_grad=True) - w = torch.rand((N,), device=dev, dtype=DTYPE_FP32, requires_grad=True) - dy = torch.randn((16, N), device=dev, dtype=DTYPE_FP32) - y = rmsnorm(x, w) - y.backward(dy) - torch.cuda.synchronize(dev) - ref = x.detach() / torch.sqrt((x.detach() ** 2).mean(1, keepdim=True) + EPS) * w.detach() - err = (y.detach() - ref).abs().max().item() - print(f" {dev}: out err={err:.3e}, dx finite={torch.isfinite(x.grad).all().item()}") - assert err < 1e-4 and torch.isfinite(x.grad).all() + for M, N, dtype, out_atol, grad_atol in ( + (16, 256, DTYPE_FP32, 1e-4, 1e-3), + (512, 256, DTYPE_BF16, 2e-1, 2e-1), + ): + expected_path = "atomic" if M < 512 else "two_stage" + for dev in ("cuda:0", "cuda:1"): + dtype_str = "f32" if dtype == DTYPE_FP32 else "bf16" + path, _ = rmsnorm_kernel_impl._select_rmsnorm_bwd_config(M, N, dtype_str, torch.device(dev)) + assert path == expected_path + x = torch.randn((M, N), device=dev, dtype=dtype, requires_grad=True) + w = torch.rand((N,), device=dev, dtype=dtype, requires_grad=True) + dy = torch.randn((M, N), device=dev, dtype=dtype) + y = rmsnorm(x, w) + y.backward(dy) + torch.cuda.synchronize(dev) + + xf = x.detach().to(DTYPE_FP32).requires_grad_(True) + wf = w.detach().to(DTYPE_FP32).requires_grad_(True) + ref = xf * torch.rsqrt(xf.square().mean(1, keepdim=True) + EPS) * wf + dx_ref, dw_ref = torch.autograd.grad(ref, (xf, wf), dy.to(DTYPE_FP32)) + + rtol = 1e-4 if dtype == DTYPE_FP32 else 1e-1 + dw_atol = 1e-2 if dtype == DTYPE_FP32 else 1.0 + torch.testing.assert_close(y.detach().to(DTYPE_FP32), ref, rtol=rtol, atol=out_atol) + torch.testing.assert_close(x.grad.to(DTYPE_FP32), dx_ref, rtol=rtol, atol=grad_atol) + torch.testing.assert_close(w.grad.to(DTYPE_FP32), dw_ref, rtol=rtol, atol=dw_atol) + print(f" M={M} N={N} {dev}: {path} forward/backward match reference") print(" -> PASSED") From de5dbc240706eda21b4c1bdf946a6aca3e156b4e Mon Sep 17 00:00:00 2001 From: Jin Pan <47354855+jhinpan@users.noreply.github.com> Date: Wed, 15 Jul 2026 21:53:55 +0000 Subject: [PATCH 2/5] [Perf] Vectorize persistent RMSNorm backward --- kernels/norm/rmsnorm_bwd_kernel.py | 277 +++++++++++++++++++---------- kernels/norm/rmsnorm_kernel.py | 21 ++- tests/kernels/test_rmsnorm.py | 39 +++- 3 files changed, 235 insertions(+), 102 deletions(-) diff --git a/kernels/norm/rmsnorm_bwd_kernel.py b/kernels/norm/rmsnorm_bwd_kernel.py index cbc914d39..933c38f92 100644 --- a/kernels/norm/rmsnorm_bwd_kernel.py +++ b/kernels/norm/rmsnorm_bwd_kernel.py @@ -14,13 +14,19 @@ import flydsl.compiler as flyc import flydsl.expr as fx from flydsl.expr import arith, const_expr, gpu, range_constexpr +from flydsl.expr.vector import ReductionOp +from flydsl.runtime.device import get_rocm_arch from kernels.common.kernels_common import atomic_add, dtype_to_elem_type from kernels.norm.rmsnorm_common import ( BLOCK_THREADS, + VEC_WIDTH, WARP_SIZE, load_scalar, + load_vec, make_single_reduction_storage, store_scalar, + store_vec, + to_elem_vec, ) # The staged path is selected by the Python wrapper only when M is large @@ -30,6 +36,12 @@ DWEIGHT_REDUCE_COLS = 64 DWEIGHT_REDUCE_ROW_LANES = 4 DWEIGHT_REDUCE_THREADS = DWEIGHT_REDUCE_COLS * DWEIGHT_REDUCE_ROW_LANES +TWO_STAGE_PARTIAL_THREADS = 512 + + +def is_rmsnorm_bwd_two_stage_vec_config(N: int, dtype_str: str) -> bool: + """Whether the staged main kernel can use full-width vec8 column I/O.""" + return dtype_str in ("f16", "bf16") and N >= TWO_STAGE_PARTIAL_THREADS * VEC_WIDTH and N % VEC_WIDTH == 0 def build_rmsnorm_bwd_module(N: int, dtype_str: str): @@ -350,10 +362,11 @@ def _build_rmsnorm_bwd_two_stage_module( ): """Build the large-M persistent backward + dweight finalizer. - The main kernel launches ``num_programs`` blocks. Each block walks rows in - a grid-stride loop, writes every assigned ``dx`` row, and accumulates one - FP32 dweight partial in registers. It then writes exactly one ``N``-wide - partial row. A second kernel reduces the ``num_programs`` partial rows and + The 512-thread main kernel walks rows in a grid-stride loop, writes every + assigned ``dx`` row, and accumulates one FP32 dweight partial in registers. + Full-width 16-bit rows use vec8 I/O and retain source, dy, and gamma across + the row reduction; other shapes keep the scalar fallback. Each program writes + one ``N``-wide partial row, then a second kernel reduces all partial rows and stores dweight directly in the parameter dtype. This replaces M competing atomic adds per dweight element with a bounded @@ -364,13 +377,19 @@ def _build_rmsnorm_bwd_two_stage_module( if num_programs <= 0: raise ValueError(f"num_programs must be positive, got {num_programs}") - RED_SLOTS = max(1, (BLOCK_THREADS + WARP_SIZE - 1) // WARP_SIZE) - NUM_COL_TILES = (N + BLOCK_THREADS - 1) // BLOCK_THREADS + RED_SLOTS = max(1, (TWO_STAGE_PARTIAL_THREADS + WARP_SIZE - 1) // WARP_SIZE) + NUM_COL_TILES = (N + TWO_STAGE_PARTIAL_THREADS - 1) // TWO_STAGE_PARTIAL_THREADS elem_bits = 32 if dtype_str == "f32" else 16 + USE_VEC = is_rmsnorm_bwd_two_stage_vec_config(N, dtype_str) + NUM_VEC_TILES = N // VEC_WIDTH if USE_VEC else 0 + NUM_VEC_ITERS = (NUM_VEC_TILES + TWO_STAGE_PARTIAL_THREADS - 1) // TWO_STAGE_PARTIAL_THREADS if USE_VEC else 0 + PARTIAL_ACC_SIZE = NUM_VEC_ITERS * VEC_WIDTH if USE_VEC else NUM_COL_TILES + arch = get_rocm_arch() if USE_VEC else "" + USE_HW_CVT_PK_BF16_F32 = (arch == "gfx950") or str(arch).startswith("gfx95") SharedStorage = make_single_reduction_storage(RED_SLOTS) DWeightReduceStorage = make_single_reduction_storage(DWEIGHT_REDUCE_THREADS) - @flyc.kernel + @flyc.kernel(known_block_size=[TWO_STAGE_PARTIAL_THREADS, 1, 1]) def rmsnorm_bwd_partial_kernel( Source: fx.Tensor, Gamma: fx.Tensor, @@ -428,7 +447,6 @@ def block_reduce_add(val): DX_buf = fx.rocdl.make_buffer_tensor(DX) DWeightPartial_buf = fx.rocdl.make_buffer_tensor(DWeightPartial) - gamma_div = fx.logical_divide(Gamma_buf, fx.make_layout(1, 1)) rstd_div = fx.logical_divide(Rstd_buf, fx.make_layout(1, 1)) partial_div = fx.logical_divide(DWeightPartial_buf, fx.make_layout(1, 1)) @@ -438,89 +456,168 @@ def block_reduce_add(val): ) copy_atom_f32 = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), 32) - # One scalar per compile-time column tile. Carrying a single Vector - # through the runtime row loop keeps the accumulator in VGPRs and gives - # scf.for one well-defined loop-carried value. - dweight_partial = fx.Vector.filled(NUM_COL_TILES, 0.0, fx.Float32) - - for row in range(fx.Int32(bid), MIn, num_programs): - row_source = fx.slice(Source_buf, (row, None)) - row_dy = fx.slice(DY_buf, (row, None)) - row_dx = fx.slice(DX_buf, (row, None)) - if const_expr(fused_add): - row_dres_out = fx.slice(DResidualOut_buf, (row, None)) - - source_div = fx.logical_divide(row_source, fx.make_layout(1, 1)) - dy_div = fx.logical_divide(row_dy, fx.make_layout(1, 1)) - dx_div = fx.logical_divide(row_dx, fx.make_layout(1, 1)) - if const_expr(fused_add): - dres_out_div = fx.logical_divide(row_dres_out, fx.make_layout(1, 1)) - - rstd = load_scalar(copy_atom_f32, fx.Float32, rstd_div, row) - - thread_acc = c_zero_f - for base in range_constexpr(0, N, BLOCK_THREADS): - idx = tid + base - is_valid = idx < N - idx_safe = is_valid.select(idx, 0) - source_e = load_scalar(copy_atom_s, elem_dtype, source_div, idx_safe) - dy_e = load_scalar(copy_atom_s, elem_dtype, dy_div, idx_safe) - g_e = load_scalar(copy_atom_s, elem_dtype, gamma_div, idx_safe) - source = source_e if dtype_str == "f32" else source_e.to(fx.Float32) - dy = dy_e if dtype_str == "f32" else dy_e.to(fx.Float32) - g = g_e if dtype_str == "f32" else g_e.to(fx.Float32) - source_hat = source * rstd - wdy = dy * g - thread_acc = thread_acc + is_valid.select(source_hat * wdy, c_zero_f) - - sum_prod = block_reduce_add(thread_acc) - c1 = sum_prod / n_float - - row_dweight = [] - for base in range_constexpr(0, N, BLOCK_THREADS): - idx = tid + base - is_valid = idx < N - idx_safe = is_valid.select(idx, 0) - source_e = load_scalar(copy_atom_s, elem_dtype, source_div, idx_safe) - dy_e = load_scalar(copy_atom_s, elem_dtype, dy_div, idx_safe) - g_e = load_scalar(copy_atom_s, elem_dtype, gamma_div, idx_safe) - source = source_e if dtype_str == "f32" else source_e.to(fx.Float32) - dy = dy_e if dtype_str == "f32" else dy_e.to(fx.Float32) - g = g_e if dtype_str == "f32" else g_e.to(fx.Float32) - source_hat = source * rstd - wdy = dy * g - dx = (wdy - source_hat * c1) * rstd + if const_expr(USE_VEC): + copy_atom_vec = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), elem_bits) + gamma_vec_div = fx.logical_divide(Gamma_buf, fx.make_layout(VEC_WIDTH, 1)) + gamma_local = [] + for tile_i in range_constexpr(NUM_VEC_ITERS): + vec_idx = tid + tile_i * TWO_STAGE_PARTIAL_THREADS + is_valid = vec_idx < NUM_VEC_TILES + vec_idx_safe = is_valid.select(vec_idx, 0) + gamma_local.append(load_vec(copy_atom_vec, VEC_WIDTH, elem_dtype, gamma_vec_div, vec_idx_safe)) + + dweight_partial = fx.Vector.filled(PARTIAL_ACC_SIZE, 0.0, fx.Float32) + for row in range(fx.Int32(bid), MIn, num_programs): + row_source = fx.slice(Source_buf, (row, None)) + row_dy = fx.slice(DY_buf, (row, None)) + row_dx = fx.slice(DX_buf, (row, None)) if const_expr(fused_add): - dres_e = load_scalar(copy_atom_s, elem_dtype, dres_out_div, idx_safe) - dres = dres_e if dtype_str == "f32" else dres_e.to(fx.Float32) - dx = dx + dres + row_dres_out = fx.slice(DResidualOut_buf, (row, None)) - if idx < N: - dx_e = dx if dtype_str == "f32" else dx.to(elem_dtype) - store_scalar(copy_atom_s, elem_dtype, elem_dtype, dx_div, idx, dx_e) - - row_dweight.append(is_valid.select(dy * source_hat, c_zero_f)) - - dweight_partial = dweight_partial + fx.Vector.from_elements(row_dweight, fx.Float32) - - # c1 lives in shared storage. All waves must finish consuming it - # before a faster wave enters the next row and overwrites s_red. - gpu.barrier() + source_div = fx.logical_divide(row_source, fx.make_layout(VEC_WIDTH, 1)) + dy_div = fx.logical_divide(row_dy, fx.make_layout(VEC_WIDTH, 1)) + dx_div = fx.logical_divide(row_dx, fx.make_layout(VEC_WIDTH, 1)) + if const_expr(fused_add): + dres_out_div = fx.logical_divide(row_dres_out, fx.make_layout(VEC_WIDTH, 1)) + + rstd = load_scalar(copy_atom_f32, fx.Float32, rstd_div, row) + thread_acc = c_zero_f + source_local = [] + dy_local = [] + for tile_i in range_constexpr(NUM_VEC_ITERS): + vec_idx = tid + tile_i * TWO_STAGE_PARTIAL_THREADS + is_valid = vec_idx < NUM_VEC_TILES + vec_idx_safe = is_valid.select(vec_idx, 0) + source_e = load_vec(copy_atom_vec, VEC_WIDTH, elem_dtype, source_div, vec_idx_safe) + dy_e = load_vec(copy_atom_vec, VEC_WIDTH, elem_dtype, dy_div, vec_idx_safe) + source_local.append(source_e) + dy_local.append(dy_e) + source_hat = source_e.to(fx.Float32) * rstd + wdy = dy_e.to(fx.Float32) * gamma_local[tile_i].to(fx.Float32) + prod = (source_hat * wdy).reduce(ReductionOp.ADD, fastmath=fm_fast) + thread_acc = thread_acc + is_valid.select(prod, c_zero_f) + + sum_prod = block_reduce_add(thread_acc) + c1 = sum_prod / n_float + + row_dweight = [] + for tile_i in range_constexpr(NUM_VEC_ITERS): + vec_idx = tid + tile_i * TWO_STAGE_PARTIAL_THREADS + is_valid = vec_idx < NUM_VEC_TILES + vec_idx_safe = is_valid.select(vec_idx, 0) + source_hat = source_local[tile_i].to(fx.Float32) * rstd + dy = dy_local[tile_i].to(fx.Float32) + wdy = dy * gamma_local[tile_i].to(fx.Float32) + dx = (wdy - source_hat * c1) * rstd + if const_expr(fused_add): + dres = load_vec( + copy_atom_vec, + VEC_WIDTH, + elem_dtype, + dres_out_div, + vec_idx_safe, + ).to(fx.Float32) + dx = dx + dres + + if vec_idx < NUM_VEC_TILES: + dx_e = to_elem_vec(dtype_str, elem_dtype, USE_HW_CVT_PK_BF16_F32, dx) + store_vec(copy_atom_vec, VEC_WIDTH, elem_dtype, dx_e, dx_div, vec_idx) + + dw = dy * source_hat + for lane in range_constexpr(VEC_WIDTH): + row_dweight.append(is_valid.select(dw[lane], c_zero_f)) + + dweight_partial = dweight_partial + fx.Vector.from_elements(row_dweight, fx.Float32) + gpu.barrier() + + for tile_i in range_constexpr(NUM_VEC_ITERS): + vec_idx = tid + tile_i * TWO_STAGE_PARTIAL_THREADS + if vec_idx < NUM_VEC_TILES: + for lane in range_constexpr(VEC_WIDTH): + idx = vec_idx * VEC_WIDTH + lane + partial_idx = bid * N + idx + store_scalar( + copy_atom_f32, + fx.Float32, + fx.Float32, + partial_div, + partial_idx, + dweight_partial[tile_i * VEC_WIDTH + lane], + ) + else: + gamma_div = fx.logical_divide(Gamma_buf, fx.make_layout(1, 1)) + dweight_partial = fx.Vector.filled(PARTIAL_ACC_SIZE, 0.0, fx.Float32) + + for row in range(fx.Int32(bid), MIn, num_programs): + row_source = fx.slice(Source_buf, (row, None)) + row_dy = fx.slice(DY_buf, (row, None)) + row_dx = fx.slice(DX_buf, (row, None)) + if const_expr(fused_add): + row_dres_out = fx.slice(DResidualOut_buf, (row, None)) - # Every program overwrites its entire partial row, including programs - # with no assigned rows (their accumulator remains zero). - for tile_i in range_constexpr(NUM_COL_TILES): - idx = tid + tile_i * BLOCK_THREADS - if idx < N: - partial_idx = bid * N + idx - store_scalar( - copy_atom_f32, - fx.Float32, - fx.Float32, - partial_div, - partial_idx, - dweight_partial[tile_i], - ) + source_div = fx.logical_divide(row_source, fx.make_layout(1, 1)) + dy_div = fx.logical_divide(row_dy, fx.make_layout(1, 1)) + dx_div = fx.logical_divide(row_dx, fx.make_layout(1, 1)) + if const_expr(fused_add): + dres_out_div = fx.logical_divide(row_dres_out, fx.make_layout(1, 1)) + + rstd = load_scalar(copy_atom_f32, fx.Float32, rstd_div, row) + thread_acc = c_zero_f + for base in range_constexpr(0, N, TWO_STAGE_PARTIAL_THREADS): + idx = tid + base + is_valid = idx < N + idx_safe = is_valid.select(idx, 0) + source_e = load_scalar(copy_atom_s, elem_dtype, source_div, idx_safe) + dy_e = load_scalar(copy_atom_s, elem_dtype, dy_div, idx_safe) + g_e = load_scalar(copy_atom_s, elem_dtype, gamma_div, idx_safe) + source = source_e if dtype_str == "f32" else source_e.to(fx.Float32) + dy = dy_e if dtype_str == "f32" else dy_e.to(fx.Float32) + g = g_e if dtype_str == "f32" else g_e.to(fx.Float32) + source_hat = source * rstd + wdy = dy * g + thread_acc = thread_acc + is_valid.select(source_hat * wdy, c_zero_f) + + sum_prod = block_reduce_add(thread_acc) + c1 = sum_prod / n_float + row_dweight = [] + for base in range_constexpr(0, N, TWO_STAGE_PARTIAL_THREADS): + idx = tid + base + is_valid = idx < N + idx_safe = is_valid.select(idx, 0) + source_e = load_scalar(copy_atom_s, elem_dtype, source_div, idx_safe) + dy_e = load_scalar(copy_atom_s, elem_dtype, dy_div, idx_safe) + g_e = load_scalar(copy_atom_s, elem_dtype, gamma_div, idx_safe) + source = source_e if dtype_str == "f32" else source_e.to(fx.Float32) + dy = dy_e if dtype_str == "f32" else dy_e.to(fx.Float32) + g = g_e if dtype_str == "f32" else g_e.to(fx.Float32) + source_hat = source * rstd + wdy = dy * g + dx = (wdy - source_hat * c1) * rstd + if const_expr(fused_add): + dres_e = load_scalar(copy_atom_s, elem_dtype, dres_out_div, idx_safe) + dres = dres_e if dtype_str == "f32" else dres_e.to(fx.Float32) + dx = dx + dres + + if idx < N: + dx_e = dx if dtype_str == "f32" else dx.to(elem_dtype) + store_scalar(copy_atom_s, elem_dtype, elem_dtype, dx_div, idx, dx_e) + row_dweight.append(is_valid.select(dy * source_hat, c_zero_f)) + + dweight_partial = dweight_partial + fx.Vector.from_elements(row_dweight, fx.Float32) + gpu.barrier() + + for tile_i in range_constexpr(NUM_COL_TILES): + idx = tid + tile_i * TWO_STAGE_PARTIAL_THREADS + if idx < N: + partial_idx = bid * N + idx + store_scalar( + copy_atom_f32, + fx.Float32, + fx.Float32, + partial_div, + partial_idx, + dweight_partial[tile_i], + ) @flyc.kernel def rmsnorm_bwd_dweight_reduce_kernel( @@ -597,7 +694,7 @@ def launch_fused_add_rmsnorm_bwd_two_stage( ) partial_launcher.launch( grid=(num_programs, 1, 1), - block=(BLOCK_THREADS, 1, 1), + block=(TWO_STAGE_PARTIAL_THREADS, 1, 1), stream=stream, ) reduce_launcher = rmsnorm_bwd_dweight_reduce_kernel(DWeightPartial, DWeight) @@ -635,7 +732,7 @@ def launch_rmsnorm_bwd_two_stage( ) partial_launcher.launch( grid=(num_programs, 1, 1), - block=(BLOCK_THREADS, 1, 1), + block=(TWO_STAGE_PARTIAL_THREADS, 1, 1), stream=stream, ) reduce_launcher = rmsnorm_bwd_dweight_reduce_kernel(DWeightPartial, DWeight) diff --git a/kernels/norm/rmsnorm_kernel.py b/kernels/norm/rmsnorm_kernel.py index 8fe932a95..ecda28ce8 100644 --- a/kernels/norm/rmsnorm_kernel.py +++ b/kernels/norm/rmsnorm_kernel.py @@ -27,6 +27,7 @@ build_fused_add_rmsnorm_bwd_two_stage_module, build_rmsnorm_bwd_module, build_rmsnorm_bwd_two_stage_module, + is_rmsnorm_bwd_two_stage_vec_config, ) from kernels.norm.rmsnorm_common import ( BLOCK_THREADS, @@ -1460,10 +1461,10 @@ def build_fused_add_rmsnorm_smoothquant_module( _FWD_CACHE: dict = {} _BWD_CACHE: dict = {} - # The staged path needs enough rows to amortize its second launch. Wide - # and FP32 rows carry more live accumulators, so cap their program count to - # avoid excess register pressure and scratch. Keep this selector as the - # single source of truth shared by plain and fused-add wrappers. + # The staged path needs enough rows to amortize its second launch. Its + # 512-thread vec8 kernel needs fewer persistent programs than the scalar + # fallback, which trades more programs for latency hiding. Keep this + # selector as the single source of truth shared by plain and fused wrappers. _BWD_TWO_STAGE_MIN_ROWS = 512 _BWD_TWO_STAGE_MAX_N = 8192 _BWD_CU_COUNT_CACHE: dict = {} @@ -1477,14 +1478,12 @@ def _get_rmsnorm_bwd_cu_count(device): def _select_rmsnorm_bwd_config(M, N, dtype_str, device): if M >= _BWD_TWO_STAGE_MIN_ROWS and N <= _BWD_TWO_STAGE_MAX_N: - if M < 1024: - programs_per_cu = 1 - elif M < 2048 or dtype_str == "f32" or N <= 2048 or N > 4096: - programs_per_cu = 2 + num_cus = _get_rmsnorm_bwd_cu_count(device) + if is_rmsnorm_bwd_two_stage_vec_config(N, dtype_str): + num_programs = num_cus if M < 2048 else (3 * num_cus) // 2 else: - programs_per_cu = 4 - num_programs = min(M, _get_rmsnorm_bwd_cu_count(device) * programs_per_cu) - return ("two_stage", num_programs) + num_programs = num_cus if M < 1024 else 2 * num_cus + return ("two_stage", min(M, num_programs)) return ("atomic", None) def _get_fwd_compiled(x, weight, out, rstd, M, N, dtype_str, store_rstd, eps, stream): diff --git a/tests/kernels/test_rmsnorm.py b/tests/kernels/test_rmsnorm.py index 23291a23a..6aa9baf63 100644 --- a/tests/kernels/test_rmsnorm.py +++ b/tests/kernels/test_rmsnorm.py @@ -70,7 +70,9 @@ _RMSNORM_BACKWARD_CONFIGS = ( (64, 256, "f32"), # small-N path, f32 (16, 512, "bf16"), # small-N path, bf16 - (512, 2048, "f16"), # staged finalizer cast/store + (512, 4096, "f16"), # staged vec8 main path and finalizer cast/store + (513, 5000, "bf16"), # vec8 column/program tails + (513, 4097, "bf16"), # 16-bit scalar fallback above the vec8 threshold (4096, 4096, "bf16"), # model-relevant large shape (64, 2000, "f32"), # small-N path, unaligned (128, 4096, "f16"), # f16, large hidden size @@ -1499,6 +1501,41 @@ def test_rmsnorm_bwd_two_stage_multistream_workspace(): torch.testing.assert_close(dweight.to(DTYPE_FP32), dw_ref, rtol=1e-1, atol=1.0) +def test_rmsnorm_bwd_vec8_contiguous_storage_offset(): + """Contiguous tensors need not have a 16-byte-aligned storage offset.""" + device = torch.device("cuda", torch.cuda.current_device()) + M, N = 512, 4096 + + def offset_rand(shape): + numel = 1 + for dim in shape: + numel *= dim + tensor = torch.randn((numel + 1,), device=device, dtype=DTYPE_BF16)[1:].view(shape) + assert tensor.is_contiguous() and tensor.data_ptr() % 16 != 0 + return tensor + + added = offset_rand((M, N)) + weight = offset_rand((N,)) + dout = offset_rand((M, N)) + dresidual_out = offset_rand((M, N)) + rstd = torch.rsqrt(added.to(DTYPE_FP32).square().mean(1) + EPS) + dx_ref, dweight_ref, _ = _reference_rmsnorm_bwd(added, weight, dout) + + dx, dweight = rmsnorm_kernel_impl.rmsnorm_bwd(added, weight, dout, rstd) + torch.testing.assert_close(dx.to(DTYPE_FP32), dx_ref, rtol=1e-1, atol=2e-1) + torch.testing.assert_close(dweight.to(DTYPE_FP32), dweight_ref, rtol=1e-1, atol=1.0) + + dx, dresidual, dweight = rmsnorm_kernel_impl.fused_add_rmsnorm_bwd(added, weight, dout, rstd, dresidual_out) + assert dx.data_ptr() == dresidual.data_ptr() + torch.testing.assert_close( + dx.to(DTYPE_FP32), + dx_ref + dresidual_out.to(DTYPE_FP32), + rtol=1e-1, + atol=2e-1, + ) + torch.testing.assert_close(dweight.to(DTYPE_FP32), dweight_ref, rtol=1e-1, atol=1.0) + + @pytest.mark.multi_gpu def test_rmsnorm_multi_gpu(): """Compiled-fn cache must not reuse a device-0 kernel on device 1 (would fault).""" From 3d15258ad3077543369214a8ff04a7d2880de054 Mon Sep 17 00:00:00 2001 From: Jin Pan <47354855+jhinpan@users.noreply.github.com> Date: Wed, 15 Jul 2026 23:02:27 +0000 Subject: [PATCH 3/5] [Perf] Reduce RMSNorm backward launch overhead --- kernels/norm/rmsnorm_kernel.py | 153 +++++++++++++++++++-------------- tests/kernels/test_rmsnorm.py | 13 +++ 2 files changed, 101 insertions(+), 65 deletions(-) diff --git a/kernels/norm/rmsnorm_kernel.py b/kernels/norm/rmsnorm_kernel.py index ecda28ce8..8f004bc7f 100644 --- a/kernels/norm/rmsnorm_kernel.py +++ b/kernels/norm/rmsnorm_kernel.py @@ -1461,6 +1461,24 @@ def build_fused_add_rmsnorm_smoothquant_module( _FWD_CACHE: dict = {} _BWD_CACHE: dict = {} + # A cached FlyDSL callable accepts the raw cudaStream_t directly. Avoid + # constructing a Python Stream wrapper or entering a no-op device guard on + # the hot same-device path, while retaining the guard when tensors belong + # to a non-current device. + _get_current_raw_stream = getattr( + torch._C, + "_cuda_getCurrentRawStream", + lambda device_index: torch.cuda.current_stream(device_index).cuda_stream, + ) + + def _launch_cached_on_device(compiled, device, args): + device_index = device.index + if torch.cuda.current_device() == device_index: + compiled(*args, _get_current_raw_stream(device_index)) + else: + with torch.cuda.device(device_index): + compiled(*args, _get_current_raw_stream(device_index)) + # The staged path needs enough rows to amortize its second launch. Its # 512-thread vec8 kernel needs fewer persistent programs than the scalar # fallback, which trades more programs for latency hiding. Keep this @@ -1528,26 +1546,28 @@ def rmsnorm_bwd(x, weight, dout, rstd, eps=EPS): eps is not used directly here — it is already baked into `rstd` by the forward — but is accepted so callers can pass it symmetrically. """ + device = x.device + dtype = x.dtype assert x.dim() == 2, "rmsnorm_bwd expects a 2D (M, N) input" assert x.is_contiguous() and dout.is_contiguous(), "rmsnorm_bwd expects contiguous inputs" assert weight.is_contiguous(), "rmsnorm_bwd: weight must be contiguous" # Same-device/same-dtype contract as the forward: gamma and dy are read at # x's element width and the kernel launches on x's device (multi-GPU correctness). - assert weight.device == x.device == dout.device == rstd.device, "rmsnorm_bwd: inputs must share a device" - assert weight.dtype == x.dtype and dout.dtype == x.dtype, "rmsnorm_bwd: weight/dout dtype must match x" + assert weight.device == device == dout.device == rstd.device, "rmsnorm_bwd: inputs must share a device" + assert weight.dtype == dtype and dout.dtype == dtype, "rmsnorm_bwd: weight/dout dtype must match x" M, N = x.shape - dtype_str = _torch_dtype_to_str(x.dtype) + dtype_str = _torch_dtype_to_str(dtype) dx = torch.empty_like(x) - # Bind compile + launch to the tensors' device (multi-GPU correctness). - with torch.cuda.device(x.device): - stream = torch.cuda.current_stream() - path, num_programs = _select_rmsnorm_bwd_config(M, N, dtype_str, x.device) - if path == "two_stage": - dweight = torch.empty_like(weight) - partial = torch.empty((num_programs * N,), device=x.device, dtype=torch.float32) - key = (path, N, dtype_str, num_programs, x.device) - compiled = _BWD_CACHE.get(key) - if compiled is None: + path, num_programs = _select_rmsnorm_bwd_config(M, N, dtype_str, device) + if path == "two_stage": + dweight = torch.empty_like(weight) + partial = torch.empty((num_programs * N,), device=device, dtype=torch.float32) + key = (path, N, dtype_str, num_programs, device) + compiled = _BWD_CACHE.get(key) + if compiled is None: + # Compilation and code-object loading are device-context bound. + with torch.cuda.device(device): + stream = torch.cuda.current_stream() launch_fn = build_rmsnorm_bwd_two_stage_module(N, dtype_str, num_programs) compiled = flyc.compile( launch_fn, @@ -1561,21 +1581,23 @@ def rmsnorm_bwd(x, weight, dout, rstd, eps=EPS): M, stream, ) - _BWD_CACHE[key] = compiled - else: - compiled(x, weight, dout, rstd, dx, dweight, partial, M, stream) - return dx, dweight - - dweight = torch.empty((N,), device=x.device, dtype=torch.float32) - key = (path, N, dtype_str, x.device) - compiled = _BWD_CACHE.get(key) - dweight.zero_() - if compiled is None: - launch_fn = build_rmsnorm_bwd_module(N, dtype_str) - compiled = flyc.compile(launch_fn, x, weight, dout, rstd, dx, dweight, M, stream) _BWD_CACHE[key] = compiled else: - compiled(x, weight, dout, rstd, dx, dweight, M, stream) + _launch_cached_on_device(compiled, device, (x, weight, dout, rstd, dx, dweight, partial, M)) + return dx, dweight + + dweight = torch.empty((N,), device=device, dtype=torch.float32) + key = (path, N, dtype_str, device) + compiled = _BWD_CACHE.get(key) + dweight.zero_() + if compiled is None: + with torch.cuda.device(device): + stream = torch.cuda.current_stream() + launch_fn = build_rmsnorm_bwd_module(N, dtype_str) + compiled = flyc.compile(launch_fn, x, weight, dout, rstd, dx, dweight, M, stream) + _BWD_CACHE[key] = compiled + else: + _launch_cached_on_device(compiled, device, (x, weight, dout, rstd, dx, dweight, M)) return dx, dweight.to(weight.dtype) class RMSNormFunction(torch.autograd.Function): @@ -1672,32 +1694,34 @@ def fused_add_rmsnorm_bwd(added, weight, dout, rstd, dresidual_out=None, eps=EPS it is treated as zero (a zero tensor is passed to the branch-free kernel). eps is already baked into `rstd`. """ + device = added.device + dtype = added.dtype assert added.dim() == 2, "fused_add_rmsnorm_bwd expects a 2D (M, N) input" assert added.is_contiguous() and dout.is_contiguous(), "fused_add_rmsnorm_bwd expects contiguous inputs" assert ( - added.dtype == weight.dtype == dout.dtype - ), f"added/weight/dout dtypes must match, got {added.dtype}/{weight.dtype}/{dout.dtype}" + dtype == weight.dtype == dout.dtype + ), f"added/weight/dout dtypes must match, got {dtype}/{weight.dtype}/{dout.dtype}" assert ( - added.device == weight.device == dout.device == rstd.device + device == weight.device == dout.device == rstd.device ), "fused_add_rmsnorm_bwd expects all tensors on the same device" if dresidual_out is not None: assert dresidual_out.is_contiguous(), "fused_add_rmsnorm_bwd expects contiguous dresidual_out" - assert dresidual_out.dtype == added.dtype, "dresidual_out dtype must match added" - assert dresidual_out.device == added.device, "dresidual_out must be on the same device as added" + assert dresidual_out.dtype == dtype, "dresidual_out dtype must match added" + assert dresidual_out.device == device, "dresidual_out must be on the same device as added" M, N = added.shape - dtype_str = _torch_dtype_to_str(added.dtype) + dtype_str = _torch_dtype_to_str(dtype) if dresidual_out is None: dresidual_out = torch.zeros_like(added) dx = torch.empty_like(added) - with torch.cuda.device(added.device): - stream = torch.cuda.current_stream() - path, num_programs = _select_rmsnorm_bwd_config(M, N, dtype_str, added.device) - if path == "two_stage": - dweight = torch.empty_like(weight) - partial = torch.empty((num_programs * N,), device=added.device, dtype=torch.float32) - key = (path, N, dtype_str, num_programs, added.device) - compiled = _FUSED_ADD_BWD_CACHE.get(key) - if compiled is None: + path, num_programs = _select_rmsnorm_bwd_config(M, N, dtype_str, device) + if path == "two_stage": + dweight = torch.empty_like(weight) + partial = torch.empty((num_programs * N,), device=device, dtype=torch.float32) + key = (path, N, dtype_str, num_programs, device) + compiled = _FUSED_ADD_BWD_CACHE.get(key) + if compiled is None: + with torch.cuda.device(device): + stream = torch.cuda.current_stream() launch_fn = build_fused_add_rmsnorm_bwd_two_stage_module(N, dtype_str, num_programs) compiled = flyc.compile( launch_fn, @@ -1712,32 +1736,31 @@ def fused_add_rmsnorm_bwd(added, weight, dout, rstd, dresidual_out=None, eps=EPS M, stream, ) - _FUSED_ADD_BWD_CACHE[key] = compiled - else: - compiled( - added, - weight, - dout, - dresidual_out, - rstd, - dx, - dweight, - partial, - M, - stream, - ) - return dx, dx, dweight - - dweight = torch.empty((N,), device=added.device, dtype=torch.float32) - key = (path, N, dtype_str, added.device) - compiled = _FUSED_ADD_BWD_CACHE.get(key) - dweight.zero_() - if compiled is None: - launch_fn = build_fused_add_rmsnorm_bwd_module(N, dtype_str) - compiled = flyc.compile(launch_fn, added, weight, dout, dresidual_out, rstd, dx, dweight, M, stream) _FUSED_ADD_BWD_CACHE[key] = compiled else: - compiled(added, weight, dout, dresidual_out, rstd, dx, dweight, M, stream) + _launch_cached_on_device( + compiled, + device, + (added, weight, dout, dresidual_out, rstd, dx, dweight, partial, M), + ) + return dx, dx, dweight + + dweight = torch.empty((N,), device=device, dtype=torch.float32) + key = (path, N, dtype_str, device) + compiled = _FUSED_ADD_BWD_CACHE.get(key) + dweight.zero_() + if compiled is None: + with torch.cuda.device(device): + stream = torch.cuda.current_stream() + launch_fn = build_fused_add_rmsnorm_bwd_module(N, dtype_str) + compiled = flyc.compile(launch_fn, added, weight, dout, dresidual_out, rstd, dx, dweight, M, stream) + _FUSED_ADD_BWD_CACHE[key] = compiled + else: + _launch_cached_on_device( + compiled, + device, + (added, weight, dout, dresidual_out, rstd, dx, dweight, M), + ) # dx == dresidual by construction; return dx as both (aliased). return dx, dx, dweight.to(weight.dtype) diff --git a/tests/kernels/test_rmsnorm.py b/tests/kernels/test_rmsnorm.py index 6aa9baf63..afe2b792e 100644 --- a/tests/kernels/test_rmsnorm.py +++ b/tests/kernels/test_rmsnorm.py @@ -1545,6 +1545,7 @@ def test_rmsnorm_multi_gpu(): if torch.cuda.device_count() < 2: pytest.skip("needs >=2 GPUs") + torch.cuda.set_device(0) torch.manual_seed(0) for M, N, dtype, out_atol, grad_atol in ( (16, 256, DTYPE_FP32, 1e-4, 1e-3), @@ -1562,6 +1563,16 @@ def test_rmsnorm_multi_gpu(): y.backward(dy) torch.cuda.synchronize(dev) + # Autograd warmed this device's backward cache. Exercise the hot + # launch again while cuda:0 remains the caller's current device; + # the cuda:1 case must enter the guarded cross-device fallback and + # restore the caller's device afterward. + caller_device = torch.cuda.current_device() + rstd = torch.rsqrt(x.detach().to(DTYPE_FP32).square().mean(1) + EPS) + dx_cached, dw_cached = rmsnorm_kernel_impl.rmsnorm_bwd(x.detach(), w.detach(), dy, rstd) + torch.cuda.synchronize(dev) + assert torch.cuda.current_device() == caller_device + xf = x.detach().to(DTYPE_FP32).requires_grad_(True) wf = w.detach().to(DTYPE_FP32).requires_grad_(True) ref = xf * torch.rsqrt(xf.square().mean(1, keepdim=True) + EPS) * wf @@ -1572,6 +1583,8 @@ def test_rmsnorm_multi_gpu(): torch.testing.assert_close(y.detach().to(DTYPE_FP32), ref, rtol=rtol, atol=out_atol) torch.testing.assert_close(x.grad.to(DTYPE_FP32), dx_ref, rtol=rtol, atol=grad_atol) torch.testing.assert_close(w.grad.to(DTYPE_FP32), dw_ref, rtol=rtol, atol=dw_atol) + torch.testing.assert_close(dx_cached.to(DTYPE_FP32), dx_ref, rtol=rtol, atol=grad_atol) + torch.testing.assert_close(dw_cached.to(DTYPE_FP32), dw_ref, rtol=rtol, atol=dw_atol) print(f" M={M} N={N} {dev}: {path} forward/backward match reference") print(" -> PASSED") From fd6a3b8800f47ae3364392a1abcb11d4be8d6eeb Mon Sep 17 00:00:00 2001 From: Jin Pan <47354855+jhinpan@users.noreply.github.com> Date: Thu, 16 Jul 2026 21:48:50 +0000 Subject: [PATCH 4/5] [Refactor] Share RMSNorm backward vector/scalar flow --- kernels/norm/rmsnorm_bwd_kernel.py | 257 ++++++++++++----------------- 1 file changed, 108 insertions(+), 149 deletions(-) diff --git a/kernels/norm/rmsnorm_bwd_kernel.py b/kernels/norm/rmsnorm_bwd_kernel.py index 933c38f92..f42a6090b 100644 --- a/kernels/norm/rmsnorm_bwd_kernel.py +++ b/kernels/norm/rmsnorm_bwd_kernel.py @@ -378,12 +378,12 @@ def _build_rmsnorm_bwd_two_stage_module( raise ValueError(f"num_programs must be positive, got {num_programs}") RED_SLOTS = max(1, (TWO_STAGE_PARTIAL_THREADS + WARP_SIZE - 1) // WARP_SIZE) - NUM_COL_TILES = (N + TWO_STAGE_PARTIAL_THREADS - 1) // TWO_STAGE_PARTIAL_THREADS elem_bits = 32 if dtype_str == "f32" else 16 USE_VEC = is_rmsnorm_bwd_two_stage_vec_config(N, dtype_str) - NUM_VEC_TILES = N // VEC_WIDTH if USE_VEC else 0 - NUM_VEC_ITERS = (NUM_VEC_TILES + TWO_STAGE_PARTIAL_THREADS - 1) // TWO_STAGE_PARTIAL_THREADS if USE_VEC else 0 - PARTIAL_ACC_SIZE = NUM_VEC_ITERS * VEC_WIDTH if USE_VEC else NUM_COL_TILES + IO_WIDTH = VEC_WIDTH if USE_VEC else 1 + NUM_IO_TILES = (N + IO_WIDTH - 1) // IO_WIDTH + NUM_IO_ITERS = (NUM_IO_TILES + TWO_STAGE_PARTIAL_THREADS - 1) // TWO_STAGE_PARTIAL_THREADS + PARTIAL_ACC_SIZE = NUM_IO_ITERS * IO_WIDTH arch = get_rocm_arch() if USE_VEC else "" USE_HW_CVT_PK_BF16_F32 = (arch == "gfx950") or str(arch).startswith("gfx95") SharedStorage = make_single_reduction_storage(RED_SLOTS) @@ -450,173 +450,132 @@ def block_reduce_add(val): rstd_div = fx.logical_divide(Rstd_buf, fx.make_layout(1, 1)) partial_div = fx.logical_divide(DWeightPartial_buf, fx.make_layout(1, 1)) - copy_atom_s = fx.make_copy_atom( - fx.rocdl.BufferCopy16b() if elem_bits <= 16 else fx.rocdl.BufferCopy32b(), + copy_atom_io = fx.make_copy_atom( + ( + fx.rocdl.BufferCopy128b() + if USE_VEC + else (fx.rocdl.BufferCopy16b() if elem_bits <= 16 else fx.rocdl.BufferCopy32b()) + ), elem_bits, ) copy_atom_f32 = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), 32) + gamma_div = fx.logical_divide(Gamma_buf, fx.make_layout(IO_WIDTH, 1)) + gamma_local = [] if const_expr(USE_VEC): - copy_atom_vec = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), elem_bits) - gamma_vec_div = fx.logical_divide(Gamma_buf, fx.make_layout(VEC_WIDTH, 1)) - gamma_local = [] - for tile_i in range_constexpr(NUM_VEC_ITERS): - vec_idx = tid + tile_i * TWO_STAGE_PARTIAL_THREADS - is_valid = vec_idx < NUM_VEC_TILES - vec_idx_safe = is_valid.select(vec_idx, 0) - gamma_local.append(load_vec(copy_atom_vec, VEC_WIDTH, elem_dtype, gamma_vec_div, vec_idx_safe)) - - dweight_partial = fx.Vector.filled(PARTIAL_ACC_SIZE, 0.0, fx.Float32) - for row in range(fx.Int32(bid), MIn, num_programs): - row_source = fx.slice(Source_buf, (row, None)) - row_dy = fx.slice(DY_buf, (row, None)) - row_dx = fx.slice(DX_buf, (row, None)) - if const_expr(fused_add): - row_dres_out = fx.slice(DResidualOut_buf, (row, None)) - - source_div = fx.logical_divide(row_source, fx.make_layout(VEC_WIDTH, 1)) - dy_div = fx.logical_divide(row_dy, fx.make_layout(VEC_WIDTH, 1)) - dx_div = fx.logical_divide(row_dx, fx.make_layout(VEC_WIDTH, 1)) - if const_expr(fused_add): - dres_out_div = fx.logical_divide(row_dres_out, fx.make_layout(VEC_WIDTH, 1)) - - rstd = load_scalar(copy_atom_f32, fx.Float32, rstd_div, row) - thread_acc = c_zero_f - source_local = [] - dy_local = [] - for tile_i in range_constexpr(NUM_VEC_ITERS): - vec_idx = tid + tile_i * TWO_STAGE_PARTIAL_THREADS - is_valid = vec_idx < NUM_VEC_TILES - vec_idx_safe = is_valid.select(vec_idx, 0) - source_e = load_vec(copy_atom_vec, VEC_WIDTH, elem_dtype, source_div, vec_idx_safe) - dy_e = load_vec(copy_atom_vec, VEC_WIDTH, elem_dtype, dy_div, vec_idx_safe) + for tile_i in range_constexpr(NUM_IO_ITERS): + io_idx = tid + tile_i * TWO_STAGE_PARTIAL_THREADS + is_valid = io_idx < NUM_IO_TILES + io_idx_safe = is_valid.select(io_idx, 0) + gamma_local.append(load_vec(copy_atom_io, IO_WIDTH, elem_dtype, gamma_div, io_idx_safe)) + + dweight_partial = fx.Vector.filled(PARTIAL_ACC_SIZE, 0.0, fx.Float32) + for row in range(fx.Int32(bid), MIn, num_programs): + row_source = fx.slice(Source_buf, (row, None)) + row_dy = fx.slice(DY_buf, (row, None)) + row_dx = fx.slice(DX_buf, (row, None)) + if const_expr(fused_add): + row_dres_out = fx.slice(DResidualOut_buf, (row, None)) + + source_div = fx.logical_divide(row_source, fx.make_layout(IO_WIDTH, 1)) + dy_div = fx.logical_divide(row_dy, fx.make_layout(IO_WIDTH, 1)) + dx_div = fx.logical_divide(row_dx, fx.make_layout(IO_WIDTH, 1)) + if const_expr(fused_add): + dres_out_div = fx.logical_divide(row_dres_out, fx.make_layout(IO_WIDTH, 1)) + + rstd = load_scalar(copy_atom_f32, fx.Float32, rstd_div, row) + thread_acc = c_zero_f + source_local = [] + dy_local = [] + for tile_i in range_constexpr(NUM_IO_ITERS): + io_idx = tid + tile_i * TWO_STAGE_PARTIAL_THREADS + is_valid = io_idx < NUM_IO_TILES + io_idx_safe = is_valid.select(io_idx, 0) + if const_expr(USE_VEC): + source_e = load_vec(copy_atom_io, IO_WIDTH, elem_dtype, source_div, io_idx_safe) + dy_e = load_vec(copy_atom_io, IO_WIDTH, elem_dtype, dy_div, io_idx_safe) source_local.append(source_e) dy_local.append(dy_e) - source_hat = source_e.to(fx.Float32) * rstd - wdy = dy_e.to(fx.Float32) * gamma_local[tile_i].to(fx.Float32) - prod = (source_hat * wdy).reduce(ReductionOp.ADD, fastmath=fm_fast) - thread_acc = thread_acc + is_valid.select(prod, c_zero_f) - - sum_prod = block_reduce_add(thread_acc) - c1 = sum_prod / n_float - - row_dweight = [] - for tile_i in range_constexpr(NUM_VEC_ITERS): - vec_idx = tid + tile_i * TWO_STAGE_PARTIAL_THREADS - is_valid = vec_idx < NUM_VEC_TILES - vec_idx_safe = is_valid.select(vec_idx, 0) - source_hat = source_local[tile_i].to(fx.Float32) * rstd - dy = dy_local[tile_i].to(fx.Float32) - wdy = dy * gamma_local[tile_i].to(fx.Float32) - dx = (wdy - source_hat * c1) * rstd - if const_expr(fused_add): - dres = load_vec( - copy_atom_vec, - VEC_WIDTH, - elem_dtype, - dres_out_div, - vec_idx_safe, - ).to(fx.Float32) - dx = dx + dres - - if vec_idx < NUM_VEC_TILES: - dx_e = to_elem_vec(dtype_str, elem_dtype, USE_HW_CVT_PK_BF16_F32, dx) - store_vec(copy_atom_vec, VEC_WIDTH, elem_dtype, dx_e, dx_div, vec_idx) - - dw = dy * source_hat - for lane in range_constexpr(VEC_WIDTH): - row_dweight.append(is_valid.select(dw[lane], c_zero_f)) + gamma_e = gamma_local[tile_i] + else: + source_e = load_scalar(copy_atom_io, elem_dtype, source_div, io_idx_safe) + dy_e = load_scalar(copy_atom_io, elem_dtype, dy_div, io_idx_safe) + gamma_e = load_scalar(copy_atom_io, elem_dtype, gamma_div, io_idx_safe) - dweight_partial = dweight_partial + fx.Vector.from_elements(row_dweight, fx.Float32) - gpu.barrier() - - for tile_i in range_constexpr(NUM_VEC_ITERS): - vec_idx = tid + tile_i * TWO_STAGE_PARTIAL_THREADS - if vec_idx < NUM_VEC_TILES: - for lane in range_constexpr(VEC_WIDTH): - idx = vec_idx * VEC_WIDTH + lane - partial_idx = bid * N + idx - store_scalar( - copy_atom_f32, - fx.Float32, - fx.Float32, - partial_div, - partial_idx, - dweight_partial[tile_i * VEC_WIDTH + lane], - ) - else: - gamma_div = fx.logical_divide(Gamma_buf, fx.make_layout(1, 1)) - dweight_partial = fx.Vector.filled(PARTIAL_ACC_SIZE, 0.0, fx.Float32) - - for row in range(fx.Int32(bid), MIn, num_programs): - row_source = fx.slice(Source_buf, (row, None)) - row_dy = fx.slice(DY_buf, (row, None)) - row_dx = fx.slice(DX_buf, (row, None)) - if const_expr(fused_add): - row_dres_out = fx.slice(DResidualOut_buf, (row, None)) + source = source_e if dtype_str == "f32" else source_e.to(fx.Float32) + dy = dy_e if dtype_str == "f32" else dy_e.to(fx.Float32) + gamma = gamma_e if dtype_str == "f32" else gamma_e.to(fx.Float32) + source_hat = source * rstd + wdy = dy * gamma + prod = source_hat * wdy + if const_expr(USE_VEC): + prod = prod.reduce(ReductionOp.ADD, fastmath=fm_fast) + + thread_acc = thread_acc + is_valid.select(prod, c_zero_f) + + sum_prod = block_reduce_add(thread_acc) + c1 = sum_prod / n_float + row_dweight = [] + for tile_i in range_constexpr(NUM_IO_ITERS): + io_idx = tid + tile_i * TWO_STAGE_PARTIAL_THREADS + is_valid = io_idx < NUM_IO_TILES + io_idx_safe = is_valid.select(io_idx, 0) + if const_expr(USE_VEC): + source_e = source_local[tile_i] + dy_e = dy_local[tile_i] + gamma_e = gamma_local[tile_i] + else: + source_e = load_scalar(copy_atom_io, elem_dtype, source_div, io_idx_safe) + dy_e = load_scalar(copy_atom_io, elem_dtype, dy_div, io_idx_safe) + gamma_e = load_scalar(copy_atom_io, elem_dtype, gamma_div, io_idx_safe) + + source = source_e if dtype_str == "f32" else source_e.to(fx.Float32) + dy = dy_e if dtype_str == "f32" else dy_e.to(fx.Float32) + gamma = gamma_e if dtype_str == "f32" else gamma_e.to(fx.Float32) + source_hat = source * rstd + wdy = dy * gamma - source_div = fx.logical_divide(row_source, fx.make_layout(1, 1)) - dy_div = fx.logical_divide(row_dy, fx.make_layout(1, 1)) - dx_div = fx.logical_divide(row_dx, fx.make_layout(1, 1)) + dx = (wdy - source_hat * c1) * rstd if const_expr(fused_add): - dres_out_div = fx.logical_divide(row_dres_out, fx.make_layout(1, 1)) - - rstd = load_scalar(copy_atom_f32, fx.Float32, rstd_div, row) - thread_acc = c_zero_f - for base in range_constexpr(0, N, TWO_STAGE_PARTIAL_THREADS): - idx = tid + base - is_valid = idx < N - idx_safe = is_valid.select(idx, 0) - source_e = load_scalar(copy_atom_s, elem_dtype, source_div, idx_safe) - dy_e = load_scalar(copy_atom_s, elem_dtype, dy_div, idx_safe) - g_e = load_scalar(copy_atom_s, elem_dtype, gamma_div, idx_safe) - source = source_e if dtype_str == "f32" else source_e.to(fx.Float32) - dy = dy_e if dtype_str == "f32" else dy_e.to(fx.Float32) - g = g_e if dtype_str == "f32" else g_e.to(fx.Float32) - source_hat = source * rstd - wdy = dy * g - thread_acc = thread_acc + is_valid.select(source_hat * wdy, c_zero_f) - - sum_prod = block_reduce_add(thread_acc) - c1 = sum_prod / n_float - row_dweight = [] - for base in range_constexpr(0, N, TWO_STAGE_PARTIAL_THREADS): - idx = tid + base - is_valid = idx < N - idx_safe = is_valid.select(idx, 0) - source_e = load_scalar(copy_atom_s, elem_dtype, source_div, idx_safe) - dy_e = load_scalar(copy_atom_s, elem_dtype, dy_div, idx_safe) - g_e = load_scalar(copy_atom_s, elem_dtype, gamma_div, idx_safe) - source = source_e if dtype_str == "f32" else source_e.to(fx.Float32) - dy = dy_e if dtype_str == "f32" else dy_e.to(fx.Float32) - g = g_e if dtype_str == "f32" else g_e.to(fx.Float32) - source_hat = source * rstd - wdy = dy * g - dx = (wdy - source_hat * c1) * rstd - if const_expr(fused_add): - dres_e = load_scalar(copy_atom_s, elem_dtype, dres_out_div, idx_safe) + if const_expr(USE_VEC): + dres = load_vec(copy_atom_io, IO_WIDTH, elem_dtype, dres_out_div, io_idx_safe).to(fx.Float32) + else: + dres_e = load_scalar(copy_atom_io, elem_dtype, dres_out_div, io_idx_safe) dres = dres_e if dtype_str == "f32" else dres_e.to(fx.Float32) - dx = dx + dres + dx = dx + dres - if idx < N: + if io_idx < NUM_IO_TILES: + if const_expr(USE_VEC): + dx_e = to_elem_vec(dtype_str, elem_dtype, USE_HW_CVT_PK_BF16_F32, dx) + store_vec(copy_atom_io, IO_WIDTH, elem_dtype, dx_e, dx_div, io_idx) + else: dx_e = dx if dtype_str == "f32" else dx.to(elem_dtype) - store_scalar(copy_atom_s, elem_dtype, elem_dtype, dx_div, idx, dx_e) - row_dweight.append(is_valid.select(dy * source_hat, c_zero_f)) + store_scalar(copy_atom_io, elem_dtype, elem_dtype, dx_div, io_idx, dx_e) - dweight_partial = dweight_partial + fx.Vector.from_elements(row_dweight, fx.Float32) - gpu.barrier() + dw = dy * source_hat + if const_expr(USE_VEC): + for lane in range_constexpr(IO_WIDTH): + row_dweight.append(is_valid.select(dw[lane], c_zero_f)) + else: + row_dweight.append(is_valid.select(dw, c_zero_f)) + + dweight_partial = dweight_partial + fx.Vector.from_elements(row_dweight, fx.Float32) + gpu.barrier() - for tile_i in range_constexpr(NUM_COL_TILES): - idx = tid + tile_i * TWO_STAGE_PARTIAL_THREADS - if idx < N: + for tile_i in range_constexpr(NUM_IO_ITERS): + io_idx = tid + tile_i * TWO_STAGE_PARTIAL_THREADS + if io_idx < NUM_IO_TILES: + for lane in range_constexpr(IO_WIDTH): + idx = io_idx * IO_WIDTH + lane partial_idx = bid * N + idx + partial_value = dweight_partial[tile_i * IO_WIDTH + lane] store_scalar( copy_atom_f32, fx.Float32, fx.Float32, partial_div, partial_idx, - dweight_partial[tile_i], + partial_value, ) @flyc.kernel From 629e0ee8bf22a2521c8ebf588e4e7906cce61fe3 Mon Sep 17 00:00:00 2001 From: Jin Pan <47354855+jhinpan@users.noreply.github.com> Date: Thu, 16 Jul 2026 21:59:41 +0000 Subject: [PATCH 5/5] [Refactor] Reuse RMSNorm backward compiled launcher cache --- kernels/norm/rmsnorm_kernel.py | 110 ++++++++++++++++++++------------- tests/kernels/test_rmsnorm.py | 16 ++++- 2 files changed, 82 insertions(+), 44 deletions(-) diff --git a/kernels/norm/rmsnorm_kernel.py b/kernels/norm/rmsnorm_kernel.py index 8f004bc7f..d14332e61 100644 --- a/kernels/norm/rmsnorm_kernel.py +++ b/kernels/norm/rmsnorm_kernel.py @@ -1453,11 +1453,13 @@ def build_fused_add_rmsnorm_smoothquant_module( # Python wrappers + autograd (quack-aligned). PR 1: plain rmsnorm. # ===================================================================== if torch is not None: + from kernels.common.tensor_shim import _run_compiled from kernels.norm.rmsnorm_common import torch_dtype_to_str as _torch_dtype_to_str - # Compiled-fn caches. Keys include device: a compiled function is bound to - # the device/context it was built on, so reusing it on another GPU faults. - # eps is a compile-time kernel constant, so it is part of the fwd key too. + # Forward caches retain CompiledFunctions; eps is a compile-time kernel + # constant, so it is part of the key. Backward caches instead retain one + # persistent JitFunction launcher per structural specialization and device; + # _run_compiled owns the launcher's single CompiledFunction in ``._cf``. _FWD_CACHE: dict = {} _BWD_CACHE: dict = {} @@ -1471,13 +1473,13 @@ def build_fused_add_rmsnorm_smoothquant_module( lambda device_index: torch.cuda.current_stream(device_index).cuda_stream, ) - def _launch_cached_on_device(compiled, device, args): + def _run_compiled_on_device(launcher, device, args): device_index = device.index if torch.cuda.current_device() == device_index: - compiled(*args, _get_current_raw_stream(device_index)) + _run_compiled(launcher, *args, _get_current_raw_stream(device_index)) else: with torch.cuda.device(device_index): - compiled(*args, _get_current_raw_stream(device_index)) + _run_compiled(launcher, *args, _get_current_raw_stream(device_index)) # The staged path needs enough rows to amortize its second launch. Its # 512-thread vec8 kernel needs fewer persistent programs than the scalar @@ -1563,14 +1565,15 @@ def rmsnorm_bwd(x, weight, dout, rstd, eps=EPS): dweight = torch.empty_like(weight) partial = torch.empty((num_programs * N,), device=device, dtype=torch.float32) key = (path, N, dtype_str, num_programs, device) - compiled = _BWD_CACHE.get(key) - if compiled is None: - # Compilation and code-object loading are device-context bound. + launcher = _BWD_CACHE.get(key) + if launcher is None: + # The builder has an arch-dependent vec conversion choice, and + # compilation/code-object loading are device-context bound. with torch.cuda.device(device): - stream = torch.cuda.current_stream() - launch_fn = build_rmsnorm_bwd_two_stage_module(N, dtype_str, num_programs) - compiled = flyc.compile( - launch_fn, + device_index = device.index + launcher = build_rmsnorm_bwd_two_stage_module(N, dtype_str, num_programs) + _run_compiled( + launcher, x, weight, dout, @@ -1579,25 +1582,35 @@ def rmsnorm_bwd(x, weight, dout, rstd, eps=EPS): dweight, partial, M, - stream, + _get_current_raw_stream(device_index), ) - _BWD_CACHE[key] = compiled + _BWD_CACHE[key] = launcher else: - _launch_cached_on_device(compiled, device, (x, weight, dout, rstd, dx, dweight, partial, M)) + _run_compiled_on_device(launcher, device, (x, weight, dout, rstd, dx, dweight, partial, M)) return dx, dweight dweight = torch.empty((N,), device=device, dtype=torch.float32) - key = (path, N, dtype_str, device) - compiled = _BWD_CACHE.get(key) + key = (path, N, dtype_str, num_programs, device) + launcher = _BWD_CACHE.get(key) dweight.zero_() - if compiled is None: + if launcher is None: with torch.cuda.device(device): - stream = torch.cuda.current_stream() - launch_fn = build_rmsnorm_bwd_module(N, dtype_str) - compiled = flyc.compile(launch_fn, x, weight, dout, rstd, dx, dweight, M, stream) - _BWD_CACHE[key] = compiled + device_index = device.index + launcher = build_rmsnorm_bwd_module(N, dtype_str) + _run_compiled( + launcher, + x, + weight, + dout, + rstd, + dx, + dweight, + M, + _get_current_raw_stream(device_index), + ) + _BWD_CACHE[key] = launcher else: - _launch_cached_on_device(compiled, device, (x, weight, dout, rstd, dx, dweight, M)) + _run_compiled_on_device(launcher, device, (x, weight, dout, rstd, dx, dweight, M)) return dx, dweight.to(weight.dtype) class RMSNormFunction(torch.autograd.Function): @@ -1718,13 +1731,13 @@ def fused_add_rmsnorm_bwd(added, weight, dout, rstd, dresidual_out=None, eps=EPS dweight = torch.empty_like(weight) partial = torch.empty((num_programs * N,), device=device, dtype=torch.float32) key = (path, N, dtype_str, num_programs, device) - compiled = _FUSED_ADD_BWD_CACHE.get(key) - if compiled is None: + launcher = _FUSED_ADD_BWD_CACHE.get(key) + if launcher is None: with torch.cuda.device(device): - stream = torch.cuda.current_stream() - launch_fn = build_fused_add_rmsnorm_bwd_two_stage_module(N, dtype_str, num_programs) - compiled = flyc.compile( - launch_fn, + device_index = device.index + launcher = build_fused_add_rmsnorm_bwd_two_stage_module(N, dtype_str, num_programs) + _run_compiled( + launcher, added, weight, dout, @@ -1734,30 +1747,41 @@ def fused_add_rmsnorm_bwd(added, weight, dout, rstd, dresidual_out=None, eps=EPS dweight, partial, M, - stream, + _get_current_raw_stream(device_index), ) - _FUSED_ADD_BWD_CACHE[key] = compiled + _FUSED_ADD_BWD_CACHE[key] = launcher else: - _launch_cached_on_device( - compiled, + _run_compiled_on_device( + launcher, device, (added, weight, dout, dresidual_out, rstd, dx, dweight, partial, M), ) return dx, dx, dweight dweight = torch.empty((N,), device=device, dtype=torch.float32) - key = (path, N, dtype_str, device) - compiled = _FUSED_ADD_BWD_CACHE.get(key) + key = (path, N, dtype_str, num_programs, device) + launcher = _FUSED_ADD_BWD_CACHE.get(key) dweight.zero_() - if compiled is None: + if launcher is None: with torch.cuda.device(device): - stream = torch.cuda.current_stream() - launch_fn = build_fused_add_rmsnorm_bwd_module(N, dtype_str) - compiled = flyc.compile(launch_fn, added, weight, dout, dresidual_out, rstd, dx, dweight, M, stream) - _FUSED_ADD_BWD_CACHE[key] = compiled + device_index = device.index + launcher = build_fused_add_rmsnorm_bwd_module(N, dtype_str) + _run_compiled( + launcher, + added, + weight, + dout, + dresidual_out, + rstd, + dx, + dweight, + M, + _get_current_raw_stream(device_index), + ) + _FUSED_ADD_BWD_CACHE[key] = launcher else: - _launch_cached_on_device( - compiled, + _run_compiled_on_device( + launcher, device, (added, weight, dout, dresidual_out, rstd, dx, dweight, M), ) diff --git a/tests/kernels/test_rmsnorm.py b/tests/kernels/test_rmsnorm.py index afe2b792e..5fbb1e8e9 100644 --- a/tests/kernels/test_rmsnorm.py +++ b/tests/kernels/test_rmsnorm.py @@ -1424,20 +1424,28 @@ def test_rmsnorm_bwd_two_stage_cache_reuse_across_m(monkeypatch, fused_add): dtype = DTYPE_BF16 N = 512 first_m, second_m = 1024, 1025 - path, _ = rmsnorm_kernel_impl._select_rmsnorm_bwd_config(first_m, N, "bf16", device) + path, num_programs = rmsnorm_kernel_impl._select_rmsnorm_bwd_config(first_m, N, "bf16", device) assert path == "two_stage" builder_name = "build_fused_add_rmsnorm_bwd_two_stage_module" if fused_add else "build_rmsnorm_bwd_two_stage_module" cache = rmsnorm_kernel_impl._FUSED_ADD_BWD_CACHE if fused_add else rmsnorm_kernel_impl._BWD_CACHE + key = (path, N, "bf16", num_programs, device) original_builder = getattr(rmsnorm_kernel_impl, builder_name) build_count = 0 + run_compiled_count = 0 def counted_builder(*args, **kwargs): nonlocal build_count build_count += 1 return original_builder(*args, **kwargs) + def counted_run_compiled(*args, **kwargs): + nonlocal run_compiled_count + run_compiled_count += 1 + return _run_compiled(*args, **kwargs) + monkeypatch.setattr(rmsnorm_kernel_impl, builder_name, counted_builder) + monkeypatch.setattr(rmsnorm_kernel_impl, "_run_compiled", counted_run_compiled) saved_cache = cache.copy() cache.clear() @@ -1463,8 +1471,13 @@ def run(M, seed): try: run(first_m, 11) + launcher = cache[key] + compiled = launcher._cf run(second_m, 12) assert build_count == 1 + assert run_compiled_count == 2 + assert cache[key] is launcher + assert launcher._cf is compiled finally: cache.clear() cache.update(saved_cache) @@ -1562,6 +1575,7 @@ def test_rmsnorm_multi_gpu(): y = rmsnorm(x, w) y.backward(dy) torch.cuda.synchronize(dev) + assert torch.cuda.current_device() == 0 # Autograd warmed this device's backward cache. Exercise the hot # launch again while cuda:0 remains the caller's current device;