From b28a0ca4f482391e75bca379db2705f0240f4ba5 Mon Sep 17 00:00:00 2001 From: Vaibhav Jindal Date: Wed, 26 Aug 2026 17:14:47 +0000 Subject: [PATCH] [cuTeDSL] Make SwiGLU CUDA-graph capturable + harmonize plain SiLU-Mul tests (#1408) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Two related changes to the plain SiLU-Mul (SwiGLU: `silu(a*gate)*b*down`) kernels — a CUDA-graph capturability fix for the CuTeDSL backend, and harmonized test coverage across all three backends (Triton / CuTeDSL / cuTile). No change to the Triton or cuTile kernels; the fused-GEMM MLP kernels are untouched. ## 1. CuTeDSL SwiGLU: CUDA-graph capturability fix The CuTeDSL SwiGLU launched all four kernels (vectorized fast path + predicated fallback, forward and backward) without threading PyTorch's current CUDA stream, so they ran on the default/null stream. As a result the op was **not CUDA-graph capturable** — graph replay produced stale/empty results (verified: graph-vs-eager max diff ~10-15, should be 0). Fix mirrors `rms_norm.py` (the existing TVM-FFI + stream reference): - add `_cute_stream()` (a cached `cuda.CUstream` over `torch.cuda.current_stream()`), - add `stream: cuda.CUstream = None` to each `@cute.jit` launcher and pass `stream=stream` into every `.launch(...)`, - for the TVM-FFI-compiled paths, compile against `make_fake_stream(use_tvm_ffi_env_stream=True)` so the kernel tracks the caller's env stream (the non-TVM-FFI `_dyn` path passes a real `_cute_stream()`). Numerics and gate/down semantics are unchanged. After the fix, CUDA-graph capture + replay is **bit-identical to eager** on all paths (fast/predicated, fwd/bwd), and the existing CuTeDSL suite still passes. ## 2. Harmonized plain SiLU-Mul tests across the three backends Each backend test file (`test_swiglu.py`, `test_swiglu_cutedsl.py`, `test_cutile_swiglu.py`) now defines a **byte-identical shared block** (`_SILU_SHAPES`, `_SILU_MULTIPLIERS`, dtype lists, an fp32 reference, tolerances — verified same md5 across files) and a **single parallel set** of tests exercising each backend's own SiLU-Mul Function: - correctness vs PyTorch (fwd+bwd) over aligned / predicated / tiny / odd / non-power-of-2 "cliff" shapes (incl. Qwen2.5 11008 / 13824 / 18944) x fp32/fp16/bf16, - gate/down multipliers, - default-multiplier backward-compat, - shape flexibility (1D-4D), - CUDA-graph capturability (guards the fix in part 1), - parity vs the Triton kernel (CuTeDSL and cuTile files). The set is intentionally non-redundant: the cliff sizes live in the shared `_SILU_SHAPES` (so every backend covers them once), the earlier cuTile-specific cliff-sweep tests were folded into that shared list, and Triton's multiplier coverage stays in its existing `test_correctness_silumul_with_multipliers` rather than a duplicate. Existing backend-specific tests are preserved: Triton's HF MLP-module + DTensor tests, and CuTeDSL's low-level `swiglu_forward` functional test. ## Results (B200) - cuTile: **68 passed** - CuTeDSL: **70 passed** - Triton: plain SiLU-Mul suite **48 passed** (adds correctness-vs-PyTorch, shape-flexibility, and CUDA-graph tests; multipliers already covered by the existing test). The only Triton failures are the **pre-existing fused-GEMM `MixtralExperts`** tests — unchanged from `main` and out of scope for plain SiLU-Mul. - The new CUDA-graph tests pass on all three backends; lint (`ruff check` + `format --check`) clean; the shared `_SILU_SHAPES` block is byte-identical across the three files. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: vaibhavjindal Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/liger_kernel/ops/cutedsl/ops/swiglu.py | 72 ++++- test/transformers/test_cutile_swiglu.py | 323 ++++++++++++++------- test/transformers/test_swiglu.py | 149 ++++++++++ test/transformers/test_swiglu_cutedsl.py | 273 ++++++++++------- 4 files changed, 590 insertions(+), 227 deletions(-) diff --git a/src/liger_kernel/ops/cutedsl/ops/swiglu.py b/src/liger_kernel/ops/cutedsl/ops/swiglu.py index 7fefea7d..ef2a8972 100644 --- a/src/liger_kernel/ops/cutedsl/ops/swiglu.py +++ b/src/liger_kernel/ops/cutedsl/ops/swiglu.py @@ -54,12 +54,14 @@ try: except ImportError: pass +import cuda.bindings.driver as cuda import cutlass import cutlass.cute as cute import cutlass.cute.arch as carch import cutlass.cute.math as cute_math from cutlass.cute.runtime import from_dlpack +from cutlass.cute.runtime import make_fake_stream from liger_kernel.ops.cutedsl.ops.utils import make_fake_tensor from liger_kernel.ops.cutedsl.ops.utils import torch2cute_dtype_map @@ -431,7 +433,7 @@ def _swiglu_bwd_vec_kernel( # --------------------------------------------------------------------------- def _make_fwd_vec(vec: int): @cute.jit - def fwd(mA, mB, mC, gate_mult: cutlass.Float32): + def fwd(mA, mB, mC, gate_mult: cutlass.Float32, stream: cuda.CUstream = None): copy_atom = cute.make_copy_atom( cute.nvgpu.CopyUniversalOp(), mA.element_type, num_bits_per_copy=vec * mA.element_type.width ) @@ -447,6 +449,7 @@ def _make_fwd_vec(vec: int): _swiglu_fwd_vec_kernel(gA, gB, gC, tiled_copy, gate_mult).launch( grid=[cute.size(gC, mode=[1]), 1, 1], block=[_NUM_THREADS, 1, 1], + stream=stream, ) return fwd @@ -454,7 +457,7 @@ def _make_fwd_vec(vec: int): def _make_bwd_vec(vec: int): @cute.jit - def bwd(mDC, mA, mB, gate_mult: cutlass.Float32): + def bwd(mDC, mA, mB, gate_mult: cutlass.Float32, stream: cuda.CUstream = None): copy_atom = cute.make_copy_atom( cute.nvgpu.CopyUniversalOp(), mA.element_type, num_bits_per_copy=vec * mA.element_type.width ) @@ -470,6 +473,7 @@ def _make_bwd_vec(vec: int): _swiglu_bwd_vec_kernel(gDC, gA, gB, tiled_copy, gate_mult).launch( grid=[cute.size(gA, mode=[1]), 1, 1], block=[_NUM_THREADS, 1, 1], + stream=stream, ) return bwd @@ -480,7 +484,7 @@ def _make_bwd_vec(vec: int): # --------------------------------------------------------------------------- def _make_fwd(vec: int, predicated: bool, packed_math: bool): @cute.jit - def fwd(mA, mB, mC, gate_mult: cutlass.Float32): + def fwd(mA, mB, mC, gate_mult: cutlass.Float32, stream: cuda.CUstream = None): thr_layout = cute.make_layout(_NUM_THREADS, stride=vec) val_layout = cute.make_layout(vec, stride=1) tiler, tv_layout = cute.make_layout_tv(thr_layout, val_layout) @@ -494,6 +498,7 @@ def _make_fwd(vec: int, predicated: bool, packed_math: bool): _swiglu_fwd_kernel(gA, gB, gC, cC, mC.shape, thr_layout, val_layout, gate_mult, predicated, packed_math).launch( grid=[cute.size(gC, mode=[1]), 1, 1], block=[cute.size(tv_layout, mode=[0]), 1, 1], + stream=stream, ) return fwd @@ -501,7 +506,7 @@ def _make_fwd(vec: int, predicated: bool, packed_math: bool): def _make_bwd(vec: int, predicated: bool, packed_math: bool): @cute.jit - def bwd(mDC, mA, mB, gate_mult: cutlass.Float32): + def bwd(mDC, mA, mB, gate_mult: cutlass.Float32, stream: cuda.CUstream = None): thr_layout = cute.make_layout(_NUM_THREADS, stride=vec) val_layout = cute.make_layout(vec, stride=1) tiler, tv_layout = cute.make_layout_tv(thr_layout, val_layout) @@ -517,6 +522,7 @@ def _make_bwd(vec: int, predicated: bool, packed_math: bool): ).launch( grid=[cute.size(gA, mode=[1]), 1, 1], block=[cute.size(tv_layout, mode=[0]), 1, 1], + stream=stream, ) return bwd @@ -577,14 +583,37 @@ def _dyn(t: torch.Tensor): return from_dlpack(t.detach()).mark_layout_dynamic() +# Cache the ``cuda.CUstream`` wrapper keyed on torch's raw stream handle so we don't +# rebuild it every launch. The kernels MUST run on PyTorch's *current* stream (not the +# default/null stream) so the op is CUDA-graph capturable and preserves ordering under +# multi-stream execution (pipeline parallelism, grad checkpointing). This is the +# non-TVM-FFI (``_dyn``) counterpart to TVM-FFI's env-stream mechanism; mirrors +# rms_norm.py / rope.py. +_stream_cache: dict = {} + + +def _cute_stream(): + raw = torch.cuda.current_stream().cuda_stream + s = _stream_cache.get(raw) + if s is None: + s = cuda.CUstream(raw) + _stream_cache[raw] = s + return s + + class _DynCaller: - """Wraps a non-TVM-FFI compiled function so callers always pass raw tensors.""" + """Wraps a non-TVM-FFI compiled function so callers always pass raw tensors. + + The compiled function bakes a trailing ``stream`` parameter (the ``@cute.jit`` + launcher's last arg); we source PyTorch's current stream here so the caller's + invocation signature stays identical to the TVM-FFI env-stream path. + """ def __init__(self, compiled): self._compiled = compiled def __call__(self, a, b, c, gm): - return self._compiled(_dyn(a), _dyn(b), _dyn(c), gm) + return self._compiled(_dyn(a), _dyn(b), _dyn(c), gm, _cute_stream()) def _get_compiled(kind: str, ref: torch.Tensor, vec: int, predicated: bool, packed_math: bool): @@ -597,12 +626,17 @@ def _get_compiled(kind: str, ref: torch.Tensor, vec: int, predicated: bool, pack if _TVM_FFI_PRESENT: # TVM-FFI: 1-D sym_int fake tensor; PyTorch tensors passed directly — no # from_dlpack overhead. Divisibility=1 since the predicated path makes no - # alignment guarantee beyond element size. + # alignment guarantee beyond element size. ``use_tvm_ffi_env_stream=True`` + # drops the stream from the FFI signature: the compiled fn runs on the + # caller's env stream (torch.cuda.current_stream()), so the launch is + # CUDA-graph capturable with no per-call stream marshalling. cute_dtype = torch2cute_dtype_map[ref.dtype] fake = make_fake_tensor(cute_dtype, (cute.sym_int(),), 1) - fn = cute.compile(maker, fake, fake, fake, gm, options="--enable-tvm-ffi") + fn = cute.compile( + maker, fake, fake, fake, gm, make_fake_stream(use_tvm_ffi_env_stream=True), options="--enable-tvm-ffi" + ) else: - compiled = cute.compile(maker, _dyn(ref), _dyn(ref), _dyn(ref), gm) + compiled = cute.compile(maker, _dyn(ref), _dyn(ref), _dyn(ref), gm, _cute_stream()) fn = _DynCaller(compiled) _COMPILE_CACHE[key] = fn return fn @@ -629,9 +663,25 @@ def _get_compiled_vec(kind: str, dtype: torch.dtype, vec: int, device_index: int return make_fake_tensor(cute_dtype, (cute.sym_int(), inner), vec) if kind == "fwd": - fn = cute.compile(_make_fwd_vec(vec), fake(), fake(), fake(), gm, options="--enable-tvm-ffi") + fn = cute.compile( + _make_fwd_vec(vec), + fake(), + fake(), + fake(), + gm, + make_fake_stream(use_tvm_ffi_env_stream=True), + options="--enable-tvm-ffi", + ) else: - fn = cute.compile(_make_bwd_vec(vec), fake(), fake(), fake(), gm, options="--enable-tvm-ffi") + fn = cute.compile( + _make_bwd_vec(vec), + fake(), + fake(), + fake(), + gm, + make_fake_stream(use_tvm_ffi_env_stream=True), + options="--enable-tvm-ffi", + ) _VEC_COMPILE_CACHE[key] = fn return fn diff --git a/test/transformers/test_cutile_swiglu.py b/test/transformers/test_cutile_swiglu.py index 538851fc..ed57f8bb 100644 --- a/test/transformers/test_cutile_swiglu.py +++ b/test/transformers/test_cutile_swiglu.py @@ -2,9 +2,7 @@ import importlib.util import pytest import torch -import torch.nn.functional as F -from test.utils import assert_verbose_allclose from test.utils import supports_bfloat16 @@ -18,8 +16,10 @@ def _has_cuda_tile(): _CUDA_TILE_AVAILABLE = _has_cuda_tile() if _CUDA_TILE_AVAILABLE: from liger_kernel.ops.cutile.ops.swiglu import LigerSiLUMulFunction as CuTileSiLUMulFunction + from liger_kernel.ops.swiglu import LigerSiLUMulFunction as TritonSiLUMulFunction else: CuTileSiLUMulFunction = None + TritonSiLUMulFunction = None pytestmark = [ pytest.mark.skipif(not torch.cuda.is_available(), reason="cuTile SwiGLU requires CUDA"), @@ -27,121 +27,222 @@ pytestmark = [ ] -def _reference_swiglu(a, b, gate_multiplier, down_multiplier): - # Compute in fp32 then cast back, matching the kernel's internal fp32 math. - return (F.silu(a.float() * gate_multiplier) * b.float() * down_multiplier).to(a.dtype) +# --------------------------------------------------------------------------- +# Shared plain SiLU-Mul coverage (kept identical across triton/cutedsl/cutile test files) +# --------------------------------------------------------------------------- +_SILU_SHAPES = [ + (4, 2048), # 2D, power-of-2 aligned + (2, 256, 512), # 3D, small + (6, 42, 431), # non-aligned / predicated path, odd width + (1, 1, 7), # tiny + (3, 1023), # odd width + (4, 11008), # non-power-of-2 "cliff" (Qwen2.5-7B) + (2, 3, 13824), # non-power-of-2 "cliff" (Qwen2.5-3B), 3D + (2, 18944), # non-power-of-2 "cliff" (Qwen2.5-14B) +] +_SILU_MULTIPLIERS = [(1.0, 1.0), (1.5, 0.75), (0.7, 1.3)] + +# fp32, fp16, bf16 (bf16 guarded) — used by the harmonized SiLU-Mul tests below. +_SILU_DTYPES = [ + pytest.param(torch.float32, id="fp32"), + pytest.param(torch.float16, id="fp16"), + pytest.param( + torch.bfloat16, + marks=pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"), + id="bf16", + ), +] +# Multiplier sweep runs on the fp32 + bf16 subset (fp16 adds no path coverage here). +_SILU_MULT_DTYPES = [ + pytest.param(torch.float32, id="fp32"), + pytest.param( + torch.bfloat16, + marks=pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"), + id="bf16", + ), +] -# n_cols chosen to cover: power-of-2 (aligned) sizes, the non-power-of-2 -# intermediate sizes that hit the old bounds-checking cliff (Qwen2.5 11008 / -# 13824 / 18944, Gemma-ish 4608), and small / odd widths that exercise the full -# power-of-2 remainder ladder (down to width 1). +def _ref_silu_mul(a, b, gate=1.0, down=1.0): + """Ground-truth SiLU-Mul reference (silu(a * gate) * b * down) computed in fp32.""" + return torch.nn.functional.silu(a.float() * gate) * b.float() * down + + +def _silu_tol(dtype): + """Harmonized (atol, rtol) for the shared plain SiLU-Mul tests.""" + if dtype == torch.float32: + return 2e-4, 1e-4 + return 1e-2, 1e-2 # fp16 / bf16 + + +# --------------------------------------------------------------------------- +# Plain SiLU-Mul tests (parallel across triton / cutedsl / cutile). +# --------------------------------------------------------------------------- @pytest.mark.flaky(reruns=3, reruns_delay=2) -@pytest.mark.parametrize( - "n_cols", - [ - 512, # < base: pure remainder ladder - 2048, # exactly one base tile - 4096, # aligned, multiple base tiles - 8192, # aligned - 4608, # Gemma-ish (non-power-of-2) - 11008, # Qwen2.5-7B - 13824, # Qwen2.5-3B - 18944, # Qwen2.5-14B - 1023, # odd: 512+256+...+1 ladder - 431, # odd remainder - 7, # tiny (widths 4+2+1) - 1, # single column - ], -) -@pytest.mark.parametrize("n_rows", [4, 73]) -@pytest.mark.parametrize( - "dtype, atol, rtol", - [ - (torch.float32, 2e-4, 1e-5), - pytest.param( - torch.bfloat16, - 1e-1, - 2e-2, - marks=pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"), - ), - (torch.float16, 1e-2, 1e-2), - ], -) -def test_cutile_swiglu_correctness(n_cols, n_rows, dtype, atol, rtol): +@pytest.mark.parametrize("shape", _SILU_SHAPES) +@pytest.mark.parametrize("dtype", _SILU_DTYPES) +def test_cutile_correctness(shape, dtype): + """Forward + backward vs the fp32 SiLU-Mul reference across shapes/dtypes.""" torch.manual_seed(0) - a = torch.randn(n_rows, n_cols, device="cuda", dtype=dtype) - b = torch.randn(n_rows, n_cols, device="cuda", dtype=dtype) - grad_output = torch.randn(n_rows, n_cols, device="cuda", dtype=dtype) - - ref_a = a.clone().requires_grad_(True) - ref_b = b.clone().requires_grad_(True) - ref_out = _reference_swiglu(ref_a, ref_b, 1.0, 1.0) - ref_out.backward(grad_output) - - cutile_a = a.clone().requires_grad_(True) - cutile_b = b.clone().requires_grad_(True) - cutile_out = CuTileSiLUMulFunction.apply(cutile_a, cutile_b) - cutile_out.backward(grad_output.clone()) - - assert_verbose_allclose(cutile_out, ref_out, atol=atol, rtol=rtol) - assert_verbose_allclose(cutile_a.grad, ref_a.grad, atol=atol, rtol=rtol, max_print=20) - assert_verbose_allclose(cutile_b.grad, ref_b.grad, atol=atol, rtol=rtol, max_print=20) - - -@pytest.mark.flaky(reruns=3, reruns_delay=2) -@pytest.mark.parametrize("n_cols", [4096, 13824, 18944]) -@pytest.mark.parametrize( - "gate_multiplier, down_multiplier", - [ - (1.0, 1.0), - (0.5, 2.0), - (2.0, 0.5), - ], -) -def test_cutile_swiglu_multipliers(n_cols, gate_multiplier, down_multiplier): - dtype, atol, rtol = torch.float32, 2e-4, 1e-5 - torch.manual_seed(0) - a = torch.randn(16, n_cols, device="cuda", dtype=dtype) - b = torch.randn(16, n_cols, device="cuda", dtype=dtype) - grad_output = torch.randn(16, n_cols, device="cuda", dtype=dtype) - - ref_a = a.clone().requires_grad_(True) - ref_b = b.clone().requires_grad_(True) - ref_out = _reference_swiglu(ref_a, ref_b, gate_multiplier, down_multiplier) - ref_out.backward(grad_output) - - cutile_a = a.clone().requires_grad_(True) - cutile_b = b.clone().requires_grad_(True) - cutile_out = CuTileSiLUMulFunction.apply(cutile_a, cutile_b, gate_multiplier, down_multiplier) - cutile_out.backward(grad_output.clone()) - - assert_verbose_allclose(cutile_out, ref_out, atol=atol, rtol=rtol) - assert_verbose_allclose(cutile_a.grad, ref_a.grad, atol=atol, rtol=rtol, max_print=20) - assert_verbose_allclose(cutile_b.grad, ref_b.grad, atol=atol, rtol=rtol, max_print=20) - - -@pytest.mark.flaky(reruns=3, reruns_delay=2) -@pytest.mark.parametrize("n_cols", [13824, 18944]) -def test_cutile_swiglu_3d_input(n_cols): - dtype, atol, rtol = torch.float32, 2e-4, 1e-5 - torch.manual_seed(0) - shape = (2, 3, n_cols) + atol, rtol = _silu_tol(dtype) a = torch.randn(*shape, device="cuda", dtype=dtype) b = torch.randn(*shape, device="cuda", dtype=dtype) - grad_output = torch.randn(*shape, device="cuda", dtype=dtype) + grad = torch.randn(*shape, device="cuda", dtype=dtype) - ref_a = a.clone().requires_grad_(True) - ref_b = b.clone().requires_grad_(True) - ref_out = _reference_swiglu(ref_a, ref_b, 1.0, 1.0) - ref_out.backward(grad_output) + ref_a = a.float().detach().requires_grad_(True) + ref_b = b.float().detach().requires_grad_(True) + ref_out = _ref_silu_mul(ref_a, ref_b) + ref_out.backward(grad.float()) - cutile_a = a.clone().requires_grad_(True) - cutile_b = b.clone().requires_grad_(True) - cutile_out = CuTileSiLUMulFunction.apply(cutile_a, cutile_b) - cutile_out.backward(grad_output.clone()) + test_a = a.clone().detach().requires_grad_(True) + test_b = b.clone().detach().requires_grad_(True) + test_out = CuTileSiLUMulFunction.apply(test_a, test_b) + test_out.backward(grad.clone()) - assert cutile_out.shape == shape - assert_verbose_allclose(cutile_out, ref_out, atol=atol, rtol=rtol) - assert_verbose_allclose(cutile_a.grad, ref_a.grad, atol=atol, rtol=rtol, max_print=20) - assert_verbose_allclose(cutile_b.grad, ref_b.grad, atol=atol, rtol=rtol, max_print=20) + torch.testing.assert_close(test_out.float(), ref_out.float(), atol=atol, rtol=rtol) + torch.testing.assert_close(test_a.grad.float(), ref_a.grad.float(), atol=atol, rtol=rtol) + torch.testing.assert_close(test_b.grad.float(), ref_b.grad.float(), atol=atol, rtol=rtol) + + +@pytest.mark.flaky(reruns=3, reruns_delay=2) +@pytest.mark.parametrize("shape", [(2, 3, 13824), (2, 256, 512)]) +@pytest.mark.parametrize("gate_multiplier, down_multiplier", _SILU_MULTIPLIERS) +@pytest.mark.parametrize("dtype", _SILU_MULT_DTYPES) +def test_cutile_multipliers(shape, gate_multiplier, down_multiplier, dtype): + """Forward + backward with gate/down multipliers vs the fp32 reference.""" + torch.manual_seed(0) + atol, rtol = _silu_tol(dtype) + a = torch.randn(*shape, device="cuda", dtype=dtype) + b = torch.randn(*shape, device="cuda", dtype=dtype) + grad = torch.randn(*shape, device="cuda", dtype=dtype) + + ref_a = a.float().detach().requires_grad_(True) + ref_b = b.float().detach().requires_grad_(True) + ref_out = _ref_silu_mul(ref_a, ref_b, gate_multiplier, down_multiplier) + ref_out.backward(grad.float()) + + test_a = a.clone().detach().requires_grad_(True) + test_b = b.clone().detach().requires_grad_(True) + test_out = CuTileSiLUMulFunction.apply(test_a, test_b, gate_multiplier, down_multiplier) + test_out.backward(grad.clone()) + + torch.testing.assert_close(test_out.float(), ref_out.float(), atol=atol, rtol=rtol) + torch.testing.assert_close(test_a.grad.float(), ref_a.grad.float(), atol=atol, rtol=rtol) + torch.testing.assert_close(test_b.grad.float(), ref_b.grad.float(), atol=atol, rtol=rtol) + + +def test_cutile_default_multipliers_backward_compat(): + """Calling without multipliers must equal calling with (1.0, 1.0).""" + torch.manual_seed(0) + a = torch.randn(4, 16, 32, device="cuda", dtype=torch.float32) + b = torch.randn(4, 16, 32, device="cuda", dtype=torch.float32) + grad = torch.randn(4, 16, 32, device="cuda", dtype=torch.float32) + + a1 = a.clone().detach().requires_grad_(True) + b1 = b.clone().detach().requires_grad_(True) + a2 = a.clone().detach().requires_grad_(True) + b2 = b.clone().detach().requires_grad_(True) + + y_default = CuTileSiLUMulFunction.apply(a1, b1) + y_explicit = CuTileSiLUMulFunction.apply(a2, b2, 1.0, 1.0) + + torch.testing.assert_close(y_default, y_explicit) + + y_default.backward(grad.clone()) + y_explicit.backward(grad.clone()) + + torch.testing.assert_close(a1.grad, a2.grad) + torch.testing.assert_close(b1.grad, b2.grad) + + +@pytest.mark.parametrize("shape", [(512,), (123,), (2, 256, 512), (3, 5, 7), (2, 4, 8, 16)]) +def test_cutile_shape_flexibility(shape): + """Arbitrary >=1-D shapes: forward + backward vs the fp32 reference; shape preserved.""" + torch.manual_seed(0) + dtype = torch.float32 + atol, rtol = _silu_tol(dtype) + a = torch.randn(*shape, device="cuda", dtype=dtype) + b = torch.randn(*shape, device="cuda", dtype=dtype) + grad = torch.randn(*shape, device="cuda", dtype=dtype) + + ref_a = a.clone().detach().requires_grad_(True) + ref_b = b.clone().detach().requires_grad_(True) + ref_out = _ref_silu_mul(ref_a, ref_b) + ref_out.backward(grad) + + test_a = a.clone().detach().requires_grad_(True) + test_b = b.clone().detach().requires_grad_(True) + test_out = CuTileSiLUMulFunction.apply(test_a, test_b) + assert test_out.shape == a.shape + test_out.backward(grad.clone()) + + torch.testing.assert_close(test_out.float(), ref_out.float(), atol=atol, rtol=rtol) + torch.testing.assert_close(test_a.grad.float(), ref_a.grad.float(), atol=atol, rtol=rtol) + torch.testing.assert_close(test_b.grad.float(), ref_b.grad.float(), atol=atol, rtol=rtol) + + +@pytest.mark.parametrize("shape, dtype", [((73, 13824), torch.bfloat16), ((4, 2048), torch.float32)]) +def test_cutile_cuda_graph(shape, dtype): + """CUDA-graph capture + replay of the forward must be bit-identical to eager.""" + if dtype == torch.bfloat16 and not supports_bfloat16(): + pytest.skip("bfloat16 not supported on this GPU") + torch.manual_seed(0) + static_a = torch.randn(*shape, device="cuda", dtype=dtype) + static_b = torch.randn(*shape, device="cuda", dtype=dtype) + + def run(): + with torch.no_grad(): + return CuTileSiLUMulFunction.apply(static_a, static_b) + + # Warm up (build / cache the kernel) on a side stream before capture. + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + for _ in range(3): + run() + torch.cuda.current_stream().wait_stream(s) + + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + static_out = run() + + # Replay on fresh data copied into the static input buffers. + fresh_a = torch.randn(*shape, device="cuda", dtype=dtype) + fresh_b = torch.randn(*shape, device="cuda", dtype=dtype) + static_a.copy_(fresh_a) + static_b.copy_(fresh_b) + g.replay() + torch.cuda.synchronize() + graph_out = static_out.clone() + + with torch.no_grad(): + eager_out = CuTileSiLUMulFunction.apply(fresh_a, fresh_b) + + max_abs_diff = (graph_out.float() - eager_out.float()).abs().max().item() + assert max_abs_diff == 0.0, f"graph vs eager mismatch: max_abs_diff={max_abs_diff}" + + +@pytest.mark.flaky(reruns=3, reruns_delay=2) +@pytest.mark.parametrize("shape", _SILU_SHAPES) +@pytest.mark.parametrize("dtype", _SILU_DTYPES) +def test_cutile_matches_triton(shape, dtype): + """cuTile Function must match the Triton Function (output + grads).""" + torch.manual_seed(0) + atol, rtol = _silu_tol(dtype) + a = torch.randn(*shape, device="cuda", dtype=dtype) + b = torch.randn(*shape, device="cuda", dtype=dtype) + grad = torch.randn(*shape, device="cuda", dtype=dtype) + + tri_a = a.clone().detach().requires_grad_(True) + tri_b = b.clone().detach().requires_grad_(True) + tri_out = TritonSiLUMulFunction.apply(tri_a, tri_b) + tri_out.backward(grad.clone()) + + test_a = a.clone().detach().requires_grad_(True) + test_b = b.clone().detach().requires_grad_(True) + test_out = CuTileSiLUMulFunction.apply(test_a, test_b) + test_out.backward(grad.clone()) + + torch.testing.assert_close(test_out.float(), tri_out.float(), atol=atol, rtol=rtol) + torch.testing.assert_close(test_a.grad.float(), tri_a.grad.float(), atol=atol, rtol=rtol) + torch.testing.assert_close(test_b.grad.float(), tri_b.grad.float(), atol=atol, rtol=rtol) diff --git a/test/transformers/test_swiglu.py b/test/transformers/test_swiglu.py index 919b35e5..0146ae39 100644 --- a/test/transformers/test_swiglu.py +++ b/test/transformers/test_swiglu.py @@ -46,6 +46,54 @@ PHI3_CONFIG = Phi3Config( SLEEP_SECONDS = 0.1 +# --------------------------------------------------------------------------- +# Shared plain SiLU-Mul coverage (kept identical across triton/cutedsl/cutile test files) +# --------------------------------------------------------------------------- +_SILU_SHAPES = [ + (4, 2048), # 2D, power-of-2 aligned + (2, 256, 512), # 3D, small + (6, 42, 431), # non-aligned / predicated path, odd width + (1, 1, 7), # tiny + (3, 1023), # odd width + (4, 11008), # non-power-of-2 "cliff" (Qwen2.5-7B) + (2, 3, 13824), # non-power-of-2 "cliff" (Qwen2.5-3B), 3D + (2, 18944), # non-power-of-2 "cliff" (Qwen2.5-14B) +] +_SILU_MULTIPLIERS = [(1.0, 1.0), (1.5, 0.75), (0.7, 1.3)] + +# fp32, fp16, bf16 (bf16 guarded) — used by the harmonized SiLU-Mul tests below. +_SILU_DTYPES = [ + pytest.param(torch.float32, id="fp32"), + pytest.param(torch.float16, id="fp16"), + pytest.param( + torch.bfloat16, + marks=pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"), + id="bf16", + ), +] +# Multiplier sweep runs on the fp32 + bf16 subset (fp16 adds no path coverage here). +_SILU_MULT_DTYPES = [ + pytest.param(torch.float32, id="fp32"), + pytest.param( + torch.bfloat16, + marks=pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"), + id="bf16", + ), +] + + +def _ref_silu_mul(a, b, gate=1.0, down=1.0): + """Ground-truth SiLU-Mul reference (silu(a * gate) * b * down) computed in fp32.""" + return torch.nn.functional.silu(a.float() * gate) * b.float() * down + + +def _silu_tol(dtype): + """Harmonized (atol, rtol) for the shared plain SiLU-Mul tests.""" + if dtype == torch.float32: + return 2e-4, 1e-4 + return 1e-2, 1e-2 # fp16 / bf16 + + @pytest.mark.parametrize( "bsz, seq_len, hidden_size, intermediate_size", [ @@ -728,3 +776,104 @@ def test_swiglu_blackwell_tiled_matches_original(monkeypatch, n_rows, n_cols, ga torch.testing.assert_close(c_tiled, c_ref, rtol=0, atol=0) torch.testing.assert_close(da_tiled, da_ref, rtol=0, atol=0) torch.testing.assert_close(db_tiled, db_ref, rtol=0, atol=0) + + +# --------------------------------------------------------------------------- +# Harmonized plain SiLU-Mul tests (parallel across triton/cutedsl/cutile). +# These exercise the Triton LigerSiLUMulFunction directly against the fp32 +# reference; the cutedsl/cutile suites mirror them (and additionally compare +# against this Triton Function via ``test__matches_triton``). +# --------------------------------------------------------------------------- +@pytest.mark.flaky(reruns=3, reruns_delay=2) +@pytest.mark.parametrize("shape", _SILU_SHAPES) +@pytest.mark.parametrize("dtype", _SILU_DTYPES) +def test_triton_silumul_correctness(shape, dtype): + """Forward + backward vs the fp32 SiLU-Mul reference across shapes/dtypes.""" + torch.manual_seed(0) + atol, rtol = _silu_tol(dtype) + a = torch.randn(*shape, device=device, dtype=dtype) + b = torch.randn(*shape, device=device, dtype=dtype) + grad = torch.randn(*shape, device=device, dtype=dtype) + + ref_a = a.float().detach().requires_grad_(True) + ref_b = b.float().detach().requires_grad_(True) + ref_out = _ref_silu_mul(ref_a, ref_b) + ref_out.backward(grad.float()) + + test_a = a.clone().detach().requires_grad_(True) + test_b = b.clone().detach().requires_grad_(True) + test_out = LigerSiLUMulFunction.apply(test_a, test_b) + test_out.backward(grad.clone()) + + torch.testing.assert_close(test_out.float(), ref_out.float(), atol=atol, rtol=rtol) + torch.testing.assert_close(test_a.grad.float(), ref_a.grad.float(), atol=atol, rtol=rtol) + torch.testing.assert_close(test_b.grad.float(), ref_b.grad.float(), atol=atol, rtol=rtol) + + +@pytest.mark.parametrize("shape", [(512,), (123,), (2, 256, 512), (3, 5, 7), (2, 4, 8, 16)]) +def test_triton_silumul_shape_flexibility(shape): + """Arbitrary >=1-D shapes: forward + backward vs the fp32 reference; shape preserved.""" + torch.manual_seed(0) + dtype = torch.float32 + atol, rtol = _silu_tol(dtype) + a = torch.randn(*shape, device=device, dtype=dtype) + b = torch.randn(*shape, device=device, dtype=dtype) + grad = torch.randn(*shape, device=device, dtype=dtype) + + ref_a = a.clone().detach().requires_grad_(True) + ref_b = b.clone().detach().requires_grad_(True) + ref_out = _ref_silu_mul(ref_a, ref_b) + ref_out.backward(grad) + + test_a = a.clone().detach().requires_grad_(True) + test_b = b.clone().detach().requires_grad_(True) + test_out = LigerSiLUMulFunction.apply(test_a, test_b) + assert test_out.shape == a.shape + test_out.backward(grad.clone()) + + torch.testing.assert_close(test_out.float(), ref_out.float(), atol=atol, rtol=rtol) + torch.testing.assert_close(test_a.grad.float(), ref_a.grad.float(), atol=atol, rtol=rtol) + torch.testing.assert_close(test_b.grad.float(), ref_b.grad.float(), atol=atol, rtol=rtol) + + +@pytest.mark.parametrize("shape, dtype", [((73, 13824), torch.bfloat16), ((4, 2048), torch.float32)]) +def test_triton_silumul_cuda_graph(shape, dtype): + """CUDA-graph capture + replay of the forward must be bit-identical to eager.""" + if dtype == torch.bfloat16 and not supports_bfloat16(): + pytest.skip("bfloat16 not supported on this GPU") + if not torch.cuda.is_available(): + pytest.skip("CUDA graph capture requires a CUDA device") + torch.manual_seed(0) + static_a = torch.randn(*shape, device=device, dtype=dtype) + static_b = torch.randn(*shape, device=device, dtype=dtype) + + def run(): + with torch.no_grad(): + return LigerSiLUMulFunction.apply(static_a, static_b) + + # Warm up (compile / autotune) on a side stream before capture. + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + for _ in range(3): + run() + torch.cuda.current_stream().wait_stream(s) + + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + static_out = run() + + # Replay on fresh data copied into the static input buffers. + fresh_a = torch.randn(*shape, device=device, dtype=dtype) + fresh_b = torch.randn(*shape, device=device, dtype=dtype) + static_a.copy_(fresh_a) + static_b.copy_(fresh_b) + g.replay() + torch.cuda.synchronize() + graph_out = static_out.clone() + + with torch.no_grad(): + eager_out = LigerSiLUMulFunction.apply(fresh_a, fresh_b) + + max_abs_diff = (graph_out.float() - eager_out.float()).abs().max().item() + assert max_abs_diff == 0.0, f"graph vs eager mismatch: max_abs_diff={max_abs_diff}" diff --git a/test/transformers/test_swiglu_cutedsl.py b/test/transformers/test_swiglu_cutedsl.py index 3bed4094..def9b960 100644 --- a/test/transformers/test_swiglu_cutedsl.py +++ b/test/transformers/test_swiglu_cutedsl.py @@ -45,105 +45,131 @@ def _tol(dtype): return 1e-2, 1e-2 -def _torch_silu_mul_ref(a, b, gate_multiplier=1.0, down_multiplier=1.0): - """Pure-PyTorch reference for silu(a * gate_mult) * b * down_mult.""" - return torch.nn.functional.silu(a * gate_multiplier) * b * down_multiplier +# --------------------------------------------------------------------------- +# Shared plain SiLU-Mul coverage (kept identical across triton/cutedsl/cutile test files) +# --------------------------------------------------------------------------- +_SILU_SHAPES = [ + (4, 2048), # 2D, power-of-2 aligned + (2, 256, 512), # 3D, small + (6, 42, 431), # non-aligned / predicated path, odd width + (1, 1, 7), # tiny + (3, 1023), # odd width + (4, 11008), # non-power-of-2 "cliff" (Qwen2.5-7B) + (2, 3, 13824), # non-power-of-2 "cliff" (Qwen2.5-3B), 3D + (2, 18944), # non-power-of-2 "cliff" (Qwen2.5-14B) +] +_SILU_MULTIPLIERS = [(1.0, 1.0), (1.5, 0.75), (0.7, 1.3)] + +# fp32, fp16, bf16 (bf16 guarded) — used by the harmonized SiLU-Mul tests below. +_SILU_DTYPES = [ + pytest.param(torch.float32, id="fp32"), + pytest.param(torch.float16, id="fp16"), + pytest.param( + torch.bfloat16, + marks=pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"), + id="bf16", + ), +] +# Multiplier sweep runs on the fp32 + bf16 subset (fp16 adds no path coverage here). +_SILU_MULT_DTYPES = [ + pytest.param(torch.float32, id="fp32"), + pytest.param( + torch.bfloat16, + marks=pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"), + id="bf16", + ), +] -@pytest.mark.parametrize( - "shape", - [ - (4096, 11008), # llama, tile-aligned - (2, 256, 512), - (6, 42, 431), # non-tile-aligned (exercises predicated kernel) - (1, 1, 7), # tiny, predicated - (3, 1023), # odd number of columns - ], -) -@pytest.mark.parametrize( - "dtype", - [ - pytest.param( - torch.bfloat16, - marks=pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"), - ), - torch.float16, - torch.float32, - ], -) -@pytest.mark.parametrize("gate_mult, down_mult", [(1.0, 1.0), (1.5, 0.75)]) -def test_cutedsl_matches_triton(shape, dtype, gate_mult, down_mult): +def _ref_silu_mul(a, b, gate=1.0, down=1.0): + """Ground-truth SiLU-Mul reference (silu(a * gate) * b * down) computed in fp32.""" + return torch.nn.functional.silu(a.float() * gate) * b.float() * down + + +def _silu_tol(dtype): + """Harmonized (atol, rtol) for the shared plain SiLU-Mul tests.""" + if dtype == torch.float32: + return 2e-4, 1e-4 + return 1e-2, 1e-2 # fp16 / bf16 + + +@pytest.mark.flaky(reruns=3, reruns_delay=2) +@pytest.mark.parametrize("shape", _SILU_SHAPES) +@pytest.mark.parametrize("dtype", _SILU_DTYPES) +def test_cutedsl_correctness(shape, dtype): + """Forward + backward vs the fp32 SiLU-Mul reference across shapes/dtypes.""" torch.manual_seed(0) - atol, rtol = _tol(dtype) + atol, rtol = _silu_tol(dtype) a = torch.randn(*shape, device=device, dtype=dtype) b = torch.randn(*shape, device=device, dtype=dtype) + grad = torch.randn(*shape, device=device, dtype=dtype) - a_ref = a.clone().requires_grad_(True) - b_ref = b.clone().requires_grad_(True) - a_cut = a.clone().requires_grad_(True) - b_cut = b.clone().requires_grad_(True) + ref_a = a.float().detach().requires_grad_(True) + ref_b = b.float().detach().requires_grad_(True) + ref_out = _ref_silu_mul(ref_a, ref_b) + ref_out.backward(grad.float()) - c_ref = LigerSiLUMulFunction.apply(a_ref, b_ref, gate_mult, down_mult) - c_cut = LigerSiLUMulCuteDSLFunction.apply(a_cut, b_cut, gate_mult, down_mult) + test_a = a.clone().detach().requires_grad_(True) + test_b = b.clone().detach().requires_grad_(True) + test_out = LigerSiLUMulCuteDSLFunction.apply(test_a, test_b) + test_out.backward(grad.clone()) - torch.testing.assert_close(c_cut.float(), c_ref.float(), atol=atol, rtol=rtol) - - grad = torch.randn_like(c_ref) - c_ref.backward(grad) - c_cut.backward(grad.clone()) - - torch.testing.assert_close(a_cut.grad.float(), a_ref.grad.float(), atol=atol, rtol=rtol) - torch.testing.assert_close(b_cut.grad.float(), b_ref.grad.float(), atol=atol, rtol=rtol) + torch.testing.assert_close(test_out.float(), ref_out.float(), atol=atol, rtol=rtol) + torch.testing.assert_close(test_a.grad.float(), ref_a.grad.float(), atol=atol, rtol=rtol) + torch.testing.assert_close(test_b.grad.float(), ref_b.grad.float(), atol=atol, rtol=rtol) -@pytest.mark.parametrize( - "shape", - [ - (2, 8, 1024), - (9, 7, 41), - ], -) -@pytest.mark.parametrize( - "gate_mult, down_mult", - [ - (1.0, 1.0), - (0.7, 1.3), - (1.5, 0.5), - ], -) -@pytest.mark.parametrize( - "dtype", - [ - torch.float32, - pytest.param( - torch.bfloat16, - marks=pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"), - ), - ], -) -def test_cutedsl_matches_pytorch(shape, gate_mult, down_mult, dtype): - """Independent ground-truth check vs pure PyTorch (catches shared CuteDSL/Triton bugs).""" +@pytest.mark.flaky(reruns=3, reruns_delay=2) +@pytest.mark.parametrize("shape", [(2, 3, 13824), (2, 256, 512)]) +@pytest.mark.parametrize("gate_multiplier, down_multiplier", _SILU_MULTIPLIERS) +@pytest.mark.parametrize("dtype", _SILU_MULT_DTYPES) +def test_cutedsl_multipliers(shape, gate_multiplier, down_multiplier, dtype): + """Forward + backward with gate/down multipliers vs the fp32 reference.""" torch.manual_seed(0) - atol, rtol = _tol(dtype) + atol, rtol = _silu_tol(dtype) a = torch.randn(*shape, device=device, dtype=dtype) b = torch.randn(*shape, device=device, dtype=dtype) + grad = torch.randn(*shape, device=device, dtype=dtype) - a_ref = a.clone().requires_grad_(True) - b_ref = b.clone().requires_grad_(True) - a_cut = a.clone().requires_grad_(True) - b_cut = b.clone().requires_grad_(True) + ref_a = a.float().detach().requires_grad_(True) + ref_b = b.float().detach().requires_grad_(True) + ref_out = _ref_silu_mul(ref_a, ref_b, gate_multiplier, down_multiplier) + ref_out.backward(grad.float()) - y_ref = _torch_silu_mul_ref(a_ref, b_ref, gate_mult, down_mult) - y_cut = LigerSiLUMulCuteDSLFunction.apply(a_cut, b_cut, gate_mult, down_mult) + test_a = a.clone().detach().requires_grad_(True) + test_b = b.clone().detach().requires_grad_(True) + test_out = LigerSiLUMulCuteDSLFunction.apply(test_a, test_b, gate_multiplier, down_multiplier) + test_out.backward(grad.clone()) - torch.testing.assert_close(y_cut.float(), y_ref.float(), atol=atol, rtol=rtol) + torch.testing.assert_close(test_out.float(), ref_out.float(), atol=atol, rtol=rtol) + torch.testing.assert_close(test_a.grad.float(), ref_a.grad.float(), atol=atol, rtol=rtol) + torch.testing.assert_close(test_b.grad.float(), ref_b.grad.float(), atol=atol, rtol=rtol) - grad = torch.randn_like(y_ref) - y_ref.backward(grad) - y_cut.backward(grad.clone()) - torch.testing.assert_close(a_cut.grad.float(), a_ref.grad.float(), atol=atol, rtol=rtol) - torch.testing.assert_close(b_cut.grad.float(), b_ref.grad.float(), atol=atol, rtol=rtol) +@pytest.mark.flaky(reruns=3, reruns_delay=2) +@pytest.mark.parametrize("shape", _SILU_SHAPES) +@pytest.mark.parametrize("dtype", _SILU_DTYPES) +def test_cutedsl_matches_triton(shape, dtype): + """CuteDSL Function must match the Triton Function (output + grads).""" + torch.manual_seed(0) + atol, rtol = _silu_tol(dtype) + a = torch.randn(*shape, device=device, dtype=dtype) + b = torch.randn(*shape, device=device, dtype=dtype) + grad = torch.randn(*shape, device=device, dtype=dtype) + + tri_a = a.clone().detach().requires_grad_(True) + tri_b = b.clone().detach().requires_grad_(True) + tri_out = LigerSiLUMulFunction.apply(tri_a, tri_b) + tri_out.backward(grad.clone()) + + cut_a = a.clone().detach().requires_grad_(True) + cut_b = b.clone().detach().requires_grad_(True) + cut_out = LigerSiLUMulCuteDSLFunction.apply(cut_a, cut_b) + cut_out.backward(grad.clone()) + + torch.testing.assert_close(cut_out.float(), tri_out.float(), atol=atol, rtol=rtol) + torch.testing.assert_close(cut_a.grad.float(), tri_a.grad.float(), atol=atol, rtol=rtol) + torch.testing.assert_close(cut_b.grad.float(), tri_b.grad.float(), atol=atol, rtol=rtol) @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) @@ -183,36 +209,73 @@ def test_cutedsl_default_multipliers_backward_compat(): torch.testing.assert_close(b1.grad, b2.grad) -@pytest.mark.parametrize( - "shape", - [ - (128 * 4,), # 1-D, tile-aligned for fp32 (vec=4, tile=512) - (123,), # 1-D, unaligned -> predicated - (3, 5, 7), # 3-D, unaligned - (2, 4, 8, 16), # 4-D, exercises shape-flatten path - ], -) +@pytest.mark.parametrize("shape", [(512,), (123,), (2, 256, 512), (3, 5, 7), (2, 4, 8, 16)]) def test_cutedsl_shape_flexibility(shape): - """Forward + backward work for arbitrary >=1-D shapes.""" + """Arbitrary >=1-D shapes: forward + backward vs the fp32 reference; shape preserved.""" + torch.manual_seed(0) dtype = torch.float32 - atol, rtol = _tol(dtype) + atol, rtol = _silu_tol(dtype) a = torch.randn(*shape, device=device, dtype=dtype) b = torch.randn(*shape, device=device, dtype=dtype) + grad = torch.randn(*shape, device=device, dtype=dtype) - a_ref = a.clone().requires_grad_(True) - b_ref = b.clone().requires_grad_(True) - a_cut = a.clone().requires_grad_(True) - b_cut = b.clone().requires_grad_(True) + ref_a = a.clone().detach().requires_grad_(True) + ref_b = b.clone().detach().requires_grad_(True) + ref_out = _ref_silu_mul(ref_a, ref_b) + ref_out.backward(grad) - y_ref = _torch_silu_mul_ref(a_ref, b_ref) - y_cut = LigerSiLUMulCuteDSLFunction.apply(a_cut, b_cut) + test_a = a.clone().detach().requires_grad_(True) + test_b = b.clone().detach().requires_grad_(True) + test_out = LigerSiLUMulCuteDSLFunction.apply(test_a, test_b) + assert test_out.shape == a.shape + test_out.backward(grad.clone()) - assert y_cut.shape == y_ref.shape - torch.testing.assert_close(y_cut, y_ref, atol=atol, rtol=rtol) + torch.testing.assert_close(test_out.float(), ref_out.float(), atol=atol, rtol=rtol) + torch.testing.assert_close(test_a.grad.float(), ref_a.grad.float(), atol=atol, rtol=rtol) + torch.testing.assert_close(test_b.grad.float(), ref_b.grad.float(), atol=atol, rtol=rtol) - grad = torch.randn_like(y_ref) - y_ref.backward(grad) - y_cut.backward(grad.clone()) - torch.testing.assert_close(a_cut.grad, a_ref.grad, atol=atol, rtol=rtol) - torch.testing.assert_close(b_cut.grad, b_ref.grad, atol=atol, rtol=rtol) +@pytest.mark.parametrize("shape, dtype", [((73, 13824), torch.bfloat16), ((4, 2048), torch.float32)]) +def test_cutedsl_cuda_graph(shape, dtype): + """CUDA-graph capture + replay of the forward must be bit-identical to eager. + + Guards the CuteDSL stream fix: the compiled kernel must run on PyTorch's + *current* (capture) stream, not the default/null stream, or replay would + read stale data / mis-order and diverge from eager. + """ + if dtype == torch.bfloat16 and not supports_bfloat16(): + pytest.skip("bfloat16 not supported on this GPU") + torch.manual_seed(0) + static_a = torch.randn(*shape, device=device, dtype=dtype) + static_b = torch.randn(*shape, device=device, dtype=dtype) + + def run(): + with torch.no_grad(): + return LigerSiLUMulCuteDSLFunction.apply(static_a, static_b) + + # Warm up (compile) on a side stream before capture. + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + for _ in range(3): + run() + torch.cuda.current_stream().wait_stream(s) + + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + static_out = run() + + # Replay on fresh data copied into the static input buffers. + fresh_a = torch.randn(*shape, device=device, dtype=dtype) + fresh_b = torch.randn(*shape, device=device, dtype=dtype) + static_a.copy_(fresh_a) + static_b.copy_(fresh_b) + g.replay() + torch.cuda.synchronize() + graph_out = static_out.clone() + + with torch.no_grad(): + eager_out = LigerSiLUMulCuteDSLFunction.apply(fresh_a, fresh_b) + + max_abs_diff = (graph_out.float() - eager_out.float()).abs().max().item() + assert max_abs_diff == 0.0, f"graph vs eager mismatch: max_abs_diff={max_abs_diff}"