perf(lfm2): fuse gating and short convolution on SM90 (#37622)

python/sglang/kernels/ops/mamba/lfm_short_conv.py

@@ -0,0 +1,369 @@
+"""Gate-fused BF16 short convolution used by LFM2 on CUDA.
+
+LFM2 materializes ``B * x``, transposes it for the generic causal-conv kernel,
+then materializes ``C * conv(B * x)``. These kernels preserve both BF16
+materialization points while operating directly on the three projection views.
+"""
+
+from __future__ import annotations
+
+import os
+
+import torch
+import triton
+import triton.language as tl
+
+_DISABLE_LFM_FUSED_CONV = os.getenv("SGLANG_DISABLE_LFM_FUSED_CONV", "0") == "1"
+
+
+@triton.jit
+def _lfm_short_conv_prefill_kernel(
+    b_ptr,
+    c_ptr,
+    x_ptr,
+    weight_ptr,
+    state_ptr,
+    query_start_loc_ptr,
+    cache_indices_ptr,
+    has_initial_state_ptr,
+    out_ptr,
+    b_stride_t: tl.constexpr,
+    c_stride_t: tl.constexpr,
+    x_stride_t: tl.constexpr,
+    weight_stride_d: tl.constexpr,
+    state_stride_slot: tl.constexpr,
+    state_stride_d: tl.constexpr,
+    state_stride_w: tl.constexpr,
+    out_stride_t: tl.constexpr,
+    blocks_per_seq,
+    dim: tl.constexpr,
+    PAD_SLOT_ID: tl.constexpr,
+    BLOCK_T: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+):
+    tile = tl.program_id(0)
+    seq = tile // blocks_per_seq
+    token_block = tile - seq * blocks_per_seq
+    start = tl.load(query_start_loc_ptr + seq)
+    end = tl.load(query_start_loc_ptr + seq + 1)
+    t = start + token_block * BLOCK_T + tl.arange(0, BLOCK_T)
+    d = tl.program_id(1) * BLOCK_D + tl.arange(0, BLOCK_D)
+    tt = t[:, None]
+    dd = d[None, :]
+    valid_d = dd < dim
+    valid_t = tt < end
+
+    slot = tl.load(cache_indices_ptr + seq)
+    active = slot != PAD_SLOT_ID
+    has_initial = tl.load(has_initial_state_ptr + seq).to(tl.int1)
+    state_base = slot * state_stride_slot + dd * state_stride_d
+
+    token0 = tt - 2
+    in_sequence0 = (token0 >= start) & (token0 < end)
+    in_mask0 = valid_d & in_sequence0 & active
+    b0 = tl.load(b_ptr + token0 * b_stride_t + dd, mask=in_mask0, other=0.0).to(
+        tl.float32
+    )
+    x0 = tl.load(x_ptr + token0 * x_stride_t + dd, mask=in_mask0, other=0.0).to(
+        tl.float32
+    )
+    new_bx0 = (b0 * x0).to(tl.bfloat16)
+    state_pos0 = token0 - start + 2
+    old0 = tl.load(
+        state_ptr + state_base + state_pos0 * state_stride_w,
+        mask=(
+            valid_d
+            & active
+            & has_initial
+            & (state_pos0 >= 0)
+            & (state_pos0 < 2)
+            & ~in_sequence0
+        ),
+        other=0.0,
+    )
+    bx0 = tl.where(in_sequence0, new_bx0, old0).to(tl.bfloat16)
+
+    token1 = tt - 1
+    in_sequence1 = (token1 >= start) & (token1 < end)
+    in_mask1 = valid_d & in_sequence1 & active
+    b1 = tl.load(b_ptr + token1 * b_stride_t + dd, mask=in_mask1, other=0.0).to(
+        tl.float32
+    )
+    x1 = tl.load(x_ptr + token1 * x_stride_t + dd, mask=in_mask1, other=0.0).to(
+        tl.float32
+    )
+    new_bx1 = (b1 * x1).to(tl.bfloat16)
+    state_pos1 = token1 - start + 2
+    old1 = tl.load(
+        state_ptr + state_base + state_pos1 * state_stride_w,
+        mask=(
+            valid_d
+            & active
+            & has_initial
+            & (state_pos1 >= 0)
+            & (state_pos1 < 2)
+            & ~in_sequence1
+        ),
+        other=0.0,
+    )
+    bx1 = tl.where(in_sequence1, new_bx1, old1).to(tl.bfloat16)
+
+    in_mask2 = valid_d & valid_t
+    b2 = tl.load(b_ptr + tt * b_stride_t + dd, mask=in_mask2, other=0.0).to(tl.float32)
+    x2 = tl.load(x_ptr + tt * x_stride_t + dd, mask=in_mask2, other=0.0).to(tl.float32)
+    bx2 = (b2 * x2).to(tl.bfloat16)
+
+    w0 = tl.load(weight_ptr + dd * weight_stride_d, mask=valid_d, other=0.0).to(
+        tl.float32
+    )
+    w1 = tl.load(weight_ptr + dd * weight_stride_d + 1, mask=valid_d, other=0.0).to(
+        tl.float32
+    )
+    w2 = tl.load(weight_ptr + dd * weight_stride_d + 2, mask=valid_d, other=0.0).to(
+        tl.float32
+    )
+    conv = tl.zeros((BLOCK_T, BLOCK_D), dtype=tl.float32)
+    conv += bx0.to(tl.float32) * w0
+    conv += bx1.to(tl.float32) * w1
+    conv += bx2.to(tl.float32) * w2
+    conv = conv.to(tl.bfloat16)
+    c = tl.load(
+        c_ptr + tt * c_stride_t + dd,
+        mask=valid_d & valid_t,
+        other=0.0,
+    ).to(tl.float32)
+    active_y = (c * conv.to(tl.float32)).to(tl.bfloat16)
+    pad_y = (c * bx2.to(tl.float32)).to(tl.bfloat16)
+    y = tl.where(active, active_y, pad_y)
+    tl.store(
+        out_ptr + tt * out_stride_t + dd,
+        y,
+        mask=valid_d & valid_t,
+    )
+
+    # Exactly one token tile per sequence commits the final two gated inputs.
+    final_prev = tl.sum(tl.where(tt == end - 1, bx1.to(tl.float32), 0.0), axis=0).to(
+        tl.bfloat16
+    )
+    final_cur = tl.sum(tl.where(tt == end - 1, bx2.to(tl.float32), 0.0), axis=0).to(
+        tl.bfloat16
+    )
+    final_block = (token_block * BLOCK_T <= end - start - 1) & (
+        end - start - 1 < (token_block + 1) * BLOCK_T
+    )
+    final_mask = valid_d & active & final_block
+    tl.store(
+        state_ptr + state_base,
+        final_prev[None, :],
+        mask=final_mask,
+    )
+    tl.store(
+        state_ptr + state_base + state_stride_w,
+        final_cur[None, :],
+        mask=final_mask,
+    )
+
+
+@triton.jit
+def _lfm_short_conv_decode_kernel(
+    b_ptr,
+    c_ptr,
+    x_ptr,
+    weight_ptr,
+    state_ptr,
+    cache_indices_ptr,
+    out_ptr,
+    b_stride_t: tl.constexpr,
+    c_stride_t: tl.constexpr,
+    x_stride_t: tl.constexpr,
+    weight_stride_d: tl.constexpr,
+    state_stride_slot: tl.constexpr,
+    state_stride_d: tl.constexpr,
+    state_stride_w: tl.constexpr,
+    out_stride_t: tl.constexpr,
+    dim: tl.constexpr,
+    PAD_SLOT_ID: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+):
+    token = tl.program_id(0)
+    d = tl.program_id(1) * BLOCK_D + tl.arange(0, BLOCK_D)
+    valid_d = d < dim
+    slot = tl.load(cache_indices_ptr + token)
+    active = slot != PAD_SLOT_ID
+    state_base = slot * state_stride_slot + d * state_stride_d
+    state_mask = valid_d & active
+
+    prev0 = tl.load(state_ptr + state_base, mask=state_mask, other=0.0).to(tl.float32)
+    prev1 = tl.load(
+        state_ptr + state_base + state_stride_w,
+        mask=state_mask,
+        other=0.0,
+    ).to(tl.float32)
+    w0 = tl.load(weight_ptr + d * weight_stride_d, mask=valid_d, other=0.0).to(
+        tl.float32
+    )
+    w1 = tl.load(weight_ptr + d * weight_stride_d + 1, mask=valid_d, other=0.0).to(
+        tl.float32
+    )
+    w2 = tl.load(weight_ptr + d * weight_stride_d + 2, mask=valid_d, other=0.0).to(
+        tl.float32
+    )
+    b = tl.load(b_ptr + token * b_stride_t + d, mask=valid_d, other=0.0).to(tl.float32)
+    x = tl.load(x_ptr + token * x_stride_t + d, mask=valid_d, other=0.0).to(tl.float32)
+    bx_bf16 = (b * x).to(tl.bfloat16)
+    bx = bx_bf16.to(tl.float32)
+    conv = tl.zeros((BLOCK_D,), dtype=tl.float32)
+    conv += prev0 * w0
+    conv += prev1 * w1
+    conv += bx * w2
+    conv_bf16 = conv.to(tl.bfloat16)
+    c = tl.load(c_ptr + token * c_stride_t + d, mask=valid_d, other=0.0).to(tl.float32)
+    active_y = (c * conv_bf16.to(tl.float32)).to(tl.bfloat16)
+    pad_y = (c * bx_bf16.to(tl.float32)).to(tl.bfloat16)
+    y = tl.where(active, active_y, pad_y)
+
+    tl.store(out_ptr + token * out_stride_t + d, y, mask=valid_d)
+    tl.store(state_ptr + state_base, prev1.to(tl.bfloat16), mask=state_mask)
+    tl.store(state_ptr + state_base + state_stride_w, bx_bf16, mask=state_mask)
+
+
+def can_use_fused_lfm_short_conv(
+    b: torch.Tensor,
+    c: torch.Tensor,
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor | None,
+    state: torch.Tensor,
+) -> bool:
+    """Return whether the exact LFM2 BF16 width-3 contract is satisfied."""
+    return (
+        not _DISABLE_LFM_FUSED_CONV
+        and b.is_cuda
+        and torch.version.hip is None
+        and torch.cuda.get_device_capability(b.device) == (9, 0)
+        and c.is_cuda
+        and x.is_cuda
+        and weight.is_cuda
+        and state.is_cuda
+        and b.device == c.device == x.device == weight.device == state.device
+        and b.dtype == c.dtype == x.dtype == torch.bfloat16
+        and b.shape == c.shape == x.shape
+        and b.ndim == 2
+        and b.stride(1) == c.stride(1) == x.stride(1) == 1
+        and weight.dtype == b.dtype
+        and weight.ndim == 2
+        and weight.shape == (b.shape[1], 3)
+        and weight.stride(1) == 1
+        and bias is None
+        and state.dtype == b.dtype
+        and state.ndim == 3
+        and state.shape[1:] == (b.shape[1], 2)
+    )
+
+
+def can_dispatch_fused_lfm_short_conv(
+    b: torch.Tensor,
+    c: torch.Tensor,
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor | None,
+    state: torch.Tensor,
+) -> bool:
+    """Return whether this call is in the end-to-end-qualified serving domain."""
+    return (
+        b.ndim == 2
+        and b.shape[1] == 2048
+        and can_use_fused_lfm_short_conv(b, c, x, weight, bias, state)
+    )
+
+
+def fused_lfm_short_conv_prefill(
+    b: torch.Tensor,
+    c: torch.Tensor,
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    state: torch.Tensor,
+    query_start_loc: torch.Tensor,
+    cache_indices: torch.Tensor,
+    has_initial_state: torch.Tensor,
+    max_seq_len: int,
+) -> torch.Tensor:
+    """Compute and cache ``C * conv(B * x)`` for packed prefill tokens."""
+    tokens, dim = b.shape
+    out = torch.empty((tokens, dim), dtype=b.dtype, device=b.device)
+    block_t = 32
+    block_d = 64
+    blocks_per_seq = triton.cdiv(max_seq_len, block_t)
+    num_sequences = query_start_loc.numel() - 1
+    _lfm_short_conv_prefill_kernel[
+        (num_sequences * blocks_per_seq, triton.cdiv(dim, block_d))
+    ](
+        b,
+        c,
+        x,
+        weight,
+        state,
+        query_start_loc,
+        cache_indices,
+        has_initial_state,
+        out,
+        b.stride(0),
+        c.stride(0),
+        x.stride(0),
+        weight.stride(0),
+        state.stride(0),
+        state.stride(1),
+        state.stride(2),
+        out.stride(0),
+        blocks_per_seq,
+        dim,
+        -1,
+        BLOCK_T=block_t,
+        BLOCK_D=block_d,
+        num_warps=4,
+    )
+    return out
+
+
+def fused_lfm_short_conv_decode(
+    b: torch.Tensor,
+    c: torch.Tensor,
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    state: torch.Tensor,
+    cache_indices: torch.Tensor,
+) -> torch.Tensor:
+    """Compute, cache, and gate one decode token per active request."""
+    tokens, dim = b.shape
+    out = torch.empty((tokens, dim), dtype=b.dtype, device=b.device)
+    block_d = 256
+    _lfm_short_conv_decode_kernel[(tokens, triton.cdiv(dim, block_d))](
+        b,
+        c,
+        x,
+        weight,
+        state,
+        cache_indices,
+        out,
+        b.stride(0),
+        c.stride(0),
+        x.stride(0),
+        weight.stride(0),
+        state.stride(0),
+        state.stride(1),
+        state.stride(2),
+        out.stride(0),
+        dim,
+        -1,
+        BLOCK_D=block_d,
+        num_warps=8,
+    )
+    return out
+
+
+__all__ = [
+    "can_dispatch_fused_lfm_short_conv",
+    "can_use_fused_lfm_short_conv",
+    "fused_lfm_short_conv_decode",
+    "fused_lfm_short_conv_prefill",
+]

python/sglang/srt/models/lfm2_moe.py

@@ -17,6 +17,11 @@ from typing import Iterable, List, Optional, Set, Tuple
 import torch
 from torch import nn
 
+from sglang.kernels.ops.mamba.lfm_short_conv import (
+    can_dispatch_fused_lfm_short_conv,
+    fused_lfm_short_conv_decode,
+    fused_lfm_short_conv_prefill,
+)
 from sglang.srt.configs.lfm2_moe import Lfm2MoeConfig
 from sglang.srt.distributed import get_pp_group
 from sglang.srt.layers.activation import SiluAndMul
@@ -342,9 +347,26 @@ class Lfm2MoeShortConv(nn.Module):
 
         proj, _ = self.in_proj(hidden_states)
         B_gate, C_gate, x = proj.chunk(3, dim=-1)
-        Bx = B_gate * x
+        use_fused = can_dispatch_fused_lfm_short_conv(
+            B_gate,
+            C_gate,
+            x,
+            self.conv_weight,
+            self.conv_bias,
+            conv_state,
+        )
 
-        if forward_batch.forward_mode.is_decode():
+        if forward_batch.forward_mode.is_decode() and use_fused:
+            gated_conv_out = fused_lfm_short_conv_decode(
+                B_gate,
+                C_gate,
+                x,
+                self.conv_weight,
+                conv_state,
+                meta.cache_indices,
+            )
+        elif forward_batch.forward_mode.is_decode():
+            Bx = B_gate * x
             conv_out = causal_conv1d_update(
                 Bx,
                 conv_state,
@@ -353,11 +375,32 @@ class Lfm2MoeShortConv(nn.Module):
                 activation=None,
                 conv_state_indices=meta.cache_indices,
             )
+            gated_conv_out = C_gate * conv_out
         elif forward_batch.forward_mode.is_target_verify():
+            Bx = B_gate * x
             conv_out = shortconv_target_verify(
                 self, Bx, meta, forward_batch.spec_info.draft_token_num, "LFM2-MoE"
             )
+            gated_conv_out = C_gate * conv_out
+        elif (
+            use_fused
+            and meta.query_start_loc is not None
+            and meta.has_initial_state is not None
+            and forward_batch.extend_seq_lens_cpu
+        ):
+            gated_conv_out = fused_lfm_short_conv_prefill(
+                B_gate,
+                C_gate,
+                x,
+                self.conv_weight,
+                conv_state,
+                meta.query_start_loc,
+                meta.cache_indices,
+                meta.has_initial_state,
+                max(forward_batch.extend_seq_lens_cpu),
+            )
         else:
+            Bx = B_gate * x
             Bx_t = Bx.transpose(0, 1).contiguous()
             conv_out = causal_conv1d_fn(
                 Bx_t,
@@ -369,8 +412,9 @@ class Lfm2MoeShortConv(nn.Module):
                 conv_states=conv_state,
                 activation=None,
             ).transpose(0, 1)
+            gated_conv_out = C_gate * conv_out
 
-        output, _ = self.out_proj(C_gate * conv_out)
+        output, _ = self.out_proj(gated_conv_out)
         return output
 
 

test/registered/kernels/ops/mamba/test_lfm_short_conv.py

@@ -0,0 +1,242 @@
+"""Correctness coverage for the gate-fused LFM2 short convolution."""
+
+import unittest
+
+import torch
+
+from sglang.kernels.ops.mamba.causal_conv1d_triton import (
+    causal_conv1d_fn,
+    causal_conv1d_update,
+)
+from sglang.kernels.ops.mamba.lfm_short_conv import (
+    can_dispatch_fused_lfm_short_conv,
+    can_use_fused_lfm_short_conv,
+    fused_lfm_short_conv_decode,
+    fused_lfm_short_conv_prefill,
+)
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=15, stage="base-b", runner_config="1-gpu-large")
+
+PAD_SLOT_ID = -1
+requires_sm90 = unittest.skipUnless(
+    torch.cuda.is_available() and torch.cuda.get_device_capability() == (9, 0),
+    "requires SM90 (Hopper)",
+)
+
+
+def _inputs(tokens: int, dim: int, slots: int, seed: int):
+    generator = torch.Generator(device="cuda").manual_seed(seed)
+
+    def make(shape):
+        return torch.randn(
+            shape,
+            device="cuda",
+            dtype=torch.bfloat16,
+            generator=generator,
+        )
+
+    return (
+        make((tokens, dim)),
+        make((tokens, dim)),
+        make((tokens, dim)),
+        make((dim, 3)),
+        make((slots, dim, 2)),
+    )
+
+
+def _reference_prefill(
+    b,
+    c,
+    x,
+    weight,
+    state,
+    query_start_loc,
+    cache_indices,
+    has_initial_state,
+    seq_lens,
+):
+    bx = b * x
+    conv = causal_conv1d_fn(
+        bx.transpose(0, 1).contiguous(),
+        weight,
+        None,
+        conv_states=state,
+        query_start_loc=query_start_loc,
+        seq_lens_cpu=seq_lens,
+        cache_indices=cache_indices,
+        has_initial_state=has_initial_state,
+        activation=None,
+    ).transpose(0, 1)
+    return c * conv
+
+
+def _reference_decode(b, c, x, weight, state, cache_indices):
+    conv = causal_conv1d_update(
+        b * x,
+        state,
+        weight,
+        None,
+        activation=None,
+        conv_state_indices=cache_indices,
+    )
+    return c * conv
+
+
+@requires_sm90
+class TestLFMShortConv(CustomTestCase):
+    def test_prefill_matches_reference_for_variable_sequences(self):
+        seq_lens = [1, 2, 33, 65]
+        tokens = sum(seq_lens)
+        b, c, x, weight, initial_state = _inputs(tokens, 2051, 3, seed=1)
+        query_start_loc = torch.tensor(
+            [0, 1, 3, 36, tokens], device="cuda", dtype=torch.int32
+        )
+        cache_indices = torch.tensor(
+            [0, 1, PAD_SLOT_ID, 2], device="cuda", dtype=torch.int32
+        )
+        has_initial_state = torch.tensor(
+            [False, True, True, False], device="cuda", dtype=torch.bool
+        )
+
+        expected_state = initial_state.clone()
+        expected = _reference_prefill(
+            b,
+            c,
+            x,
+            weight,
+            expected_state,
+            query_start_loc,
+            cache_indices,
+            has_initial_state,
+            seq_lens,
+        )
+        actual_state = initial_state.clone()
+        actual = fused_lfm_short_conv_prefill(
+            b,
+            c,
+            x,
+            weight,
+            actual_state,
+            query_start_loc,
+            cache_indices,
+            has_initial_state,
+            max(seq_lens),
+        )
+
+        active_tokens = torch.cat(
+            [
+                torch.full(
+                    (seq_len,),
+                    cache_index != PAD_SLOT_ID,
+                    device="cuda",
+                    dtype=torch.bool,
+                )
+                for seq_len, cache_index in zip(seq_lens, cache_indices.tolist())
+            ]
+        )
+        torch.testing.assert_close(
+            actual[active_tokens].float(),
+            expected[active_tokens].float(),
+            rtol=0.02,
+            atol=0.05,
+        )
+        self.assertTrue(torch.equal(actual_state, expected_state))
+
+    def test_decode_matches_reference_and_preserves_padded_slots(self):
+        tokens, dim, slots = 5, 2051, 4
+        b, c, x, weight, initial_state = _inputs(tokens, dim, slots, seed=2)
+        cache_indices = torch.tensor(
+            [2, PAD_SLOT_ID, 0, 3, PAD_SLOT_ID],
+            device="cuda",
+            dtype=torch.int32,
+        )
+
+        expected_state = initial_state.clone()
+        expected = _reference_decode(b, c, x, weight, expected_state, cache_indices)
+        actual_state = initial_state.clone()
+        actual = fused_lfm_short_conv_decode(
+            b, c, x, weight, actual_state, cache_indices
+        )
+
+        active_tokens = cache_indices != PAD_SLOT_ID
+        torch.testing.assert_close(
+            actual[active_tokens].float(),
+            expected[active_tokens].float(),
+            rtol=0.02,
+            atol=0.05,
+        )
+        self.assertTrue(torch.equal(actual_state, expected_state))
+        self.assertTrue(torch.equal(actual_state[1], initial_state[1]))
+
+    def test_decode_cuda_graph_replay(self):
+        tokens, dim, slots = 4, 2048, 4
+        b, c, x, weight, initial_state = _inputs(tokens, dim, slots, seed=3)
+        cache_indices = torch.arange(tokens, device="cuda", dtype=torch.int32)
+
+        # Compile before capture, then prove the captured call reads live inputs and
+        # writes the caller-owned state/output on replay.
+        fused_lfm_short_conv_decode(
+            b, c, x, weight, initial_state.clone(), cache_indices
+        )
+        graph_state = initial_state.clone()
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            actual = fused_lfm_short_conv_decode(
+                b, c, x, weight, graph_state, cache_indices
+            )
+
+        b.add_(torch.tensor(0.25, device="cuda", dtype=b.dtype))
+        x.mul_(torch.tensor(0.75, device="cuda", dtype=x.dtype))
+        expected_state = initial_state.clone()
+        expected = _reference_decode(b, c, x, weight, expected_state, cache_indices)
+        graph_state.copy_(initial_state)
+        graph.replay()
+        torch.cuda.synchronize()
+
+        torch.testing.assert_close(
+            actual.float(), expected.float(), rtol=0.02, atol=0.05
+        )
+        self.assertTrue(torch.equal(graph_state, expected_state))
+
+    def test_dispatch_is_narrow(self):
+        b, c, x, weight, state = _inputs(4, 257, 4, seed=4)
+        self.assertTrue(can_use_fused_lfm_short_conv(b, c, x, weight, None, state))
+        self.assertFalse(
+            can_dispatch_fused_lfm_short_conv(b, c, x, weight, None, state)
+        )
+
+        b_2048, c_2048, x_2048, weight_2048, state_2048 = _inputs(4, 2048, 4, seed=5)
+        self.assertTrue(
+            can_dispatch_fused_lfm_short_conv(
+                b_2048, c_2048, x_2048, weight_2048, None, state_2048
+            )
+        )
+        self.assertFalse(
+            can_use_fused_lfm_short_conv(
+                b, c, x, weight, torch.zeros(257, device="cuda"), state
+            )
+        )
+        self.assertFalse(
+            can_use_fused_lfm_short_conv(
+                b, c, x, weight[:, :2].contiguous(), None, state
+            )
+        )
+        self.assertFalse(
+            can_use_fused_lfm_short_conv(b, c.float(), x, weight, None, state)
+        )
+        self.assertFalse(
+            can_use_fused_lfm_short_conv(
+                b.transpose(0, 1),
+                c.transpose(0, 1),
+                x.transpose(0, 1),
+                weight,
+                None,
+                state,
+            )
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()