quack-kernels 0.1.9__py3-none-any.whl → 0.1.11__py3-none-any.whl

quack/fast_math.py

@@ -0,0 +1,97 @@
+# Copyright (c) 2025, Tri Dao.
+
+from typing import Tuple
+
+import cutlass
+import cutlass.cute as cute
+from cutlass import Int32, Uint32
+from cutlass.cutlass_dsl import T, dsl_user_op
+from cutlass._mlir.dialects import llvm
+
+
+@cute.jit
+def clz(x: Int32) -> Int32:
+    # for i in cutlass.range_constexpr(32):
+    #     if (1 << (31 - i)) & x:
+    #         return Int32(i)
+    # return Int32(32)
+    # Early exit is not supported yet
+    res = Int32(32)
+    done = False
+    for i in cutlass.range(32):
+        if ((1 << (31 - i)) & x) and not done:
+            res = Int32(i)
+            done = True
+    return res
+
+
+def find_log2(x: Int32) -> Int32:
+    a: Int32 = Int32(31 - clz(x))
+    return a + ((x & (x - 1)) != 0)  # Round up, add 1 if not a power of 2.
+
+
+@dsl_user_op
+def umulhi(a: Int32, b: Int32, *, loc=None, ip=None) -> Uint32:
+    return Uint32(
+        llvm.inline_asm(
+            T.i32(),
+            [Int32(a).ir_value(loc=loc, ip=ip), Int32(b).ir_value(loc=loc, ip=ip)],
+            "mul.hi.u32 $0, $1, $2;",
+            "=r,r,r",
+            has_side_effects=False,
+            is_align_stack=False,
+            asm_dialect=llvm.AsmDialect.AD_ATT,
+        )
+    )
+
+
+class FastDivmod:
+    def __init__(
+        self, divisor: Int32, multipler: Uint32, shift_right: Uint32, *, loc=None, ip=None
+    ):
+        self.divisor = divisor
+        self.multiplier = multipler
+        self.shift_right = shift_right
+        self._loc = loc
+
+    # called by host
+    @staticmethod
+    def create(divisor: Int32, *, loc=None, ip=None) -> "FastDivmod":
+        """Construct the FastDivmod object, in host code.
+        This precomputes some values based on the divisor and is computationally expensive.
+        """
+        p = Uint32(31 + find_log2(divisor))
+        divisor_u32 = Uint32(divisor)
+        multiplier = Uint32(((cutlass.Uint64(1) << p) + divisor_u32 - 1) // divisor_u32)
+        shift_right = Uint32(p - 32)
+        return FastDivmod(divisor, multiplier, shift_right, loc=loc, ip=ip)
+
+    @cute.jit
+    def div(self, dividend: Int32) -> Int32:
+        return (
+            Int32(umulhi(dividend, self.multiplier) >> self.shift_right)
+            if self.divisor != 1
+            else dividend
+        )
+
+    def divmod(self, dividend: Int32) -> Tuple[Int32, Int32]:
+        quotient = self.div(dividend)
+        remainder = dividend - quotient * self.divisor
+        return quotient, remainder
+
+    def __extract_mlir_values__(self):
+        values, self._values_pos = [], []
+        for obj in [self.divisor, self.multiplier, self.shift_right]:
+            obj_values = cutlass.extract_mlir_values(obj)
+            values += obj_values
+            self._values_pos.append(len(obj_values))
+        return values
+
+    def __new_from_mlir_values__(self, values):
+        obj_list = []
+        for obj, n_items in zip(
+            [self.divisor, self.multiplier, self.shift_right], self._values_pos
+        ):
+            obj_list.append(cutlass.new_from_mlir_values(obj, values[:n_items]))
+            values = values[n_items:]
+        return FastDivmod(*(tuple(obj_list)), loc=self._loc)

quack/gemm_config.py

@@ -0,0 +1,61 @@
+# Copyright (C) 2025, Tri Dao.
+import itertools
+from typing import Optional
+from pydantic import BaseModel
+
+
+class GemmConfig(BaseModel, frozen=True):
+    tile_m: int = 256
+    tile_n: int = 128
+    cluster_m: int = 2
+    cluster_n: int = 1
+    swap_ab: bool = False
+    pingpong: bool = False
+    raster_order: int = 2
+    max_swizzle_size: int = 1
+
+
+def get_all_configs(
+    epilogue: Optional[str],
+    tune_pingpong=True,
+    tune_raster_order=True,
+) -> list[GemmConfig]:
+    tile_n_vals = [128, 144, 160, 176, 192, 208]
+    tile_mn_vals = [(256, tile_n) for tile_n in tile_n_vals]
+    if epilogue in ["swiglu"]:
+        tile_mn_vals = [(m, n) for m, n in tile_mn_vals if n % 32 == 0]
+    cluster = [(1, 1), (1, 2), (2, 1)]
+    # cluster = [(1, 2), (2, 1)]
+    if epilogue in ["lse"]:
+        cluster = [(1, 2), (2, 1)]
+    swap_ab_vals = [False, True]
+    if epilogue in ["lse", "swiglu"]:
+        swap_ab_vals = [False]
+    pingpong_vals = [False, True] if tune_pingpong else [False]
+    raster_swizzle = (
+        [(0, 1)]
+        if not tune_raster_order
+        else [(1, 1), (1, 2), (1, 4), (1, 8), (2, 1), (2, 2), (2, 4), (2, 8)]
+    )
+    return [
+        GemmConfig(
+            tile_m=tile_m if not pingpong else 128,
+            tile_n=tile_n,
+            cluster_m=cluster_m,
+            cluster_n=cluster_n,
+            swap_ab=swap_ab,
+            pingpong=pingpong,
+            raster_order=raster_order,
+            max_swizzle_size=max_swizzle_size,
+        )
+        for (tile_m, tile_n), (cluster_m, cluster_n), swap_ab, pingpong, (
+            raster_order,
+            max_swizzle_size,
+        ) in itertools.product(
+            tile_mn_vals,
+            cluster,
+            swap_ab_vals,
+            pingpong_vals,
+            raster_swizzle,
+        )
+    ]

quack/gemm_interface.py

@@ -0,0 +1,321 @@
+# Copyright (c) 2025, Tri Dao
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+
+from quack.gemm_config import GemmConfig, get_all_configs
+
+from quack.autotuner import autotune, AutotuneConfig
+from quack.lse import logsumexp
+
+
+def gemm_swiglu_out_ref(
+    A: Tensor, B: Tensor, out: Optional[Tensor], store_preact: bool
+) -> (Tensor, Tensor):
+    preact = torch.mm(A, B)
+    out_ = F.silu(preact[..., ::2]) * preact[..., 1::2]
+    if out is not None:
+        out.copy_(out_)
+    else:
+        out = out_
+    if not store_preact:
+        preact = None
+    return out, preact
+
+
+@autotune(
+    configs=[AutotuneConfig(config=c) for c in get_all_configs(epilogue=None)], key=["sm_carveout"]
+)
+def gemm_tuned(
+    A: Tensor,
+    B: Tensor,
+    sm_carveout: int = 0,
+    config: Optional[GemmConfig] = None,
+) -> (Tensor, Optional[Tensor]):
+    if config is None:
+        config = GemmConfig(
+            tile_m=256,
+            tile_n=192,
+            cluster_m=2,
+            cluster_n=1,
+            pingpong=False,
+            raster_order=2,
+            max_swizzle_size=1,
+        )
+    out = torch.ops.quack.gemm_impl.default(
+        A if not config.swap_ab else B.T,
+        B if not config.swap_ab else A.T,
+        sm_carveout,
+        config.tile_m,
+        config.tile_n,
+        config.cluster_m,
+        config.cluster_n,
+        not config.swap_ab,  # C_rowmajor
+        config.pingpong,
+        config.raster_order,
+        config.max_swizzle_size,
+    )
+    return out if not config.swap_ab else out.T
+
+
+@torch.library.custom_op("quack::gemm", mutates_args=(), device_types="cuda")
+def gemm(A: Tensor, B: Tensor, sm_carveout: int = 0) -> Tensor:
+    return gemm_tuned(A, B, sm_carveout)
+
+
+@torch.library.register_fake("quack::gemm")
+def gemm_ref(A: Tensor, B: Tensor, sm_carveout: int = 0) -> Tensor:
+    return torch.mm(A, B)
+
+
+@autotune(configs=[AutotuneConfig(config=c) for c in get_all_configs("add")])
+def gemm_add_tuned(
+    A: Tensor,
+    B: Tensor,
+    C: Tensor,
+    config: Optional[GemmConfig] = None,
+) -> (Tensor, Optional[Tensor]):
+    if config is None:
+        config = GemmConfig(
+            tile_m=256,
+            tile_n=192,
+            cluster_m=2,
+            cluster_n=1,
+            pingpong=False,
+            raster_order=2,
+            max_swizzle_size=1,
+        )
+    out = torch.ops.quack.gemm_add_impl.default(
+        A if not config.swap_ab else B.T,
+        B if not config.swap_ab else A.T,
+        C if not config.swap_ab else C.T,
+        config.tile_m,
+        config.tile_n,
+        config.cluster_m,
+        config.cluster_n,
+        config.pingpong,
+        config.raster_order,
+        config.max_swizzle_size,
+    )
+    return out if not config.swap_ab else out.T
+
+
+@torch.library.custom_op("quack::gemm_add", mutates_args=(), device_types="cuda")
+def gemm_add(A: Tensor, B: Tensor, C: Tensor) -> Tensor:
+    return gemm_add_tuned(A, B, C)
+
+
+@torch.library.register_fake("quack::gemm_add")
+def gemm_add_ref(A: Tensor, B: Tensor, C: Tensor) -> Tensor:
+    return C + torch.mm(A, B)
+
+
+@torch.library.custom_op("quack::gemm_add_t", mutates_args=(), device_types="cuda")
+def gemm_t_add(A: Tensor, B: Tensor, C: Tensor) -> Tensor:
+    return gemm_add_tuned(A, B.T, C)
+
+
+@torch.library.register_fake("quack::gemm_add_t")
+def gemm_t_add_ref(A: Tensor, B: Tensor, C: Tensor) -> Tensor:
+    return gemm_add_ref(A, B.T, C)
+
+
+@autotune(
+    configs=[AutotuneConfig(config=c) for c in get_all_configs("swiglu")], key=["store_preact"]
+)
+def gemm_swiglu_tuned(
+    A: Tensor,
+    B: Tensor,
+    store_preact: bool = True,
+    config: Optional[GemmConfig] = None,
+) -> (Tensor, Optional[Tensor]):
+    if config is None:
+        config = GemmConfig(
+            tile_m=256,
+            tile_n=192,
+            cluster_m=2,
+            cluster_n=1,
+            pingpong=False,
+            raster_order=2,
+            max_swizzle_size=1,
+        )
+    # out, preact
+    return torch.ops.quack.gemm_swiglu_impl.default(
+        A,
+        B,
+        store_preact,
+        config.tile_m,
+        config.tile_n,
+        config.cluster_m,
+        config.cluster_n,
+        config.pingpong,
+        config.raster_order,
+        config.max_swizzle_size,
+    )
+
+
+# Specifying the schema manually here since torch.library._infer_schema doesn't work when return
+# type is a tuple of Tensor
+@torch.library.custom_op(
+    "quack::gemm_swiglu",
+    mutates_args=(),
+    device_types="cuda",
+    schema="(Tensor A, Tensor B, bool store_preact) -> (Tensor, Tensor)",
+)
+def gemm_swiglu(A: Tensor, B: Tensor, store_preact: bool = True) -> (Tensor, Tensor):
+    return gemm_swiglu_tuned(A, B, store_preact=store_preact)
+
+
+@torch.library.register_fake("quack::gemm_swiglu")
+def gemm_swiglu_ref(A: Tensor, B: Tensor, store_preact: bool) -> (Tensor, Tensor):
+    return gemm_swiglu_out_ref(A, B, None, store_preact)
+
+
+# @torch.library.custom_op("quack::gemm_swiglu_t", mutates_args=(), device_types="cuda",
+#                          schema="(Tensor A, Tensor B, bool store_preact) -> (Tensor, Tensor)")
+# def gemm_swiglu_t(A: Tensor, B: Tensor, store_preact: bool = True) -> (Tensor, Tensor):
+#     return gemm_swiglu_tuned(A, B.T, store_preact=store_preact)
+
+
+# @torch.library.register_fake("quack::gemm_swiglu_t")
+# def gemm_swiglu_t_ref(A: Tensor, B: Tensor, store_preact: bool) -> (Tensor, Tensor):
+#     return gemm_swiglu_ref(A, B.T, store_preact)
+
+
+@autotune(
+    configs=[AutotuneConfig(config=c) for c in get_all_configs("dswiglu")], key=["sm_carveout"]
+)
+def gemm_dswiglu_tuned(
+    A: Tensor,
+    B: Tensor,
+    preact: Tensor,
+    sm_carveout: int = 0,
+    config: Optional[GemmConfig] = None,
+) -> (Tensor, Tensor):
+    if config is None:
+        config = GemmConfig(
+            tile_m=128,
+            tile_n=192,
+            cluster_m=2,
+            cluster_n=1,
+            pingpong=True,
+            raster_order=2,
+            max_swizzle_size=1,
+        )
+    out, postact = torch.ops.quack.gemm_dswiglu_impl.default(
+        A if not config.swap_ab else B.T,
+        B if not config.swap_ab else A.T,
+        preact if not config.swap_ab else preact.T,
+        sm_carveout,
+        config.tile_m,
+        config.tile_n,
+        config.cluster_m,
+        config.cluster_n,
+        not config.swap_ab,  # C_rowmajor
+        config.pingpong,
+        config.raster_order,
+        config.max_swizzle_size,
+    )
+    return (out, postact) if not config.swap_ab else (out.T, postact.T)
+
+
+# Specifying the schema manually here since torch.library._infer_schema doesn't work when return
+# type is a tuple of Tensor
+@torch.library.custom_op(
+    "quack::gemm_dswiglu",
+    mutates_args=(),
+    device_types="cuda",
+    schema="(Tensor A, Tensor B, Tensor preact, int sm_carveout=0) -> (Tensor, Tensor)",
+)
+def gemm_dswiglu(A: Tensor, B: Tensor, preact: Tensor, sm_carveout: int = 0) -> (Tensor, Tensor):
+    return gemm_dswiglu_tuned(A, B, preact, sm_carveout)
+
+
+@torch.library.register_fake("quack::gemm_dswiglu")
+def gemm_dswiglu_ref(
+    A: Tensor, B: Tensor, preact: Tensor, sm_carveout: int = 0
+) -> (Tensor, Tensor):
+    # A: (M, K), B: (K, N), preact: (M, 2 * N)
+    dout = torch.mm(A, B)
+    p0, p1 = preact[..., ::2], preact[..., 1::2]
+    sigmoid = torch.sigmoid(p0)
+    silu = F.silu(p0)
+    postact = silu * p1
+    d0 = sigmoid * (1 + p0 * (1 - sigmoid)) * p1 * dout
+    d1 = F.silu(p0) * dout
+    out = torch.stack([d0, d1], dim=-1).reshape(d0.shape[:-1] + (2 * d0.shape[-1],))
+    return out, postact
+
+
+@autotune(configs=[AutotuneConfig(config=c) for c in get_all_configs("lse")])
+def gemm_lse_tuned(
+    A: Tensor,
+    B: Tensor,
+    softcap: float = 0.0,
+    config: Optional[GemmConfig] = None,
+) -> (Tensor, Tensor):
+    if config is None:
+        config = GemmConfig(
+            tile_m=256,
+            tile_n=192,
+            cluster_m=2,
+            cluster_n=1,
+            pingpong=False,
+            raster_order=2,
+            max_swizzle_size=1,
+        )
+    out, lse_partial = torch.ops.quack.gemm_lse_impl.default(
+        A,
+        B,
+        None,  # bias
+        softcap,
+        config.tile_m,
+        config.tile_n,
+        config.cluster_m,
+        config.cluster_n,
+        config.pingpong,
+        config.raster_order,
+        config.max_swizzle_size,
+    )
+    lse = logsumexp(lse_partial)
+    return out, lse
+
+
+@torch.library.custom_op(
+    "quack::gemm_lse",
+    mutates_args=(),
+    device_types="cuda",
+    schema="(Tensor A, Tensor B, float softcap=0.0) -> (Tensor, Tensor)",
+)
+def gemm_lse(A: Tensor, B: Tensor, softcap: float = 0.0) -> (Tensor, Tensor):
+    return gemm_lse_tuned(A, B, softcap)
+
+
+@torch.library.register_fake("quack::gemm_lse")
+def gemm_lse_ref(A: Tensor, B: Tensor, softcap: float = 0.0) -> (Tensor, Tensor):
+    # A: (M, K), B: (K, N)
+    out = torch.mm(A, B)
+    if softcap > 0:
+        out_fp32 = torch.tanh(out.to(torch.float32) / softcap) * softcap
+        out = out_fp32.to(out.dtype)
+    else:
+        out_fp32 = out.to(torch.float32)
+    lse = torch.logsumexp(out_fp32, dim=-1)
+    return out, lse
+
+
+@torch.library.custom_op(
+    "quack::gemm_lse_t",
+    mutates_args=(),
+    device_types="cuda",
+    schema="(Tensor A, Tensor B, float softcap=0.0) -> (Tensor, Tensor)",
+)
+def gemm_lse_t(A: Tensor, B: Tensor, softcap: float = 0.0) -> (Tensor, Tensor):
+    return gemm_lse_tuned(A, B.T, softcap)
+
+
+@torch.library.register_fake("quack::gemm_lse_t")
+def gemm_lse_t_ref(A: Tensor, B: Tensor, softcap: float = 0.0) -> (Tensor, Tensor):
+    return gemm_lse_ref(A, B.T, softcap)

quack/linear.py

@@ -0,0 +1,176 @@
+# Copyright (c) 2025, Tri Dao
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+from torch.amp import custom_fwd, custom_bwd
+
+
+from gemm_cublas import gemm as gemm_cb, gemm_add_ as gemm_add_cb_
+# from gemm_cublas.interface import gemm_tuned as gemm_cb, gemm_add_tuned_ as gemm_add_cb_
+
+from quack import gemm, gemm_lse  # TODO: implement these
+
+
+def linear_fwd_convert_type(*tensors):
+    autocast_dtype = torch.get_autocast_dtype("cuda")
+    if torch.is_autocast_enabled():
+        tensors = tuple(t.to(dtype=autocast_dtype) for t in tensors)
+    return tensors
+
+
+def linear_fwd_postprocess(ctx, x, weight, weight_og, needs_input_grad, needs_weight_grad):
+    if not needs_input_grad:
+        weight, weight_og = None, None
+    if not needs_weight_grad:
+        x = None
+    ctx.save_for_backward(x, weight, weight_og if ctx.fuse_grad_accum else None)
+
+
+def linear_bwd_compute_input_grad(ctx, dout, weight, use_tuned_gemm=True, sm_carveout=0):
+    if ctx.needs_input_grad[0]:
+        assert weight is not None
+        # return gemm(dout, weight) if use_tuned_gemm else (dout @ weight)
+        return (
+            gemm(dout, weight, sm_carveout=sm_carveout)
+            if use_tuned_gemm
+            else gemm_cb(dout, weight, sm_carveout=sm_carveout)
+        )
+    else:
+        return None
+
+
+def linear_bwd_compute_weight_grad(ctx, dout, x, weight_og, sm_carveout=0):
+    if ctx.needs_input_grad[1]:
+        assert x is not None
+        x = x.reshape(-1, x.shape[-1])
+        # fuse_grad_accum is not compatible with torch.compile
+        if not ctx.fuse_grad_accum or weight_og.grad is None or torch.compiler.is_compiling():
+            dweight = gemm_cb(dout.T, x, out_dtype=ctx.weight_dtype, sm_carveout=sm_carveout)
+        else:
+            # print("Using fuse grad accum in Linear", dout.shape, x.shape, weight_og.grad.shape)
+            gemm_add_cb_(dout.T, x, weight_og.grad, sm_carveout=sm_carveout)
+            dweight = weight_og.grad
+            weight_og.grad = None  # So that pytorch doesn't add dweight to weight_og.grad again
+    else:
+        dweight = None
+    return dweight
+
+
+class LinearFunc(torch.autograd.Function):
+    @staticmethod
+    @custom_fwd(device_type="cuda")
+    def forward(ctx, x, weight, fuse_grad_accum=False):
+        """
+        x: (..., in_features)
+        weight: (out_features, in_features)
+        out: (..., out_features)
+        """
+        ctx.weight_dtype = weight.dtype
+        ctx.fuse_grad_accum = fuse_grad_accum
+        weight_og = weight
+        x, weight = linear_fwd_convert_type(x, weight)
+        batch_shape = x.shape[:-1]
+        x = x.reshape(-1, x.shape[-1])
+        # out = F.linear(x, weight)
+        out = gemm(x, weight.T)
+        linear_fwd_postprocess(
+            ctx,
+            x,
+            weight,
+            weight_og,
+            needs_input_grad=ctx.needs_input_grad[0],
+            needs_weight_grad=ctx.needs_input_grad[1],
+        )
+        return out.reshape(*batch_shape, out.shape[-1])
+
+    @staticmethod
+    @custom_bwd(device_type="cuda")
+    def backward(ctx, dout):
+        """
+        dout: (..., out_features)
+        """
+        x, weight, weight_og = ctx.saved_tensors  # weight_og is None if not ctx.fuse_grad_accum
+        batch_shape = dout.shape[:-1]
+        dout = dout.reshape(-1, dout.shape[-1])
+        dx = linear_bwd_compute_input_grad(ctx, dout, weight, use_tuned_gemm=True)
+        dx = dx.reshape(*batch_shape, dx.shape[-1]) if dx is not None else None
+        dweight = linear_bwd_compute_weight_grad(ctx, dout, x, weight_og)
+        return dx, dweight, None
+
+
+def linear_func(x, weight, fuse_grad_accum=False):
+    return LinearFunc.apply(x, weight, fuse_grad_accum)
+
+
+class LinearLSEFunc(torch.autograd.Function):
+    @staticmethod
+    @custom_fwd(device_type="cuda")
+    def forward(ctx, x, weight, fuse_grad_accum=False):
+        """
+        x: (..., in_features)
+        weight: (out_features, in_features)
+        out: (..., out_features)
+        """
+        needs_weight_grad = weight.requires_grad
+        needs_input_grad = x.requires_grad
+        ctx.weight_dtype = weight.dtype
+        ctx.fuse_grad_accum = fuse_grad_accum
+        weight_og = weight
+        x, weight = linear_fwd_convert_type(x, weight)
+        batch_shape = x.shape[:-1]
+        x = x.reshape(-1, x.shape[-1])
+        out, lse = gemm_lse(x, weight.T)
+        lse = lse.reshape(*batch_shape)
+        linear_fwd_postprocess(ctx, x, weight, weight_og, needs_weight_grad, needs_input_grad)
+        ctx.mark_non_differentiable(lse)
+        return out.reshape(*batch_shape, out.shape[-1]), lse
+
+    @staticmethod
+    @custom_bwd(device_type="cuda")
+    def backward(ctx, dout, dlse_ignored):
+        """
+        dout: (..., out_features)
+        """
+        x, weight, weight_og = ctx.saved_tensors  # weight_og is None if not ctx.fuse_grad_accum
+        batch_shape = dout.shape[:-1]
+        dout = dout.reshape(-1, dout.shape[-1])
+        # cuBLAS seems faster for this so we just use it instead of cutlass gemm
+        dx = linear_bwd_compute_input_grad(ctx, dout, weight, use_tuned_gemm=False)
+        dx = dx.reshape(*batch_shape, dx.shape[-1]) if dx is not None else None
+        dweight = linear_bwd_compute_weight_grad(ctx, dout, x, weight_og)
+        return dx, dweight, None
+
+
+def linear_lse_func(x, weight, fuse_grad_accum=False):
+    return LinearLSEFunc.apply(x, weight, fuse_grad_accum)
+
+
+class Linear(nn.Linear):
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        bias: bool = False,
+        device=None,
+        dtype=None,
+        fuse_grad_accum: bool = False,
+    ) -> None:
+        super().__init__(in_features, out_features, bias=bias, device=device, dtype=dtype)
+        self.fuse_grad_accum = fuse_grad_accum
+
+    def forward(self, input: Tensor) -> Tensor:
+        if self.bias is None and input.is_cuda:
+            return linear_func(input, self.weight, fuse_grad_accum=self.fuse_grad_accum)
+        else:
+            return F.linear(input, self.weight, self.bias)
+
+
+class LinearLSE(Linear):
+    def forward(self, input: Tensor) -> Tensor:
+        if self.bias is None and input.is_cuda:
+            return linear_lse_func(input, self.weight, fuse_grad_accum=self.fuse_grad_accum)
+        else:
+            out = F.linear(input, self.weight, self.bias)
+            lse = torch.logsumexp(out, dim=-1)
+            return out, lse