From 69a04179c60c25eb10c5d7de5ac03f4ddb9073e9 Mon Sep 17 00:00:00 2001 From: pearblossom <3364870135@qq.com> Date: Sun, 9 Aug 2026 19:32:26 +0800 Subject: [PATCH 1/6] feat: add LigerMLP module Signed-off-by: pearblossom <3364870135@qq.com> --- benchmark/scripts/benchmark_mlp.py | 110 ++++ src/liger_kernel/ops/__init__.py | 1 + src/liger_kernel/ops/mlp.py | 632 ++++++++++++++++++++++ src/liger_kernel/transformers/__init__.py | 2 + src/liger_kernel/transformers/mlp.py | 20 + test/transformers/test_mlp.py | 200 +++++++ 6 files changed, 965 insertions(+) create mode 100644 benchmark/scripts/benchmark_mlp.py create mode 100644 src/liger_kernel/ops/mlp.py create mode 100644 src/liger_kernel/transformers/mlp.py create mode 100644 test/transformers/test_mlp.py diff --git a/benchmark/scripts/benchmark_mlp.py b/benchmark/scripts/benchmark_mlp.py new file mode 100644 index 000000000..5c53987a8 --- /dev/null +++ b/benchmark/scripts/benchmark_mlp.py @@ -0,0 +1,110 @@ +import torch + +from benchmark_model_configs import MODEL_REGISTRY +from benchmark_model_configs import build_model_config_sweep +from benchmark_model_configs import build_token_length_sweep +from benchmark_model_configs import get_benchmark_model_config +from transformers.models.llama.configuration_llama import LlamaConfig +from transformers.models.llama.modeling_llama import LlamaMLP +from utils import SingleBenchmarkRunInput +from utils import build_memory_bench_fn +from utils import build_speed_bench_fn +from utils import parse_benchmark_script_args +from utils import run_benchmarks + +from liger_kernel.transformers.mlp import LigerMLP +from liger_kernel.utils import infer_device + +device = infer_device() + + +def setup_mlp(input: SingleBenchmarkRunInput): + """Create input tensor and SwiGLU MLP layer from benchmark config.""" + cfg = input.extra_benchmark_config + if isinstance(input.x, str): + model_cfg = MODEL_REGISTRY[input.x] + seq_len = cfg["seq_len"] + hidden_size = model_cfg.hidden_size + intermediate_size = model_cfg.intermediate_size + dtype = model_cfg.dtype + else: + seq_len = input.x + hidden_size = cfg["hidden_size"] + intermediate_size = cfg["intermediate_size"] + dtype = cfg["dtype"] + + llama_config = LlamaConfig( + hidden_size=hidden_size, + intermediate_size=intermediate_size, + hidden_act=cfg["hidden_act"], + ) + x = torch.randn( + cfg["bsz"], + seq_len, + hidden_size, + device=device, + dtype=dtype, + requires_grad=True, + ) + if input.kernel_provider == "liger": + layer = LigerMLP(config=llama_config).to(device).to(dtype) + elif input.kernel_provider == "huggingface": + layer = LlamaMLP(config=llama_config).to(device).to(dtype) + else: + raise ValueError(f"Invalid provider: {input.kernel_provider} for MLP") + return x, layer + + +if __name__ == "__main__": + args = parse_benchmark_script_args() + + if args.sweep_mode == "model_config": + common_configs = build_model_config_sweep( + kernel_name="mlp", + setup_fn=setup_mlp, + model_keys=["hidden_size", "intermediate_size", "dtype"], + probe_provider="huggingface", + extra_configs={ + "bsz": 1, + "hidden_act": "silu", + }, + probe_dim="T", + bt=args.bt, + overwrite=args.overwrite, + ) + else: + model = get_benchmark_model_config(args.model) + probe_seq_len = 1024 + + common_configs = build_token_length_sweep( + kernel_name="mlp", + probe_x=probe_seq_len, + model=model, + setup_fn=setup_mlp, + model_keys=["hidden_size", "intermediate_size", "dtype"], + extra_configs={ + "bsz": 1, + "hidden_act": "silu", + }, + scale_dim="T", + x_label="total tokens", + probe_provider="huggingface", + overwrite=args.overwrite, + ) + + common_configs["kernel_providers"] = ["huggingface", "liger"] + + run_benchmarks( + bench_test_fn=build_speed_bench_fn(setup_mlp), + kernel_operation_modes=["forward", "backward", "full"], + metric_name="speed", + metric_unit="ms", + **common_configs, + ) + run_benchmarks( + bench_test_fn=build_memory_bench_fn(setup_mlp), + kernel_operation_modes=["full", "forward", "backward"], + metric_name="memory", + metric_unit="MB", + **common_configs, + ) diff --git a/src/liger_kernel/ops/__init__.py b/src/liger_kernel/ops/__init__.py index fc4fed124..1bc228882 100644 --- a/src/liger_kernel/ops/__init__.py +++ b/src/liger_kernel/ops/__init__.py @@ -70,6 +70,7 @@ from liger_kernel.ops.mhc import LigerMHCCoeffsFunction # noqa: F401 from liger_kernel.ops.mhc import LigerMHCPostResFunction # noqa: F401 from liger_kernel.ops.mhc import LigerMHCPreFunction # noqa: F401 +from liger_kernel.ops.mlp import LigerMLPFunction # noqa: F401 from liger_kernel.ops.modulated_rms_norm import LigerModulatedRMSNormFunction # noqa: F401 from liger_kernel.ops.modulated_rms_norm import modulated_rms_norm_backward # noqa: F401 from liger_kernel.ops.modulated_rms_norm import modulated_rms_norm_forward # noqa: F401 diff --git a/src/liger_kernel/ops/mlp.py b/src/liger_kernel/ops/mlp.py new file mode 100644 index 000000000..dc5f79918 --- /dev/null +++ b/src/liger_kernel/ops/mlp.py @@ -0,0 +1,632 @@ +""" +Fused SwiGLU MLP: forward/backward orchestration + autograd.Function. + +Shapes +------ +input : [B, S, dim] +gate_weight: [hidden_dim, dim] +up_weight : [hidden_dim, dim] +down_weight: [dim, hidden_dim] + +1. The forward kernel fuses gate_proj + up_proj + SiLU + gating into a single +Triton kernel (down_proj stays a plain cuBLAS GEMM). +2. The backward kernels fuse the SiLU/gating gradient (dG, dU) into one kernel, +and dI = dG @ Wg + dU @ Wu into another, avoiding an extra HBM round trip for the +intermediate dA tensor. +3. G and U are cached from the forward pass instead of being recomputed in +the backward pass (recompute-vs-store trade-off). +""" + +import torch +import torch.nn.functional as F +import triton +import triton.language as tl + +from triton.tools.tensor_descriptor import TensorDescriptor + + +@triton.jit +def _l2_swizzle(pid, M, N, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, GROUP_SIZE_M: tl.constexpr): + num_pid_m = tl.cdiv(M, BLOCK_M) + num_pid_n = tl.cdiv(N, BLOCK_N) + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + base_pid_m = group_id * GROUP_SIZE_M + effective_group_size_m = tl.minimum(num_pid_m - base_pid_m, GROUP_SIZE_M) + pid_in_group = pid % num_pid_in_group + pid_m = pid_in_group % effective_group_size_m + base_pid_m + pid_n = pid_in_group // effective_group_size_m + return pid_m, pid_n + + +# ===================================================================== +# Forward Training +# ===================================================================== +def _host_descriptor_pre_hook_fwd(nargs): + BLOCK_M = nargs["BLOCK_M"] + BLOCK_N = nargs["BLOCK_N"] + BLOCK_K = nargs["BLOCK_K"] + + nargs["desc_input"].block_shape = [1, BLOCK_M, BLOCK_K] + nargs["desc_Wg"].block_shape = [BLOCK_N, BLOCK_K] + nargs["desc_Wu"].block_shape = [BLOCK_N, BLOCK_K] + nargs["desc_A"].block_shape = [1, BLOCK_M, BLOCK_N] + nargs["desc_G"].block_shape = [1, BLOCK_M, BLOCK_N] + nargs["desc_U"].block_shape = [1, BLOCK_M, BLOCK_N] + + +@triton.autotune( + configs=[ + triton.Config( + {"BLOCK_M": BM, "BLOCK_N": BN, "BLOCK_K": BK, "GROUP_SIZE_M": GS}, + num_warps=w, + num_stages=s, + pre_hook=_host_descriptor_pre_hook_fwd, + ) + for BM in [64, 128] + for BN in [128] + for BK in [64, 128] + for GS in [1, 8] # GROUP_SIZE_M = 1 equal to no L2 swizzle. + for w in [4, 8] + for s in [2, 3] + ], + key=["dim", "hidden_dim", "bucket_M"], +) +# grid = [ceil(S/Block_M) * ceil(hidden_dim/Block_N), B] +@triton.jit +def _swiglu_kernel_forward( + # ======= TensorDescriptor ======= + desc_input, + desc_Wg, + desc_Wu, + desc_A, + desc_G, + desc_U, + # ============ shape ============ + S, + dim, + hidden_dim, + bucket_M, + # =========== meta data =========== + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, +): + pid = tl.program_id(0) + b = tl.program_id(1) + # L2 swizzle + pid_m, pid_n = _l2_swizzle(pid, S, hidden_dim, BLOCK_M, BLOCK_N, GROUP_SIZE_M) + + # Compute the starting positions for input, gate and up + M_start = pid_m * BLOCK_M + N_start = pid_n * BLOCK_N + + # Compute input @ gate^T and input @ up^T + acc_gate = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + acc_up = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + for k in range(0, dim, BLOCK_K): + input_block = desc_input.load([b, M_start, k]) # [1, BLOCK_M, BLOCK_K] + input_block = tl.reshape(input_block, [BLOCK_M, BLOCK_K]) # [BLOCK_M, BLOCK_K] + gate_block = desc_Wg.load([N_start, k]) # [BLOCK_N, BLOCK_K] + acc_gate = tl.dot(input_block, tl.trans(gate_block), acc=acc_gate) + up_block = desc_Wu.load([N_start, k]) # [BLOCK_N, BLOCK_K] + acc_up = tl.dot(input_block, tl.trans(up_block), acc=acc_up) + + # Compute A_block + A_block = acc_gate * tl.sigmoid(acc_gate) * acc_up # [BLOCK_M, BLOCK_N] + + # Write back + dtype = desc_input.dtype + A_block = tl.reshape(A_block, [1, BLOCK_M, BLOCK_N]) + desc_A.store([b, M_start, N_start], A_block.to(dtype)) + G_block = tl.reshape(acc_gate, [1, BLOCK_M, BLOCK_N]) + desc_G.store([b, M_start, N_start], G_block.to(dtype)) + U_block = tl.reshape(acc_up, [1, BLOCK_M, BLOCK_N]) + desc_U.store([b, M_start, N_start], U_block.to(dtype)) + + +# ===================================================================== +# Forward Inference +# ===================================================================== +""" No need to save GU during inference.""" + + +def _host_descriptor_pre_hook_fwd_inference(nargs): + BLOCK_M = nargs["BLOCK_M"] + BLOCK_N = nargs["BLOCK_N"] + BLOCK_K = nargs["BLOCK_K"] + + nargs["desc_input"].block_shape = [1, BLOCK_M, BLOCK_K] + nargs["desc_Wg"].block_shape = [BLOCK_N, BLOCK_K] + nargs["desc_Wu"].block_shape = [BLOCK_N, BLOCK_K] + nargs["desc_A"].block_shape = [1, BLOCK_M, BLOCK_N] + + +@triton.autotune( + configs=[ + triton.Config( + {"BLOCK_M": BM, "BLOCK_N": BN, "BLOCK_K": BK, "GROUP_SIZE_M": GS}, + num_warps=w, + num_stages=s, + pre_hook=_host_descriptor_pre_hook_fwd_inference, + ) + for BM in [32, 64] + for BN in [128] + for BK in [64, 128] + for GS in [1, 4] # GROUP_SIZE_M = 1 equal to no L2 swizzle. + for w in [4, 8] + for s in [2, 3] + ], + key=["dim", "hidden_dim", "bucket_M"], +) +# grid = [ceil(S/Block_M) * ceil(hidden_dim/Block_N), B] +@triton.jit +def _swiglu_kernel_forward_inference( + # ======= TensorDescriptor ======= + desc_input, + desc_Wg, + desc_Wu, + desc_A, + # ============ shape ============ + S, + dim, + hidden_dim, + bucket_M, + # =========== meta data =========== + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, +): + pid = tl.program_id(0) + b = tl.program_id(1) + # L2 swizzle + pid_m, pid_n = _l2_swizzle(pid, S, hidden_dim, BLOCK_M, BLOCK_N, GROUP_SIZE_M) + + # Compute the starting positions for input, gate and up + M_start = pid_m * BLOCK_M + N_start = pid_n * BLOCK_N + + # Compute input @ gate^T and input @ up^T + acc_gate = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + acc_up = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + for k in range(0, dim, BLOCK_K): + input_block = desc_input.load([b, M_start, k]) # [1, BLOCK_M, BLOCK_K] + input_block = tl.reshape(input_block, [BLOCK_M, BLOCK_K]) # [BLOCK_M, BLOCK_K] + gate_block = desc_Wg.load([N_start, k]) # [BLOCK_N, BLOCK_K] + acc_gate = tl.dot(input_block, tl.trans(gate_block), acc=acc_gate) + up_block = desc_Wu.load([N_start, k]) # [BLOCK_N, BLOCK_K] + acc_up = tl.dot(input_block, tl.trans(up_block), acc=acc_up) + + # Compute A_block + A_block = acc_gate * tl.sigmoid(acc_gate) * acc_up # [BLOCK_M, BLOCK_N] + + # Write back + dtype = desc_input.dtype + A_block = tl.reshape(A_block, [1, BLOCK_M, BLOCK_N]) + desc_A.store([b, M_start, N_start], A_block.to(dtype)) + + +# ===================================================================== +# Backward dU and dG +# ===================================================================== +""" +Note: +- We don't use autotune here, +- because autotune could break the in-place overwrite operation. +- Instead, we directly use the optimal configuration for the forward kernel. +""" + + +# grid = [ceil(S/Block_M) * ceil(hidden_dim/Block_N), B] +@triton.jit +def _swiglu_kernel_backward_dGU( + # ======= TensorDescriptor ======= + desc_dO, + desc_Wd, + desc_G, + desc_U, + desc_dG, + desc_dU, + # ============ shape ============ + S, + dim, + hidden_dim, + # =========== meta data =========== + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, +): + pid = tl.program_id(0) + b = tl.program_id(1) + # L2 swizzle + pid_m, pid_n = _l2_swizzle(pid, S, hidden_dim, BLOCK_M, BLOCK_N, GROUP_SIZE_M) + + # Compute the starting positions for input, gate and up + M_start = pid_m * BLOCK_M + N_start = pid_n * BLOCK_N + + acc_A = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + for k in range(0, dim, BLOCK_K): + dO_block = desc_dO.load([b, M_start, k]) # [1, BLOCK_M, BLOCK_K] + dO_block = tl.reshape(dO_block, [BLOCK_M, BLOCK_K]) # [BLOCK_M, BLOCK_K] + down_block = desc_Wd.load([k, N_start]) # [BLOCK_K, BLOCK_N] + acc_A = tl.dot(dO_block, down_block, acc=acc_A) # [BLOCK_M, BLOCK_N] + + # Load G and U + G_block = desc_G.load([b, M_start, N_start]) # [1, BLOCK_M, BLOCK_N] + G_block = tl.reshape(G_block, [BLOCK_M, BLOCK_N]) # [BLOCK_M, BLOCK_N] + U_block = desc_U.load([b, M_start, N_start]) # [1, BLOCK_M, BLOCK_N] + U_block = tl.reshape(U_block, [BLOCK_M, BLOCK_N]) # [BLOCK_M, BLOCK_N] + + # Compute and save dU, dG + G_block = G_block.to(tl.float32) + U_block = U_block.to(tl.float32) + sigmoid_G = tl.sigmoid(G_block) + silu_G = G_block * sigmoid_G + dtype = desc_G.dtype + + dU = acc_A * silu_G # [BLOCK_M, BLOCK_N] + dU = tl.reshape(dU, [1, BLOCK_M, BLOCK_N]) + desc_dU.store([b, M_start, N_start], dU.to(dtype)) + + dG = acc_A * U_block * (sigmoid_G + silu_G - sigmoid_G * silu_G) # [BLOCK_M, BLOCK_N] + dG = tl.reshape(dG, [1, BLOCK_M, BLOCK_N]) + desc_dG.store([b, M_start, N_start], dG.to(dtype)) + + +# ===================================================================== +# Backward dI +# ===================================================================== +def _host_descriptor_pre_hook_bwd_dI(nargs): + BLOCK_M = nargs["BLOCK_M"] + BLOCK_N = nargs["BLOCK_N"] + BLOCK_K = nargs["BLOCK_K"] + + nargs["desc_dG"].block_shape = [1, BLOCK_M, BLOCK_K] + nargs["desc_Wg"].block_shape = [BLOCK_K, BLOCK_N] + nargs["desc_dU"].block_shape = [1, BLOCK_M, BLOCK_K] + nargs["desc_Wu"].block_shape = [BLOCK_K, BLOCK_N] + nargs["desc_dI"].block_shape = [1, BLOCK_M, BLOCK_N] + + +@triton.autotune( + configs=[ + triton.Config( + {"BLOCK_M": BM, "BLOCK_N": BN, "BLOCK_K": BK, "GROUP_SIZE_M": GS}, + num_warps=w, + num_stages=s, + pre_hook=_host_descriptor_pre_hook_bwd_dI, + ) + for BM in [64, 128] + for BN in [128] + for BK in [64, 128] + for GS in [1, 8] # GROUP_SIZE_M = 1 equal to no L2 swizzle. + for w in [4, 8] + for s in [2, 3] + ], + key=["dim", "hidden_dim", "bucket_M"], +) +# grid = [ceil(S/Block_M) * ceil(dim/Block_N), B] +@triton.jit +def _swiglu_kernel_backward_dI( + # ======= TensorDescriptor ======= + desc_dG, + desc_Wg, + desc_dU, + desc_Wu, + desc_dI, + # ============ shape ============ + S, + dim, + hidden_dim, + bucket_M, + # =========== meta data =========== + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, +): + pid = tl.program_id(0) + b = tl.program_id(1) + # L2 swizzle + pid_m, pid_n = _l2_swizzle(pid, S, dim, BLOCK_M, BLOCK_N, GROUP_SIZE_M) + + # Compute the starting positions + M_start = pid_m * BLOCK_M + N_start = pid_n * BLOCK_N + + # Compute dI = dG @ gate_weight + dU @ up_weight + # [B, S, hidden_dim] @ [hidden_dim, dim] + ... @ ... = [B, S, dim] + acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + for k in range(0, hidden_dim, BLOCK_K): + dG_block = desc_dG.load([b, M_start, k]) # [1, BLOCK_M, BLOCK_K] + dG_block = tl.reshape(dG_block, [BLOCK_M, BLOCK_K]) # [BLOCK_M, BLOCK_K] + gate_weight_block = desc_Wg.load([k, N_start]) # [BLOCK_K, BLOCK_N] + acc = tl.dot(dG_block, gate_weight_block, acc=acc) + + dU_block = desc_dU.load([b, M_start, k]) # [1, BLOCK_M, BLOCK_K] + dU_block = tl.reshape(dU_block, [BLOCK_M, BLOCK_K]) # [BLOCK_M, BLOCK_K] + up_weight_block = desc_Wu.load([k, N_start]) # [BLOCK_K, BLOCK_N] + acc = tl.dot(dU_block, up_weight_block, acc=acc) + + # Save dI + dI = acc # [BLOCK_M, BLOCK_N] + dI = tl.reshape(dI, [1, BLOCK_M, BLOCK_N]) + desc_dI.store([b, M_start, N_start], dI.to(desc_dI.dtype)) + + +def get_bucket_m(M, threshold=4096, bucket_size=128): + """ + Returns the globally unique bucket index mapped to M = B * S. + - [0, 256) : Fine-grained bucketing with a granularity of 32, occupying buckets 0 to 7. + - [256, threshold): Coarse-grained bucketing with a granularity of 128. + - >= threshold : All fall into the last fixed large bucket. + """ + # The first 256 lengths are bucketed with a granularity of 32, consuming 256 // 32 = 8 buckets. + offset = 256 // 32 + + if M >= threshold: + # All extra-long sequences are placed into the final bucket: 8 + 4096//128 = 40 + return offset + threshold // bucket_size + elif M < 256: + # Fine-grained range. + return M // 32 + + # Coarse-grained range, must add the previous offset to avoid index conflicts with the [0, 256) range. + return offset + M // bucket_size + + +DTYPE_SUPPORT_LIST = (torch.float16, torch.bfloat16, torch.float32) + + +def _ensure_tma_compatible(t: torch.Tensor) -> torch.Tensor: + """ + Check if tensor meets TMA descriptor requirements. + If not, make it contiguous. + 1. Last dimension must be contiguous (stride[-1] == 1). + 2. All other dimensions must be 16-byte aligned. + """ + strides = t.stride() + elem_bytes = t.element_size() + if strides[-1] != 1: + return t.contiguous() + for s in strides[:-1]: + if (s * elem_bytes) % 16 != 0: + return t.contiguous() + return t + + +def _check_inputs(input, gate_weight, up_weight, down_weight): + if input.dtype not in DTYPE_SUPPORT_LIST: + raise TypeError(f"flash_swiglu_mlp only supports {DTYPE_SUPPORT_LIST}, got {input.dtype}.") + + dtypes = {t.dtype for t in [input, gate_weight, up_weight, down_weight]} + if len(dtypes) != 1: + raise TypeError(f"All tensors must share the same dtype, got {dtypes}.") + + input = _ensure_tma_compatible(input) + gate_weight = _ensure_tma_compatible(gate_weight) + up_weight = _ensure_tma_compatible(up_weight) + down_weight = _ensure_tma_compatible(down_weight) + + return input, gate_weight, up_weight, down_weight + + +def swiglu_forward( + input: torch.Tensor, + gate_weight: torch.Tensor, + up_weight: torch.Tensor, + down_weight: torch.Tensor, +) -> torch.Tensor: + """Training-mode forward: computes O, and returns AGU.""" + input, gate_weight, up_weight, down_weight = _check_inputs(input, gate_weight, up_weight, down_weight) + + B, S, dim = input.shape + hidden_dim = gate_weight.shape[0] + bucket_M = get_bucket_m(B * S) + # Optimization: merge 3 torch.empty calls into 1 to reduce allocator overhead. + AGU = torch.empty(B, S, 3 * hidden_dim, device=input.device, dtype=input.dtype) + # Slice out 3 independent 3D views. + A = AGU[..., :hidden_dim] + G = AGU[..., hidden_dim : 2 * hidden_dim] + U = AGU[..., 2 * hidden_dim :] + + # Dummy block + dummy_block_3D = [1, 1, 1] + dummy_block_2D = [1, 1] + + # Create TensorDescriptor + desc_input = TensorDescriptor.from_tensor(tensor=input, block_shape=dummy_block_3D, padding="zero") + desc_Wg = TensorDescriptor.from_tensor(tensor=gate_weight, block_shape=dummy_block_2D, padding="zero") + desc_Wu = TensorDescriptor.from_tensor(tensor=up_weight, block_shape=dummy_block_2D, padding="zero") + desc_A = TensorDescriptor.from_tensor(tensor=A, block_shape=dummy_block_3D, padding="zero") + desc_G = TensorDescriptor.from_tensor(tensor=G, block_shape=dummy_block_3D, padding="zero") + desc_U = TensorDescriptor.from_tensor(tensor=U, block_shape=dummy_block_3D, padding="zero") + + grid = lambda META: ( + triton.cdiv(S, META["BLOCK_M"]) * triton.cdiv(hidden_dim, META["BLOCK_N"]), + B, + ) + + _swiglu_kernel_forward[grid]( + desc_input, + desc_Wg, + desc_Wu, + desc_A, + desc_G, + desc_U, + S, + dim, + hidden_dim, + bucket_M, + ) + fwd_best_config = _swiglu_kernel_forward.best_config + + out = F.linear(A, down_weight) + + return out, AGU, fwd_best_config + + +def swiglu_forward_inference( + input: torch.Tensor, + gate_weight: torch.Tensor, + up_weight: torch.Tensor, + down_weight: torch.Tensor, +) -> torch.Tensor: + """Inference-mode forward.""" + input, gate_weight, up_weight, down_weight = _check_inputs(input, gate_weight, up_weight, down_weight) + B, S, dim = input.shape + hidden_dim = gate_weight.shape[0] + bucket_M = get_bucket_m(B * S) + A = torch.empty((B, S, hidden_dim), device=input.device, dtype=input.dtype) + + # Dummy block + dummy_block_3D = [1, 1, 1] + dummy_block_2D = [1, 1] + + # Create TensorDescriptor + desc_input = TensorDescriptor.from_tensor(tensor=input, block_shape=dummy_block_3D, padding="zero") + desc_Wg = TensorDescriptor.from_tensor(tensor=gate_weight, block_shape=dummy_block_2D, padding="zero") + desc_Wu = TensorDescriptor.from_tensor(tensor=up_weight, block_shape=dummy_block_2D, padding="zero") + desc_A = TensorDescriptor.from_tensor(tensor=A, block_shape=dummy_block_3D, padding="zero") + + grid = lambda META: ( + triton.cdiv(S, META["BLOCK_M"]) * triton.cdiv(hidden_dim, META["BLOCK_N"]), + B, + ) + + _swiglu_kernel_forward_inference[grid]( + desc_input, + desc_Wg, + desc_Wu, + desc_A, + S, + dim, + hidden_dim, + bucket_M, + ) + + out = F.linear(A, down_weight) + + return out + + +def swiglu_backward( + dO: torch.Tensor, + input: torch.Tensor, + gate_weight: torch.Tensor, + up_weight: torch.Tensor, + down_weight: torch.Tensor, + AGU: torch.Tensor, + fwd_best_config, +) -> torch.Tensor: + dO = _ensure_tma_compatible(dO) + B, S, dim = input.shape + hidden_dim = gate_weight.shape[0] + bucket_M = get_bucket_m(B * S) + + A = AGU[..., :hidden_dim] + G = AGU[..., hidden_dim : 2 * hidden_dim] + U = AGU[..., 2 * hidden_dim :] + + # ========================== Compute dWd =================================== + dWd = dO.reshape(-1, dim).T @ A.reshape(-1, hidden_dim) # [dim, hidden_dim] + + # ================= Compute dU and dG by triton kernel ===================== + # Reuse the memory of GU to store dGU, in order to save memory. + dG, dU = G, U + + # Use forward best config. + BLOCK_M = fwd_best_config.kwargs["BLOCK_M"] + BLOCK_N = fwd_best_config.kwargs["BLOCK_N"] + BLOCK_K = fwd_best_config.kwargs["BLOCK_K"] + GROUP_SIZE_M = fwd_best_config.kwargs["GROUP_SIZE_M"] + + desc_dO = TensorDescriptor.from_tensor(tensor=dO, block_shape=[1, BLOCK_M, BLOCK_K], padding="zero") + desc_Wd = TensorDescriptor.from_tensor(tensor=down_weight, block_shape=[BLOCK_K, BLOCK_N], padding="zero") + desc_G = TensorDescriptor.from_tensor(tensor=G, block_shape=[1, BLOCK_M, BLOCK_N], padding="zero") + desc_U = TensorDescriptor.from_tensor(tensor=U, block_shape=[1, BLOCK_M, BLOCK_N], padding="zero") + desc_dG = TensorDescriptor.from_tensor(tensor=dG, block_shape=[1, BLOCK_M, BLOCK_N], padding="zero") + desc_dU = TensorDescriptor.from_tensor(tensor=dU, block_shape=[1, BLOCK_M, BLOCK_N], padding="zero") + + grid_dGU = lambda META: ( + triton.cdiv(S, META["BLOCK_M"]) * triton.cdiv(hidden_dim, META["BLOCK_N"]), + B, + ) + + _swiglu_kernel_backward_dGU[grid_dGU]( + desc_dO, + desc_Wd, + desc_G, + desc_U, + desc_dG, + desc_dU, + S, + dim, + hidden_dim, + BLOCK_M=BLOCK_M, + BLOCK_N=BLOCK_N, + BLOCK_K=BLOCK_K, + GROUP_SIZE_M=GROUP_SIZE_M, + num_warps=fwd_best_config.num_warps, + num_stages=fwd_best_config.num_stages, + ) + + # ================== Compute dWug by a single large GEMM =================== + dGU = AGU[..., hidden_dim:] + dWug = dGU.reshape(-1, 2 * hidden_dim).T @ input.reshape(-1, dim) # [2*hidden_dim, dim] + dWg, dWu = dWug[:hidden_dim], dWug[hidden_dim:] # [hidden_dim, dim] + + # ====================== Compute dI by triton kernel ======================= + dummy_block_3D = [1, 1, 1] + dummy_block_2D = [1, 1] + + dI = torch.empty((B, S, dim), device=input.device, dtype=input.dtype) + desc_dI = TensorDescriptor.from_tensor(tensor=dI, block_shape=dummy_block_3D, padding="zero") + desc_Wg = TensorDescriptor.from_tensor(tensor=gate_weight, block_shape=dummy_block_2D, padding="zero") + desc_Wu = TensorDescriptor.from_tensor(tensor=up_weight, block_shape=dummy_block_2D, padding="zero") + + grid_dI = lambda META: ( + triton.cdiv(S, META["BLOCK_M"]) * triton.cdiv(dim, META["BLOCK_N"]), + B, + ) + + _swiglu_kernel_backward_dI[grid_dI]( + desc_dG, + desc_Wg, + desc_dU, + desc_Wu, + desc_dI, + S, + dim, + hidden_dim, + bucket_M, + ) + + return dI, dWg, dWu, dWd + + +class LigerMLPFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, input, gate_weight, up_weight, down_weight): + assert input.is_cuda and gate_weight.is_cuda and up_weight.is_cuda and down_weight.is_cuda + # Note: + # PyTorch automatically disables global gradient computation during + # the forward pass, so torch.is_grad_enabled() cannot be used. + if any(ctx.needs_input_grad): + out, AGU, fwd_best_config = swiglu_forward(input, gate_weight, up_weight, down_weight) + ctx.save_for_backward(input, gate_weight, up_weight, down_weight, AGU) + ctx.fwd_best_config = fwd_best_config + return out + else: + # Inference mode: no GU saving. + return swiglu_forward_inference(input, gate_weight, up_weight, down_weight) + + @staticmethod + def backward(ctx, dO): + input, gate_weight, up_weight, down_weight, AGU = ctx.saved_tensors + dI, dWg, dWu, dWd = swiglu_backward(dO, input, gate_weight, up_weight, down_weight, AGU, ctx.fwd_best_config) + return dI, dWg, dWu, dWd diff --git a/src/liger_kernel/transformers/__init__.py b/src/liger_kernel/transformers/__init__.py index 26bdef91b..27f48a776 100644 --- a/src/liger_kernel/transformers/__init__.py +++ b/src/liger_kernel/transformers/__init__.py @@ -15,6 +15,7 @@ from liger_kernel.transformers.llama4_rope import liger_llama4_text_rotary_pos_emb # noqa: F401 from liger_kernel.transformers.llama4_rope import liger_llama4_vision_rotary_pos_emb # noqa: F401 from liger_kernel.transformers.mhc import LigerMHC # noqa: F401 +from liger_kernel.transformers.mlp import LigerMLP # noqa: F401 from liger_kernel.transformers.modulated_rms_norm import LigerModulatedRMSNorm # noqa: F401 from liger_kernel.transformers.multi_token_attention import LigerMultiTokenAttention # noqa: F401 from liger_kernel.transformers.poly_norm import LigerPolyNorm # noqa: F401 @@ -200,6 +201,7 @@ def __getattr__(name: str): "LigerMultiTokenAttention", "LigerSoftmax", "LigerSparsemax", + "LigerMLP", ] # Add transformer-dependent symbols only if available diff --git a/src/liger_kernel/transformers/mlp.py b/src/liger_kernel/transformers/mlp.py new file mode 100644 index 000000000..cec442600 --- /dev/null +++ b/src/liger_kernel/transformers/mlp.py @@ -0,0 +1,20 @@ +import torch +import torch.nn as nn + +from liger_kernel.ops import LigerMLPFunction + + +class LigerMLP(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) + if config.hidden_act not in ["silu", "swish"]: + raise ValueError(f"Activation function {config.hidden_act} not supported.") + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return LigerMLPFunction.apply(x, self.gate_proj.weight, self.up_proj.weight, self.down_proj.weight) diff --git a/test/transformers/test_mlp.py b/test/transformers/test_mlp.py new file mode 100644 index 000000000..a65029255 --- /dev/null +++ b/test/transformers/test_mlp.py @@ -0,0 +1,200 @@ +import pytest +import torch + +from test.utils import supports_bfloat16 +from transformers.models.llama.configuration_llama import LlamaConfig +from transformers.models.llama.modeling_llama import LlamaMLP + +from liger_kernel.transformers.mlp import LigerMLP +from liger_kernel.utils import infer_device + +device = infer_device() + + +@pytest.mark.parametrize( + "bsz, seq_len, hidden_size, intermediate_size", + [ + (2, 256, 256, 512), + # weird shapes (TMA requires every row stride to be 16-byte aligned, + # so hidden/intermediate sizes are kept multiples of 8) + (6, 42, 128, 432), + (3, 37, 96, 264), + ], +) +@pytest.mark.parametrize( + "dtype, atol, rtol", + [ + # fp32: Triton (fp32 accum) vs cuBLAS (TF32) diverge through 3 chained + # GEMMs; activations reach ~1e4 magnitude, so atol must scale accordingly + (torch.float32, 2e2, 2e-2), + pytest.param( + torch.bfloat16, + 1e4, + 1e-2, + marks=pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"), + ), + (torch.float16, 1e2, 1e-2), + ], +) +def test_correctness_llamamlp(bsz, seq_len, hidden_size, intermediate_size, dtype, atol, rtol): + config = LlamaConfig( + hidden_size=hidden_size, + intermediate_size=intermediate_size, + hidden_act="silu", + num_attention_heads=1, + ) + + _input = torch.randn(bsz, seq_len, hidden_size, device=device, dtype=dtype) + + x1 = _input.clone().requires_grad_(True) + x2 = _input.clone().requires_grad_(True) + + # initialize weights directly in their final (contiguous) orientation + G = torch.randn(intermediate_size, hidden_size, device=device, dtype=dtype) + U = torch.randn(intermediate_size, hidden_size, device=device, dtype=dtype) + D = torch.randn(hidden_size, intermediate_size, device=device, dtype=dtype) + + llama_mlp = LlamaMLP(config=config).to(device).to(dtype) + llama_mlp.gate_proj.weight.data = G + llama_mlp.up_proj.weight.data = U + llama_mlp.down_proj.weight.data = D + + liger_mlp = LigerMLP(config=config).to(device).to(dtype) + liger_mlp.gate_proj.weight.data = G + liger_mlp.up_proj.weight.data = U + liger_mlp.down_proj.weight.data = D + + y1 = llama_mlp(x1) + y2 = liger_mlp(x2) + + assert torch.allclose(y1, y2, atol=atol, rtol=rtol) + + dy = torch.randn_like(y1) + + y1.backward(dy.clone(), retain_graph=True) + y2.backward(dy.clone(), retain_graph=True) + + assert torch.allclose( + llama_mlp.gate_proj.weight.grad, + liger_mlp.gate_proj.weight.grad, + atol=atol, + rtol=rtol, + ) + assert torch.allclose( + llama_mlp.up_proj.weight.grad, + liger_mlp.up_proj.weight.grad, + atol=atol, + rtol=rtol, + ) + assert torch.allclose( + llama_mlp.down_proj.weight.grad, + liger_mlp.down_proj.weight.grad, + atol=atol, + rtol=rtol, + ) + + assert torch.allclose(x1.grad, x2.grad, atol=atol, rtol=rtol) + + +@pytest.mark.parametrize( + "bsz, seq_len, hidden_size, intermediate_size", + [ + (2, 256, 256, 512), + # weird shapes (TMA requires every row stride to be 16-byte aligned, + # so hidden/intermediate sizes are kept multiples of 8) + (6, 42, 128, 432), + (3, 37, 96, 264), + ], +) +@pytest.mark.parametrize( + "dtype, atol, rtol", + [ + # fp32: Triton (fp32 accum) vs cuBLAS (TF32) diverge through 3 chained + # GEMMs; activations reach ~1e4 magnitude, so atol must scale accordingly + (torch.float32, 2e2, 2e-2), + pytest.param( + torch.bfloat16, + 1e4, + 1e-2, + marks=pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"), + ), + (torch.float16, 1e2, 1e-2), + ], +) +def test_correctness_inference_mode(bsz, seq_len, hidden_size, intermediate_size, dtype, atol, rtol): + """Exercise the inference forward path (no GU saved for backward).""" + config = LlamaConfig( + hidden_size=hidden_size, + intermediate_size=intermediate_size, + hidden_act="silu", + num_attention_heads=1, + ) + + _input = torch.randn(bsz, seq_len, hidden_size, device=device, dtype=dtype) + + # initialize weights directly in their final (contiguous) orientation + G = torch.randn(intermediate_size, hidden_size, device=device, dtype=dtype) + U = torch.randn(intermediate_size, hidden_size, device=device, dtype=dtype) + D = torch.randn(hidden_size, intermediate_size, device=device, dtype=dtype) + + llama_mlp = LlamaMLP(config=config).to(device).to(dtype) + llama_mlp.gate_proj.weight.data = G + llama_mlp.up_proj.weight.data = U + llama_mlp.down_proj.weight.data = D + + liger_mlp = LigerMLP(config=config).to(device).to(dtype) + liger_mlp.gate_proj.weight.data = G + liger_mlp.up_proj.weight.data = U + liger_mlp.down_proj.weight.data = D + + with torch.no_grad(): + y1 = llama_mlp(_input) + y2 = liger_mlp(_input) + + assert torch.allclose(y1, y2, atol=atol, rtol=rtol) + + +@pytest.mark.parametrize( + "hidden_act", + [ + "gelu", + "relu", + "tanh", + ], +) +def test_invalid_hidden_act_raises(hidden_act): + config = LlamaConfig( + hidden_size=128, + intermediate_size=256, + hidden_act=hidden_act, + ) + with pytest.raises(ValueError): + LigerMLP(config=config) + + +@pytest.mark.parametrize("hidden_act", ["silu", "swish"]) +def test_supported_hidden_act(hidden_act): + config = LlamaConfig( + hidden_size=128, + intermediate_size=256, + hidden_act=hidden_act, + ) + LigerMLP(config=config) + + +@pytest.mark.xfail( + reason="TMA descriptors require 3 * intermediate_size * elem_bytes to be 16-byte aligned; " + "intermediate_size=431 with bf16 (stride 2586 bytes) violates this", + raises=AssertionError, + strict=True, +) +def test_misaligned_intermediate_size_not_supported(): + config = LlamaConfig( + hidden_size=128, + intermediate_size=431, + hidden_act="silu", + num_attention_heads=1, + ) + liger_mlp = LigerMLP(config=config).to(device).to(torch.bfloat16) + x = torch.randn(2, 8, 128, device=device, dtype=torch.bfloat16, requires_grad=True) + liger_mlp(x) From 6aee4dd8f90f522c1c5ef0d930c6860182d751b6 Mon Sep 17 00:00:00 2001 From: pearblossom <3364870135@qq.com> Date: Mon, 10 Aug 2026 14:57:17 +0800 Subject: [PATCH 2/6] fix: address review feedback Signed-off-by: pearblossom <3364870135@qq.com> --- src/liger_kernel/ops/mlp.py | 106 ++++++++++++++++++++-------------- test/transformers/test_mlp.py | 55 +++++++++--------- 2 files changed, 91 insertions(+), 70 deletions(-) diff --git a/src/liger_kernel/ops/mlp.py b/src/liger_kernel/ops/mlp.py index dc5f79918..bf3c0d70b 100644 --- a/src/liger_kernel/ops/mlp.py +++ b/src/liger_kernel/ops/mlp.py @@ -391,11 +391,24 @@ def _ensure_tma_compatible(t: torch.Tensor) -> torch.Tensor: """ strides = t.stride() elem_bytes = t.element_size() - if strides[-1] != 1: - return t.contiguous() for s in strides[:-1]: if (s * elem_bytes) % 16 != 0: - return t.contiguous() + raise ValueError( + f"Tensor does not meet TMA descriptor requirements:" + f"last stride must be 1, and all stride*elem_size must be multiple of 16." + f"Got strides={strides}, elem_size={elem_bytes}" + ) + if strides[-1] != 1: + t_contig = t.contiguous() + # Must check again. + new_strides = t_contig.stride() + for s in new_strides[:-1]: + if (s * elem_bytes) % 16 != 0: + raise ValueError( + f"Tensor after .contiguous() still not TMA-compatible: " + f"last stride must be 1, and all stride*elem_size must be multiple of 16." + ) + return t_contig return t @@ -415,15 +428,12 @@ def _check_inputs(input, gate_weight, up_weight, down_weight): return input, gate_weight, up_weight, down_weight -def swiglu_forward( +def gemm_swiglu( input: torch.Tensor, gate_weight: torch.Tensor, up_weight: torch.Tensor, - down_weight: torch.Tensor, ) -> torch.Tensor: - """Training-mode forward: computes O, and returns AGU.""" - input, gate_weight, up_weight, down_weight = _check_inputs(input, gate_weight, up_weight, down_weight) - + """Training-mode forward: computes and returns AGU.""" B, S, dim = input.shape hidden_dim = gate_weight.shape[0] bucket_M = get_bucket_m(B * S) @@ -465,19 +475,15 @@ def swiglu_forward( ) fwd_best_config = _swiglu_kernel_forward.best_config - out = F.linear(A, down_weight) + return AGU, fwd_best_config - return out, AGU, fwd_best_config - -def swiglu_forward_inference( +def gemm_swiglu_inference( input: torch.Tensor, gate_weight: torch.Tensor, up_weight: torch.Tensor, - down_weight: torch.Tensor, ) -> torch.Tensor: """Inference-mode forward.""" - input, gate_weight, up_weight, down_weight = _check_inputs(input, gate_weight, up_weight, down_weight) B, S, dim = input.shape hidden_dim = gate_weight.shape[0] bucket_M = get_bucket_m(B * S) @@ -509,33 +515,21 @@ def swiglu_forward_inference( bucket_M, ) - out = F.linear(A, down_weight) - - return out + return A -def swiglu_backward( +def swiglu_backward_dGU( dO: torch.Tensor, - input: torch.Tensor, - gate_weight: torch.Tensor, - up_weight: torch.Tensor, down_weight: torch.Tensor, AGU: torch.Tensor, fwd_best_config, ) -> torch.Tensor: - dO = _ensure_tma_compatible(dO) - B, S, dim = input.shape - hidden_dim = gate_weight.shape[0] - bucket_M = get_bucket_m(B * S) + B, S, dim = dO.shape + hidden_dim = down_weight.shape[1] - A = AGU[..., :hidden_dim] G = AGU[..., hidden_dim : 2 * hidden_dim] U = AGU[..., 2 * hidden_dim :] - # ========================== Compute dWd =================================== - dWd = dO.reshape(-1, dim).T @ A.reshape(-1, hidden_dim) # [dim, hidden_dim] - - # ================= Compute dU and dG by triton kernel ===================== # Reuse the memory of GU to store dGU, in order to save memory. dG, dU = G, U @@ -575,16 +569,23 @@ def swiglu_backward( num_stages=fwd_best_config.num_stages, ) - # ================== Compute dWug by a single large GEMM =================== - dGU = AGU[..., hidden_dim:] - dWug = dGU.reshape(-1, 2 * hidden_dim).T @ input.reshape(-1, dim) # [2*hidden_dim, dim] - dWg, dWu = dWug[:hidden_dim], dWug[hidden_dim:] # [hidden_dim, dim] - # ====================== Compute dI by triton kernel ======================= +def swiglu_backward_dI( + gate_weight: torch.Tensor, + up_weight: torch.Tensor, + AGU: torch.Tensor, +) -> torch.Tensor: + B, S, _ = AGU.shape + hidden_dim, dim = gate_weight.shape + bucket_M = get_bucket_m(B * S) + dI = torch.empty((B, S, dim), device=gate_weight.device, dtype=gate_weight.dtype) + dG = AGU[..., hidden_dim : 2 * hidden_dim] + dU = AGU[..., 2 * hidden_dim :] + dummy_block_3D = [1, 1, 1] dummy_block_2D = [1, 1] - - dI = torch.empty((B, S, dim), device=input.device, dtype=input.dtype) + desc_dG = TensorDescriptor.from_tensor(tensor=dG, block_shape=dummy_block_3D, padding="zero") + desc_dU = TensorDescriptor.from_tensor(tensor=dU, block_shape=dummy_block_3D, padding="zero") desc_dI = TensorDescriptor.from_tensor(tensor=dI, block_shape=dummy_block_3D, padding="zero") desc_Wg = TensorDescriptor.from_tensor(tensor=gate_weight, block_shape=dummy_block_2D, padding="zero") desc_Wu = TensorDescriptor.from_tensor(tensor=up_weight, block_shape=dummy_block_2D, padding="zero") @@ -606,27 +607,48 @@ def swiglu_backward( bucket_M, ) - return dI, dWg, dWu, dWd + return dI class LigerMLPFunction(torch.autograd.Function): @staticmethod def forward(ctx, input, gate_weight, up_weight, down_weight): assert input.is_cuda and gate_weight.is_cuda and up_weight.is_cuda and down_weight.is_cuda + input, gate_weight, up_weight, down_weight = _check_inputs(input, gate_weight, up_weight, down_weight) # Note: # PyTorch automatically disables global gradient computation during # the forward pass, so torch.is_grad_enabled() cannot be used. if any(ctx.needs_input_grad): - out, AGU, fwd_best_config = swiglu_forward(input, gate_weight, up_weight, down_weight) + AGU, fwd_best_config = gemm_swiglu(input, gate_weight, up_weight) ctx.save_for_backward(input, gate_weight, up_weight, down_weight, AGU) ctx.fwd_best_config = fwd_best_config - return out + hidden_dim = gate_weight.shape[0] + return F.linear(AGU[..., :hidden_dim], down_weight) else: # Inference mode: no GU saving. - return swiglu_forward_inference(input, gate_weight, up_weight, down_weight) + A = gemm_swiglu_inference(input, gate_weight, up_weight) + return F.linear(A, down_weight) @staticmethod def backward(ctx, dO): + dO = _ensure_tma_compatible(dO) input, gate_weight, up_weight, down_weight, AGU = ctx.saved_tensors - dI, dWg, dWu, dWd = swiglu_backward(dO, input, gate_weight, up_weight, down_weight, AGU, ctx.fwd_best_config) + hidden_dim, dim = gate_weight.shape + + # Compute dWd. + A = AGU[..., :hidden_dim] + dWd = dO.reshape(-1, dim).T @ A.reshape(-1, hidden_dim) # [dim, hidden_dim] + + # Compute dU and dG by triton kernel. + # The original G and U will be overwritten by dG and dU. + swiglu_backward_dGU(dO, down_weight, AGU, ctx.fwd_best_config) + + # Compute dWug by a single large GEMM. + dGU = AGU[..., hidden_dim:] + dWug = dGU.reshape(-1, 2 * hidden_dim).T @ input.reshape(-1, dim) # [2*hidden_dim, dim] + dWg, dWu = dWug[:hidden_dim], dWug[hidden_dim:] # [hidden_dim, dim] + + # Compute dI by triton kernel. + dI = swiglu_backward_dI(gate_weight, up_weight, AGU) + return dI, dWg, dWu, dWd diff --git a/test/transformers/test_mlp.py b/test/transformers/test_mlp.py index a65029255..98ff9b76c 100644 --- a/test/transformers/test_mlp.py +++ b/test/transformers/test_mlp.py @@ -10,6 +10,24 @@ device = infer_device() +FP32_DIFF_THRESHOLD = 1e-5 + + +def calc_diff(x: torch.Tensor, y: torch.Tensor): + x, y = x.double(), y.double() + denominator = (x * x + y * y).sum() + if denominator == 0: # Which means that all elements in x and y are 0 + return 0.0 + sim = 2 * (x * y).sum() / denominator + return 1 - sim + + +def _assert_close(x: torch.Tensor, y: torch.Tensor, dtype: torch.dtype, atol, rtol): + if dtype == torch.float32: + assert calc_diff(x, y) < FP32_DIFF_THRESHOLD + else: + assert torch.allclose(x, y, atol=atol, rtol=rtol) + @pytest.mark.parametrize( "bsz, seq_len, hidden_size, intermediate_size", @@ -24,9 +42,7 @@ @pytest.mark.parametrize( "dtype, atol, rtol", [ - # fp32: Triton (fp32 accum) vs cuBLAS (TF32) diverge through 3 chained - # GEMMs; activations reach ~1e4 magnitude, so atol must scale accordingly - (torch.float32, 2e2, 2e-2), + (torch.float32, None, None), pytest.param( torch.bfloat16, 1e4, @@ -67,33 +83,18 @@ def test_correctness_llamamlp(bsz, seq_len, hidden_size, intermediate_size, dtyp y1 = llama_mlp(x1) y2 = liger_mlp(x2) - assert torch.allclose(y1, y2, atol=atol, rtol=rtol) + _assert_close(y1, y2, dtype, atol, rtol) dy = torch.randn_like(y1) y1.backward(dy.clone(), retain_graph=True) y2.backward(dy.clone(), retain_graph=True) - assert torch.allclose( - llama_mlp.gate_proj.weight.grad, - liger_mlp.gate_proj.weight.grad, - atol=atol, - rtol=rtol, - ) - assert torch.allclose( - llama_mlp.up_proj.weight.grad, - liger_mlp.up_proj.weight.grad, - atol=atol, - rtol=rtol, - ) - assert torch.allclose( - llama_mlp.down_proj.weight.grad, - liger_mlp.down_proj.weight.grad, - atol=atol, - rtol=rtol, - ) + _assert_close(llama_mlp.gate_proj.weight.grad, liger_mlp.gate_proj.weight.grad, dtype, atol, rtol) + _assert_close(llama_mlp.up_proj.weight.grad, liger_mlp.up_proj.weight.grad, dtype, atol, rtol) + _assert_close(llama_mlp.down_proj.weight.grad, liger_mlp.down_proj.weight.grad, dtype, atol, rtol) - assert torch.allclose(x1.grad, x2.grad, atol=atol, rtol=rtol) + _assert_close(x1.grad, x2.grad, dtype, atol, rtol) @pytest.mark.parametrize( @@ -109,9 +110,7 @@ def test_correctness_llamamlp(bsz, seq_len, hidden_size, intermediate_size, dtyp @pytest.mark.parametrize( "dtype, atol, rtol", [ - # fp32: Triton (fp32 accum) vs cuBLAS (TF32) diverge through 3 chained - # GEMMs; activations reach ~1e4 magnitude, so atol must scale accordingly - (torch.float32, 2e2, 2e-2), + (torch.float32, None, None), pytest.param( torch.bfloat16, 1e4, @@ -151,7 +150,7 @@ def test_correctness_inference_mode(bsz, seq_len, hidden_size, intermediate_size y1 = llama_mlp(_input) y2 = liger_mlp(_input) - assert torch.allclose(y1, y2, atol=atol, rtol=rtol) + _assert_close(y1, y2, dtype, atol, rtol) @pytest.mark.parametrize( @@ -185,7 +184,7 @@ def test_supported_hidden_act(hidden_act): @pytest.mark.xfail( reason="TMA descriptors require 3 * intermediate_size * elem_bytes to be 16-byte aligned; " "intermediate_size=431 with bf16 (stride 2586 bytes) violates this", - raises=AssertionError, + raises=ValueError, strict=True, ) def test_misaligned_intermediate_size_not_supported(): From bc75948ed86f12d1c10d19e0910b14a06b065b98 Mon Sep 17 00:00:00 2001 From: pearblossom <3364870135@qq.com> Date: Mon, 10 Aug 2026 22:12:14 +0800 Subject: [PATCH 3/6] =?UTF-8?q?Add=20store=5Fpreact=20flag=20in=20gemm=5Fs?= =?UTF-8?q?wiglu=20(replaces=20gemm=5Fswiglu=5Finference);=20update=20test?= =?UTF-8?q?s=20to=20cosine=E2=80=91similarity=20for=20all=20precision=20va?= =?UTF-8?q?lidations.?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: pearblossom <3364870135@qq.com> --- src/liger_kernel/ops/mlp.py | 126 ++++++++++++++-------------------- test/transformers/test_mlp.py | 41 +++++------ 2 files changed, 70 insertions(+), 97 deletions(-) diff --git a/src/liger_kernel/ops/mlp.py b/src/liger_kernel/ops/mlp.py index bf3c0d70b..051b76845 100644 --- a/src/liger_kernel/ops/mlp.py +++ b/src/liger_kernel/ops/mlp.py @@ -405,8 +405,8 @@ def _ensure_tma_compatible(t: torch.Tensor) -> torch.Tensor: for s in new_strides[:-1]: if (s * elem_bytes) % 16 != 0: raise ValueError( - f"Tensor after .contiguous() still not TMA-compatible: " - f"last stride must be 1, and all stride*elem_size must be multiple of 16." + "Tensor after .contiguous() still not TMA-compatible: " + "last stride must be 1, and all stride*elem_size must be multiple of 16." ) return t_contig return t @@ -432,62 +432,26 @@ def gemm_swiglu( input: torch.Tensor, gate_weight: torch.Tensor, up_weight: torch.Tensor, -) -> torch.Tensor: - """Training-mode forward: computes and returns AGU.""" + store_preact: bool = True, +): + """ + Fused gate/up projection + SwiGLU. + store_preact=True (training): also stores G and U for backward, returns (AGU, fwd_best_config). + store_preact=False (inference): returns (A, None). + """ B, S, dim = input.shape hidden_dim = gate_weight.shape[0] bucket_M = get_bucket_m(B * S) - # Optimization: merge 3 torch.empty calls into 1 to reduce allocator overhead. - AGU = torch.empty(B, S, 3 * hidden_dim, device=input.device, dtype=input.dtype) - # Slice out 3 independent 3D views. - A = AGU[..., :hidden_dim] - G = AGU[..., hidden_dim : 2 * hidden_dim] - U = AGU[..., 2 * hidden_dim :] - - # Dummy block - dummy_block_3D = [1, 1, 1] - dummy_block_2D = [1, 1] - - # Create TensorDescriptor - desc_input = TensorDescriptor.from_tensor(tensor=input, block_shape=dummy_block_3D, padding="zero") - desc_Wg = TensorDescriptor.from_tensor(tensor=gate_weight, block_shape=dummy_block_2D, padding="zero") - desc_Wu = TensorDescriptor.from_tensor(tensor=up_weight, block_shape=dummy_block_2D, padding="zero") - desc_A = TensorDescriptor.from_tensor(tensor=A, block_shape=dummy_block_3D, padding="zero") - desc_G = TensorDescriptor.from_tensor(tensor=G, block_shape=dummy_block_3D, padding="zero") - desc_U = TensorDescriptor.from_tensor(tensor=U, block_shape=dummy_block_3D, padding="zero") - - grid = lambda META: ( - triton.cdiv(S, META["BLOCK_M"]) * triton.cdiv(hidden_dim, META["BLOCK_N"]), - B, - ) - - _swiglu_kernel_forward[grid]( - desc_input, - desc_Wg, - desc_Wu, - desc_A, - desc_G, - desc_U, - S, - dim, - hidden_dim, - bucket_M, - ) - fwd_best_config = _swiglu_kernel_forward.best_config - - return AGU, fwd_best_config - -def gemm_swiglu_inference( - input: torch.Tensor, - gate_weight: torch.Tensor, - up_weight: torch.Tensor, -) -> torch.Tensor: - """Inference-mode forward.""" - B, S, dim = input.shape - hidden_dim = gate_weight.shape[0] - bucket_M = get_bucket_m(B * S) - A = torch.empty((B, S, hidden_dim), device=input.device, dtype=input.dtype) + if store_preact: + # Optimization: merge 3 torch.empty calls into 1 to reduce allocator overhead. + AGU = torch.empty(B, S, 3 * hidden_dim, device=input.device, dtype=input.dtype) + # Slice out 3 independent 3D views. + A = AGU[..., :hidden_dim] + G = AGU[..., hidden_dim : 2 * hidden_dim] + U = AGU[..., 2 * hidden_dim :] + else: + A = torch.empty((B, S, hidden_dim), device=input.device, dtype=input.dtype) # Dummy block dummy_block_3D = [1, 1, 1] @@ -504,18 +468,34 @@ def gemm_swiglu_inference( B, ) - _swiglu_kernel_forward_inference[grid]( - desc_input, - desc_Wg, - desc_Wu, - desc_A, - S, - dim, - hidden_dim, - bucket_M, - ) - - return A + if store_preact: + desc_G = TensorDescriptor.from_tensor(tensor=G, block_shape=dummy_block_3D, padding="zero") + desc_U = TensorDescriptor.from_tensor(tensor=U, block_shape=dummy_block_3D, padding="zero") + _swiglu_kernel_forward[grid]( + desc_input, + desc_Wg, + desc_Wu, + desc_A, + desc_G, + desc_U, + S, + dim, + hidden_dim, + bucket_M, + ) + return AGU, _swiglu_kernel_forward.best_config + else: + _swiglu_kernel_forward_inference[grid]( + desc_input, + desc_Wg, + desc_Wu, + desc_A, + S, + dim, + hidden_dim, + bucket_M, + ) + return A, None def swiglu_backward_dGU( @@ -618,16 +598,16 @@ def forward(ctx, input, gate_weight, up_weight, down_weight): # Note: # PyTorch automatically disables global gradient computation during # the forward pass, so torch.is_grad_enabled() cannot be used. - if any(ctx.needs_input_grad): - AGU, fwd_best_config = gemm_swiglu(input, gate_weight, up_weight) - ctx.save_for_backward(input, gate_weight, up_weight, down_weight, AGU) + store_preact = any(ctx.needs_input_grad) + out, fwd_best_config = gemm_swiglu(input, gate_weight, up_weight, store_preact) + if store_preact: + ctx.save_for_backward(input, gate_weight, up_weight, down_weight, out) ctx.fwd_best_config = fwd_best_config hidden_dim = gate_weight.shape[0] - return F.linear(AGU[..., :hidden_dim], down_weight) + return F.linear(out[..., :hidden_dim], down_weight) else: # Inference mode: no GU saving. - A = gemm_swiglu_inference(input, gate_weight, up_weight) - return F.linear(A, down_weight) + return F.linear(out, down_weight) @staticmethod def backward(ctx, dO): @@ -647,7 +627,7 @@ def backward(ctx, dO): dGU = AGU[..., hidden_dim:] dWug = dGU.reshape(-1, 2 * hidden_dim).T @ input.reshape(-1, dim) # [2*hidden_dim, dim] dWg, dWu = dWug[:hidden_dim], dWug[hidden_dim:] # [hidden_dim, dim] - + # Compute dI by triton kernel. dI = swiglu_backward_dI(gate_weight, up_weight, AGU) diff --git a/test/transformers/test_mlp.py b/test/transformers/test_mlp.py index 98ff9b76c..d12531e5d 100644 --- a/test/transformers/test_mlp.py +++ b/test/transformers/test_mlp.py @@ -10,7 +10,7 @@ device = infer_device() -FP32_DIFF_THRESHOLD = 1e-5 +DIFF_THRESHOLD = 1e-5 def calc_diff(x: torch.Tensor, y: torch.Tensor): @@ -22,11 +22,8 @@ def calc_diff(x: torch.Tensor, y: torch.Tensor): return 1 - sim -def _assert_close(x: torch.Tensor, y: torch.Tensor, dtype: torch.dtype, atol, rtol): - if dtype == torch.float32: - assert calc_diff(x, y) < FP32_DIFF_THRESHOLD - else: - assert torch.allclose(x, y, atol=atol, rtol=rtol) +def _assert_close(x: torch.Tensor, y: torch.Tensor): + assert calc_diff(x, y) < DIFF_THRESHOLD @pytest.mark.parametrize( @@ -40,19 +37,17 @@ def _assert_close(x: torch.Tensor, y: torch.Tensor, dtype: torch.dtype, atol, rt ], ) @pytest.mark.parametrize( - "dtype, atol, rtol", + "dtype", [ - (torch.float32, None, None), + torch.float32, pytest.param( torch.bfloat16, - 1e4, - 1e-2, marks=pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"), ), - (torch.float16, 1e2, 1e-2), + torch.float16, ], ) -def test_correctness_llamamlp(bsz, seq_len, hidden_size, intermediate_size, dtype, atol, rtol): +def test_correctness_llamamlp(bsz, seq_len, hidden_size, intermediate_size, dtype): config = LlamaConfig( hidden_size=hidden_size, intermediate_size=intermediate_size, @@ -83,18 +78,18 @@ def test_correctness_llamamlp(bsz, seq_len, hidden_size, intermediate_size, dtyp y1 = llama_mlp(x1) y2 = liger_mlp(x2) - _assert_close(y1, y2, dtype, atol, rtol) + _assert_close(y1, y2) dy = torch.randn_like(y1) y1.backward(dy.clone(), retain_graph=True) y2.backward(dy.clone(), retain_graph=True) - _assert_close(llama_mlp.gate_proj.weight.grad, liger_mlp.gate_proj.weight.grad, dtype, atol, rtol) - _assert_close(llama_mlp.up_proj.weight.grad, liger_mlp.up_proj.weight.grad, dtype, atol, rtol) - _assert_close(llama_mlp.down_proj.weight.grad, liger_mlp.down_proj.weight.grad, dtype, atol, rtol) + _assert_close(llama_mlp.gate_proj.weight.grad, liger_mlp.gate_proj.weight.grad) + _assert_close(llama_mlp.up_proj.weight.grad, liger_mlp.up_proj.weight.grad) + _assert_close(llama_mlp.down_proj.weight.grad, liger_mlp.down_proj.weight.grad) - _assert_close(x1.grad, x2.grad, dtype, atol, rtol) + _assert_close(x1.grad, x2.grad) @pytest.mark.parametrize( @@ -108,19 +103,17 @@ def test_correctness_llamamlp(bsz, seq_len, hidden_size, intermediate_size, dtyp ], ) @pytest.mark.parametrize( - "dtype, atol, rtol", + "dtype", [ - (torch.float32, None, None), + torch.float32, pytest.param( torch.bfloat16, - 1e4, - 1e-2, marks=pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"), ), - (torch.float16, 1e2, 1e-2), + torch.float16, ], ) -def test_correctness_inference_mode(bsz, seq_len, hidden_size, intermediate_size, dtype, atol, rtol): +def test_correctness_inference_mode(bsz, seq_len, hidden_size, intermediate_size, dtype): """Exercise the inference forward path (no GU saved for backward).""" config = LlamaConfig( hidden_size=hidden_size, @@ -150,7 +143,7 @@ def test_correctness_inference_mode(bsz, seq_len, hidden_size, intermediate_size y1 = llama_mlp(_input) y2 = liger_mlp(_input) - _assert_close(y1, y2, dtype, atol, rtol) + _assert_close(y1, y2) @pytest.mark.parametrize( From 279614fdf91729337294f2df66a8e1e783b7aa75 Mon Sep 17 00:00:00 2001 From: pearblossom <3364870135@qq.com> Date: Fri, 14 Aug 2026 20:08:26 +0800 Subject: [PATCH 4/6] add support for Falcon H1 and replace LigerSwiGLUMLP with LigerMLP Signed-off-by: pearblossom <3364870135@qq.com> --- src/liger_kernel/ops/mlp.py | 74 ++++++-- src/liger_kernel/transformers/mlp.py | 35 ++++ src/liger_kernel/transformers/monkey_patch.py | 17 +- test/transformers/test_mlp.py | 162 +++++++++++++++++- test/transformers/test_monkey_patch.py | 30 +++- 5 files changed, 297 insertions(+), 21 deletions(-) diff --git a/src/liger_kernel/ops/mlp.py b/src/liger_kernel/ops/mlp.py index 051b76845..c159571c3 100644 --- a/src/liger_kernel/ops/mlp.py +++ b/src/liger_kernel/ops/mlp.py @@ -92,6 +92,7 @@ def _swiglu_kernel_forward( BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, GROUP_SIZE_M: tl.constexpr, + gate_multiplier: tl.constexpr, ): pid = tl.program_id(0) b = tl.program_id(1) @@ -114,7 +115,12 @@ def _swiglu_kernel_forward( acc_up = tl.dot(input_block, tl.trans(up_block), acc=acc_up) # Compute A_block - A_block = acc_gate * tl.sigmoid(acc_gate) * acc_up # [BLOCK_M, BLOCK_N] + # Compile-time branch elimination, no runtime overhead. + if gate_multiplier == 1.0: + A_block = acc_gate * tl.sigmoid(acc_gate) * acc_up # [BLOCK_M, BLOCK_N] + else: + acc_gate = acc_gate * gate_multiplier + A_block = acc_gate * tl.sigmoid(acc_gate) * acc_up # [BLOCK_M, BLOCK_N] # Write back dtype = desc_input.dtype @@ -178,6 +184,7 @@ def _swiglu_kernel_forward_inference( BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, GROUP_SIZE_M: tl.constexpr, + gate_multiplier: tl.constexpr, ): pid = tl.program_id(0) b = tl.program_id(1) @@ -200,7 +207,12 @@ def _swiglu_kernel_forward_inference( acc_up = tl.dot(input_block, tl.trans(up_block), acc=acc_up) # Compute A_block - A_block = acc_gate * tl.sigmoid(acc_gate) * acc_up # [BLOCK_M, BLOCK_N] + # Compile-time branch elimination, no runtime overhead. + if gate_multiplier == 1.0: + A_block = acc_gate * tl.sigmoid(acc_gate) * acc_up # [BLOCK_M, BLOCK_N] + else: + acc_gate = acc_gate * gate_multiplier + A_block = acc_gate * tl.sigmoid(acc_gate) * acc_up # [BLOCK_M, BLOCK_N] # Write back dtype = desc_input.dtype @@ -238,6 +250,8 @@ def _swiglu_kernel_backward_dGU( BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, GROUP_SIZE_M: tl.constexpr, + gate_multiplier: tl.constexpr, + down_multiplier: tl.constexpr, ): pid = tl.program_id(0) b = tl.program_id(1) @@ -262,17 +276,26 @@ def _swiglu_kernel_backward_dGU( U_block = tl.reshape(U_block, [BLOCK_M, BLOCK_N]) # [BLOCK_M, BLOCK_N] # Compute and save dU, dG + # NOTE: on Falcon H1 path, the forward kernel already stores G pre-multiplied + # by gate_multiplier, so do NOT multiply it again here. G_block = G_block.to(tl.float32) U_block = U_block.to(tl.float32) + sigmoid_G = tl.sigmoid(G_block) silu_G = G_block * sigmoid_G dtype = desc_G.dtype - dU = acc_A * silu_G # [BLOCK_M, BLOCK_N] + if gate_multiplier == 1.0 and down_multiplier == 1.0: + dU = acc_A * silu_G # [BLOCK_M, BLOCK_N] + dG = acc_A * U_block * (sigmoid_G + silu_G - sigmoid_G * silu_G) # [BLOCK_M, BLOCK_N] + else: + dU = acc_A * down_multiplier * silu_G # [BLOCK_M, BLOCK_N] + dG = ( + acc_A * down_multiplier * U_block * (sigmoid_G + silu_G - sigmoid_G * silu_G) * gate_multiplier + ) # [BLOCK_M, BLOCK_N] + dU = tl.reshape(dU, [1, BLOCK_M, BLOCK_N]) desc_dU.store([b, M_start, N_start], dU.to(dtype)) - - dG = acc_A * U_block * (sigmoid_G + silu_G - sigmoid_G * silu_G) # [BLOCK_M, BLOCK_N] dG = tl.reshape(dG, [1, BLOCK_M, BLOCK_N]) desc_dG.store([b, M_start, N_start], dG.to(dtype)) @@ -432,6 +455,7 @@ def gemm_swiglu( input: torch.Tensor, gate_weight: torch.Tensor, up_weight: torch.Tensor, + gate_multiplier: float, store_preact: bool = True, ): """ @@ -482,6 +506,7 @@ def gemm_swiglu( dim, hidden_dim, bucket_M, + gate_multiplier=gate_multiplier, ) return AGU, _swiglu_kernel_forward.best_config else: @@ -494,6 +519,7 @@ def gemm_swiglu( dim, hidden_dim, bucket_M, + gate_multiplier=gate_multiplier, ) return A, None @@ -503,6 +529,8 @@ def swiglu_backward_dGU( down_weight: torch.Tensor, AGU: torch.Tensor, fwd_best_config, + gate_multiplier: float, + down_multiplier: float, ) -> torch.Tensor: B, S, dim = dO.shape hidden_dim = down_weight.shape[1] @@ -545,6 +573,8 @@ def swiglu_backward_dGU( BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K, GROUP_SIZE_M=GROUP_SIZE_M, + gate_multiplier=gate_multiplier, + down_multiplier=down_multiplier, num_warps=fwd_best_config.num_warps, num_stages=fwd_best_config.num_stages, ) @@ -592,36 +622,56 @@ def swiglu_backward_dI( class LigerMLPFunction(torch.autograd.Function): @staticmethod - def forward(ctx, input, gate_weight, up_weight, down_weight): + def forward( + ctx, + input, + gate_weight, + up_weight, + down_weight, + gate_multiplier=1.0, + down_multiplier=1.0, + ): assert input.is_cuda and gate_weight.is_cuda and up_weight.is_cuda and down_weight.is_cuda input, gate_weight, up_weight, down_weight = _check_inputs(input, gate_weight, up_weight, down_weight) # Note: # PyTorch automatically disables global gradient computation during # the forward pass, so torch.is_grad_enabled() cannot be used. store_preact = any(ctx.needs_input_grad) - out, fwd_best_config = gemm_swiglu(input, gate_weight, up_weight, store_preact) + out, fwd_best_config = gemm_swiglu(input, gate_weight, up_weight, gate_multiplier, store_preact) if store_preact: ctx.save_for_backward(input, gate_weight, up_weight, down_weight, out) ctx.fwd_best_config = fwd_best_config + ctx.gate_multiplier = gate_multiplier + ctx.down_multiplier = down_multiplier hidden_dim = gate_weight.shape[0] - return F.linear(out[..., :hidden_dim], down_weight) + if down_multiplier == 1.0: + return F.linear(out[..., :hidden_dim], down_weight) + else: + return F.linear(out[..., :hidden_dim], down_weight) * down_multiplier else: # Inference mode: no GU saving. - return F.linear(out, down_weight) + if down_multiplier == 1.0: + return F.linear(out, down_weight) + else: + return F.linear(out, down_weight) * down_multiplier @staticmethod def backward(ctx, dO): dO = _ensure_tma_compatible(dO) input, gate_weight, up_weight, down_weight, AGU = ctx.saved_tensors + gate_multiplier, down_multiplier = ctx.gate_multiplier, ctx.down_multiplier hidden_dim, dim = gate_weight.shape # Compute dWd. A = AGU[..., :hidden_dim] - dWd = dO.reshape(-1, dim).T @ A.reshape(-1, hidden_dim) # [dim, hidden_dim] + if down_multiplier == 1.0: + dWd = dO.reshape(-1, dim).T @ A.reshape(-1, hidden_dim) # [dim, hidden_dim] + else: + dWd = dO.reshape(-1, dim).T @ A.reshape(-1, hidden_dim) * down_multiplier # Compute dU and dG by triton kernel. # The original G and U will be overwritten by dG and dU. - swiglu_backward_dGU(dO, down_weight, AGU, ctx.fwd_best_config) + swiglu_backward_dGU(dO, down_weight, AGU, ctx.fwd_best_config, gate_multiplier, down_multiplier) # Compute dWug by a single large GEMM. dGU = AGU[..., hidden_dim:] @@ -631,4 +681,4 @@ def backward(ctx, dO): # Compute dI by triton kernel. dI = swiglu_backward_dI(gate_weight, up_weight, AGU) - return dI, dWg, dWu, dWd + return dI, dWg, dWu, dWd, None, None diff --git a/src/liger_kernel/transformers/mlp.py b/src/liger_kernel/transformers/mlp.py index cec442600..3086f0bad 100644 --- a/src/liger_kernel/transformers/mlp.py +++ b/src/liger_kernel/transformers/mlp.py @@ -18,3 +18,38 @@ def __init__(self, config): def forward(self, x: torch.Tensor) -> torch.Tensor: return LigerMLPFunction.apply(x, self.gate_proj.weight, self.up_proj.weight, self.down_proj.weight) + + +class LigerFalconH1MLP(nn.Module): + """ + Patch FalconH1MLP to use LigerMLPFunction with gate / down multipliers. + Falcon H1's MLP block pre-scales the gate pre-activation and post-scales the + down projection output: + y = down_proj(silu(gate_proj(x) * gate_mult) * up_proj(x)) * down_mult + https://github.com/huggingface/transformers/blob/main/src/transformers/models/falcon_h1/modeling_falcon_h1.py + """ + + def __init__(self, config): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size + bias = getattr(config, "mlp_bias", False) + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=bias) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=bias) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=bias) + if config.hidden_act not in ["silu", "swish"]: + raise ValueError(f"Activation function {config.hidden_act} not supported.") + gate_multiplier, down_multiplier = config.mlp_multipliers + self.gate_multiplier = gate_multiplier + self.down_multiplier = down_multiplier + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return LigerMLPFunction.apply( + x, + self.gate_proj.weight, + self.up_proj.weight, + self.down_proj.weight, + float(self.gate_multiplier), + float(self.down_multiplier), + ) diff --git a/src/liger_kernel/transformers/monkey_patch.py b/src/liger_kernel/transformers/monkey_patch.py index 4d33d7e41..5f8b27fd4 100755 --- a/src/liger_kernel/transformers/monkey_patch.py +++ b/src/liger_kernel/transformers/monkey_patch.py @@ -37,7 +37,12 @@ from liger_kernel.transformers.swiglu import LigerBlockSparseTop2MLP from liger_kernel.transformers.swiglu import LigerExperts from liger_kernel.transformers.swiglu import LigerPhi3SwiGLUMLP -from liger_kernel.transformers.swiglu import LigerSwiGLUMLP + +USE_FLASH_SWIGLU = True +if USE_FLASH_SWIGLU: + from liger_kernel.transformers.mlp import LigerMLP as LigerSwiGLUMLP +else: + from liger_kernel.transformers.swiglu import LigerSwiGLUMLP try: import peft @@ -465,6 +470,11 @@ def apply_liger_kernel_to_llama4( from liger_kernel.transformers.model.llama4 import lce_forward as llama4_lce_forward + # NOTE: Llama4's text MLP receives a 2D [T, hidden] input, which the fused LigerMLP + # kernels do not support, so we use the original LigerSwiGLUMLP implementation + # here instead of LigerMLP. + from liger_kernel.transformers.swiglu import LigerSwiGLUMLP + if rope: from liger_kernel.transformers.llama4_rope import apply_liger_llama4_rope_full @@ -2911,7 +2921,10 @@ def apply_liger_kernel_to_falcon_h1( from transformers.models.falcon_h1 import modeling_falcon_h1 from transformers.models.falcon_h1.modeling_falcon_h1 import FalconH1Model - from liger_kernel.transformers.swiglu import LigerFalconH1SwiGLUMLP + if USE_FLASH_SWIGLU: + from liger_kernel.transformers.mlp import LigerFalconH1MLP as LigerFalconH1SwiGLUMLP + else: + from liger_kernel.transformers.swiglu import LigerFalconH1SwiGLUMLP if rope: logger.info("Apply liger rotary pos emb.") diff --git a/test/transformers/test_mlp.py b/test/transformers/test_mlp.py index d12531e5d..cd23d351d 100644 --- a/test/transformers/test_mlp.py +++ b/test/transformers/test_mlp.py @@ -5,12 +5,43 @@ from transformers.models.llama.configuration_llama import LlamaConfig from transformers.models.llama.modeling_llama import LlamaMLP +from liger_kernel.transformers.mlp import LigerFalconH1MLP from liger_kernel.transformers.mlp import LigerMLP from liger_kernel.utils import infer_device +try: + from transformers.models.falcon_h1.configuration_falcon_h1 import FalconH1Config + from transformers.models.falcon_h1.modeling_falcon_h1 import FalconH1MLP + + FALCON_H1_AVAILABLE = True +except ImportError: + FALCON_H1_AVAILABLE = False + device = infer_device() DIFF_THRESHOLD = 1e-5 +# Why the FalconH1 tests use a looser tolerance only for bf16. +# +# The reference implementation and the fused kernel round at different places: +# the reference rounds every intermediate result to the compute dtype (the +# gate/up GEMM outputs, the scaled gate g * gate_mult, the silu output, the +# silu(g) * up product...), while the kernel keeps everything in fp32 until the +# very end. +# +# Without multipliers there are fewer such differences (one less intermediate +# rounding), so bf16 stays at ~8e-6 and fits within 1e-5. With multipliers, the +# extra g * gate_mult product adds one more rounding step, pushing bf16 to +# ~1.02e-5, just over 1e-5 - hence the 2e-5 relaxation for bf16 only. fp32 +# (~1.1e-6) and fp16 (~1.6e-7) are far below 1e-5 either way and keep the +# strict threshold, which also guards against formula errors. +FALCON_H1_BF16_DIFF_THRESHOLD = 2e-5 + + +def falcon_h1_threshold(dtype): + return FALCON_H1_BF16_DIFF_THRESHOLD if dtype == torch.bfloat16 else DIFF_THRESHOLD + + +falcon_h1_unavailable = pytest.mark.skipif(not FALCON_H1_AVAILABLE, reason="falcon_h1 module not available") def calc_diff(x: torch.Tensor, y: torch.Tensor): @@ -22,8 +53,8 @@ def calc_diff(x: torch.Tensor, y: torch.Tensor): return 1 - sim -def _assert_close(x: torch.Tensor, y: torch.Tensor): - assert calc_diff(x, y) < DIFF_THRESHOLD +def _assert_close(x: torch.Tensor, y: torch.Tensor, threshold: float = DIFF_THRESHOLD): + assert calc_diff(x, y) < threshold @pytest.mark.parametrize( @@ -190,3 +221,130 @@ def test_misaligned_intermediate_size_not_supported(): liger_mlp = LigerMLP(config=config).to(device).to(torch.bfloat16) x = torch.randn(2, 8, 128, device=device, dtype=torch.bfloat16, requires_grad=True) liger_mlp(x) + + +@falcon_h1_unavailable +@pytest.mark.parametrize( + "bsz, seq_len, hidden_size, intermediate_size", + [ + (2, 256, 256, 512), + # weird shapes (TMA requires every row stride to be 16-byte aligned, + # so hidden/intermediate sizes are kept multiples of 8) + (6, 42, 128, 432), + (3, 37, 96, 264), + ], +) +@pytest.mark.parametrize( + "dtype", + [ + torch.float32, + pytest.param( + torch.bfloat16, + marks=pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"), + ), + torch.float16, + ], +) +@pytest.mark.parametrize("mlp_multipliers", [[1.0, 1.0], [1.7, 0.85]]) +def test_correctness_falconh1mlp(bsz, seq_len, hidden_size, intermediate_size, dtype, mlp_multipliers): + config = FalconH1Config( + hidden_size=hidden_size, + intermediate_size=intermediate_size, + hidden_act="silu", + mlp_bias=False, + mlp_multipliers=mlp_multipliers, + num_attention_heads=1, + num_key_value_heads=1, + ) + + _input = torch.randn(bsz, seq_len, hidden_size, device=device, dtype=dtype) + + x1 = _input.clone().requires_grad_(True) + x2 = _input.clone().requires_grad_(True) + + G = torch.randn(intermediate_size, hidden_size, device=device, dtype=dtype) + U = torch.randn(intermediate_size, hidden_size, device=device, dtype=dtype) + D = torch.randn(hidden_size, intermediate_size, device=device, dtype=dtype) + + falcon_mlp = FalconH1MLP(config=config).to(device).to(dtype) + falcon_mlp.gate_proj.weight.data = G + falcon_mlp.up_proj.weight.data = U + falcon_mlp.down_proj.weight.data = D + + liger_mlp = LigerFalconH1MLP(config=config).to(device).to(dtype) + liger_mlp.gate_proj.weight.data = G + liger_mlp.up_proj.weight.data = U + liger_mlp.down_proj.weight.data = D + + y1 = falcon_mlp(x1) + y2 = liger_mlp(x2) + + threshold = falcon_h1_threshold(dtype) + _assert_close(y1, y2, threshold) + + dy = torch.randn_like(y1) + + y1.backward(dy.clone(), retain_graph=True) + y2.backward(dy.clone(), retain_graph=True) + + _assert_close(falcon_mlp.gate_proj.weight.grad, liger_mlp.gate_proj.weight.grad, threshold) + _assert_close(falcon_mlp.up_proj.weight.grad, liger_mlp.up_proj.weight.grad, threshold) + _assert_close(falcon_mlp.down_proj.weight.grad, liger_mlp.down_proj.weight.grad, threshold) + + _assert_close(x1.grad, x2.grad, threshold) + + +@falcon_h1_unavailable +@pytest.mark.parametrize( + "bsz, seq_len, hidden_size, intermediate_size", + [ + (2, 256, 256, 512), + (6, 42, 128, 432), + (3, 37, 96, 264), + ], +) +@pytest.mark.parametrize( + "dtype", + [ + torch.float32, + pytest.param( + torch.bfloat16, + marks=pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"), + ), + torch.float16, + ], +) +@pytest.mark.parametrize("mlp_multipliers", [[1.0, 1.0], [1.7, 0.85]]) +def test_correctness_falconh1mlp_inference_mode(bsz, seq_len, hidden_size, intermediate_size, dtype, mlp_multipliers): + """Exercise the inference forward path (no GU saved for backward).""" + config = FalconH1Config( + hidden_size=hidden_size, + intermediate_size=intermediate_size, + hidden_act="silu", + mlp_bias=False, + mlp_multipliers=mlp_multipliers, + num_attention_heads=1, + num_key_value_heads=1, + ) + + _input = torch.randn(bsz, seq_len, hidden_size, device=device, dtype=dtype) + + G = torch.randn(intermediate_size, hidden_size, device=device, dtype=dtype) + U = torch.randn(intermediate_size, hidden_size, device=device, dtype=dtype) + D = torch.randn(hidden_size, intermediate_size, device=device, dtype=dtype) + + falcon_mlp = FalconH1MLP(config=config).to(device).to(dtype) + falcon_mlp.gate_proj.weight.data = G + falcon_mlp.up_proj.weight.data = U + falcon_mlp.down_proj.weight.data = D + + liger_mlp = LigerFalconH1MLP(config=config).to(device).to(dtype) + liger_mlp.gate_proj.weight.data = G + liger_mlp.up_proj.weight.data = U + liger_mlp.down_proj.weight.data = D + + with torch.no_grad(): + y1 = falcon_mlp(_input) + y2 = liger_mlp(_input) + + _assert_close(y1, y2, falcon_h1_threshold(dtype)) diff --git a/test/transformers/test_monkey_patch.py b/test/transformers/test_monkey_patch.py index 25099f7be..657386980 100755 --- a/test/transformers/test_monkey_patch.py +++ b/test/transformers/test_monkey_patch.py @@ -22,7 +22,8 @@ from liger_kernel.transformers import LigerPhi3SwiGLUMLP from liger_kernel.transformers import LigerQwen3MoeSwiGLUMLP from liger_kernel.transformers import LigerRMSNorm -from liger_kernel.transformers import LigerSwiGLUMLP + +# from liger_kernel.transformers import LigerSwiGLUMLP from liger_kernel.transformers import monkey_patch from liger_kernel.transformers.layer_norm import LigerLayerNorm from liger_kernel.transformers.model.falcon_h1 import lce_forward as falcon_h1_lce_forward @@ -41,9 +42,15 @@ from liger_kernel.transformers.model.qwen3_next import lce_forward as qwen3_next_lce_forward from liger_kernel.transformers.model.smollm3 import lce_forward as smolllm3_lce_forward from liger_kernel.transformers.monkey_patch import MODEL_TYPE_TO_APPLY_LIGER_FN +from liger_kernel.transformers.monkey_patch import USE_FLASH_SWIGLU from liger_kernel.transformers.monkey_patch import _apply_liger_kernel from liger_kernel.transformers.monkey_patch import _apply_liger_kernel_to_instance +if USE_FLASH_SWIGLU: + from liger_kernel.transformers.mlp import LigerMLP as LigerSwiGLUMLP +else: + from liger_kernel.transformers.swiglu import LigerSwiGLUMLP + # We only support transformers >= 4.52.0 transformer_version = version.parse(transformers.__version__) MIN_SUPPORTED_TRANSFORMERS_VERSION = version.parse("4.52.0") @@ -1340,6 +1347,10 @@ def test_apply_liger_kernel_to_instance_for_llama4_for_causal_lm(): with patch("transformers.models.llama4.modeling_llama4"): from transformers.models.llama4.modeling_llama4 import Llama4ForCausalLM + # Compare against the class actually bound by the patch: llama4 keeps the + # original swiglu implementation because its MLPs receive 2D input. + from liger_kernel.transformers.swiglu import LigerSwiGLUMLP + # Instantiate a dummy model config = transformers.models.llama4.configuration_llama4.Llama4TextConfig( dtype=torch.bfloat16, @@ -1391,6 +1402,10 @@ def test_apply_liger_kernel_to_instance_for_llama4_for_conditional_generation(): with patch("transformers.models.llama4.modeling_llama4"): from transformers.models.llama4.modeling_llama4 import Llama4ForConditionalGeneration + # Compare against the class actually bound by the patch: llama4 keeps the + # original swiglu implementation because its MLPs receive 2D input. + from liger_kernel.transformers.swiglu import LigerSwiGLUMLP + # Instantiate a dummy model config = transformers.models.llama4.configuration_llama4.Llama4Config( dtype=torch.bfloat16, @@ -3194,8 +3209,9 @@ def test_apply_liger_kernel_to_instance_for_qwen3_next(): for expert in layer.mlp.experts: assert inspect.getsource(expert.forward) != inspect.getsource(LigerQwen3MoeSwiGLUMLP.forward) if hasattr(layer.mlp, "shared_expert"): + # Compare against the class actually bound by the patch assert inspect.getsource(layer.mlp.shared_expert.forward) != inspect.getsource( - LigerSwiGLUMLP.forward + LigerQwen3MoeSwiGLUMLP.forward ) else: assert inspect.getsource(layer.mlp.forward) != inspect.getsource(LigerQwen3MoeSwiGLUMLP.forward) @@ -3217,8 +3233,9 @@ def test_apply_liger_kernel_to_instance_for_qwen3_next(): for expert in layer.mlp.experts: assert inspect.getsource(expert.forward) == inspect.getsource(LigerQwen3MoeSwiGLUMLP.forward) if hasattr(layer.mlp, "shared_expert"): + # Compare against the class actually bound by the patch assert inspect.getsource(layer.mlp.shared_expert.forward) == inspect.getsource( - LigerSwiGLUMLP.forward + LigerQwen3MoeSwiGLUMLP.forward ) else: assert inspect.getsource(layer.mlp.forward) == inspect.getsource(LigerQwen3MoeSwiGLUMLP.forward) @@ -3643,6 +3660,7 @@ def test_apply_liger_kernel_to_instance_for_hunyuan_v1_dense(): # Ensure any monkey patching is cleaned up for subsequent tests with patch("transformers.models.hunyuan_v1_dense.modeling_hunyuan_v1_dense"): from liger_kernel.transformers.model.hunyuan_v1 import lce_forward as hunyuan_v1_dense_lce_forward + from liger_kernel.transformers.swiglu import LigerHunyuanV1SwiGLUMLP # Instantiate a dummy model config = transformers.models.hunyuan_v1_dense.configuration_hunyuan_v1_dense.HunYuanDenseV1Config( @@ -3660,7 +3678,8 @@ def test_apply_liger_kernel_to_instance_for_hunyuan_v1_dense(): assert inspect.getsource(dummy_model_instance.forward) != inspect.getsource(hunyuan_v1_dense_lce_forward) assert inspect.getsource(dummy_model_instance.model.norm.forward) != inspect.getsource(LigerRMSNorm.forward) for layer in dummy_model_instance.model.layers: - assert inspect.getsource(layer.mlp.forward) != inspect.getsource(LigerSwiGLUMLP.forward) + # Compare against the class actually bound by the patch + assert inspect.getsource(layer.mlp.forward) != inspect.getsource(LigerHunyuanV1SwiGLUMLP.forward) assert inspect.getsource(layer.input_layernorm.forward) != inspect.getsource(LigerRMSNorm.forward) assert inspect.getsource(layer.post_attention_layernorm.forward) != inspect.getsource(LigerRMSNorm.forward) @@ -3671,7 +3690,8 @@ def test_apply_liger_kernel_to_instance_for_hunyuan_v1_dense(): assert inspect.getsource(dummy_model_instance.forward) == inspect.getsource(hunyuan_v1_dense_lce_forward) assert inspect.getsource(dummy_model_instance.model.norm.forward) == inspect.getsource(LigerRMSNorm.forward) for layer in dummy_model_instance.model.layers: - assert inspect.getsource(layer.mlp.forward) == inspect.getsource(LigerSwiGLUMLP.forward) + # Compare against the class actually bound by the patch + assert inspect.getsource(layer.mlp.forward) == inspect.getsource(LigerHunyuanV1SwiGLUMLP.forward) assert inspect.getsource(layer.input_layernorm.forward) == inspect.getsource(LigerRMSNorm.forward) assert inspect.getsource(layer.post_attention_layernorm.forward) == inspect.getsource(LigerRMSNorm.forward) From bb43631154518a0b68217bf08d7e6af997a3ebe5 Mon Sep 17 00:00:00 2001 From: ma-jh <3364870135@qq.com> Date: Sat, 15 Aug 2026 10:30:11 +0800 Subject: [PATCH 5/6] Revert monkey patch changes, keep only core kernel modifications Signed-off-by: ma-jh <3364870135@qq.com> --- src/liger_kernel/transformers/__init__.py | 2 ++ src/liger_kernel/transformers/monkey_patch.py | 19 ++--------- test/transformers/test_monkey_patch.py | 32 ++++--------------- 3 files changed, 11 insertions(+), 42 deletions(-) diff --git a/src/liger_kernel/transformers/__init__.py b/src/liger_kernel/transformers/__init__.py index 27f48a776..0bad8baad 100644 --- a/src/liger_kernel/transformers/__init__.py +++ b/src/liger_kernel/transformers/__init__.py @@ -15,6 +15,7 @@ from liger_kernel.transformers.llama4_rope import liger_llama4_text_rotary_pos_emb # noqa: F401 from liger_kernel.transformers.llama4_rope import liger_llama4_vision_rotary_pos_emb # noqa: F401 from liger_kernel.transformers.mhc import LigerMHC # noqa: F401 +from liger_kernel.transformers.mlp import LigerFalconH1MLP # noqa: F401 from liger_kernel.transformers.mlp import LigerMLP # noqa: F401 from liger_kernel.transformers.modulated_rms_norm import LigerModulatedRMSNorm # noqa: F401 from liger_kernel.transformers.multi_token_attention import LigerMultiTokenAttention # noqa: F401 @@ -202,6 +203,7 @@ def __getattr__(name: str): "LigerSoftmax", "LigerSparsemax", "LigerMLP", + "LigerFalconH1MLP", ] # Add transformer-dependent symbols only if available diff --git a/src/liger_kernel/transformers/monkey_patch.py b/src/liger_kernel/transformers/monkey_patch.py index 5f8b27fd4..83de3a076 100755 --- a/src/liger_kernel/transformers/monkey_patch.py +++ b/src/liger_kernel/transformers/monkey_patch.py @@ -37,12 +37,7 @@ from liger_kernel.transformers.swiglu import LigerBlockSparseTop2MLP from liger_kernel.transformers.swiglu import LigerExperts from liger_kernel.transformers.swiglu import LigerPhi3SwiGLUMLP - -USE_FLASH_SWIGLU = True -if USE_FLASH_SWIGLU: - from liger_kernel.transformers.mlp import LigerMLP as LigerSwiGLUMLP -else: - from liger_kernel.transformers.swiglu import LigerSwiGLUMLP +from liger_kernel.transformers.swiglu import LigerSwiGLUMLP try: import peft @@ -470,11 +465,6 @@ def apply_liger_kernel_to_llama4( from liger_kernel.transformers.model.llama4 import lce_forward as llama4_lce_forward - # NOTE: Llama4's text MLP receives a 2D [T, hidden] input, which the fused LigerMLP - # kernels do not support, so we use the original LigerSwiGLUMLP implementation - # here instead of LigerMLP. - from liger_kernel.transformers.swiglu import LigerSwiGLUMLP - if rope: from liger_kernel.transformers.llama4_rope import apply_liger_llama4_rope_full @@ -2921,10 +2911,7 @@ def apply_liger_kernel_to_falcon_h1( from transformers.models.falcon_h1 import modeling_falcon_h1 from transformers.models.falcon_h1.modeling_falcon_h1 import FalconH1Model - if USE_FLASH_SWIGLU: - from liger_kernel.transformers.mlp import LigerFalconH1MLP as LigerFalconH1SwiGLUMLP - else: - from liger_kernel.transformers.swiglu import LigerFalconH1SwiGLUMLP + from liger_kernel.transformers.swiglu import LigerFalconH1SwiGLUMLP if rope: logger.info("Apply liger rotary pos emb.") @@ -3660,4 +3647,4 @@ def _apply_liger_kernel_to_instance(model: PreTrainedModel, **kwargs) -> None: f"Applying Liger kernels to model instance with model type: {model_type} with kwargs: {applicable_kwargs}" ) - apply_fn(model=model, **applicable_kwargs) + apply_fn(model=model, **applicable_kwargs) \ No newline at end of file diff --git a/test/transformers/test_monkey_patch.py b/test/transformers/test_monkey_patch.py index 657386980..09c80e823 100755 --- a/test/transformers/test_monkey_patch.py +++ b/test/transformers/test_monkey_patch.py @@ -22,8 +22,7 @@ from liger_kernel.transformers import LigerPhi3SwiGLUMLP from liger_kernel.transformers import LigerQwen3MoeSwiGLUMLP from liger_kernel.transformers import LigerRMSNorm - -# from liger_kernel.transformers import LigerSwiGLUMLP +from liger_kernel.transformers import LigerSwiGLUMLP from liger_kernel.transformers import monkey_patch from liger_kernel.transformers.layer_norm import LigerLayerNorm from liger_kernel.transformers.model.falcon_h1 import lce_forward as falcon_h1_lce_forward @@ -42,15 +41,9 @@ from liger_kernel.transformers.model.qwen3_next import lce_forward as qwen3_next_lce_forward from liger_kernel.transformers.model.smollm3 import lce_forward as smolllm3_lce_forward from liger_kernel.transformers.monkey_patch import MODEL_TYPE_TO_APPLY_LIGER_FN -from liger_kernel.transformers.monkey_patch import USE_FLASH_SWIGLU from liger_kernel.transformers.monkey_patch import _apply_liger_kernel from liger_kernel.transformers.monkey_patch import _apply_liger_kernel_to_instance -if USE_FLASH_SWIGLU: - from liger_kernel.transformers.mlp import LigerMLP as LigerSwiGLUMLP -else: - from liger_kernel.transformers.swiglu import LigerSwiGLUMLP - # We only support transformers >= 4.52.0 transformer_version = version.parse(transformers.__version__) MIN_SUPPORTED_TRANSFORMERS_VERSION = version.parse("4.52.0") @@ -1347,10 +1340,6 @@ def test_apply_liger_kernel_to_instance_for_llama4_for_causal_lm(): with patch("transformers.models.llama4.modeling_llama4"): from transformers.models.llama4.modeling_llama4 import Llama4ForCausalLM - # Compare against the class actually bound by the patch: llama4 keeps the - # original swiglu implementation because its MLPs receive 2D input. - from liger_kernel.transformers.swiglu import LigerSwiGLUMLP - # Instantiate a dummy model config = transformers.models.llama4.configuration_llama4.Llama4TextConfig( dtype=torch.bfloat16, @@ -1402,10 +1391,6 @@ def test_apply_liger_kernel_to_instance_for_llama4_for_conditional_generation(): with patch("transformers.models.llama4.modeling_llama4"): from transformers.models.llama4.modeling_llama4 import Llama4ForConditionalGeneration - # Compare against the class actually bound by the patch: llama4 keeps the - # original swiglu implementation because its MLPs receive 2D input. - from liger_kernel.transformers.swiglu import LigerSwiGLUMLP - # Instantiate a dummy model config = transformers.models.llama4.configuration_llama4.Llama4Config( dtype=torch.bfloat16, @@ -3209,9 +3194,8 @@ def test_apply_liger_kernel_to_instance_for_qwen3_next(): for expert in layer.mlp.experts: assert inspect.getsource(expert.forward) != inspect.getsource(LigerQwen3MoeSwiGLUMLP.forward) if hasattr(layer.mlp, "shared_expert"): - # Compare against the class actually bound by the patch assert inspect.getsource(layer.mlp.shared_expert.forward) != inspect.getsource( - LigerQwen3MoeSwiGLUMLP.forward + LigerSwiGLUMLP.forward ) else: assert inspect.getsource(layer.mlp.forward) != inspect.getsource(LigerQwen3MoeSwiGLUMLP.forward) @@ -3233,9 +3217,8 @@ def test_apply_liger_kernel_to_instance_for_qwen3_next(): for expert in layer.mlp.experts: assert inspect.getsource(expert.forward) == inspect.getsource(LigerQwen3MoeSwiGLUMLP.forward) if hasattr(layer.mlp, "shared_expert"): - # Compare against the class actually bound by the patch assert inspect.getsource(layer.mlp.shared_expert.forward) == inspect.getsource( - LigerQwen3MoeSwiGLUMLP.forward + LigerSwiGLUMLP.forward ) else: assert inspect.getsource(layer.mlp.forward) == inspect.getsource(LigerQwen3MoeSwiGLUMLP.forward) @@ -3660,7 +3643,6 @@ def test_apply_liger_kernel_to_instance_for_hunyuan_v1_dense(): # Ensure any monkey patching is cleaned up for subsequent tests with patch("transformers.models.hunyuan_v1_dense.modeling_hunyuan_v1_dense"): from liger_kernel.transformers.model.hunyuan_v1 import lce_forward as hunyuan_v1_dense_lce_forward - from liger_kernel.transformers.swiglu import LigerHunyuanV1SwiGLUMLP # Instantiate a dummy model config = transformers.models.hunyuan_v1_dense.configuration_hunyuan_v1_dense.HunYuanDenseV1Config( @@ -3678,8 +3660,7 @@ def test_apply_liger_kernel_to_instance_for_hunyuan_v1_dense(): assert inspect.getsource(dummy_model_instance.forward) != inspect.getsource(hunyuan_v1_dense_lce_forward) assert inspect.getsource(dummy_model_instance.model.norm.forward) != inspect.getsource(LigerRMSNorm.forward) for layer in dummy_model_instance.model.layers: - # Compare against the class actually bound by the patch - assert inspect.getsource(layer.mlp.forward) != inspect.getsource(LigerHunyuanV1SwiGLUMLP.forward) + assert inspect.getsource(layer.mlp.forward) != inspect.getsource(LigerSwiGLUMLP.forward) assert inspect.getsource(layer.input_layernorm.forward) != inspect.getsource(LigerRMSNorm.forward) assert inspect.getsource(layer.post_attention_layernorm.forward) != inspect.getsource(LigerRMSNorm.forward) @@ -3690,8 +3671,7 @@ def test_apply_liger_kernel_to_instance_for_hunyuan_v1_dense(): assert inspect.getsource(dummy_model_instance.forward) == inspect.getsource(hunyuan_v1_dense_lce_forward) assert inspect.getsource(dummy_model_instance.model.norm.forward) == inspect.getsource(LigerRMSNorm.forward) for layer in dummy_model_instance.model.layers: - # Compare against the class actually bound by the patch - assert inspect.getsource(layer.mlp.forward) == inspect.getsource(LigerHunyuanV1SwiGLUMLP.forward) + assert inspect.getsource(layer.mlp.forward) == inspect.getsource(LigerSwiGLUMLP.forward) assert inspect.getsource(layer.input_layernorm.forward) == inspect.getsource(LigerRMSNorm.forward) assert inspect.getsource(layer.post_attention_layernorm.forward) == inspect.getsource(LigerRMSNorm.forward) @@ -3738,4 +3718,4 @@ def test_apply_liger_kernel_to_instance_for_nemotron(): try: print(dummy_model_instance) except Exception as e: - pytest.fail(f"An exception occured in extra_expr: {type(e).__name__} - {e}") + pytest.fail(f"An exception occured in extra_expr: {type(e).__name__} - {e}") \ No newline at end of file From 85d07e74ae3735ec486f524b333eb6e4c0068b8a Mon Sep 17 00:00:00 2001 From: ma-jh <3364870135@qq.com> Date: Sat, 15 Aug 2026 10:32:26 +0800 Subject: [PATCH 6/6] Revert monkey patch changes Signed-off-by: ma-jh <3364870135@qq.com> --- src/liger_kernel/transformers/monkey_patch.py | 2 +- test/transformers/test_monkey_patch.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/liger_kernel/transformers/monkey_patch.py b/src/liger_kernel/transformers/monkey_patch.py index 83de3a076..4d33d7e41 100755 --- a/src/liger_kernel/transformers/monkey_patch.py +++ b/src/liger_kernel/transformers/monkey_patch.py @@ -3647,4 +3647,4 @@ def _apply_liger_kernel_to_instance(model: PreTrainedModel, **kwargs) -> None: f"Applying Liger kernels to model instance with model type: {model_type} with kwargs: {applicable_kwargs}" ) - apply_fn(model=model, **applicable_kwargs) \ No newline at end of file + apply_fn(model=model, **applicable_kwargs) diff --git a/test/transformers/test_monkey_patch.py b/test/transformers/test_monkey_patch.py index 09c80e823..25099f7be 100755 --- a/test/transformers/test_monkey_patch.py +++ b/test/transformers/test_monkey_patch.py @@ -3718,4 +3718,4 @@ def test_apply_liger_kernel_to_instance_for_nemotron(): try: print(dummy_model_instance) except Exception as e: - pytest.fail(f"An exception occured in extra_expr: {type(e).__name__} - {e}") \ No newline at end of file + pytest.fail(f"An exception occured in extra_expr: {type(e).__name__} - {e}")