diff --git a/modeling/transformers/scripts/benchmark_hf_model.sh b/modeling/transformers/scripts/benchmark_hf_model.sh index b8e310e3..da0ba165 100755 --- a/modeling/transformers/scripts/benchmark_hf_model.sh +++ b/modeling/transformers/scripts/benchmark_hf_model.sh @@ -22,7 +22,7 @@ usage() { cat < torch.Tensor: + """Depthwise causal conv1d for the prefill path (drop-in for ``causal_conv1d_fn``). + + Args: + x: ``(B=1, D, L)`` unpadded input (``Bx`` in the LFM2 short-conv block). + weight: ``(D, K=3)`` depthwise conv weights (``conv.weight.view(D, K)``). + bias: optional ``(D,)`` bias (LFM2-8B-A1B uses ``conv_bias=False`` -> None). + activation: accepted for signature compatibility; must be ``None`` (LFM2 + does not fuse an activation into the conv). + seq_idx: accepted for signature compatibility. Packed-sequence boundaries + are not supported by this fused path; only ``None`` is handled. + + Returns: + ``(B=1, D, L)`` conv output. + """ + assert activation is None, "LFM2 short-conv fuses no activation; activation must be None" + assert seq_idx is None, "lfm2_causal_conv1d_fn_cutile does not support packed sequences (seq_idx)" + + B, D, L = x.shape + assert B == 1, "lfm2_causal_conv1d_fn_cutile only supports B=1" + K = weight.shape[1] + assert K == 3, f"expected kernel_size 3, got {K}" + + x_2d = x.squeeze(0).contiguous() # (D, L) + x_padded = F.pad(x_2d, (K - 1, 0)) # (D, L + K - 1), left pad only + w = weight.contiguous() + output = torch.empty(D, L, dtype=x.dtype, device=x.device) + + # NOTE: the sequence length is deliberately *not* passed as a ct.Constant -- + # it would become part of the JIT specialization key and force a fresh cuTile + # compile for every distinct prompt length (a long-prompt stall). Bounds are + # handled by `check_bounds` on the gathers/scatter, so one compiled kernel + # serves all lengths. + BLOCK_T = 256 + grid = (D, (L + BLOCK_T - 1) // BLOCK_T) + ct.launch( + torch.cuda.current_stream(), + grid, + _causal_conv1d_prefill_kernel, + (x_padded, w, output, BLOCK_T), + ) + + out = output.unsqueeze(0) # (1, D, L) + if bias is not None: + out = out + bias.view(1, -1, 1) + return out diff --git a/src/tilegym/transformers/lfm2_moe/kernels/causal_conv1d_update.py b/src/tilegym/transformers/lfm2_moe/kernels/causal_conv1d_update.py new file mode 100644 index 00000000..9d663b81 --- /dev/null +++ b/src/tilegym/transformers/lfm2_moe/kernels/causal_conv1d_update.py @@ -0,0 +1,125 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: MIT + +"""LFM2-MoE depthwise causal conv1d decode-update cuTile kernel. + +Replacement for the module-level ``causal_conv1d_update`` called by +``Lfm2MoeShortConv.forward`` on the single-token cached-decode path. + +LFM2 stores a full ``K``-wide (``conv_L_cache = 3``) rolling window in the +conv-state cache. For a new input ``x`` and state ``[s0, s1, s2]`` the update +rolls the window (dropping the oldest ``s0``) and computes the conv over the +new window ``[s1, s2, x]``: + + out = s1 * w0 + s2 * w1 + x * w2 + +and writes the rolled window back into ``conv_state`` **in place**. No +activation is fused (LFM2 gates externally). This mirrors the pure-torch +``Lfm2MoeShortConv.slow_forward`` decode branch +(``sum(update_conv_state(x) * weight, dim=-1)``). +""" + +import cuda.tile as ct +import torch + +ConstInt = ct.Constant[int] + + +@ct.kernel +def _causal_conv1d_update_kernel( + x, # (D,) + conv_state, # (D, K=3), updated in place + weight, # (D, K=3) + output, # (D,) + BLOCK_D: ConstInt, +): + bid = ct.bid(0) + d_start = bid * BLOCK_D + offs = ct.arange(BLOCK_D, dtype=ct.int32) + d_idx = d_start + offs + + s0 = ct.astype(ct.gather(conv_state, (d_idx, 0), check_bounds=True), ct.float32) + s1 = ct.astype(ct.gather(conv_state, (d_idx, 1), check_bounds=True), ct.float32) + s2 = ct.astype(ct.gather(conv_state, (d_idx, 2), check_bounds=True), ct.float32) + xv = ct.astype(ct.gather(x, (d_idx,), check_bounds=True), ct.float32) + + w0 = ct.astype(ct.gather(weight, (d_idx, 0), check_bounds=True), ct.float32) + w1 = ct.astype(ct.gather(weight, (d_idx, 1), check_bounds=True), ct.float32) + w2 = ct.astype(ct.gather(weight, (d_idx, 2), check_bounds=True), ct.float32) + + # Roll the window (drop s0) then convolve over [s1, s2, x]. + result = s1 * w0 + s2 * w1 + xv * w2 + + ct.scatter(output, (d_idx,), ct.astype(result, output.dtype), check_bounds=True) + # Shift state in place: [s0, s1, s2] -> [s1, s2, x] + ct.scatter(conv_state, (d_idx, 0), ct.astype(s1, conv_state.dtype), check_bounds=True) + ct.scatter(conv_state, (d_idx, 1), ct.astype(s2, conv_state.dtype), check_bounds=True) + ct.scatter(conv_state, (d_idx, 2), ct.astype(xv, conv_state.dtype), check_bounds=True) + + +def lfm2_causal_conv1d_update_cutile( + x: torch.Tensor, + conv_state: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None = None, + activation=None, +) -> torch.Tensor: + """Depthwise causal conv1d decode-update (drop-in for ``causal_conv1d_update``). + + The call contract differs across transformers releases: + + - transformers >= 5.13-style ``Lfm2MoeShortConv.forward`` passes the *3D* + ``hidden_states`` ``(B, D, 1)`` and multiplies the result with the 3D + gate ``C`` (``y = C * hidden_states``) -- the output must stay 3D. + - Older releases (e.g. 5.10.x, ``cuda_kernels_forward``) pass + ``Bx.squeeze(-1)`` -- a *2D* ``(B, D)`` input -- and unsqueeze the + result themselves. + + Both conventions are accepted; the returned tensor keeps the input's rank. + + Args: + x: ``(B=1, D, L=1)`` or ``(B=1, D)`` current-timestep input (``B * x`` + after the in-proj chunk). + conv_state: ``(B=1, D, K=3)`` rolling window cache, updated **in place**. + weight: ``(D, K=3)`` depthwise conv weights. + bias: optional ``(D,)`` bias (LFM2-8B-A1B uses ``conv_bias=False`` -> None). + activation: accepted for signature compatibility; must be ``None``. + + Returns: + Conv output for the current timestep, same rank as ``x`` + (``(B=1, D, 1)`` or ``(B=1, D)``). + """ + assert activation is None, "LFM2 short-conv fuses no activation; activation must be None" + + keep_3d = x.dim() == 3 + if keep_3d: + B, D, L = x.shape + assert L == 1, f"lfm2_causal_conv1d_update_cutile is the single-token decode path, got L={L}" + else: + assert x.dim() == 2, f"expected x of rank 2 or 3, got shape {tuple(x.shape)}" + B, D = x.shape + assert B == 1, "lfm2_causal_conv1d_update_cutile only supports B=1" + K = weight.shape[1] + assert K == 3, f"expected kernel_size 3, got {K}" + + x_1d = x.reshape(D).contiguous() # (D,) + cs = conv_state.squeeze(0) # (D, 3) view -> mutated in place + w = weight.contiguous() + output = torch.empty(D, dtype=x.dtype, device=x.device) + + BLOCK_D = 256 + grid = ((D + BLOCK_D - 1) // BLOCK_D,) + ct.launch( + torch.cuda.current_stream(), + grid, + _causal_conv1d_update_kernel, + (x_1d, cs, w, output, BLOCK_D), + ) + + # Match the input rank (5.13+ passes 3D and gates directly; <= 5.10 passes + # 2D and unsqueezes at the call site). + out = output.view(1, D, 1) if keep_3d else output.unsqueeze(0) + if bias is not None: + out = out + (bias.view(1, -1, 1) if keep_3d else bias.view(1, -1)) + return out diff --git a/src/tilegym/transformers/lfm2_moe/modeling_lfm2_moe.py b/src/tilegym/transformers/lfm2_moe/modeling_lfm2_moe.py new file mode 100644 index 00000000..3e65b33e --- /dev/null +++ b/src/tilegym/transformers/lfm2_moe/modeling_lfm2_moe.py @@ -0,0 +1,173 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: MIT + +"""TileGym replacement modules for `transformers.models.lfm2_moe.modeling_lfm2_moe`. + +The MoE block and the dense SwiGLU MLP are replaced here; RoPE / RMSNorm / +attention are patched elsewhere (registry-level / class-level) by +`apply_tilegym_kernel_to_lfm2_moe`, and the hybrid short-convolution operator +layers are routed through the fused cuTile kernels in `kernels/` when +`use_cutile=True`. + +`Lfm2MoeSparseMoeBlockTileGym` keeps the exact same nested-parameter layout as +the stock `Lfm2MoeSparseMoeBlock` (`self.experts = Lfm2MoeExperts(...)`, +`self.gate = Lfm2MoeTopKRouter(...)`, and the optional `self.expert_bias` +buffer) so HuggingFace `state_dict` loading works unchanged. Forward replaces +the per-expert Python loop in `Lfm2MoeExperts` with TileGym's batched +`fused_moe` kernel. + +Weight-layout compatibility notes (verified against HF LFM2-MoE source): + +- HF `self.experts.gate_up_proj`: shape ``(E, 2*I, H)``. The first ``I`` rows + along axis 1 are the **gate** projection, the second ``I`` rows are the + **up** projection — confirmed by HF's + ``linear(x, gate_up_proj[e]).chunk(2, dim=-1)`` which produces + ``(gate, up)`` in that order. +- HF `self.experts.down_proj`: shape ``(E, H, I)``. +- TileGym `fused_moe(w1, w2)` expects ``w1: (E, 2*I, H)`` with the standard + ``silu_and_mul`` ordering ``silu(x[:, :I]) * x[:, I:]`` (i.e. ``[gate, up]``) + and ``w2: (E, H, I)`` — identical to HF, so the parameters are passed through + with **no merge / no reorder**. + +Routing semantics (differ from OLMoE — reproduced inline from +`Lfm2MoeTopKRouter.forward`): + +- LFM2-MoE routes with a **sigmoid** (not softmax): ``sigmoid(logits)``. +- When ``use_expert_bias`` is set, a per-expert bias is added *only to select* + the top-k experts; the gathered weights are the un-biased sigmoid values + (DeepSeek-V3 style). The bias lives in the ``expert_bias`` buffer. +- ``norm_topk_prob`` divides the top-k weights by ``(sum + 1e-6)``. +- ``routed_scaling_factor`` multiplies the weights. + +All of the above is applied in the wrapper before calling ``fused_moe`` — the +kernel has no norm/scaling/bias arguments and multiplies the (already final) +routing weights into the down-projection output un-normalized. +""" + +import torch +import torch.nn.functional as F +from torch import nn + +from tilegym.ops import fused_moe +from tilegym.ops import silu_and_mul + + +class Lfm2MoeMLPTileGym(nn.Module): + """Drop-in replacement for the dense ``Lfm2MoeMLP`` (used in the first + ``num_dense_layers`` layers) that fuses the SiLU-and-mul activation via + TileGym's ``silu_and_mul`` kernel. + + LFM2-MoE names the SwiGLU projections ``w1`` (gate), ``w3`` (up) and + ``w2`` (down) — not the usual ``gate_proj``/``up_proj``/``down_proj`` — so + the generic ``get_swiglu_module`` helpers are not state_dict compatible. + This class keeps those exact ``nn.Linear`` attribute names so the + HuggingFace ``state_dict`` loads with ``strict=True``. Follows the + ``Phi3MLPTileGym`` precedent (nn.Linear projections + fused activation). + """ + + def __init__(self, config, intermediate_size: int | None = None): + super().__init__() + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size if intermediate_size is None else intermediate_size + self.w1 = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) # gate + self.w3 = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) # up + self.w2 = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) # down + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gate = self.w1(x) + up = self.w3(x) + # silu_and_mul(cat([gate, up])) == silu(gate) * up, matching the stock + # forward ``w2(F.silu(w1(x)) * w3(x))``. + return self.w2(silu_and_mul(torch.cat([gate, up], dim=-1))) + + +class Lfm2MoeSparseMoeBlockTileGym(nn.Module): + """Drop-in replacement for ``Lfm2MoeSparseMoeBlock`` that routes the expert + compute through TileGym's batched ``fused_moe`` kernel. + + The nested submodule layout (``self.experts``, ``self.gate``) and the + optional ``expert_bias`` buffer are kept identical to the stock class so the + HuggingFace state_dict loads with ``strict=True``. + """ + + def __init__(self, config): + super().__init__() + # Import here so the module import is cheap and doesn't run HF init + # at TileGym import time. + from transformers.models.lfm2_moe.modeling_lfm2_moe import Lfm2MoeExperts + + self.experts = Lfm2MoeExperts(config) + try: + # transformers >= 5.13-style: dedicated router module. + from transformers.models.lfm2_moe.modeling_lfm2_moe import Lfm2MoeTopKRouter + + self.gate = Lfm2MoeTopKRouter(config) + except ImportError: + # Older releases (e.g. 5.10.x) inline the router as a plain + # `nn.Linear(hidden_size, num_experts)` named `gate`. Both layouts + # expose the same state_dict key `gate.weight` of shape (E, H), + # and `_route` only needs `.weight`, so they are interchangeable. + self.gate = nn.Linear(config.hidden_size, config.num_experts, bias=False) + self.use_expert_bias = config.use_expert_bias + if self.use_expert_bias: + # Match the stock buffer exactly (name / dtype / shape) so strict + # state_dict loading succeeds. + self.register_buffer("expert_bias", torch.zeros(config.num_experts, dtype=torch.float32)) + + # Cache router metadata for convenience. + self.top_k = config.num_experts_per_tok + self.num_experts = config.num_experts + self.norm_topk_prob = config.norm_topk_prob + self.routed_scaling_factor = config.routed_scaling_factor + self.hidden_size = config.hidden_size + + def _route(self, hidden_flat: torch.Tensor): + """Reproduce ``Lfm2MoeTopKRouter.forward`` inline. + + Returns ``(topk_weights, topk_indices)`` where: + - ``topk_weights`` is the final (bias-selected, normalized, scaled) + routing weight, cast back to ``hidden_flat.dtype``. + - ``topk_indices`` is ``torch.long`` (output of ``torch.topk``). + """ + # gate.weight is (num_experts, hidden_size); F.linear handles the + # transpose and the matmul is tiny, so no cuTile matmul is needed. + router_logits = F.linear(hidden_flat, self.gate.weight) + routing_weights = router_logits.sigmoid() + + if self.use_expert_bias: + # Bias is used only to *select* the experts; the returned weights + # are the un-biased sigmoid values gathered at the selected indices. + scores_for_routing = routing_weights + self.expert_bias + _, topk_indices = torch.topk(scores_for_routing, self.top_k, dim=-1) + topk_weights = torch.gather(routing_weights, dim=1, index=topk_indices).type_as(router_logits) + else: + topk_weights, topk_indices = torch.topk(routing_weights, self.top_k, dim=-1) + + if self.norm_topk_prob: + topk_weights = topk_weights / (topk_weights.sum(dim=-1, keepdim=True) + 1e-6) + topk_weights = topk_weights * self.routed_scaling_factor + + return topk_weights.to(hidden_flat.dtype), topk_indices + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + batch_size, sequence_length, hidden_dim = hidden_states.shape + hidden_flat = hidden_states.reshape(-1, hidden_dim).contiguous() + + topk_weights, topk_indices = self._route(hidden_flat) + + # TileGym's fused_moe expects (M, H) input, (E, 2I, H) w1, (E, H, I) w2. + # ``topk_indices`` from torch.topk is int64; cast to int32 for the + # kernel which uses 32-bit indices internally. + out_flat = fused_moe( + hidden_flat, + w1=self.experts.gate_up_proj, + w2=self.experts.down_proj, + topk_weights=topk_weights, + topk_ids=topk_indices.to(torch.int32), + ) + + # Match the dtype contract of the stock block, which returns a single + # tensor (not a tuple). + out_flat = out_flat.to(hidden_states.dtype) + return out_flat.view(batch_size, sequence_length, hidden_dim) diff --git a/src/tilegym/transformers/monkey_patch.py b/src/tilegym/transformers/monkey_patch.py index 172b43eb..ccc90f9d 100644 --- a/src/tilegym/transformers/monkey_patch.py +++ b/src/tilegym/transformers/monkey_patch.py @@ -15,6 +15,7 @@ from tilegym.ops import get_swiglu_module from tilegym.transformers.deepseek2.modeling_deepseek import DeepseekV2MoETileGym from tilegym.transformers.deepseek2.modeling_deepseek import tilegym_deepseek_v2_forward +from tilegym.transformers.lfm2_moe.modeling_lfm2_moe import Lfm2MoeSparseMoeBlockTileGym from tilegym.transformers.phi3.modeling_phi3 import Phi3MLPTileGym from tilegym.transformers.phi3.modeling_phi3 import get_fmha_phi3_interface @@ -552,6 +553,104 @@ def apply_tilegym_kernel_to_olmo3( logger.info("Patched Olmo3DecoderLayer.forward with fused residual_add+RMSNorm") +def apply_tilegym_kernel_to_lfm2_moe( + rope: bool = True, + rms_norm: bool = True, + attn: bool = True, + moe: bool = True, + swiglu: bool = True, + conv: bool = True, + model: PreTrainedModel = None, + use_cutile: bool = False, +) -> None: + """ + Apply TileGym kernels to replace original implementation in HuggingFace LFM2-MoE models + (e.g. LiquidAI/LFM2-8B-A1B). + + LFM2-MoE is a hybrid Mixture-of-Experts model. For LFM2-8B-A1B (24 layers): the + token mixer is a short convolution in 18 layers and grouped-query attention in 6 + layers; the channel mixer is a dense MLP in the first `num_dense_layers` (2) layers + and a sparse MoE block (32 experts, top-4) in the remaining 22 layers. + Per-component compute this patch replaces: + - Llama-style RMSNorm (operator/ffn/embedding norms + per-head Q/K norms) + - Standard RoPE (full rotation, rope_theta=1e6) + - GQA FMHA for the `full_attention` layers (32 Q heads, 8 KV heads, head_dim=64) + - Lfm2MoeSparseMoeBlock with Lfm2MoeTopKRouter (sigmoid routing + expert-bias + selection, norm_topk_prob, routed_scaling_factor) and Lfm2MoeExperts whose + gate_up_proj (E, 2I, H) and down_proj (E, H, I) are already stacked in the format + TileGym's fused_moe expects. + - Dense SwiGLU MLP (`Lfm2MoeMLP`, w1/w3/w2) with a fused silu-and-mul path. + - Short-convolution operator layers via fused cuTile depthwise causal conv1d + kernels (prefill + decode-update, kernel_size=3, no activation). Requires + `use_cutile=True` since these are cuTile kernels. + + Args: + rope (bool): Patch `apply_rotary_pos_emb`. Default True. + rms_norm (bool): Patch `Lfm2MoeRMSNorm`. Default True. + attn (bool): Patch `ALL_ATTENTION_FUNCTIONS["sdpa"]` with FMHA. Default True. + moe (bool): Patch `Lfm2MoeSparseMoeBlock` with the TileGym fused-MoE variant. Default True. + swiglu (bool): Patch the dense `Lfm2MoeMLP` with a fused silu-and-mul MLP. Default True. + conv (bool): Patch the short-conv path with fused cuTile causal conv1d kernels. + Only takes effect when `use_cutile=True` (cuTile kernels). Default True. + model (PreTrainedModel): Unused; present for API symmetry with sibling helpers. + use_cutile (bool): Switch the TileGym backend to cuTile. Default False. + """ + import transformers + from packaging import version + + # Degrade gracefully on transformers installs that predate + # `transformers.models.lfm2_moe` rather than crashing the dispatcher. + if version.parse(transformers.__version__) < version.parse("4.57.0"): + logger.warning("LFM2-MoE support requires a transformers release that ships `models.lfm2_moe`") + return + + logger.info("--------------------------------") + logger.info("apply_tilegym_kernel_to_lfm2_moe") + logger.info("--------------------------------") + + try: + from transformers.models.lfm2_moe import modeling_lfm2_moe + except ImportError: + logger.warning("transformers.models.lfm2_moe is not available in this transformers install") + return + + if use_cutile: + set_backend("cutile") + + if rope: + # `apply_rotary_pos_emb` is a module-level global resolved at call time in + # Lfm2MoeAttention.forward, so reassigning it takes effect like the other + # models. (A user-invoked kernels-hub `kernelize()` exchange could shadow + # this, but that is not the default generate path.) + modeling_lfm2_moe.apply_rotary_pos_emb = get_apply_rope_func(model="llama") + if rms_norm: + modeling_lfm2_moe.Lfm2MoeRMSNorm = get_rms_norm_module() + if attn: + from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS + + ALL_ATTENTION_FUNCTIONS["sdpa"] = get_fmha_interface() + if moe: + modeling_lfm2_moe.Lfm2MoeSparseMoeBlock = Lfm2MoeSparseMoeBlockTileGym + if swiglu: + from tilegym.transformers.lfm2_moe.modeling_lfm2_moe import Lfm2MoeMLPTileGym + + modeling_lfm2_moe.Lfm2MoeMLP = Lfm2MoeMLPTileGym + + if use_cutile and conv: + # Fused cuTile depthwise causal conv1d for the short-conv operator layers. + # `causal_conv1d_fn` / `causal_conv1d_update` are read as module-level + # globals in Lfm2MoeShortConv.forward at call time, so patching the + # globals routes both the prefill and the single-token decode-update + # through the TileGym kernels. `is_fast_path_available` only exists on + # some transformers versions; flipping it on is harmless where absent. + from tilegym.transformers.lfm2_moe.kernels.causal_conv1d_prefill import lfm2_causal_conv1d_fn_cutile + from tilegym.transformers.lfm2_moe.kernels.causal_conv1d_update import lfm2_causal_conv1d_update_cutile + + modeling_lfm2_moe.causal_conv1d_fn = lfm2_causal_conv1d_fn_cutile + modeling_lfm2_moe.causal_conv1d_update = lfm2_causal_conv1d_update_cutile + modeling_lfm2_moe.is_fast_path_available = True + + MODEL_TYPE_TO_APPLY_TILEGYM_FN = { "llama": apply_tilegym_kernel_to_llama, "deepseek_v2": apply_tilegym_kernel_to_deepseek_v2, @@ -563,6 +662,7 @@ def apply_tilegym_kernel_to_olmo3( "phi3": apply_tilegym_kernel_to_phi3, "olmo3": apply_tilegym_kernel_to_olmo3, "olmoe": apply_tilegym_kernel_to_olmoe, + "lfm2_moe": apply_tilegym_kernel_to_lfm2_moe, } diff --git a/tests/kernel_inventory/runtime_inputs.yaml b/tests/kernel_inventory/runtime_inputs.yaml index 5f6b6649..137cf8d1 100644 --- a/tests/kernel_inventory/runtime_inputs.yaml +++ b/tests/kernel_inventory/runtime_inputs.yaml @@ -14,3 +14,5 @@ cases: mutates: [q, k] src/tilegym/transformers/qwen3_5::definition::qwen3_5_causal_conv1d_update_silu: mutates: [conv_state] + src/tilegym/transformers/lfm2_moe::definition::lfm2_moe_causal_conv1d_update: + mutates: [conv_state] diff --git a/tests/ops/test_causal_conv1d_lfm2.py b/tests/ops/test_causal_conv1d_lfm2.py new file mode 100644 index 00000000..c92ae72f --- /dev/null +++ b/tests/ops/test_causal_conv1d_lfm2.py @@ -0,0 +1,137 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: MIT + +"""Correctness tests for the LFM2-MoE fused cuTile causal conv1d kernels. + +These wrappers are not backend-dispatched ops, so they are exercised directly +(rather than via the ``PyTestCase`` harness) against a pure-torch reference that +matches ``Lfm2MoeShortConv``'s stock semantics. Requires a CUDA device with +cuTile; skipped otherwise. +""" + +import pytest +import torch +import torch.nn.functional as F + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="LFM2 conv kernels require CUDA/cuTile") + + +def _ref_prefill(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + """nn.Conv1d(groups=D, padding=K-1)(x)[..., :L] — depthwise causal conv.""" + B, D, L = x.shape + K = weight.shape[1] + xp = F.pad(x, (K - 1, 0)) + w = weight.view(D, 1, K) + return F.conv1d(xp, w, groups=D) + + +def _ref_update(x: torch.Tensor, conv_state: torch.Tensor, weight: torch.Tensor): + """Roll the K-wide window and convolve; returns (out, rolled_state). + + Mirrors HF's fallback ``causal_conv1d_update`` on the single-token decode + path: ``x`` is 3D ``(B, D, 1)`` and the output keeps that rank. + + The multiply-accumulate is done in float32 to match the kernel under test + (which upcasts to f32 internally, like the CUDA causal-conv1d package) -- + a pure-bf16 reference rounds each product to bf16 before summing and is + *less* accurate than the kernel, showing up as spurious 1-2 ulp diffs. + """ + rolled = torch.cat([conv_state[..., 1:], x], dim=-1) + out = (rolled.float() * weight.float().unsqueeze(0)).sum(-1, keepdim=True).to(x.dtype) + return out, rolled + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("D, L", [(2048, 16), (2048, 128), (256, 7), (300, 33)]) +def test_op_prefill(D, L, dtype): + from tilegym.transformers.lfm2_moe.kernels.causal_conv1d_prefill import lfm2_causal_conv1d_fn_cutile + + torch.manual_seed(0) + device = torch.device("cuda") + x = torch.randn(1, D, L, dtype=dtype, device=device) + weight = torch.randn(D, 3, dtype=dtype, device=device) + + out = lfm2_causal_conv1d_fn_cutile(x, weight) + ref = _ref_prefill(x, weight) + + assert out.shape == ref.shape == (1, D, L) + atol = 1e-4 if dtype == torch.float32 else 2e-2 + torch.testing.assert_close(out.float(), ref.float(), atol=atol, rtol=1e-3) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("D", [2048, 256, 300]) +def test_op_update(D, dtype): + """Exercise the kernel with the exact 3D call contract of + ``Lfm2MoeShortConv.forward``'s cached-decode branch: x is ``(B, D, 1)`` + and the output must be ``(B, D, 1)`` (it is multiplied with the 3D gate + ``C`` at the call site).""" + from tilegym.transformers.lfm2_moe.kernels.causal_conv1d_update import lfm2_causal_conv1d_update_cutile + + torch.manual_seed(0) + device = torch.device("cuda") + x = torch.randn(1, D, 1, dtype=dtype, device=device) + conv_state = torch.randn(1, D, 3, dtype=dtype, device=device) + weight = torch.randn(D, 3, dtype=dtype, device=device) + + ref_out, ref_state = _ref_update(x, conv_state, weight) + + cs = conv_state.clone() + out = lfm2_causal_conv1d_update_cutile(x, cs, weight) + + assert out.shape == (1, D, 1) + atol = 1e-4 if dtype == torch.float32 else 2e-2 + # output matches reference + torch.testing.assert_close(out.float(), ref_out.float(), atol=atol, rtol=1e-3) + # conv_state was rolled in place to [s1, s2, x] + torch.testing.assert_close(cs.float(), ref_state.float(), atol=atol, rtol=1e-3) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_op_update_2d_legacy_call(dtype): + """Older transformers (<= 5.10.x, `cuda_kernels_forward`) call the update + with a 2D ``(B, D)`` input (``Bx.squeeze(-1)``) and unsqueeze the result + at the call site. The wrapper must accept that convention too.""" + from tilegym.transformers.lfm2_moe.kernels.causal_conv1d_update import lfm2_causal_conv1d_update_cutile + + torch.manual_seed(0) + device = torch.device("cuda") + D = 512 + x = torch.randn(1, D, 1, dtype=dtype, device=device) + conv_state = torch.randn(1, D, 3, dtype=dtype, device=device) + weight = torch.randn(D, 3, dtype=dtype, device=device) + + ref_out, ref_state = _ref_update(x, conv_state, weight) + + cs = conv_state.clone() + out = lfm2_causal_conv1d_update_cutile(x.squeeze(-1), cs, weight) # 2D legacy call + + assert out.shape == (1, D) # rank follows the input + atol = 1e-4 if dtype == torch.float32 else 2e-2 + torch.testing.assert_close(out.float(), ref_out.squeeze(-1).float(), atol=atol, rtol=1e-3) + torch.testing.assert_close(cs.float(), ref_state.float(), atol=atol, rtol=1e-3) + + +def test_op_prefill_matches_update_stepwise(): + """A full-sequence prefill should equal stepping the update kernel token by + token from a zero-initialized K-wide state (end-to-end conv consistency).""" + from tilegym.transformers.lfm2_moe.kernels.causal_conv1d_prefill import lfm2_causal_conv1d_fn_cutile + from tilegym.transformers.lfm2_moe.kernels.causal_conv1d_update import lfm2_causal_conv1d_update_cutile + + torch.manual_seed(0) + device = torch.device("cuda") + D, L = 512, 12 + x = torch.randn(1, D, L, dtype=torch.float32, device=device) + weight = torch.randn(D, 3, dtype=torch.float32, device=device) + + prefill = lfm2_causal_conv1d_fn_cutile(x, weight) # (1, D, L) + + conv_state = torch.zeros(1, D, 3, dtype=torch.float32, device=device) + step_outs = [] + for t in range(L): + # Keep the 3D (1, D, 1) call shape used by the HF decode path. + step_outs.append(lfm2_causal_conv1d_update_cutile(x[:, :, t : t + 1], conv_state, weight)) + stepwise = torch.cat(step_outs, dim=-1) # (1, D, L) + + torch.testing.assert_close(prefill, stepwise, atol=1e-4, rtol=1e-3)