From 3c0c2e9ff6c59a792e3e000c5f5f94eaffb6e7a6 Mon Sep 17 00:00:00 2001 From: iamanishx Date: Mon, 17 Aug 2026 22:26:39 +0530 Subject: [PATCH 1/3] lfm 2 moe Signed-off-by: iamanishx --- .../scripts/benchmark_hf_model.sh | 10 +- .../src/tilegym_hf_bench/tilegym_patch.py | 5 + src/tilegym/transformers/__init__.py | 1 + src/tilegym/transformers/lfm2_moe/__init__.py | 3 + .../lfm2_moe_causal_conv1d_prefill.json | 58 +++++++ .../lfm2_moe_causal_conv1d_update.json | 67 +++++++ .../lfm2_moe_causal_conv1d_prefill.json | 25 +++ .../lfm2_moe_causal_conv1d_update.json | 25 +++ .../transformers/lfm2_moe/kernels/__init__.py | 3 + .../lfm2_moe/kernels/causal_conv1d_prefill.py | 98 +++++++++++ .../lfm2_moe/kernels/causal_conv1d_update.py | 111 ++++++++++++ .../lfm2_moe/modeling_lfm2_moe.py | 164 ++++++++++++++++++ src/tilegym/transformers/monkey_patch.py | 100 +++++++++++ .../kernel_inventory/kernel_runtime_utils.py | 1 + tests/ops/test_causal_conv1d_lfm2.py | 107 ++++++++++++ 15 files changed, 777 insertions(+), 1 deletion(-) create mode 100644 src/tilegym/transformers/lfm2_moe/__init__.py create mode 100644 src/tilegym/transformers/lfm2_moe/kernel_definitions/lfm2_moe_causal_conv1d_prefill.json create mode 100644 src/tilegym/transformers/lfm2_moe/kernel_definitions/lfm2_moe_causal_conv1d_update.json create mode 100644 src/tilegym/transformers/lfm2_moe/kernel_solutions/lfm2_moe_causal_conv1d_prefill.json create mode 100644 src/tilegym/transformers/lfm2_moe/kernel_solutions/lfm2_moe_causal_conv1d_update.json create mode 100644 src/tilegym/transformers/lfm2_moe/kernels/__init__.py create mode 100644 src/tilegym/transformers/lfm2_moe/kernels/causal_conv1d_prefill.py create mode 100644 src/tilegym/transformers/lfm2_moe/kernels/causal_conv1d_update.py create mode 100644 src/tilegym/transformers/lfm2_moe/modeling_lfm2_moe.py create mode 100644 tests/ops/test_causal_conv1d_lfm2.py 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) + + 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, L, 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..b894a70a --- /dev/null +++ b/src/tilegym/transformers/lfm2_moe/kernels/causal_conv1d_update.py @@ -0,0 +1,111 @@ +# 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``). + + Matches the call contract of ``Lfm2MoeShortConv.forward``'s single-token + cached-decode branch, which passes the *3D* ``hidden_states`` tensor and + multiplies the result with the 3D gate ``C`` (``y = C * hidden_states``). + + Args: + x: ``(B=1, D, L=1)`` current-timestep input (``B * x`` after the + in-proj chunk, still 3D at the call site). + 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: + ``(B=1, D, 1)`` conv output for the current timestep (same rank as the + input, like the HF reference implementation). + """ + assert activation is None, "LFM2 short-conv fuses no activation; activation must be None" + + B, D, L = x.shape + assert B == 1, "lfm2_causal_conv1d_update_cutile only supports B=1" + assert L == 1, f"lfm2_causal_conv1d_update_cutile is the single-token decode path, got L={L}" + 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), + ) + + out = output.view(1, D, 1) # (1, D, 1), matching the HF fallback's output rank + if bias is not None: + out = out + bias.view(1, -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..d65328cf --- /dev/null +++ b/src/tilegym/transformers/lfm2_moe/modeling_lfm2_moe.py @@ -0,0 +1,164 @@ +# 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 + from transformers.models.lfm2_moe.modeling_lfm2_moe import Lfm2MoeTopKRouter + + self.experts = Lfm2MoeExperts(config) + self.gate = Lfm2MoeTopKRouter(config) + 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/kernel_runtime_utils.py b/tests/kernel_inventory/kernel_runtime_utils.py index eea6031a..69b96bc6 100644 --- a/tests/kernel_inventory/kernel_runtime_utils.py +++ b/tests/kernel_inventory/kernel_runtime_utils.py @@ -53,6 +53,7 @@ "olmo3_dual_rms_norm": ("q", "k"), "olmoe_dual_rms_norm": ("q", "k"), "qwen3_5_causal_conv1d_update_silu": ("conv_state",), + "lfm2_moe_causal_conv1d_update": ("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..79a6645c --- /dev/null +++ b/tests/ops/test_causal_conv1d_lfm2.py @@ -0,0 +1,107 @@ +# 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. + """ + rolled = torch.cat([conv_state[..., 1:], x], dim=-1) + out = (rolled * weight.unsqueeze(0)).sum(-1, keepdim=True) + 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_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_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) + + +def test_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) From a6b0f1860c8c002b03867de0ab37f13ee0a5a4c8 Mon Sep 17 00:00:00 2001 From: iamanishx Date: Sun, 23 Aug 2026 11:20:27 +0530 Subject: [PATCH 2/3] fix(lfm2_moe): match HF conv API drift, fix JIT specialization, satisfy kernel inventory - causal_conv1d_update_cutile: accept both decode call conventions and return the matching rank. transformers 5.13+ calls it from Lfm2MoeShortConv.forward with a 3D (B, D, 1) tensor and gates the result directly, while 5.10.x calls it from cuda_kernels_forward with Bx.squeeze(-1), a 2D (B, D) tensor. The previous 2D-only unpacking raised ValueError on the 5.13+ path. - Lfm2MoeSparseMoeBlockTileGym: fall back to nn.Linear when Lfm2MoeTopKRouter is not importable. transformers 5.10.x inlines the router inside Lfm2MoeSparseMoeBlock. Both layouts expose the same gate.weight state_dict key of shape (num_experts, hidden_size), so strict loading works either way. - causal_conv1d_prefill: drop the unused sequence-length ct.Constant. It was part of the JIT specialization key, so every distinct prompt length forced a fresh cuTile compile for no benefit. Bounds are already handled by check_bounds on the gathers and scatter. - kernel_definitions: point Definition.reference at pinned transformers permalinks with line anchors and align the reference run signatures with the solution entry points, both required by tests/kernel_inventory. - tests: compute the bf16 update reference in float32 to match the kernel's f32 accumulation, and add coverage for the legacy 2D decode call. Validated on RTX PRO 6000 Blackwell (SM120): tests/ops/test_causal_conv1d_lfm2.py, 17 passed. tests/kernel_inventory -k lfm2, 2 passed. Signed-off-by: iamanishx --- .../lfm2_moe_causal_conv1d_prefill.json | 12 +++---- .../lfm2_moe_causal_conv1d_update.json | 14 ++++---- .../lfm2_moe/kernels/causal_conv1d_prefill.py | 8 +++-- .../lfm2_moe/kernels/causal_conv1d_update.py | 36 +++++++++++++------ .../lfm2_moe/modeling_lfm2_moe.py | 13 +++++-- tests/ops/test_causal_conv1d_lfm2.py | 32 ++++++++++++++++- 6 files changed, 86 insertions(+), 29 deletions(-) diff --git a/src/tilegym/transformers/lfm2_moe/kernel_definitions/lfm2_moe_causal_conv1d_prefill.json b/src/tilegym/transformers/lfm2_moe/kernel_definitions/lfm2_moe_causal_conv1d_prefill.json index 61e9c975..1b9aeaad 100644 --- a/src/tilegym/transformers/lfm2_moe/kernel_definitions/lfm2_moe_causal_conv1d_prefill.json +++ b/src/tilegym/transformers/lfm2_moe/kernel_definitions/lfm2_moe_causal_conv1d_prefill.json @@ -21,19 +21,19 @@ }, "description": "LFM2-MoE prefill-path depthwise causal conv1d (kernel_size=3, no activation). Input is unpadded; the entry point left-pads by K-1 internally.", "inputs": { - "x": { + "weight": { "dtype": "bfloat16", "shape": [ - "B", "D", - "T" + "K" ] }, - "weight": { + "x": { "dtype": "bfloat16", "shape": [ + "B", "D", - "K" + "T" ] } }, @@ -49,7 +49,7 @@ ] } }, - "reference": "# Source: https://github.com/huggingface/transformers/blob/main/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py (Lfm2MoeShortConv, causal_conv1d_fn path)\nimport torch\nimport torch.nn.functional as F\n\ndef run(x, weight):\n B, D, L = x.shape\n K = weight.shape[1]\n xp = F.pad(x, (K - 1, 0))\n w = weight.view(D, 1, K)\n return F.conv1d(xp, w, groups=D)", + "reference": "# Source: https://github.com/huggingface/transformers/blob/5eddc12edfaf8cafde8c9bae4ccb12f8a139b4f9/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py#L385-L405\nimport torch\nimport torch.nn.functional as F\n\ndef run(x, weight, bias=None, activation=None, seq_idx=None):\n # LFM2 short-conv fuses no activation and does not use packed sequences,\n # so `activation` and `seq_idx` are always None at the call site.\n _, hidden_size, seq_len = x.shape\n padding = weight.shape[-1] - 1\n out = F.conv1d(\n x.to(weight.dtype),\n weight=weight.unsqueeze(1),\n bias=bias,\n padding=padding,\n groups=hidden_size,\n )[:, :, :seq_len]\n return out.to(x.dtype)", "tags": [ "model:lfm2_moe", "stage:prefill", diff --git a/src/tilegym/transformers/lfm2_moe/kernel_definitions/lfm2_moe_causal_conv1d_update.json b/src/tilegym/transformers/lfm2_moe/kernel_definitions/lfm2_moe_causal_conv1d_update.json index befdc08f..cd2c39ad 100644 --- a/src/tilegym/transformers/lfm2_moe/kernel_definitions/lfm2_moe_causal_conv1d_update.json +++ b/src/tilegym/transformers/lfm2_moe/kernel_definitions/lfm2_moe_causal_conv1d_update.json @@ -22,27 +22,27 @@ }, "description": "LFM2-MoE decode-path depthwise causal conv1d update (kernel_size=3, no activation). The conv_state input holds the full K-wide window and is rolled/updated in place.", "inputs": { - "x": { + "conv_state": { "dtype": "bfloat16", "shape": [ "B", "D", - "T" + "K" ] }, - "conv_state": { + "weight": { "dtype": "bfloat16", "shape": [ - "B", "D", "K" ] }, - "weight": { + "x": { "dtype": "bfloat16", "shape": [ + "B", "D", - "K" + "T" ] } }, @@ -58,7 +58,7 @@ ] } }, - "reference": "# Source: https://github.com/huggingface/transformers/blob/main/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py (Lfm2MoeShortConv, causal_conv1d_update path)\nimport torch\n\ndef run(x, conv_state, weight, bias=None, activation=None):\n rolled = torch.cat([conv_state[..., 1:], x], dim=-1)\n out = (rolled * weight.unsqueeze(0)).sum(-1, keepdim=True)\n conv_state[..., 0] = conv_state[..., 1]\n conv_state[..., 1] = conv_state[..., 2]\n conv_state[..., 2] = x[..., 0]\n return out", + "reference": "# Source: https://github.com/huggingface/transformers/blob/5eddc12edfaf8cafde8c9bae4ccb12f8a139b4f9/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py#L365-L382\nimport torch\nimport torch.nn.functional as F\n\ndef run(x, conv_state, weight, bias=None, activation=None):\n # `conv_state` holds the full K-wide window and is rolled in place.\n # LFM2 short-conv fuses no activation, so `activation` is always None.\n _, hidden_size, seq_len = x.shape\n state_len = conv_state.shape[-1]\n hidden_states_new = torch.cat([conv_state, x], dim=-1).to(weight.dtype)\n conv_state.copy_(hidden_states_new[:, :, -state_len:])\n out = F.conv1d(hidden_states_new, weight.unsqueeze(1), bias, padding=0, groups=hidden_size)\n return out[:, :, -seq_len:].to(x.dtype)", "tags": [ "model:lfm2_moe", "stage:decode", diff --git a/src/tilegym/transformers/lfm2_moe/kernels/causal_conv1d_prefill.py b/src/tilegym/transformers/lfm2_moe/kernels/causal_conv1d_prefill.py index fc6c1f09..8a7f3886 100644 --- a/src/tilegym/transformers/lfm2_moe/kernels/causal_conv1d_prefill.py +++ b/src/tilegym/transformers/lfm2_moe/kernels/causal_conv1d_prefill.py @@ -26,7 +26,6 @@ def _causal_conv1d_prefill_kernel( x, # (D, T_padded) left-padded by K-1 weight, # (D, K=3) output, # (D, T) - T: ConstInt, BLOCK_T: ConstInt, ): bid_d = ct.bid(0) @@ -83,13 +82,18 @@ def lfm2_causal_conv1d_fn_cutile( 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, L, BLOCK_T), + (x_padded, w, output, BLOCK_T), ) out = output.unsqueeze(0) # (1, D, L) diff --git a/src/tilegym/transformers/lfm2_moe/kernels/causal_conv1d_update.py b/src/tilegym/transformers/lfm2_moe/kernels/causal_conv1d_update.py index b894a70a..9d663b81 100644 --- a/src/tilegym/transformers/lfm2_moe/kernels/causal_conv1d_update.py +++ b/src/tilegym/transformers/lfm2_moe/kernels/causal_conv1d_update.py @@ -67,27 +67,39 @@ def lfm2_causal_conv1d_update_cutile( ) -> torch.Tensor: """Depthwise causal conv1d decode-update (drop-in for ``causal_conv1d_update``). - Matches the call contract of ``Lfm2MoeShortConv.forward``'s single-token - cached-decode branch, which passes the *3D* ``hidden_states`` tensor and - multiplies the result with the 3D gate ``C`` (``y = C * hidden_states``). + 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)`` current-timestep input (``B * x`` after the - in-proj chunk, still 3D at the call site). + 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: - ``(B=1, D, 1)`` conv output for the current timestep (same rank as the - input, like the HF reference implementation). + 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" - B, D, L = x.shape + 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" - assert L == 1, f"lfm2_causal_conv1d_update_cutile is the single-token decode path, got L={L}" K = weight.shape[1] assert K == 3, f"expected kernel_size 3, got {K}" @@ -105,7 +117,9 @@ def lfm2_causal_conv1d_update_cutile( (x_1d, cs, w, output, BLOCK_D), ) - out = output.view(1, D, 1) # (1, D, 1), matching the HF fallback's output rank + # 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) + 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 index d65328cf..3e65b33e 100644 --- a/src/tilegym/transformers/lfm2_moe/modeling_lfm2_moe.py +++ b/src/tilegym/transformers/lfm2_moe/modeling_lfm2_moe.py @@ -96,10 +96,19 @@ def __init__(self, config): # 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 - from transformers.models.lfm2_moe.modeling_lfm2_moe import Lfm2MoeTopKRouter self.experts = Lfm2MoeExperts(config) - self.gate = Lfm2MoeTopKRouter(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 diff --git a/tests/ops/test_causal_conv1d_lfm2.py b/tests/ops/test_causal_conv1d_lfm2.py index 79a6645c..33e2d20d 100644 --- a/tests/ops/test_causal_conv1d_lfm2.py +++ b/tests/ops/test_causal_conv1d_lfm2.py @@ -31,9 +31,14 @@ def _ref_update(x: torch.Tensor, conv_state: torch.Tensor, weight: torch.Tensor) 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 * weight.unsqueeze(0)).sum(-1, keepdim=True) + out = (rolled.float() * weight.float().unsqueeze(0)).sum(-1, keepdim=True).to(x.dtype) return out, rolled @@ -83,6 +88,31 @@ def test_update(D, dtype): torch.testing.assert_close(cs.float(), ref_state.float(), atol=atol, rtol=1e-3) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_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_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 b176ea496a8d2b17d379d83db5650730041fed4d Mon Sep 17 00:00:00 2001 From: iamanishx Date: Sun, 23 Aug 2026 11:32:57 +0530 Subject: [PATCH 3/3] test(lfm2_moe): rename conv tests to test_op* so CI ops job selects them The CI ops job runs pytest with -k test_op, so the previous names were silently deselected. Signed-off-by: iamanishx --- tests/ops/test_causal_conv1d_lfm2.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/ops/test_causal_conv1d_lfm2.py b/tests/ops/test_causal_conv1d_lfm2.py index 33e2d20d..c92ae72f 100644 --- a/tests/ops/test_causal_conv1d_lfm2.py +++ b/tests/ops/test_causal_conv1d_lfm2.py @@ -44,7 +44,7 @@ def _ref_update(x: torch.Tensor, conv_state: torch.Tensor, weight: torch.Tensor) @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) @pytest.mark.parametrize("D, L", [(2048, 16), (2048, 128), (256, 7), (300, 33)]) -def test_prefill(D, L, dtype): +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) @@ -62,7 +62,7 @@ def test_prefill(D, L, dtype): @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) @pytest.mark.parametrize("D", [2048, 256, 300]) -def test_update(D, dtype): +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 @@ -89,7 +89,7 @@ def test_update(D, dtype): @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -def test_update_2d_legacy_call(dtype): +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.""" @@ -113,7 +113,7 @@ def test_update_2d_legacy_call(dtype): torch.testing.assert_close(cs.float(), ref_state.float(), atol=atol, rtol=1e-3) -def test_prefill_matches_update_stepwise(): +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