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

Add SM100 fused linear SwiGLU (#1340)

Browse Source 
## Summary

Adds a unified CuTe DSL SwiGLU implementation on top of the current
`main`
SM100 GEMM:

- a differentiable elementwise forward/backward replacement for
  `LigerSiLUMulFunction`;
- an exact-SM100 fused-linear forward path that computes gate and up
projections
in one grouped GEMM and applies SwiGLU in the epilogue without
materializing
either projection, reading a plainly **stacked** (or **separate**)
gate/up
  weight directly with no pre-packing and no extra weight copy;
- safe fallback to two linear projections plus the elementwise CuTe DSL
kernel
  for unsupported architectures, dtypes, and K alignment.

This PR is rebased directly onto `main`; the former dependency on #1339
has
been removed.

## Implementation

- The fused path consumes a plainly **stacked** `[2N, K]` weight (gate
rows
`[0:N]` then up rows `[N:2N]`) directly:
`run_grouped_epilogue_gemm_stacked`
rebuilds the B tensor with a nested `((32, 2, T):(K, N·K, 32·K))`
N-layout so
gate/up for the same 32 output features still land as adjacent
accumulator
fragments local to the same SM100 CTA — **without any physical repack or
a
  second weight copy**. The MMA mainloop and the FP32 TMEM epilogue are
  unchanged from the interleaved-pack version.
- The grouped custom-epilogue mode is integrated into the current
reusable
`_sm100_gemm.py` without restoring the retired FLCE row-reduction paths.
- The existing single-fragment FLCE configuration remains unchanged.
- Native fused execution requires exact compute capability `(10, 0)`,
  FP16/BF16 operands, and K divisible by 64.
- Unsupported inputs run two `torch.nn.functional.linear` calls followed
by the
CuTe DSL elementwise kernel (separate tensors are used directly here, no
copy).
- Elementwise and fused kernels launch under the input tensor's CUDA
device
guard and on PyTorch's active CUDA stream for both TVM-FFI and DLPack
calling
  paths.

## Weight API

```python
# Stacked (zero-copy): gate rows [0:N] then up rows [N:2N]
# e.g. a model's fused gate_up_proj.weight, or torch.cat([gate, up], dim=0)
h = fused_swiglu(x, gate_up_weight)

# Separate tensors: concatenated once internally for the native path
h = fused_swiglu(x, gate_weight, up_weight)
```

- The **stacked** form is read as a re-strided view of a single buffer,
so there
is **no pre-pack step and no transient 2× weight copy** (previously
incurred by
the 32-row interleave). This resolves the review feedback on the
weight-packing
  memory spike.
- **Separate** `gate`/`up` tensors are accepted for convenience; the
native path
  concatenates them once (a strided view cannot span two independent
allocations), while the fallback path consumes them with no copy. Prefer
the
  stacked form on hot paths.
- Because the weight stays a standard `[2N, K]` (or two `[N, K]`)
tensor, it is
directly usable as an `nn.Parameter` and is checkpoint-/shard-friendly.
- `output_features` is **keyword-only** and inferred as `N` when
omitted;
trailing padded columns (when `N % 32 != 0`) are trimmed automatically.
- `pack_swiglu_weights` is retained as a thin `torch.cat` convenience
helper.

## B200 optimization

NCU at `M=8192, H=4096, N=14336` classified the grouped kernel as
compute-bound:

- compute throughput: **96.07%**
- memory throughput: **62.11%**
- registers/thread: **96**
- dynamic shared memory/block: **213.25 KiB**
- achieved occupancy: **8.63%**

A manual pipeline-stage/swizzle sweep selected:

- five A/B stages at 1K tokens;
- six stages from 2K through 16K;
- four stages with the existing eight-way swizzle from 32K upward.

The large-M specialization raises normalized speedup from
`1.015x`/`0.981x`
at 32K/64K to `1.053x`/`1.052x`.

## BF16 benchmarks

NVIDIA B200, `H=4096`, `N=14336`, seven warmups, 21 interleaved samples.
The
stacked weight is read as a zero-copy view, so nothing is excluded from
latency
on the fast path (separate tensors incur a one-time `torch.cat`). Peak
memory is
measured in isolated provider processes and includes input, one weight
representation, output, and temporary PyTorch CUDA allocations.

| Tokens | Two GEMMs + SwiGLU (ms) | Fused SwiGLU (ms) | Speedup |
Baseline peak (MiB) | Fused peak (MiB) | Memory ratio |
|---:|---:|---:|---:|---:|---:|---:|
| 1,024 | 0.214208 | 0.223136 | 0.960x | 348 | 260 | 1.34x |
| 2,048 | 0.355776 | 0.356896 | 0.997x | 440 | 296 | 1.49x |
| 4,096 | 0.691104 | 0.647360 | 1.068x | 624 | 368 | 1.70x |
| 8,192 | 1.285664 | 1.217984 | 1.056x | 992 | 512 | 1.94x |
| 16,384 | 2.565120 | 2.412224 | 1.063x | 1,728 | 800 | 2.16x |
| 32,768 | 5.130048 | 4.873984 | 1.053x | 3,200 | 1,376 | 2.33x |
| 65,536 | 10.174336 | 9.672512 | 1.052x | 6,144 | 2,528 | 2.43x |

## Correctness

- CuTe DSL SwiGLU, upstream elementwise SwiGLU, SM100 GEMM, and FLCE
regression
  suites: **301 passed, 1 skipped**
- Stacked zero-copy path, separate-weight path (byte-identical to
stacked), and
weight-API validation are covered in `test/cutedsl/test_swiglu.py`
(verified
  alongside `test/cutedsl/test_sm100_gemm.py` on B200 SM100)
- Current-stream coverage: elementwise forward, elementwise backward,
and
  fused-linear forward
- FP16/BF16 fused correctness, padded output tails, exact-SM100
dispatch, FP32
fallback, arbitrary elementwise shapes, multipliers, and PyTorch/Triton
parity
- `make checkstyle`: **passed**

Validation environment:

- NVIDIA B200, SM100, driver `580.105.08`
- PyTorch `2.13.0.1+cu130`
- Triton `3.7.1`
- `nvidia-cutlass-dsl` `4.6.0`
- `apache-tvm-ffi` `0.1.12`

## Scope

- `src/liger_kernel/ops/cutedsl/ops/swiglu.py`
- `src/liger_kernel/ops/cutedsl/ops/_sm100_gemm.py`
- `test/cutedsl/test_swiglu.py`

---------

Co-authored-by: Stella Dey <[email protected]>
Co-authored-by: Copilot <[email protected]>
Co-authored-by: Vaibhav Jindal <[email protected]>
This commit is contained in:
Justin Hu
2026-09-09 23:11:17 +00:00
committed by GitHub
co-authored by Stella Dey Copilot Vaibhav Jindal
parent ad2436e80e
commit d2b5474de7
3 changed files with 555 additions and 64 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/liger_kernel/ops/cutedsl/ops/_sm100_gemm.py
+206 -64
View File
@@ -75,6 +75,12 @@ _EPI_TILE = (_CTA_M, _EPI_N)
_CLUSTER_SHAPE_MN = (2, 1)
_NUM_AB_STAGES = 6
_NUM_AB_STAGES_GROUPED_LARGE_M = 4
_NUM_AB_STAGES_GROUPED_SHORT_M = 5
_SWIZZLE_SIZES_GROUPED = (8, 4, 2)
_SWIZZLE_GROUPED_SHORT_M = 4
_M_TILES_GROUPED_LARGE_M = 256
_M_TILES_GROUPED_SHORT_M = 8
_NUM_ACC_STAGES = 2
_NUM_OUT_STAGES = 2
_NUM_TMEM_COLS = 512
@@ -91,10 +97,14 @@ _TMEM_DEALLOC_BARRIER = 2
_EPILOGUE_BARRIER = 3
K_ALIGNMENT = _MMA_K
EPILOGUE_TILE_SIZE = _EPI_N
__all__ = [
"EPILOGUE_TILE_SIZE",
"K_ALIGNMENT",
"run_epilogue_gemm",
"run_grouped_epilogue_gemm",
"run_grouped_epilogue_gemm_stacked",
]
@@ -280,6 +290,7 @@ def _kernel(
num_ab_stages: cutlass.Constexpr,
io_dtype: cutlass.Constexpr,
epilogue: cutlass.Constexpr,
epilogue_fragments: cutlass.Constexpr,
use_tma_output: cutlass.Constexpr,
):
warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx())
@@ -412,18 +423,25 @@ def _kernel(
thr_mma = tiled_mma.get_slice(mma_tile_coord_v)
t_cg_a = thr_mma.partition_A(g_a_mkl)
t_cg_b = thr_mma.partition_B(g_b_nkl)
m_c_mnl = cute.make_identity_tensor((m_a_mkl.shape[0], m_b_nkl.shape[0], 1))
m_c_mnl = cute.make_identity_tensor((m_a_mkl.shape[0], cute.size(m_b_nkl, mode=[0]), 1))
g_c_mnl = cute.local_tile(
m_c_mnl,
cute.slice_(_MMA_TILER, (None, None, 0)),
(None, None, None),
)
t_cg_c = thr_mma.partition_C(g_c_mnl)
g_out_mnl = cute.local_tile(
m_out_mnl,
cute.slice_(_MMA_TILER, (None, None, 0)),
(None, None, None),
)
if const_expr(epilogue_fragments == 2):
g_out_mnl = cute.local_tile(
m_out_mnl,
(_MMA_M, _MMA_N // 2, 1),
(None, None, None),
)
else:
g_out_mnl = cute.local_tile(
m_out_mnl,
cute.slice_(_MMA_TILER, (None, None, 0)),
(None, None, None),
)
t_cg_out = thr_mma.partition_C(g_out_mnl)
a_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape)
@@ -503,8 +521,16 @@ def _kernel(
t_rs_r_out,
t_rs_s_out,
) = _make_tma_epilogue_partitions(tidx, t_ct_acc_base, t_cg_c, s_out, io_dtype)
t_cg_out_transformed = utils.gemm.sm100.transform_partitioned_tensor_layout(t_cg_out)
t_cg_out_epi = cute.flat_divide(t_cg_out_transformed, _EPI_TILE)
if const_expr(epilogue_fragments == 2):
t_tr_r_lhs = cute.make_rmem_tensor(t_tr_r_acc.shape, Float32)
t_rs_r_lhs = tiled_copy_r2s.retile(t_tr_r_lhs)
t_cg_out_epi = cute.flat_divide(
t_cg_out[((None, None), 0, 0, 0, None, None, None)],
_EPI_TILE,
)
else:
t_cg_out_transformed = utils.gemm.sm100.transform_partitioned_tensor_layout(t_cg_out)
t_cg_out_epi = cute.flat_divide(t_cg_out_transformed, _EPI_TILE)
b_sg_s_out, b_sg_g_out_partitioned = cpasync.tma_partition(
tma_atom_out,
0,
@@ -525,7 +551,7 @@ def _kernel(
_NUM_ACC_STAGES,
)
m = m_a_mkl.shape[0]
n = m_b_nkl.shape[0]
n = cute.size(m_b_nkl, mode=[0])
while work_tile.is_valid_tile:
tile_coord = work_tile.tile_idx
@@ -547,50 +573,105 @@ def _kernel(
acc_pipeline.consumer_wait(acc_consumer_state)
t_tr_t_acc = cute.group_modes(t_tr_t_acc, 3, cute.rank(t_tr_t_acc))
fragment_count = _MMA_N // _EPI_N
for fragment_idx in cutlass.range_constexpr(fragment_count):
cute.copy(
tiled_copy_t2r,
t_tr_t_acc[(None, None, None, fragment_idx)],
t_tr_r_acc,
)
if fragment_idx == fragment_count - 1:
with cute.arch.elect_one():
acc_pipeline.consumer_release(acc_consumer_state)
acc_consumer_state.advance()
epilogue(t_rs_r_acc, t_rs_r_out)
if const_expr(not use_tma_output):
t_rs_g_c = tiled_copy_r2s.retile(t_tr_g_c[(None, None, None, fragment_idx)])
for element in cutlass.range(
cute.size(t_rs_r_acc),
unroll_full=True,
):
output_coord = t_rs_g_c[element]
output_row = output_coord[0]
output_col = output_coord[1]
if output_row < m and output_col < n:
m_out_direct_mnl[output_row, output_col, 0] = t_rs_r_out[element]
if const_expr(use_tma_output):
out_buffer = fragment_idx % _NUM_OUT_STAGES
if warp_idx == _EPILOGUE_WARPS[0]:
out_pipeline.producer_acquire()
epilogue_barrier.arrive_and_wait()
if const_expr(epilogue_fragments == 1):
fragment_count = _MMA_N // _EPI_N
for fragment_idx in cutlass.range_constexpr(fragment_count):
cute.copy(
tiled_copy_r2s,
t_rs_r_out,
t_rs_s_out[(None, None, None, out_buffer)],
tiled_copy_t2r,
t_tr_t_acc[(None, None, None, fragment_idx)],
t_tr_r_acc,
)
cute.arch.fence_proxy("async.shared", space="cta")
epilogue_barrier.arrive_and_wait()
if warp_idx == _EPILOGUE_WARPS[0]:
if fragment_idx == fragment_count - 1:
with cute.arch.elect_one():
acc_pipeline.consumer_release(acc_consumer_state)
acc_consumer_state.advance()
epilogue(t_rs_r_acc, t_rs_r_out)
if const_expr(not use_tma_output):
t_rs_g_c = tiled_copy_r2s.retile(t_tr_g_c[(None, None, None, fragment_idx)])
for element in cutlass.range(
cute.size(t_rs_r_acc),
unroll_full=True,
):
output_coord = t_rs_g_c[element]
output_row = output_coord[0]
output_col = output_coord[1]
if output_row < m and output_col < n:
m_out_direct_mnl[output_row, output_col, 0] = t_rs_r_out[element]
if const_expr(use_tma_output):
out_buffer = fragment_idx % _NUM_OUT_STAGES
if warp_idx == _EPILOGUE_WARPS[0]:
out_pipeline.producer_acquire()
epilogue_barrier.arrive_and_wait()
cute.copy(
tma_atom_out,
b_sg_s_out[(None, out_buffer)],
b_sg_g_out[(None, fragment_idx)],
tiled_copy_r2s,
t_rs_r_out,
t_rs_s_out[(None, None, None, out_buffer)],
)
out_pipeline.producer_commit()
cute.arch.fence_proxy("async.shared", space="cta")
epilogue_barrier.arrive_and_wait()
if warp_idx == _EPILOGUE_WARPS[0]:
cute.copy(
tma_atom_out,
b_sg_s_out[(None, out_buffer)],
b_sg_g_out[(None, fragment_idx)],
)
out_pipeline.producer_commit()
if const_expr(epilogue_fragments == 2):
pair_count = _MMA_N // (2 * _EPI_N)
for pair_idx in cutlass.range_constexpr(pair_count):
lhs_subtile = pair_idx * 2
cute.copy(
tiled_copy_t2r,
t_tr_t_acc[(None, None, None, lhs_subtile + 1)],
t_tr_r_acc,
)
cute.copy(
tiled_copy_t2r,
t_tr_t_acc[(None, None, None, lhs_subtile)],
t_tr_r_lhs,
)
if pair_idx == pair_count - 1:
with cute.arch.elect_one():
acc_pipeline.consumer_release(acc_consumer_state)
acc_consumer_state.advance()
epilogue(t_rs_r_lhs, t_rs_r_acc, t_rs_r_out)
if const_expr(not use_tma_output):
t_rs_g_c = tiled_copy_r2s.retile(t_tr_g_c[(None, None, None, lhs_subtile)])
for element in cutlass.range(
cute.size(t_rs_r_acc),
unroll_full=True,
):
output_coord = t_rs_g_c[element]
output_row = output_coord[0]
physical_col = output_coord[1]
physical_base = n_tile * _MMA_N + pair_idx * _EPI_N * 2
logical_col = n_tile * (_MMA_N // 2) + pair_idx * _EPI_N + physical_col - physical_base
if output_row < m and physical_col < n and logical_col < m_out_direct_mnl.shape[1]:
m_out_direct_mnl[output_row, logical_col, 0] = t_rs_r_out[element]
if const_expr(use_tma_output):
out_buffer = pair_idx % _NUM_OUT_STAGES
if warp_idx == _EPILOGUE_WARPS[0]:
out_pipeline.producer_acquire()
epilogue_barrier.arrive_and_wait()
cute.copy(
tiled_copy_r2s,
t_rs_r_out,
t_rs_s_out[(None, None, None, out_buffer)],
)
cute.arch.fence_proxy("async.shared", space="cta")
epilogue_barrier.arrive_and_wait()
if warp_idx == _EPILOGUE_WARPS[0]:
cute.copy(
tma_atom_out,
b_sg_s_out[(None, out_buffer)],
b_sg_g_out[(None, pair_idx)],
)
out_pipeline.producer_commit()
if const_expr(use_tma_output) and warp_idx == _EPILOGUE_WARPS[0]:
out_pipeline.producer_tail()
@@ -606,6 +687,8 @@ def _host(
m_out: cute.Tensor,
num_ab_stages: cutlass.Constexpr,
epilogue: cutlass.Constexpr,
epilogue_fragments: cutlass.Constexpr,
b_stacked: cutlass.Constexpr,
use_tma_output: cutlass.Constexpr,
swizzle_size: cutlass.Constexpr,
max_active_clusters: cutlass.Constexpr,
@@ -615,10 +698,27 @@ def _host(
m_a.iterator,
cute.append(m_a.layout, cute.make_layout(1)),
)
m_b_nkl = cute.make_tensor(
m_b.iterator,
cute.append(m_b.layout, cute.make_layout(1)),
)
if const_expr(b_stacked):
# ``m_b`` is a plain row-major ``[2N', K]`` stack (gate rows ``[0:N']``,
# up rows ``[N':2N']``). Present it to the grouped kernel with a nested
# N-layout so that N-tile ``t`` gathers ``[gate 0:32, up 0:32, gate
# 32:64, up 32:64, ...]`` -- byte-identical to the interleaved packed
# tensor the epilogue expects, but WITHOUT any physical repack.
b_k = cute.size(m_b, mode=[1])
b_np = cute.size(m_b, mode=[0]) // 2
b_t = b_np // _EPI_N
m_b_nkl = cute.make_tensor(
m_b.iterator,
cute.make_layout(
((_EPI_N, 2, b_t), b_k, 1),
stride=((b_k, b_np * b_k, _EPI_N * b_k), 1, 0),
),
)
else:
m_b_nkl = cute.make_tensor(
m_b.iterator,
cute.append(m_b.layout, cute.make_layout(1)),
)
m_out_mnl = cute.make_tensor(
m_out.iterator,
cute.append(m_out.layout, cute.make_layout(1)),
@@ -693,7 +793,7 @@ def _host(
_EPI_TILE,
)
m_c_mnl = cute.make_identity_tensor((m_a_mkl.shape[0], m_b_nkl.shape[0], 1))
m_c_mnl = cute.make_identity_tensor((m_a_mkl.shape[0], cute.size(m_b_nkl, mode=[0]), 1))
g_c_mnl = cute.zipped_divide(
m_c_mnl,
tiler=(_CTA_M, _MMA_N),
@@ -725,6 +825,7 @@ def _host(
num_ab_stages,
io_dtype,
epilogue,
epilogue_fragments,
use_tma_output,
).launch(
grid=grid,
@@ -811,14 +912,14 @@ def _fake_dynamic_tensor(tensor, leading_dim, assumed_align, stride_divisibility
@functools.cache
def _validate_epilogue_callback(epilogue):
def _validate_epilogue_callback(epilogue, expected_parameters):
wrapped = getattr(epilogue, "__wrapped__", None)
if wrapped is None or getattr(epilogue, "_dsl_cls", None) is None:
raise TypeError("epilogue must be a module-level function decorated with @cute.jit.")
if "<locals>" in wrapped.__qualname__ or wrapped.__closure__ or wrapped.__code__.co_freevars:
raise ValueError("epilogue must be module-level and cannot close over local variables.")
if len(inspect.signature(wrapped).parameters) != 2:
raise TypeError("epilogue must accept exactly 2 parameters.")
if len(inspect.signature(wrapped).parameters) != expected_parameters:
raise TypeError(f"epilogue must accept exactly {expected_parameters} parameters.")
return wrapped.__module__, wrapped.__qualname__, wrapped
@@ -831,6 +932,7 @@ def _validate_epilogue_signature(
aligned,
devices,
dtypes,
epilogue_fragments,
):
for name, shape, is_contiguous, is_aligned, device in zip(
("a", "b", "out"),
@@ -860,12 +962,14 @@ def _validate_epilogue_signature(
raise ValueError(f"a and b K dimensions must match, got {a_shape[1]} and {b_shape[1]}.")
if a_shape[1] % K_ALIGNMENT:
raise ValueError(f"K must be divisible by {K_ALIGNMENT}, got {a_shape[1]}.")
expected_shape = (a_shape[0], b_shape[0])
if epilogue_fragments == 2 and b_shape[0] % (2 * EPILOGUE_TILE_SIZE):
raise ValueError(f"Grouped b rows must be divisible by {2 * EPILOGUE_TILE_SIZE}, got {b_shape[0]}.")
expected_shape = (a_shape[0], b_shape[0] // epilogue_fragments)
if out_shape != expected_shape:
raise ValueError(f"out must have shape {expected_shape}, got {out_shape}.")
def _validate_epilogue_inputs(a, b, out):
def _validate_epilogue_inputs(a, b, out, epilogue_fragments):
if out.stride(1) != 1 or out.stride(0) < out.shape[1]:
raise ValueError("out must be a dense row-major tensor or row-padded row-major view.")
tensors = (a, b, out)
@@ -875,6 +979,7 @@ def _validate_epilogue_inputs(a, b, out):
tuple(tensor.data_ptr() % 16 == 0 for tensor in tensors),
tuple(tensor.device for tensor in tensors),
tuple(tensor.dtype for tensor in tensors),
epilogue_fragments,
)
@@ -889,11 +994,20 @@ def _select_epilogue_config(a):
return _NUM_AB_STAGES, 1
def _run_epilogue_gemm(a, b, out, epilogue):
_validate_epilogue_inputs(a, b, out)
epilogue_key = _validate_epilogue_callback(epilogue)
def _run_epilogue_gemm(a, b, out, epilogue, epilogue_fragments, b_stacked=False):
_validate_epilogue_inputs(a, b, out, epilogue_fragments)
epilogue_key = _validate_epilogue_callback(epilogue, epilogue_fragments + 1)
current_stream = _current_stream(a.device)
num_ab_stages, swizzle_size = _select_epilogue_config(a)
if epilogue_fragments == 2:
m_tiles = (a.shape[0] + _CTA_M - 1) // _CTA_M
if m_tiles == _M_TILES_GROUPED_SHORT_M:
num_ab_stages = _NUM_AB_STAGES_GROUPED_SHORT_M
swizzle_size = _SWIZZLE_GROUPED_SHORT_M
else:
num_ab_stages = _NUM_AB_STAGES_GROUPED_LARGE_M if m_tiles >= _M_TILES_GROUPED_LARGE_M else _NUM_AB_STAGES
swizzle_size = next((size for size in _SWIZZLE_SIZES_GROUPED if m_tiles % size == 0), 1)
else:
num_ab_stages, swizzle_size = _select_epilogue_config(a)
use_tma_output = out.stride(0) * out.element_size() % 16 == 0
max_active_clusters = _max_active_clusters(a.device.index)
@@ -903,6 +1017,8 @@ def _run_epilogue_gemm(a, b, out, epilogue):
b.dtype,
out.dtype,
epilogue_key,
epilogue_fragments,
b_stacked,
use_tma_output,
swizzle_size,
num_ab_stages,
@@ -935,6 +1051,8 @@ def _run_epilogue_gemm(a, b, out, epilogue):
m_out,
num_ab_stages,
epilogue,
epilogue_fragments,
b_stacked,
use_tma_output,
swizzle_size,
max_active_clusters,
@@ -960,4 +1078,28 @@ def _run_epilogue_gemm(a, b, out, epilogue):
def run_epilogue_gemm(a, b, out, epilogue):
"""Run ``a @ b.T`` and apply ``epilogue(accumulator, output)`` per fragment."""
with _device_guard(a.device):
_run_epilogue_gemm(a, b, out, epilogue)
_run_epilogue_gemm(a, b, out, epilogue, epilogue_fragments=1)
def run_grouped_epilogue_gemm(a, packed_b, out, epilogue):
"""Apply ``epilogue(first, second, output)`` to adjacent packed fragments.
``packed_b`` interleaves two logical matrices in 32-row tiles, allowing the
callback to fuse arbitrary binary projection math without materialization.
"""
with _device_guard(a.device):
_run_epilogue_gemm(a, packed_b, out, epilogue, epilogue_fragments=2)
def run_grouped_epilogue_gemm_stacked(a, stacked_b, out, epilogue):
"""Grouped epilogue GEMM over a plainly *stacked* ``[2N', K]`` weight.
Identical numerics to :func:`run_grouped_epilogue_gemm`, but ``stacked_b``
holds the two logical matrices as a simple concatenation (rows ``[0:N']``
then ``[N':2N']``) instead of the bespoke 32-row interleave. The kernel
reinterprets the strides on the fly (nested N-layout), so no pre-packing and
no extra weight copy is needed -- the epilogue still sees adjacent
``(first, second)`` fragments.
"""
with _device_guard(a.device):
_run_epilogue_gemm(a, stacked_b, out, epilogue, epilogue_fragments=2, b_stacked=True)
src/liger_kernel/ops/cutedsl/ops/swiglu.py
+176
View File
@@ -47,6 +47,8 @@ The public API mirrors ``liger_kernel.ops.swiglu``:
``swiglu_forward`` / ``swiglu_backward`` / ``LigerSiLUMulCuteDSLFunction``.
"""
import functools
import torch
try:
@@ -63,6 +65,9 @@ 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._sm100_gemm import EPILOGUE_TILE_SIZE
from liger_kernel.ops.cutedsl.ops._sm100_gemm import K_ALIGNMENT
from liger_kernel.ops.cutedsl.ops._sm100_gemm import run_grouped_epilogue_gemm_stacked
from liger_kernel.ops.cutedsl.ops.utils import make_fake_tensor
from liger_kernel.ops.cutedsl.ops.utils import torch2cute_dtype_map
from liger_kernel.ops.utils import ensure_contiguous
@@ -686,6 +691,177 @@ def _get_compiled_vec(kind: str, dtype: torch.dtype, vec: int, device_index: int
return fn
# ---------------------------------------------------------------------------
# Fused linear SwiGLU
# ---------------------------------------------------------------------------
@cute.jit
def _fused_swiglu_epilogue(gate, up, out):
_silu_mul_fwd_packed(gate, up, out, cutlass.Float32(1.0))
def _validate_and_stack_gate_up(gate_weight, up_weight):
"""Validate two ``[N, K]`` projections and concatenate into ``[2N, K]``.
This is the *separate weights* entry point. Because ``gate_weight`` and
``up_weight`` are independent allocations, forming the single contiguous
buffer the SM100 kernel reads requires one copy (a strided view cannot span
two allocations). Callers on a hot path should pre-stack once and pass the
stacked tensor to :func:`fused_swiglu` to avoid this per-call copy.
"""
if gate_weight.ndim != 2 or up_weight.ndim != 2:
raise ValueError("gate_weight and up_weight must both be 2D tensors.")
if gate_weight.shape != up_weight.shape:
raise ValueError(
f"gate_weight and up_weight must have identical shapes, got {gate_weight.shape} and {up_weight.shape}."
)
if gate_weight.dtype != up_weight.dtype:
raise TypeError(
f"gate_weight and up_weight must have the same dtype, got {gate_weight.dtype} and {up_weight.dtype}."
)
if gate_weight.device != up_weight.device:
raise ValueError(
f"gate_weight and up_weight must be on the same device, got {gate_weight.device} and {up_weight.device}."
)
return torch.cat((gate_weight, up_weight), dim=0)
def pack_swiglu_weights(gate_weight, up_weight):
"""Stack ``[N, K]`` gate/up weights into a plain ``[2N, K]`` tensor.
The fused SM100 kernel now consumes a *simply stacked* weight (gate rows
followed by up rows) directly, reinterpreting the strides on the fly, so no
bespoke interleave is materialized. This helper is kept for convenience and
API compatibility: it is a single ``torch.cat`` (one contiguous copy) for
callers that hold two separate tensors. Callers that already have a fused
``gate_up_proj`` weight can skip this and pass it straight to
:func:`fused_swiglu`; callers with two separate tensors can also pass them
directly as ``fused_swiglu(a, gate, up)``.
Returns ``(stacked_weight, output_features)`` where ``output_features`` is
the per-projection output dimension ``N``.
"""
return _validate_and_stack_gate_up(gate_weight, up_weight), gate_weight.shape[0]
@functools.lru_cache(maxsize=32)
def _validate_fused_swiglu_signature(a_shape, weight_shape, output_features):
if len(a_shape) != 2 or len(weight_shape) != 2:
raise ValueError("a and gate_up_weight must both be 2D tensors.")
if a_shape[1] != weight_shape[1]:
raise ValueError(f"Input and weight K dimensions must match, got {a_shape[1]} and {weight_shape[1]}.")
if weight_shape[0] % 2 != 0:
raise ValueError(f"gate_up_weight must have an even number of rows (2N), got {weight_shape[0]}.")
features = weight_shape[0] // 2
if output_features is None:
output_features = features
if not 0 < output_features <= features:
raise ValueError(f"output_features must be in [1, {features}], got {output_features}.")
return output_features
def _native_fused_swiglu_supported(a):
if a.device.type != "cuda" or a.dtype not in (torch.float16, torch.bfloat16) or a.shape[1] % K_ALIGNMENT != 0:
return False
device_id = a.device.index if a.device.index is not None else torch.cuda.current_device()
return infer_device_arch(device_id) == "blackwell" and torch.cuda.get_device_capability(a.device) == (10, 0)
def _pad_stacked_weight(gate_up_weight, features):
"""Pad each projection up to a multiple of ``EPILOGUE_TILE_SIZE``.
Returns ``(stacked, padded_features)``. When ``features`` is already a
multiple of the tile size (the common case, e.g. LLaMA/Qwen MLPs) the input
is returned unchanged -- no copy, no memory spike.
"""
padded = ((features + EPILOGUE_TILE_SIZE - 1) // EPILOGUE_TILE_SIZE) * EPILOGUE_TILE_SIZE
if padded == features:
return gate_up_weight, padded
stacked = gate_up_weight.new_zeros(2 * padded, gate_up_weight.shape[1])
stacked[:features].copy_(gate_up_weight[:features])
stacked[padded : padded + features].copy_(gate_up_weight[features:])
return stacked, padded
def _fused_swiglu_sm100(a, stacked_gate_up_weight, output_features):
padded_features = stacked_gate_up_weight.shape[0] // 2
out = torch.empty(
a.shape[0],
padded_features,
device=a.device,
dtype=a.dtype,
)
run_grouped_epilogue_gemm_stacked(
a.contiguous(),
stacked_gate_up_weight.contiguous(),
out,
_fused_swiglu_epilogue,
)
return out[:, :output_features]
def fused_swiglu(a, gate_weight, up_weight=None, *, output_features=None):
"""Fused gate/up projections + SwiGLU with an SM100 fast path.
Two calling conventions are supported:
* **Stacked (zero-copy)** -- ``fused_swiglu(a, gate_up_weight)`` where
``gate_up_weight`` is a plainly stacked ``[2N, K]`` tensor: the gate rows
``[0:N]`` followed by the up rows ``[N:2N]`` (e.g. ``torch.cat([gate,
up])`` or a model's fused ``gate_up_proj.weight``). The SM100 kernel
reinterprets the stacked strides on the fly and reads it directly with
**no extra weight copy**.
* **Separate** -- ``fused_swiglu(a, gate_weight, up_weight)`` where each is
``[N, K]``. They are concatenated into one ``[2N, K]`` buffer before the
native kernel runs (**one copy** -- a strided view cannot span two
separate allocations; the fallback path uses them directly with no copy).
Prefer the stacked form on hot paths to avoid this per-call concat.
``output_features`` (keyword-only) optionally trims trailing padded columns
of the per-projection output; it is inferred as ``N`` when omitted. Returns
``silu(x @ gate.T) * (x @ up.T)`` of shape ``[M, N]``.
"""
if a.device.type != "cuda" or gate_weight.device.type != "cuda":
raise ValueError("a and weights must be CUDA tensors.")
if up_weight is not None:
if not isinstance(up_weight, torch.Tensor):
raise TypeError(
"up_weight must be a tensor of shape [N, K]; 'output_features' is keyword-only "
"(call fused_swiglu(a, gate_up_weight, output_features=...))."
)
if up_weight.device.type != "cuda":
raise ValueError("a and weights must be CUDA tensors.")
# Separate projections: stack into one buffer (option (a) -- one copy).
gate_up_weight = _validate_and_stack_gate_up(gate_weight, up_weight)
else:
# Already stacked [2N, K] -- consumed as a zero-copy strided view.
gate_up_weight = gate_weight
if a.device != gate_up_weight.device:
raise ValueError(
f"a and gate_up_weight must be on the same device, got {a.device} and {gate_up_weight.device}."
)
if a.dtype != gate_up_weight.dtype:
raise TypeError(f"a and gate_up_weight must have the same dtype, got {a.dtype} and {gate_up_weight.dtype}.")
_validate_supported_dtype(a.dtype)
output_features = _validate_fused_swiglu_signature(
tuple(a.shape),
tuple(gate_up_weight.shape),
output_features,
)
features = gate_up_weight.shape[0] // 2
stacked, padded_features = _pad_stacked_weight(gate_up_weight, features)
if _native_fused_swiglu_supported(a):
return _fused_swiglu_sm100(a, stacked, output_features)
gate_weight = stacked[:padded_features]
up_weight = stacked[padded_features:]
gate = torch.nn.functional.linear(a, gate_weight)
up = torch.nn.functional.linear(a, up_weight)
return swiglu_forward(gate, up)[2][:, :output_features]
# ---------------------------------------------------------------------------
# Public functional API
# ---------------------------------------------------------------------------
test/cutedsl/test_swiglu.py
+173
View File
@@ -31,8 +31,13 @@ pytestmark = pytest.mark.skipif(
reason="nvidia-cutlass-dsl + CUDA GPU required for CuteDSL SwiGLU",
)
# Native fused-linear SwiGLU requires exact SM100 (Blackwell) compute capability.
sm100_available = cutedsl_available and torch.cuda.is_available() and torch.cuda.get_device_capability() == (10, 0)
if cutedsl_available:
from liger_kernel.ops.cutedsl.ops.swiglu import LigerSiLUMulCuteDSLFunction
from liger_kernel.ops.cutedsl.ops.swiglu import fused_swiglu
from liger_kernel.ops.cutedsl.ops.swiglu import pack_swiglu_weights
from liger_kernel.ops.cutedsl.ops.swiglu import swiglu_forward as cutedsl_forward
@@ -278,3 +283,171 @@ def test_cutedsl_cuda_graph(shape, dtype):
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.skipif(not sm100_available, reason="fused linear SwiGLU requires SM100")
@pytest.mark.parametrize(
"shape",
[
(33, 128, 96), # non-tile-aligned tokens
(128, 256, 130), # padded output tail
(256, 128, 130),
(1024, 128, 96),
],
)
@pytest.mark.parametrize(
"dtype",
[
torch.float16,
pytest.param(
torch.bfloat16,
marks=pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"),
),
],
)
def test_fused_linear_cutedsl_matches_pytorch(shape, dtype):
m, k, n = shape
torch.manual_seed(0)
scale = k**-0.5
x = torch.randn(m, k, device=device, dtype=dtype)
gate_weight = torch.randn(n, k, device=device, dtype=dtype) * scale
up_weight = torch.randn(n, k, device=device, dtype=dtype) * scale
gate_up_weight = torch.cat((gate_weight, up_weight), dim=0)
actual = fused_swiglu(x, gate_up_weight)
gate = torch.nn.functional.linear(x, gate_weight)
up = torch.nn.functional.linear(x, up_weight)
expected = torch.nn.functional.silu(gate) * up
assert actual.shape == (m, n)
torch.testing.assert_close(actual.float(), expected.float(), atol=0.05, rtol=0.03)
def test_fused_linear_cutedsl_fp32_matches_pytorch():
m, k, n = 4, 64, 35
x = torch.randn(m, k, device=device, dtype=torch.float32)
gate_weight = torch.randn(n, k, device=device, dtype=torch.float32)
up_weight = torch.randn(n, k, device=device, dtype=torch.float32)
gate_up_weight = torch.cat((gate_weight, up_weight), dim=0)
actual = fused_swiglu(x, gate_up_weight)
expected = torch.nn.functional.silu(torch.nn.functional.linear(x, gate_weight))
expected *= torch.nn.functional.linear(x, up_weight)
torch.testing.assert_close(actual, expected)
@pytest.mark.skipif(not sm100_available, reason="fused linear SwiGLU requires SM100")
def test_fused_linear_cutedsl_uses_current_stream():
m, k, n = 64, 64, 32
x = torch.zeros(m, k, device=device, dtype=torch.bfloat16)
gate_weight = torch.ones(n, k, device=device, dtype=torch.bfloat16) / k
up_weight = torch.ones_like(gate_weight) / k
gate_up_weight = torch.cat((gate_weight, up_weight), dim=0)
torch.cuda.synchronize()
stream = torch.cuda.Stream()
with torch.cuda.stream(stream):
torch.cuda._sleep(5_000_000)
x.fill_(1)
actual = fused_swiglu(x, gate_up_weight)
stream.synchronize()
expected = torch.nn.functional.silu(torch.nn.functional.linear(x, gate_weight))
expected *= torch.nn.functional.linear(x, up_weight)
torch.testing.assert_close(actual.float(), expected.float(), atol=0.05, rtol=0.03)
@pytest.mark.skipif(not sm100_available, reason="fused linear SwiGLU requires SM100")
def test_fused_linear_cutedsl_stacked_zero_copy_weight():
"""A plain ``torch.cat([gate, up])`` stack (i.e. a model's fused
``gate_up_proj.weight``) is consumed directly, with no pre-packing and no
extra weight copy when the output dim is already tile-aligned."""
m, k, n = 512, 256, 448 # n divisible by EPILOGUE_TILE_SIZE -> zero-copy
torch.manual_seed(0)
x = torch.randn(m, k, device=device, dtype=torch.bfloat16)
gate_weight = torch.randn(n, k, device=device, dtype=torch.bfloat16) * k**-0.5
up_weight = torch.randn(n, k, device=device, dtype=torch.bfloat16) * k**-0.5
gate_up_weight = torch.cat((gate_weight, up_weight), dim=0).contiguous()
weight_ptr = gate_up_weight.data_ptr()
actual = fused_swiglu(x, gate_up_weight)
# The kernel must not have repacked / reallocated the weight.
assert gate_up_weight.data_ptr() == weight_ptr
expected = torch.nn.functional.silu(torch.nn.functional.linear(x, gate_weight))
expected *= torch.nn.functional.linear(x, up_weight)
assert actual.shape == (m, n)
torch.testing.assert_close(actual.float(), expected.float(), atol=0.05, rtol=0.03)
def test_pack_swiglu_weights_is_a_plain_stack():
"""``pack_swiglu_weights`` is now a simple concatenation that unpacks back
into gate/up with trivial slicing (no bespoke 32-row interleave)."""
n, k = 448, 128
gate_weight = torch.randn(n, k)
up_weight = torch.randn(n, k)
stacked, output_features = pack_swiglu_weights(gate_weight, up_weight)
assert output_features == n
assert stacked.shape == (2 * n, k)
torch.testing.assert_close(stacked[:n], gate_weight)
torch.testing.assert_close(stacked[n:], up_weight)
@pytest.mark.skipif(not sm100_available, reason="fused linear SwiGLU requires SM100")
@pytest.mark.parametrize(
"shape",
[
(128, 256, 128), # tile-aligned
(256, 128, 130), # padded output tail
(33, 128, 96), # non-tile-aligned tokens
],
)
@pytest.mark.parametrize(
"dtype",
[
torch.float16,
pytest.param(
torch.bfloat16,
marks=pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"),
),
],
)
def test_fused_linear_cutedsl_separate_weights(shape, dtype):
"""``fused_swiglu(a, gate, up)`` accepts two *separate* ``[N, K]`` tensors
and matches both PyTorch and the stacked-weight call."""
m, k, n = shape
torch.manual_seed(0)
scale = k**-0.5
x = torch.randn(m, k, device=device, dtype=dtype)
gate_weight = torch.randn(n, k, device=device, dtype=dtype) * scale
up_weight = torch.randn(n, k, device=device, dtype=dtype) * scale
actual = fused_swiglu(x, gate_weight, up_weight)
expected = torch.nn.functional.silu(torch.nn.functional.linear(x, gate_weight))
expected *= torch.nn.functional.linear(x, up_weight)
assert actual.shape == (m, n)
torch.testing.assert_close(actual.float(), expected.float(), atol=0.05, rtol=0.03)
# Separate and stacked calling conventions must agree exactly.
stacked = fused_swiglu(x, torch.cat((gate_weight, up_weight), dim=0))
torch.testing.assert_close(actual, stacked)
@pytest.mark.skipif(not sm100_available, reason="fused linear SwiGLU requires SM100")
def test_fused_linear_cutedsl_separate_weights_validation():
"""Mismatched separate weights raise clear errors (not cryptic tensor-bool
ambiguity), and ``output_features`` stays keyword-only."""
k, n = 128, 96
x = torch.randn(8, k, device=device, dtype=torch.bfloat16)
gate_weight = torch.randn(n, k, device=device, dtype=torch.bfloat16) * k**-0.5
up_weight = torch.randn(n + 32, k, device=device, dtype=torch.bfloat16) * k**-0.5
with pytest.raises(ValueError, match="identical shapes"):
fused_swiglu(x, gate_weight, up_weight)
with pytest.raises(TypeError, match="same dtype"):
fused_swiglu(x, gate_weight, gate_weight.half())