linkedin/Liger-Kernel
1
0
Fork 0
mirror of https://github.com/linkedin/Liger-Kernel.git synced 2026-10-12 02:17:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

[cuTeDSL] Make SwiGLU CUDA-graph capturable + harmonize plain SiLU-Mul tests (#1408)

Browse Source 
## 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 <[email protected]>

Co-authored-by: vaibhavjindal <[email protected]>
Co-authored-by: Copilot <[email protected]>
This commit is contained in:
Vaibhav Jindal
2026-08-26 17:14:47 +00:00
committed by GitHub
co-authored by Copilot vaibhavjindal
parent 4310065d8c
commit b28a0ca4f4
4 changed files with 590 additions and 227 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/liger_kernel/ops/cutedsl/ops/swiglu.py
+61 -11
View File
@@ -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
test/transformers/test_cutile_swiglu.py
+212 -111
View File
@@ -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)
test/transformers/test_swiglu.py
+149
View File
@@ -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_<be>_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}"
test/transformers/test_swiglu_cutedsl.py
+168 -105
View File
@@ -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}"