From 31c068284edfd8026c0f722b2216c257c26ab862 Mon Sep 17 00:00:00 2001 From: Zhipeng Date: Sun, 16 Aug 2026 05:49:54 +0000 Subject: [PATCH] perf(cutedsl): add Blackwell fast path to RMSNorm RMSNorm was the only generic CuTe DSL op with no Blackwell-specific path. swiglu has packed-f32x2 SFU math and a TVM-FFI direct-call path, rope has a TMA kernel tuned on B200, and cross_entropy dispatches on infer_device_arch() == "blackwell". rms_norm.py had no arch query at all. Its kernels were not the problem. On a B200 the vector forward already beat Triton on device time (5.8us vs 7.8us at bf16 4096x2048), as did the fused backward (26.3us vs 32.4us). Wall clock said the opposite because every call first marshalled its tensors into cute.Tensor handles (from_dlpack + memref construction), costing ~48us per launch. The op was host-bound at every shape up to 16384x4096: measured wall time tracked host time (~54us), not the kernel. Two changes: * Compile the vector forward and fused backward with --enable-tvm-ffi against abstract tensors, so PyTorch tensors are passed straight to the compiled function. eps/offset stay runtime scalars, so neither forces a recompile. Forward host cost 54us -> 15us, backward 74us -> 26us. This is a calling convention, not an arch feature, so it applies wherever apache-tvm-ffi is importable; without it the original marshalling launch still runs. * Use packed-f32x2 SFU math (fma/mul/add) for the fused backward's dX/dW inner products, which issue two fp32 lanes per instruction. Gated to sm_10x and to n_cols > 4096, the 16-warp regime where the kernel is issue-bound rather than memory-bound; measurements at each width are recorded in _use_packed_math. The forward deliberately stays scalar: it is DRAM-bound (6.6 TB/s at 16384x8192) and the packed variant measured within +-0.5% at every width and dtype, so it would be complexity that buys nothing. Both halves are switchable for A/B via LIGER_RMS_FORCE_NO_FFI and LIGER_RMS_FORCE_NO_PACKED. Measured on B200, bf16, llama mode, op-level (cutedsl vs Triton): shape fwd before fwd after bwd before bwd after 1024x2048 0.65x 2.36x 0.81x 2.36x 4096x2048 0.66x 2.32x 0.81x 2.29x 4096x4096 0.69x 2.33x 0.82x 1.51x 4096x8192 0.66x 1.52x 1.08x 1.17x 16384x2048 0.67x 1.65x 1.73x 1.75x 16384x4096 0.90x 1.12x 1.75x 1.75x 16384x8192 1.24x 1.24x 1.17x 1.29x 32768x4096 1.11x 1.12x 1.84x 1.83x fp32 moves the same way (forward 0.66-1.08x -> 0.99-2.29x, backward 0.80-1.73x -> 1.15-2.58x). No shape regressed in either dtype. Adds apache-tvm-ffi to the cutedsl extra. It is optional at runtime, but the whole forward win depends on it, and the existing swiglu kernel already has the same latent dependency. Tests: three new cases cover ground the parity suite cannot reach -- it only goes up to 4096 wide, so it never compiles the packed backward, and it always runs on the default stream. The FFI path is checked bit-identical to the marshalling path (affine and non-affine, the latter pinning the dummy-W layout), the packed backward against the scalar one, and stream binding via CUDA graph capture. Both comparison tests assert that the two variants were really compiled, so they cannot degenerate into comparing a path with itself. Full cutedsl suite on B200: 579 passed, identical 33-failure set before and after (pre-existing bf16/fp32 parity gaps in cross_entropy and RMSNorm's bf16-llama dW), so this introduces no new failures. --- README.md | 2 + setup.py | 5 + src/liger_kernel/ops/cutedsl/ops/rms_norm.py | 463 ++++++++++++++++++- test/cutedsl/test_rms_norm.py | 135 ++++++ 4 files changed, 591 insertions(+), 14 deletions(-) diff --git a/README.md b/README.md index 5e3fce335..1acc11f50 100644 --- a/README.md +++ b/README.md @@ -219,6 +219,8 @@ LIGER_KERNEL_IMPL=cutedsl python your_script.py It currently provides genuine `cutlass.cute` implementations of **RMSNorm**, **cross entropy**, and **fused scaled cross entropy**. Ops without a CuTe DSL kernel transparently fall back to the default Triton kernel. +The `cutedsl` extra also pulls in `apache-tvm-ffi`, which lets compiled kernels take PyTorch tensors directly rather than marshalling each one per call. It is optional — every kernel falls back to the marshalling launch without it — but short kernels are dominated by that per-call cost, so installing it is strongly recommended. + ### Fused Scaled Cross Entropy `LigerFusedLinearScaledCrossEntropyFunction` is an additional per-token operator, not a replacement for the reduction-oriented Triton `LigerFusedLinearCrossEntropyFunction`. It takes `input[M, H]`, `weight[V, H]`, and `target[M]`, applies `logits / temperature`, and returns FP32 negative log-likelihood `[M]` plus optional differentiable vocabulary entropy `[M]` in the input dtype. Reductions remain in PyTorch, and rows whose target equals `ignore_index` contribute zero outputs and gradients. diff --git a/setup.py b/setup.py index c9f03e7f1..a680d1121 100644 --- a/setup.py +++ b/setup.py @@ -39,6 +39,11 @@ def get_optional_dependencies(): ] cutedsl_deps = [ "nvidia-cutlass-dsl>=4.6.0", + # Lets compiled CuTe DSL kernels take PyTorch tensors directly instead of + # marshalling each one through DLPack per call. The kernels fall back to + # the marshalling launch when it is absent, but on short kernels that + # per-call cost dominates: RMSNorm forward measured 53us -> 15us on B200. + "apache-tvm-ffi>=0.1.0", ] dev_deps = [ "transformers>=4.52.0", diff --git a/src/liger_kernel/ops/cutedsl/ops/rms_norm.py b/src/liger_kernel/ops/cutedsl/ops/rms_norm.py index ed32e9303..8acd0483d 100644 --- a/src/liger_kernel/ops/cutedsl/ops/rms_norm.py +++ b/src/liger_kernel/ops/cutedsl/ops/rms_norm.py @@ -37,17 +37,22 @@ import cuda.bindings.driver as cuda import cutlass import cutlass.cute as cute +import cutlass.cute.arch as carch import cutlass.utils import torch from cutlass import Float32 from cutlass import Int32 from cutlass import const_expr +from cutlass.cute.runtime import make_fake_stream from liger_kernel.ops.cutedsl.ops.rms_norm_fastpath import backward_warp_count from liger_kernel.ops.cutedsl.ops.rms_norm_fastpath import fast_path_vector_width from liger_kernel.ops.cutedsl.ops.rms_norm_fastpath import fwd_warp_count +from liger_kernel.ops.cutedsl.ops.utils import make_fake_tensor from liger_kernel.ops.cutedsl.ops.utils import to_cute_tensor +from liger_kernel.ops.cutedsl.ops.utils import torch2cute_dtype_map +from liger_kernel.utils import infer_device_arch # --------------------------------------------------------------------------- # Tuning / debug env vars are read ONCE at import. These are launch-time knobs @@ -65,12 +70,157 @@ _FORCE_SPLIT_BWD = bool(int(os.environ.get("LIGER_RMS_FORCE_SPLIT_BWD") or 0)) _AUTOTUNE_FILE = os.environ.get("LIGER_RMS_AUTOTUNE_FILE") or None _FUSED_STRIP_MULT = int(os.environ.get("LIGER_RMS_FUSED_STRIP_MULT") or 0) +# A/B switches for the two halves of the Blackwell fast path (see below). +_FORCE_NO_FFI = bool(int(os.environ.get("LIGER_RMS_FORCE_NO_FFI") or 0)) +_FORCE_NO_PACKED = bool(int(os.environ.get("LIGER_RMS_FORCE_NO_PACKED") or 0)) try: _BACKWARD_WARPS = int(os.environ.get("LIGER_RMS_BACKWARD_WARPS") or 0) or None except Exception: _BACKWARD_WARPS = None +# --------------------------------------------------------------------------- +# Blackwell fast path +# --------------------------------------------------------------------------- +# The vector kernels are short: at bf16 4096x2048 the forward runs in ~5.8us and +# the fused backward in ~26us on a B200 -- both already faster than Triton's +# device time (7.8us / 32.4us). Wall-clock told the opposite story, because the +# default launch path pays ~48us per call marshalling torch tensors into +# cute.Tensor handles (``from_dlpack`` + memref construction) before it can even +# enqueue. The op was therefore *host*-bound at every shape up to 16384x4096: +# measured wall time tracked host time (~54us), not kernel time. +# +# This path removes that tax the same way the CuTe DSL swiglu kernel does: +# +# * ``--enable-tvm-ffi`` -- the kernel is compiled once per (dtype, geometry) +# against *abstract* tensors, so PyTorch tensors are passed straight to the +# compiled function with no per-call ``from_dlpack``. This is a calling +# convention, independent of the GPU arch, so it is used wherever the +# optional ``apache-tvm-ffi`` package is importable. +# * packed-f32x2 SFU math in the fused backward -- Blackwell processes two +# fp32 lanes per instruction (``fma_packed_f32x2`` / ``mul_packed_f32x2`` / +# ``add_packed_f32x2``), halving the issued math for the dX/dW inner +# products. Gated on sm_10x, checked per launch from the *input tensor's* +# device, and narrowed further by ``_use_packed_math`` (below) to the row +# widths where it measured faster. +# +# Both halves fall back cleanly: without ``tvm_ffi`` the original marshalling +# launch runs, and off Blackwell the scalar math runs. +# +# The forward deliberately does NOT use packed math: it is DRAM-bound (measured +# 6.6 TB/s at 16384x8192, ~85% of B200 peak) and the packed variant measured +# within +-0.5% -- noise -- at every width and dtype tried, so it would be +# complexity that buys nothing. + +# Static probe: is the optional ``apache-tvm-ffi`` package importable? +try: + import tvm_ffi # noqa: F401 + + _TVM_FFI_PRESENT = True +# NOTE: must be a single exception name, not a tuple. The CuteDSL AST +# preprocessor parses this module's source and only handles ``ast.Name`` / bare +# except handlers (``handler.type.id``); a tuple ``except (A, B):`` raises +# AttributeError at compile. +except Exception: # pragma: no cover - depends on optional deps + _TVM_FFI_PRESENT = False + +# Static probe: are the packed-f32x2 SFU symbols importable? They exist in the +# CUTLASS wheel on every arch, so presence does NOT imply usability -- the +# compute capability is checked per launch by ``_is_blackwell``. +try: + _ = (carch.mul_packed_f32x2, carch.add_packed_f32x2, carch.fma_packed_f32x2) + _PACKED_OPS_PRESENT = True +except Exception: # pragma: no cover - depends on optional deps + _PACKED_OPS_PRESENT = False + + +def _is_blackwell(device=None) -> bool: + """Whether ``device`` is a **data-center Blackwell** (sm_100 / sm_103). + + Gated to the sm_10x data-center parts only -- B200 (sm_100) and B300 + (sm_103). Hopper (sm_90) and *consumer* Blackwell (RTX 50xx, sm_120) do not + get the packed math: the packed-f32x2 SFU ops are a data-center-Blackwell + feature (their Python symbols exist on every arch, so presence != usability) + and those parts run the scalar path instead. Checked per launch from the + *input tensor's* device: a process may hold tensors on GPUs of different + archs, and ``torch.cuda.set_device`` can move the current device after + import. Mirrors ``_is_blackwell`` in the CuTe DSL swiglu op. + """ + try: + if not torch.cuda.is_available(): + return False + if isinstance(device, torch.device): + device_id = device.index if device.index is not None else torch.cuda.current_device() + elif device is None: + device_id = torch.cuda.current_device() + else: + device_id = device + return infer_device_arch(device_id) in ("blackwell", "blackwell_ultra") + except Exception: # pragma: no cover - no CUDA / bad device + return False + + +def _use_packed_math(device, n_cols: int, vec: int) -> bool: + """Whether to compile the packed-f32x2 fused backward for this launch. + + Blackwell-only, and only in the wide-row regime. The backward's warp count + is width-driven (``backward_warp_count``: 16 warps above 4096, else 8 or 4), + and the win tracks that split -- above 4096 each thread holds several + register-resident vector tiles and the kernel becomes issue-bound, which is + exactly what halving the math instructions relieves. At or below 4096 the + same kernel is already memory-bound and the packed form is neutral or a + slight loss. + + Measured on a B200, fused backward device time, 16384 rows (scalar -> + packed):: + + bf16 H=1024 71.4 -> 73.2 us (+2.5%) + bf16 H=2048 84.7 -> 86.1 us (+1.7%) + bf16 H=4096 110.7 -> 107.7 us (-2.7%) + bf16 H=6144 218.7 -> 201.6 us (-7.8%) + bf16 H=8192 245.8 -> 222.0 us (-9.7%) + fp32 H=6144 251.3 -> 194.6 us (-22.6%) + fp32 H=8192 297.9 -> 296.4 us (-0.5%) + + ``vec`` must be even for the pairwise loop; every supported dtype gives an + even vector width, so this only ever rejects a hypothetical odd one. + """ + if not _PACKED_OPS_PRESENT or _FORCE_NO_PACKED or vec % 2: + return False + return n_cols > 4096 and _is_blackwell(device) + + +def _use_ffi() -> bool: + """Whether to use the TVM-FFI direct-call path (no per-call marshalling).""" + return _TVM_FFI_PRESENT and not _FORCE_NO_FFI + + +# Compiled TVM-FFI callables, keyed on everything the specialization bakes. +_ffi_compile_cache: dict = {} + + +def _ffi_fake(dtype: torch.dtype, shape, divisibility: int): + """Abstract tensor for ``cute.compile``; ``dtype`` is a torch dtype.""" + return make_fake_tensor(torch2cute_dtype_map[dtype], shape, divisibility) + + +def _ffi_fake_stream(): + """Placeholder stream that resolves to TVM-FFI's env stream at call time. + + Compiling against it drops the stream parameter from the FFI signature + entirely: the compiled function reads the caller's environment stream + instead, so the launch path never queries or marshals a stream from Python. + + That environment stream tracks ``torch.cuda.current_stream()``, so no + host-side stream plumbing is needed here -- and adding any would be pure + overhead. Verified from profiler stream ids: launching under the default + stream, and under two different ``torch.cuda.stream(...)`` contexts, places + the kernel on three correspondingly different streams, and the op captures + into a CUDA graph and replays correctly. + """ + return make_fake_stream(use_tvm_ffi_env_stream=True) + + # Lightweight debug logger controlled by env var LIGER_RMS_DEBUG def _rms_debug(msg): if _DEBUG: @@ -543,6 +693,7 @@ def _rms_norm_bwd_fused_vector_kernel( NUM_VEC_TILES: cutlass.Constexpr, NUM_THREADS: cutlass.Constexpr, NUM_WARPS: cutlass.Constexpr, + PACKED_MATH: cutlass.Constexpr = False, ): """Aligned affine backward matching Triton's persistent execution shape. @@ -625,14 +776,32 @@ def _rms_norm_bwd_fused_vector_kernel( cute.autovec_copy(gdYv[None, vec_idx], dy_frags[None, ct]) dot = Float32(0.0) + # Two independent fp32 accumulators feed the packed-f32x2 FFMA chain. + dot0 = Float32(0.0) + dot1 = Float32(0.0) for ct in cutlass.range_constexpr(NUM_VEC_TILES): vec_idx = ct * NUM_THREADS + tid if r_valid and vec_idx < n_vec: - dot = dot + ( - x_frags[None, ct].load().to(Float32) - * dy_frags[None, ct].load().to(Float32) - * (w_frags[None, ct].load().to(Float32) + offset) - ).reduce(cute.ReductionOp.ADD, Float32(0.0), 0) + if const_expr(PACKED_MATH): + xf = x_frags[None, ct] + dyf = dy_frags[None, ct] + wf = w_frags[None, ct] + for i in cutlass.range_constexpr(0, VEC, 2): + m0 = wf[i].to(Float32) + m1 = wf[i + 1].to(Float32) + m0, m1 = carch.add_packed_f32x2((m0, m1), (offset, offset)) + m0, m1 = carch.mul_packed_f32x2((dyf[i].to(Float32), dyf[i + 1].to(Float32)), (m0, m1)) + dot0, dot1 = carch.fma_packed_f32x2( + (xf[i].to(Float32), xf[i + 1].to(Float32)), (m0, m1), (dot0, dot1) + ) + else: + dot = dot + ( + x_frags[None, ct].load().to(Float32) + * dy_frags[None, ct].load().to(Float32) + * (w_frags[None, ct].load().to(Float32) + offset) + ).reduce(cute.ReductionOp.ADD, Float32(0.0), 0) + if const_expr(PACKED_MATH): + dot = dot0 + dot1 dot_total = _cta_reduce_sum_warp0(dot, sm_warp, sm_result, lane, warp, NUM_WARPS) coef = (Float32(0.0) - rstd * rstd * dot_total) / Float32(N_COLS) @@ -640,15 +809,43 @@ def _rms_norm_bwd_fused_vector_kernel( for ct in cutlass.range_constexpr(NUM_VEC_TILES): vec_idx = ct * NUM_THREADS + tid if r_valid and vec_idx < n_vec: - xf = x_frags[None, ct].load().to(Float32) - dyf = dy_frags[None, ct].load().to(Float32) - mk = dyf * (w_frags[None, ct].load().to(Float32) + offset) - dx_frag.store((rstd * (mk + coef * xf)).to(mdX.element_type)) - cute.autovec_copy(dx_frag, gdXv[None, vec_idx]) - xhat = xf * rstd - if const_expr(CASTING_MODE == _CASTING_MODE_LLAMA): - xhat = xhat.to(mX.element_type).to(Float32) - dw_acc[None, ct].store(dw_acc[None, ct].load() + dyf * xhat) + if const_expr(PACKED_MATH): + xr = x_frags[None, ct] + dyr = dy_frags[None, ct] + wr = w_frags[None, ct] + ar = dw_acc[None, ct] + for i in cutlass.range_constexpr(0, VEC, 2): + x0 = xr[i].to(Float32) + x1 = xr[i + 1].to(Float32) + dy0 = dyr[i].to(Float32) + dy1 = dyr[i + 1].to(Float32) + # m = dy * (w + offset) + m0, m1 = carch.add_packed_f32x2((wr[i].to(Float32), wr[i + 1].to(Float32)), (offset, offset)) + m0, m1 = carch.mul_packed_f32x2((dy0, dy1), (m0, m1)) + # dx = rstd * (m + coef * x) + t0, t1 = carch.fma_packed_f32x2((coef, coef), (x0, x1), (m0, m1)) + t0, t1 = carch.mul_packed_f32x2((t0, t1), (rstd, rstd)) + dx_frag[i] = t0.to(mdX.element_type) + dx_frag[i + 1] = t1.to(mdX.element_type) + # dw += dy * xhat, xhat = x * rstd + h0, h1 = carch.mul_packed_f32x2((x0, x1), (rstd, rstd)) + if const_expr(CASTING_MODE == _CASTING_MODE_LLAMA): + h0 = h0.to(mX.element_type).to(Float32) + h1 = h1.to(mX.element_type).to(Float32) + a0, a1 = carch.fma_packed_f32x2((dy0, dy1), (h0, h1), (ar[i], ar[i + 1])) + ar[i] = a0 + ar[i + 1] = a1 + cute.autovec_copy(dx_frag, gdXv[None, vec_idx]) + else: + xf = x_frags[None, ct].load().to(Float32) + dyf = dy_frags[None, ct].load().to(Float32) + mk = dyf * (w_frags[None, ct].load().to(Float32) + offset) + dx_frag.store((rstd * (mk + coef * xf)).to(mdX.element_type)) + cute.autovec_copy(dx_frag, gdXv[None, vec_idx]) + xhat = xf * rstd + if const_expr(CASTING_MODE == _CASTING_MODE_LLAMA): + xhat = xhat.to(mX.element_type).to(Float32) + dw_acc[None, ct].store(dw_acc[None, ct].load() + dyf * xhat) # Emit this strip's dW partials; the host sums the num_strips partial rows. dw_row = mdW[strip, None] @@ -790,6 +987,7 @@ def _rms_norm_bwd_fused_host( NUM_THREADS: cutlass.Constexpr, NUM_WARPS: cutlass.Constexpr, SMEM_BYTES: cutlass.Constexpr, + PACKED_MATH: cutlass.Constexpr = False, stream: cuda.CUstream = None, ): num_strips = mdW.shape[0] @@ -807,6 +1005,7 @@ def _rms_norm_bwd_fused_host( NUM_VEC_TILES, NUM_THREADS, NUM_WARPS, + PACKED_MATH, ).launch( grid=[num_strips, 1, 1], block=[NUM_THREADS, 1, 1], @@ -815,6 +1014,199 @@ def _rms_norm_bwd_fused_host( ) +# ============================================================================= +# TVM-FFI specializations (direct call, no per-call tensor marshalling) +# ============================================================================= +# Every constexpr is captured by the closure rather than passed as an argument, +# so the compiled function's runtime signature is just tensors + fp32 scalars -- +# which is what lets PyTorch tensors be handed straight to it. ``eps`` / ``offset`` +# stay runtime scalars so a change of either does not force a recompile. +def _make_fwd_vector_ffi(casting_mode, elementwise_affine, n_cols, vec, num_vec_tiles, num_warps, num_threads): + smem_bytes = (((num_warps + 1) * 4 + 15) // 16) * 16 + + @cute.jit + def fwd( + mX: cute.Tensor, + mW: cute.Tensor, + mY: cute.Tensor, + mRSTD: cute.Tensor, + eps: Float32, + offset: Float32, + stream: cuda.CUstream = None, + ): + _rms_norm_fwd_vector_kernel( + mX, + mW, + mY, + mRSTD, + eps, + offset, + casting_mode, + elementwise_affine, + n_cols, + vec, + num_vec_tiles, + num_warps, + num_threads, + ).launch( + grid=[mX.shape[0], 1, 1], + block=[num_threads, 1, 1], + smem=smem_bytes, + stream=stream, + ) + + return fwd + + +def _make_bwd_fused_ffi(casting_mode, n_cols, vec, num_vec_tiles, num_threads, num_warps, packed): + smem_bytes = (((num_warps + 1) * 4 + 15) // 16) * 16 + + @cute.jit + def bwd( + mdY: cute.Tensor, + mX: cute.Tensor, + mW: cute.Tensor, + mRSTD: cute.Tensor, + mdX: cute.Tensor, + mdW: cute.Tensor, + offset: Float32, + stream: cuda.CUstream = None, + ): + _rms_norm_bwd_fused_vector_kernel( + mdY, + mX, + mW, + mRSTD, + mdX, + mdW, + offset, + casting_mode, + n_cols, + vec, + num_vec_tiles, + num_threads, + num_warps, + packed, + ).launch( + grid=[mdW.shape[0], 1, 1], + block=[num_threads, 1, 1], + smem=smem_bytes, + stream=stream, + ) + + return bwd + + +def _get_fwd_vector_ffi( + x_dtype, + w_dtype, + device_index, + casting_mode, + elementwise_affine, + n_cols, + vec, + num_vec_tiles, + num_warps, + num_threads, +): + """Compile (once per specialization) the TVM-FFI vector forward.""" + key = ( + "fwd_vec_ffi", + x_dtype, + w_dtype, + device_index, + casting_mode, + elementwise_affine, + n_cols, + vec, + num_vec_tiles, + num_warps, + num_threads, + ) + fn = _ffi_compile_cache.get(key) + if fn is not None: + return fn + # X / Y: dynamic row count, static width, ``vec``-divisible row stride (the + # host has already validated the 16-byte base and row alignment this bakes in). + fake_x = _ffi_fake(x_dtype, (cute.sym_int(), n_cols), vec) + fake_y = _ffi_fake(x_dtype, (cute.sym_int(), n_cols), vec) + fake_rstd = _ffi_fake(torch.float32, (cute.sym_int(),), 1) + # Non-affine bakes a dummy fp32 vector for W and is handed RSTD at call time; + # ELEMENTWISE_AFFINE is constexpr-False so the kernel never reads it. + fake_w = _ffi_fake(w_dtype, (n_cols,), vec) if elementwise_affine else fake_rstd + fn = cute.compile( + _make_fwd_vector_ffi(casting_mode, elementwise_affine, n_cols, vec, num_vec_tiles, num_warps, num_threads), + fake_x, + fake_w, + fake_y, + fake_rstd, + Float32(0.0), + Float32(0.0), + _ffi_fake_stream(), + options="--enable-tvm-ffi", + ) + _ffi_compile_cache[key] = fn + if _DEBUG: + _rms_debug(f"Compiled TVM-FFI fwd kernel for key: {key}") + return fn + + +def _get_bwd_fused_ffi( + x_dtype, + dy_dtype, + w_dtype, + device_index, + casting_mode, + n_cols, + vec, + num_vec_tiles, + num_threads, + num_warps, + packed, +): + """Compile (once per specialization) the TVM-FFI fused affine backward.""" + key = ( + "bwd_fused_ffi", + x_dtype, + dy_dtype, + w_dtype, + device_index, + casting_mode, + n_cols, + vec, + num_vec_tiles, + num_threads, + num_warps, + packed, + ) + fn = _ffi_compile_cache.get(key) + if fn is not None: + return fn + fake_dy = _ffi_fake(dy_dtype, (cute.sym_int(), n_cols), vec) + fake_x = _ffi_fake(x_dtype, (cute.sym_int(), n_cols), vec) + fake_w = _ffi_fake(w_dtype, (n_cols,), vec) + fake_rstd = _ffi_fake(torch.float32, (cute.sym_int(),), 1) + fake_dx = _ffi_fake(dy_dtype, (cute.sym_int(), n_cols), vec) + # dW partials are fp32 (num_strips, n_cols); num_strips varies with the launch. + fake_dw = _ffi_fake(torch.float32, (cute.sym_int(), n_cols), 4) + fn = cute.compile( + _make_bwd_fused_ffi(casting_mode, n_cols, vec, num_vec_tiles, num_threads, num_warps, packed), + fake_dy, + fake_x, + fake_w, + fake_rstd, + fake_dx, + fake_dw, + Float32(0.0), + _ffi_fake_stream(), + options="--enable-tvm-ffi", + ) + _ffi_compile_cache[key] = fn + if _DEBUG: + _rms_debug(f"Compiled TVM-FFI fused bwd kernel for key: {key}") + return fn + + def _is_16b_row_aligned(t): """Whether a contiguous 1D/2D tensor can safely use 16-byte row vectors.""" if t is None or t.data_ptr() % 16 or t.stride(-1) != 1: @@ -850,6 +1242,27 @@ def _launch_fwd_vector(X, W, Y, RSTD, eps, offset, casting_mode, elementwise_aff The warp count (hence thread count and register-resident tiles per thread) is chosen from the hidden width by ``fwd_warp_count``; ``num_vec_tiles`` follows. """ + n_cols_ = X.shape[1] + if _use_ffi(): + num_warps = fwd_warp_count(n_cols_, vec) + num_threads = 32 * num_warps + fn = _get_fwd_vector_ffi( + X.dtype, + W.dtype if elementwise_affine else None, + X.device.index, + casting_mode, + elementwise_affine, + n_cols_, + vec, + (n_cols_ // vec + num_threads - 1) // num_threads, + num_warps, + num_threads, + ) + # Non-affine passes RSTD as the dummy W (fp32, matching the baked layout). + w_arg = W if elementwise_affine else RSTD + fn(X, w_arg, Y, RSTD, Float32(float(eps)), Float32(float(offset))) + return + stream = _cute_stream() # Cache the marshaled handles for the INPUTS (X, W) -- their addresses are stable # across steps (weights always; activations under a reused-buffer harness), so they @@ -1010,6 +1423,26 @@ def _launch_bwd_fused( num_warps, ): """Launch the aligned register-resident affine backward specialization.""" + packed = _use_packed_math(X.device, X.shape[1], vec) + if _use_ffi(): + n_cols_ = X.shape[1] + num_threads_ = 32 * num_warps + fn = _get_bwd_fused_ffi( + X.dtype, + dY.dtype, + W.dtype, + X.device.index, + casting_mode, + n_cols_, + vec, + num_vec_tiles, + num_threads_, + num_warps, + packed, + ) + fn(dY, X, W, RSTD, dX, dW_partial, Float32(float(offset))) + return + stream = _cute_stream() dy_ct = _to_cute_cached(dY, assumed_align=16) x_ct = _to_cute_cached(X, assumed_align=16) @@ -1045,6 +1478,7 @@ def _launch_bwd_fused( dY.dtype, W.dtype, casting_mode, + packed, ) if _DEBUG: _rms_debug(f"_launch_bwd_fused reload_policy={reload_policy}") @@ -1080,6 +1514,7 @@ def _launch_bwd_fused( num_threads, num_warps, smem_bytes, + packed, stream, ) if _DEBUG: diff --git a/test/cutedsl/test_rms_norm.py b/test/cutedsl/test_rms_norm.py index b71268b44..6947c616b 100644 --- a/test/cutedsl/test_rms_norm.py +++ b/test/cutedsl/test_rms_norm.py @@ -148,3 +148,138 @@ def test_rms_norm_mixed_weight_dtype_no_cache_collision(): torch.testing.assert_close(y_cd, y_tr, atol=atol, rtol=rtol) torch.testing.assert_close(dx_cd, dx_tr, atol=atol, rtol=rtol) torch.testing.assert_close(dw_cd, dw_tr, atol=atol, rtol=rtol) + + +# ============================================================================= +# Blackwell fast path (TVM-FFI direct call + packed-f32x2 fused backward) +# ============================================================================= +# These exercise machinery the parity suite above cannot reach: it only covers +# hidden widths <= 4096, so it never compiles the packed-f32x2 backward, and it +# always runs on the default stream. +def _rms_norm_mod(): + """The cutedsl rms_norm module, or skip if CUTLASS isn't installed.""" + try: + from liger_kernel.ops.cutedsl.ops import rms_norm as mod + except ImportError as exc: + pytest.skip(f"cutedsl backend not importable (cutlass.cute missing?): {exc}") + return mod + + +def _fwd_bwd(mod, X, W, dY, casting_mode="llama", offset=0.0): + """One forward + backward through the module-level functional API. + + ``dW`` is ``None`` in the non-affine case; callers zip over the names. + """ + Y, X2, RSTD, BS, NW, cm = mod.rms_norm_forward(X, W, 1e-6, offset, casting_mode, None) + dX, dW = mod.rms_norm_backward(dY, X2, W, RSTD, offset, cm, BS, NW, False, None) + return Y, dX, dW + + +@cuda_required +@pytest.mark.skipif(not _supports_bf16(), reason="bf16 needs SM80+") +@pytest.mark.parametrize("n_cols", [2048, 8192]) +@pytest.mark.parametrize("casting_mode", ["llama", "none"]) +@pytest.mark.parametrize("elementwise_affine", [True, False]) +def test_rms_norm_ffi_matches_marshalling_path(monkeypatch, n_cols, casting_mode, elementwise_affine): + """The TVM-FFI direct-call path must be bit-identical to the marshalling path. + + Both compile the same kernel; only the calling convention differs, so any + difference means the abstract-tensor layouts baked at compile time do not + describe the tensors actually handed over (wrong stride divisibility, wrong + dtype for the non-affine dummy W, ...). The non-affine case is the one that + pins the dummy-W layout: it hands RSTD over in W's slot, so the compiled + signature must expect an fp32 vector there. + """ + mod = _rms_norm_mod() + if not mod._TVM_FFI_PRESENT: + pytest.skip("apache-tvm-ffi not installed; only the marshalling path exists") + + set_seed(0) + X = torch.randn(512, n_cols, device="cuda", dtype=torch.bfloat16) + W = torch.randn(n_cols, device="cuda", dtype=torch.bfloat16) if elementwise_affine else None + dY = torch.randn(512, n_cols, device="cuda", dtype=torch.bfloat16) + + monkeypatch.setattr(mod, "_FORCE_NO_FFI", True) + ref = [None if t is None else t.clone() for t in _fwd_bwd(mod, X, W, dY, casting_mode)] + monkeypatch.setattr(mod, "_FORCE_NO_FFI", False) + got = _fwd_bwd(mod, X, W, dY, casting_mode) + + # Guard the test itself: both launch paths must really have been exercised, + # or this degenerates into comparing one path against itself. + assert any(k[0] == "fwd_vec" for k in mod._compile_cache), "marshalling path never compiled" + assert any(k[0] == "fwd_vec_ffi" for k in mod._ffi_compile_cache), "FFI path never compiled" + + for name, a, b in zip(("Y", "dX", "dW"), got, ref): + if a is None or b is None: + assert a is b is None, f"{name} present on only one launch path" + continue + assert torch.equal(a, b), f"{name} differs between the FFI and marshalling launch paths" + + +@cuda_required +@pytest.mark.skipif(not _supports_bf16(), reason="bf16 needs SM80+") +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +def test_rms_norm_packed_backward_matches_scalar(monkeypatch, dtype): + """The packed-f32x2 fused backward must agree with the scalar one. + + Only the instruction form changes -- the packed variant pairs lanes and + keeps two fp32 accumulators, so results differ at most by fp32 reassociation + (and one output rounding step), never structurally. Uses a width above the + ``_use_packed_math`` threshold so the packed kernel is the one compiled. + """ + mod = _rms_norm_mod() + if not mod._is_blackwell(torch.device("cuda")): + pytest.skip("packed-f32x2 math needs a data-center Blackwell (sm_100/sm_103)") + + set_seed(0) + n_cols = 8192 # > 4096, the _use_packed_math threshold + X = torch.randn(1024, n_cols, device="cuda", dtype=dtype) + W = torch.randn(n_cols, device="cuda", dtype=dtype) + dY = torch.randn(1024, n_cols, device="cuda", dtype=dtype) + + assert mod._use_packed_math(X.device, n_cols, 16 // X.element_size()) + monkeypatch.setattr(mod, "_FORCE_NO_PACKED", True) + ref = [t.clone() for t in _fwd_bwd(mod, X, W, dY)] + monkeypatch.setattr(mod, "_FORCE_NO_PACKED", False) + got = _fwd_bwd(mod, X, W, dY) + + # Guard the test itself: the packed flag is the last element of the fused + # backward compile keys, and both specializations must have been built. + bwd_keys = [k for k in (*mod._compile_cache, *mod._ffi_compile_cache) if k[0].startswith("bwd_fused")] + assert {k[-1] for k in bwd_keys} == {True, False}, f"only one backward variant compiled: {bwd_keys}" + + tol = 1e-5 if dtype == torch.float32 else 8e-3 + for name, a, b in zip(("Y", "dX", "dW"), got, ref): + torch.testing.assert_close(a, b, atol=tol, rtol=tol, msg=lambda m, n=name: f"{n}: {m}") + + +@cuda_required +@pytest.mark.skipif(not _supports_bf16(), reason="bf16 needs SM80+") +def test_rms_norm_runs_on_the_current_stream(): + """Kernels must follow ``torch.cuda.current_stream()``, not the default stream. + + The TVM-FFI path drops the stream argument and reads the caller's + environment stream instead; if that ever stopped tracking torch, every + launch under ``torch.cuda.stream(...)`` would silently race with the + surrounding side-stream work. A CUDA graph is the sharpest available probe: + capture only records work issued on the capturing stream, so a kernel that + escaped to the default stream would leave the poisoned output untouched. + """ + from liger_kernel.ops.cutedsl.ops.rms_norm import LigerRMSNormFunction as fn + + set_seed(0) + X = torch.randn(256, 4096, device="cuda", dtype=torch.bfloat16) + W = torch.randn(4096, device="cuda", dtype=torch.bfloat16) + + expected = fn.apply(X, W, 1e-6, 0.0, "llama", False, None).clone() + torch.cuda.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = fn.apply(X, W, 1e-6, 0.0, "llama", False, None) + captured.zero_() # poison: only a captured kernel can restore this + torch.cuda.synchronize() + graph.replay() + torch.cuda.synchronize() + + assert torch.equal(captured, expected), "kernel did not run on the capturing stream"