quack-kernels 0.1.5__tar.gz → 0.1.6__tar.gz

{quack_kernels-0.1.5/quack_kernels.egg-info → quack_kernels-0.1.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: quack-kernels
-Version: 0.1.5
+Version: 0.1.6
 Requires-Python: >=3.9
 License-File: LICENSE
 Requires-Dist: nvidia-cutlass-dsl==4.1.0.dev0

{quack_kernels-0.1.5 → quack_kernels-0.1.6}/README.md

@@ -19,6 +19,7 @@ pip install quack-kernels
 - 🦆 RMSNorm forward
 - 🦆 Softmax forward + backward
 - 🦆 Cross entropy forward + backward
+- 🦆 Layernorm forward
 
 Upcoming:
 - 🦆 RMSNorm backward
@@ -30,13 +31,23 @@ Upcoming:
 from quack import rmsnorm, softmax, cross_entropy
 ```
 
-## Caveats 🦆⚠️
+## Documentations
 
-**Tensor Size Limitation**: We currently only support tensors ≤ 4GB due to CuTe-DSL using int32 for indexing.
+[2025-07-10] We have a comprehensive
+[blogpost](media/2025-07-10-membound-sol.md) on how to get memory-bound kernels
+to speed-of-light, right in the comfort of Python thanks to the [CuTe-DSL](https://docs.nvidia.com/cutlass/media/docs/pythonDSL/cute_dsl_general/dsl_introduction.html).
 
-🦆 **Workaround**: For larger tensors, split your input tensors into chunks of
-size ≤ 4GB each. We will implement this automatic chunking in the pytorch part
-of the code in the near future, but if you need it in the meantime, we welcome contributions!
+## Performance
+
+<div align="center">
+<figure>
+  <img
+  src="media/bf16_kernel_benchmarks_single_row.svg"
+  >
+</figure>
+</div>
+
+See our [blogpost](media/2025-07-10-membound-sol.md) for the details.
 
 ## Development
 

{quack_kernels-0.1.5 → quack_kernels-0.1.6}/quack/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.1.5"
+__version__ = "0.1.6"
 
 from quack.rmsnorm import rmsnorm
 from quack.softmax import softmax

{quack_kernels-0.1.5 → quack_kernels-0.1.6}/quack/cross_entropy.py

@@ -104,7 +104,10 @@ class CrossEntropy(ReductionBase):
         shape: cute.Shape = mX.shape
         idX = cute.make_identity_tensor(shape)
         # slice for CTAs
-        gX, cX = [cute.local_tile(mT, tiler_mn, (bidx, cluster_y)) for mT in (mX, idX)]
+        # We use domain_offset_i64 to deal with tensors larger than 2^31 elements
+        mX_off = utils.domain_offset_i64((bidx * tiler_mn[0], 0), mX)
+        gX = cute.local_tile(mX_off, tiler_mn, (0, cluster_y))
+        cX = cute.local_tile(idX, tiler_mn, (bidx, cluster_y))
 
         smem = cutlass.utils.SmemAllocator()
         sX = smem.allocate_tensor(
@@ -150,7 +153,9 @@ class CrossEntropy(ReductionBase):
 
         target_logit = cute.Float32.zero
         if row < shape[0] and tXcX[0][1] == 0:
-            target_logit = cute.Float32(mX[row, target])
+            # Use Int64 for indexing to deal with large tensors
+            mX_off = utils.domain_offset_i64((row, 0), mX)
+            target_logit = cute.Float32(mX_off[0, target])
 
         threads_per_row = tv_layout.shape[0][0]
         if cutlass.const_expr(not self.online_softmax):
@@ -363,11 +368,10 @@ class CrossEntropyBackward:
         )
 
         idX = cute.make_identity_tensor(shape)
-
-        gX, gdX, cX, gTarget, gDLoss, gLse = [
-            cute.local_tile(mT, tiler_mn, (bidx, bidy))
-            for mT in (mX, mdX, idX, mTarget, mDLoss, mLSE)
-        ]
+        # We use domain_offset_i64 to deal with tensors larger than 2^31 elements
+        mX, mdX = [utils.domain_offset_i64((bidx * tiler_mn[0], 0), mT) for mT in (mX, mdX)]
+        gX, gdX = [cute.local_tile(mT, tiler_mn, (0, bidy)) for mT in (mX, mdX)]
+        cX = cute.local_tile(idX, tiler_mn, (bidx, bidy))
 
         copy_atom_load_X = cute.make_copy_atom(
             cute.nvgpu.CopyUniversalOp(), gX.element_type, num_bits_per_copy=128

quack_kernels-0.1.6/quack/layernorm.py

@@ -0,0 +1,351 @@
+# Copyright (c) 2025, Wentao Guo, Ted Zadouri, Tri Dao.
+
+
+import torch
+from typing import Optional
+
+import cuda.bindings.driver as cuda
+
+import cutlass
+import cutlass.cute as cute
+from cutlass.cute.runtime import from_dlpack
+import quack.utils as utils
+from quack.reduction_base import ReductionBase, torch2cute_dtype_map
+
+
+class LayerNorm(ReductionBase):
+    def __init__(self, dtype: cutlass.Numeric, N: int):
+        super().__init__(dtype, N, stage=2)  # 2 stages for mean and var
+        self.reload_from = None if N <= 16384 else "smem"
+        self.delay_w_load = False
+
+    def _calculate_threads_per_row(self):
+        N = self.N
+        return (
+            8
+            if N <= 64
+            else (
+                16
+                if N <= 128
+                else (32 if N <= 3072 else (64 if N <= 6144 else (128 if N <= 16384 else 256)))
+            )
+        )
+
+    def _set_cluster_n(self):
+        N = self.N
+        # cluster_n = 4 is faster and cluster_n = 2 for N=64k for some reason
+        # Similarly cluster_n = 8 is faster for N=128k
+        if cutlass.const_expr(self.dtype.width == 16):
+            cluster_n = (
+                1
+                if N <= 16 * 1024
+                else (
+                    2
+                    if N <= 32 * 1024
+                    else (4 if N <= 64 * 1024 else (8 if N <= 128 * 1024 else 16))
+                )
+            )
+        else:  # fp32
+            cluster_n = (
+                1
+                if N <= 32 * 1024
+                else (
+                    2
+                    if N <= 64 * 1024
+                    else (4 if N <= 128 * 1024 else (8 if N <= 256 * 1024 else 16))
+                )
+            )
+        self.cluster_n = cluster_n
+
+    @cute.jit
+    def __call__(
+        self,
+        mX: cute.Tensor,
+        mW: cute.Tensor,
+        mO: cute.Tensor,
+        mRstd: Optional[cute.Tensor],
+        mMean: Optional[cute.Tensor],
+        stream: cuda.CUstream,
+        eps: cutlass.Float32 = 1e-6,
+    ):
+        assert mX.element_type == self.dtype
+        assert mO.element_type == self.dtype
+        self._set_cluster_n()
+        tiler_mn, tv_layout = self._get_tv_layout()
+        num_threads = cute.size(tv_layout, mode=[0])
+        num_warps = num_threads // cute.arch.WARP_SIZE
+        mW_expanded_layout = cute.prepend(mW.layout, cute.make_layout((tiler_mn[0],), stride=(0,)))
+        mW = cute.make_tensor(mW.iterator, mW_expanded_layout)
+        if cutlass.const_expr(mRstd is not None):
+            mRstd_expanded_layout = cute.append(
+                mRstd.layout, cute.make_layout((self.N,), stride=(0,))
+            )
+            mRstd = cute.make_tensor(mRstd.iterator, mRstd_expanded_layout)
+        if cutlass.const_expr(mMean is not None):
+            mMean_expanded_layout = cute.append(
+                mMean.layout, cute.make_layout((self.N,), stride=(0,))
+            )
+            mMean = cute.make_tensor(mMean.iterator, mMean_expanded_layout)
+        self.kernel(mX, mW, mO, mRstd, mMean, eps, tv_layout, tiler_mn, self.reload_from).launch(
+            grid=[cute.ceil_div(mX.shape[0], tiler_mn[0]), self.cluster_n, 1],
+            block=[num_threads, 1, 1],
+            cluster=[1, self.cluster_n, 1] if cutlass.const_expr(self.cluster_n > 1) else None,
+            smem=self._smem_size_in_bytes(tiler_mn, num_warps),
+            stream=stream,
+        )
+
+    @cute.kernel
+    def kernel(
+        self,
+        mX: cute.Tensor,
+        mW: cute.Tensor,
+        mO: cute.Tensor,
+        mRstd: Optional[cute.Tensor],
+        mMean: Optional[cute.Tensor],
+        eps: cute.Float32,
+        tv_layout: cute.Layout,
+        tiler_mn: cute.Shape,
+        reload_from: cutlass.Constexpr = None,
+        delay_w_load: cutlass.Constexpr = False,
+    ):
+        tidx, _, _ = cute.arch.thread_idx()
+        bidx, _, _ = cute.arch.block_idx()
+        if cutlass.const_expr(self.cluster_n > 1):
+            cluster_y = cute.arch.block_idx()[1]
+        else:
+            cluster_y = cutlass.const_expr(0)
+
+        smem = cutlass.utils.SmemAllocator()
+        sX = smem.allocate_tensor(
+            mX.element_type, cute.make_ordered_layout(tiler_mn, order=(1, 0)), byte_alignment=16
+        )
+        reduction_buffer, mbar_ptr = self._allocate_reduction_buffer_and_mbar(smem, tv_layout)
+
+        shape = mX.shape
+        idX = cute.make_identity_tensor(shape)
+        # slice for CTAs
+        # We use domain_offset_i64 to deal with tensors larger than 2^31 elements
+        mX, mO = [utils.domain_offset_i64((bidx * tiler_mn[0], 0), mT) for mT in (mX, mO)]
+        gX, gO = [cute.local_tile(mT, tiler_mn, (0, cluster_y)) for mT in (mX, mO)]
+        cX = cute.local_tile(idX, tiler_mn, (bidx, cluster_y))
+        gW = cute.local_tile(mW, tiler_mn, (0, cluster_y))
+        gRstd = (
+            cute.local_tile(mRstd, tiler_mn, (bidx, cluster_y))
+            if cutlass.const_expr(mRstd is not None)
+            else None
+        )
+        gMean = (
+            cute.local_tile(mMean, tiler_mn, (bidx, cluster_y))
+            if cutlass.const_expr(mMean is not None)
+            else None
+        )
+
+        # declare the atoms which will be used later for memory copy
+        copy_atom_load_X = cute.make_copy_atom(
+            cute.nvgpu.CopyUniversalOp(), mX.element_type, num_bits_per_copy=128
+        )
+        copy_atom_load_X_async = cute.make_copy_atom(
+            cute.nvgpu.cpasync.CopyG2SOp(), mX.element_type, num_bits_per_copy=128
+        )
+        copy_atom_load_W = cute.make_copy_atom(
+            cute.nvgpu.CopyUniversalOp(), mW.element_type, num_bits_per_copy=128
+        )
+        copy_atom_store_O = cute.make_copy_atom(
+            cute.nvgpu.CopyUniversalOp(), mO.element_type, num_bits_per_copy=128
+        )
+
+        thr_copy_X = cute.make_tiled_copy(copy_atom_load_X_async, tv_layout, tiler_mn).get_slice(
+            tidx
+        )
+        thr_copy_W = cute.make_tiled_copy(copy_atom_load_W, tv_layout, tiler_mn).get_slice(tidx)
+        thr_copy_O = cute.make_tiled_copy(copy_atom_store_O, tv_layout, tiler_mn).get_slice(tidx)
+
+        tWgW = thr_copy_W.partition_S(gW)
+        tXgX = thr_copy_X.partition_S(gX)
+        tXsX = thr_copy_X.partition_D(sX)
+        tXgO = thr_copy_O.partition_D(gO)
+        tXrRstd = thr_copy_O.partition_D(gRstd) if cutlass.const_expr(mRstd is not None) else None
+        tXrMean = thr_copy_O.partition_D(gMean) if cutlass.const_expr(mMean is not None) else None
+        tXcX = thr_copy_X.partition_S(cX)[(0, None), None, None]
+
+        # allocate fragments for gmem->rmem
+        tWrW = cute.make_fragment_like(tWgW)
+        tXrW = thr_copy_X.retile(tWrW)
+        tXrX, tXrO = [cute.make_fragment_like(thr) for thr in (tXgX, tXgO)]
+
+        num_warps = cute.size(tv_layout, mode=[0]) // cute.arch.WARP_SIZE
+        self._initialize_cluster(tidx, mbar_ptr, num_warps)
+
+        tXpX = utils.predicate_k(thr_copy_X.partition_S(cX), limit=shape[1])
+        row = tXcX[0][0]
+        if row < shape[0]:
+            cute.copy(copy_atom_load_X_async, tXgX, tXsX, pred=tXpX)
+        cute.arch.cp_async_commit_group()
+
+        tWpW = utils.predicate_k(thr_copy_W.partition_S(cX), limit=shape[1])
+        if cutlass.const_expr(not delay_w_load):
+            cute.copy(copy_atom_load_W, tWgW, tWrW, pred=tWpW)
+
+        cute.arch.cp_async_wait_group(0)
+        cute.autovec_copy(tXsX, tXrX)
+        x = tXrX.load().to(cute.Float32)
+        threads_per_row = tv_layout.shape[0][0]
+        sum_x = utils.row_reduce(
+            x,
+            cute.ReductionOp.ADD,
+            threads_per_row,
+            reduction_buffer[None, None, 0],
+            mbar_ptr + 0 if cutlass.const_expr(self.cluster_n > 1) else None,
+            init_val=0.0,
+            hook_fn=cute.arch.cluster_wait if cutlass.const_expr(self.cluster_n > 1) else None,
+        )
+        mean = sum_x / shape[1]
+        if cutlass.const_expr(reload_from == "smem"):
+            cute.autovec_copy(tXsX, tXrX)
+            x = tXrX.load().to(cute.Float32)
+        elif cutlass.const_expr(reload_from == "gmem"):
+            cute.copy(copy_atom_load_X, tXgX, tXrX, pred=tXpX)
+            x = tXrX.load().to(cute.Float32)
+
+        sum_sq_x_sub_mean = utils.row_reduce(
+            (x - mean) * (x - mean),
+            cute.ReductionOp.ADD,
+            threads_per_row,
+            reduction_buffer[None, None, 1],
+            mbar_ptr + 1 if cutlass.const_expr(self.cluster_n > 1) else None,
+            init_val=0.0,
+        )
+        rstd = utils.rsqrt(sum_sq_x_sub_mean / shape[1] + eps)
+        if cutlass.const_expr(mRstd is not None):
+            # Only the thread corresponding to column 0 writes out the rstd to gmem
+            if (
+                tXcX[0][1] == 0
+                and row < shape[0]
+                and (self.cluster_n == 1 or cute.arch.block_idx_in_cluster() == 0)
+            ):
+                tXrRstd[0] = rstd
+        if cutlass.const_expr(mMean is not None):
+            # Only the thread corresponding to column 0 writes out the mean to gmem
+            if (
+                tXcX[0][1] == 0
+                and row < shape[0]
+                and (self.cluster_n == 1 or cute.arch.block_idx_in_cluster() == 0)
+            ):
+                tXrMean[0] = mean
+        if cutlass.const_expr(delay_w_load):
+            cute.copy(copy_atom_load_W, tWgW, tWrW, pred=tWpW)
+        if cutlass.const_expr(reload_from == "smem"):
+            cute.autovec_copy(tXsX, tXrX)
+            x = tXrX.load().to(cute.Float32)
+        elif cutlass.const_expr(reload_from == "gmem"):
+            cute.copy(copy_atom_load_X, tXgX, tXrX, pred=tXpX)
+            x = tXrX.load().to(cute.Float32)
+        x_hat = (x - mean) * rstd
+        w = tXrW.load().to(cute.Float32)
+        y = x_hat * w
+        tXrO.store(y.to(tXrO.element_type))
+        tOpO = utils.predicate_k(thr_copy_O.partition_S(cX), limit=shape[1])
+        if row < shape[0]:
+            cute.copy(copy_atom_store_O, tXrO, tXgO, pred=tOpO)
+
+
+def layernorm(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float = 1e-6,
+    return_rstd: bool = False,
+    return_mean: bool = False,
+) -> torch.Tensor:
+    """LayerNorm forward pass.
+
+    Args:
+        x: Input tensor of shape (M, N)
+        weight: Weight tensor of shape (N,)
+        eps: Small value for numerical stability
+        return_rstd: Whether to return the reciprocal standard deviation
+        return_mean: Whether to return the mean
+
+    Returns:
+        Normalized output tensor of same shape as x
+        If return_rstd is True, also returns rstd tensor of shape (M,)
+        If return_mean is True, also returns mean tensor of shape (M,)
+    """
+    assert x.dim() == 2, "Input must be 2D"
+    assert weight.dim() == 1, "Weight must be 1D"
+    assert x.shape[-1] == weight.shape[0], "Last dimension of input must match weight dimension"
+    assert x.is_cuda and weight.is_cuda, "Tensors must be on CUDA device"
+    assert x.dtype in [torch.float16, torch.bfloat16, torch.float32], "Unsupported dtype"
+    assert weight.dtype == torch.float32, "Weight must be float32"
+    M, N = x.shape
+    device = x.device
+    out = torch.empty_like(x)
+    rstd = torch.empty(M, device=device, dtype=torch.float32) if return_rstd else None
+    mean = torch.empty(M, device=device, dtype=torch.float32) if return_mean else None
+    dtype = torch2cute_dtype_map[x.dtype]
+    convert_from_dlpack = lambda x: (
+        from_dlpack(x.detach(), assumed_align=16).mark_compact_shape_dynamic(
+            mode=0, stride_order=(0, 1)
+        )
+    )
+    x_tensor, out_tensor = [
+        # utils.convert_from_dlpack(t, leading_dim=t.ndim - 1, divisibility=128 // dtype.width)
+        convert_from_dlpack(t)
+        for t in (x, out)
+    ]
+    weight_tensor = utils.convert_from_dlpack(
+        weight.detach(), leading_dim=0, divisibility=128 // cutlass.Float32.width
+    )
+    rstd_tensor = (
+        from_dlpack(rstd.detach(), assumed_align=4).mark_layout_dynamic(leading_dim=0)
+        if rstd is not None
+        else None
+    )
+    mean_tensor = (
+        from_dlpack(mean.detach(), assumed_align=4).mark_layout_dynamic(leading_dim=0)
+        if mean is not None
+        else None
+    )
+    current_stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
+    compile_key = (dtype, N, rstd is not None, mean is not None)
+    if compile_key not in layernorm.compile_cache:
+        rmsnorm_op = LayerNorm(dtype, N)
+        layernorm.compile_cache[compile_key] = cute.compile(
+            rmsnorm_op,
+            x_tensor,
+            weight_tensor,
+            out_tensor,
+            rstd_tensor,
+            mean_tensor,
+            current_stream,
+        )
+    layernorm.compile_cache[compile_key](
+        x_tensor, weight_tensor, out_tensor, rstd_tensor, mean_tensor, current_stream, eps
+    )
+    return (
+        (out, rstd, mean)
+        if return_mean and return_rstd
+        else (
+            (out, rstd)
+            if return_rstd and not return_mean
+            else ((out, mean) if return_mean and not return_rstd else (out))
+        )
+    )
+
+
+layernorm.compile_cache = {}
+
+
+def layernorm_ref(x: torch.Tensor, w: torch.Tensor, eps: float = 1e-6):
+    x_f32 = x.float()
+    return torch.nn.functional.layer_norm(x_f32, w.shape, w, None, eps).to(x.dtype)
+
+
+def rstd_ref(x: torch.Tensor, eps: float = 1e-6):
+    x_f32 = x.float()
+    mean = x_f32.mean(dim=-1, keepdim=True)
+    var = ((x_f32 - mean) ** 2).mean(dim=-1)
+    return 1.0 / torch.sqrt(var + eps)
+
+
+def mean_ref(x: torch.Tensor) -> torch.Tensor:
+    return x.float().mean(dim=-1)

{quack_kernels-0.1.5 → quack_kernels-0.1.6}/quack/rmsnorm.py

@@ -117,7 +117,10 @@ class RMSNorm(ReductionBase):
         shape = mX.shape
         idX = cute.make_identity_tensor(shape)
         # slice for CTAs
-        gX, gO, cX = [cute.local_tile(mT, tiler_mn, (bidx, cluster_y)) for mT in (mX, mO, idX)]
+        # We use domain_offset_i64 to deal with tensors larger than 2^31 elements
+        mX, mO = [utils.domain_offset_i64((bidx * tiler_mn[0], 0), mT) for mT in (mX, mO)]
+        gX, gO = [cute.local_tile(mT, tiler_mn, (0, cluster_y)) for mT in (mX, mO)]
+        cX = cute.local_tile(idX, tiler_mn, (bidx, cluster_y))
         gW = cute.local_tile(mW, tiler_mn, (0, cluster_y))
         gRstd = (
             cute.local_tile(mRstd, tiler_mn, (bidx, cluster_y))

{quack_kernels-0.1.5 → quack_kernels-0.1.6}/quack/softmax.py

@@ -98,7 +98,10 @@ class Softmax(ReductionBase):
         shape = mX.shape
         idX = cute.make_identity_tensor(shape)
         # slice for CTAs
-        gX, gO, cX = [cute.local_tile(mT, tiler_mn, (bidx, cluster_y)) for mT in (mX, mO, idX)]
+        # We use domain_offset_i64 to deal with tensors larger than 2^31 elements
+        mX, mO = [utils.domain_offset_i64((bidx * tiler_mn[0], 0), mT) for mT in (mX, mO)]
+        gX, gO = [cute.local_tile(mT, tiler_mn, (0, cluster_y)) for mT in (mX, mO)]
+        cX = cute.local_tile(idX, tiler_mn, (bidx, cluster_y))
 
         smem = cutlass.utils.SmemAllocator()
         sX = smem.allocate_tensor(
@@ -312,9 +315,11 @@ class SoftmaxBackward(ReductionBase):
         shape = mdY.shape
         idX = cute.make_identity_tensor(shape)
         # slice for CTAs
-        gdY, gY, gdX, cX = [
-            cute.local_tile(mT, tiler_mn, (bidx, cluster_y)) for mT in (mdY, mY, mdX, idX)
+        mdY, mY, mdX = [
+            utils.domain_offset_i64((bidx * tiler_mn[0], 0), mT) for mT in (mdY, mY, mdX)
         ]
+        gdY, gY, gdX = [cute.local_tile(mT, tiler_mn, (0, cluster_y)) for mT in (mdY, mY, mdX)]
+        cX = cute.local_tile(idX, tiler_mn, (bidx, cluster_y))
 
         smem = cutlass.utils.SmemAllocator()
         sdY = smem.allocate_tensor(

{quack_kernels-0.1.5 → quack_kernels-0.1.6}/quack/utils.py

@@ -390,3 +390,19 @@ def i64_to_f32x2(c: cutlass.Int64, *, loc=None, ip=None) -> Tuple[Float32, Float
         vector.extract(vec_f32x2, dynamic_position=[], static_position=[1], loc=loc, ip=ip)
     )
     return res0, res1
+
+
+@dsl_user_op
+def domain_offset_i64(coord: cute.Coord, tensor: cute.Tensor, *, loc=None, ip=None) -> cute.Tensor:
+    flat_coord_i64 = tuple(cutlass.Int64(c) for c in cute.flatten(coord))
+    flat_stride = cute.flatten_to_tuple(tensor.stride)
+    offset = sum(c * s for c, s in zip(flat_coord_i64, flat_stride))
+    assert isinstance(tensor.iterator, cute.Pointer)
+    # HACK: we assume that applying the offset does not change the pointer alignment
+    new_ptr = cute.make_ptr(
+        tensor.element_type,
+        tensor.iterator.toint() + offset * tensor.element_type.width // 8,
+        tensor.memspace,
+        assumed_align=tensor.iterator.max_alignment,
+    )
+    return cute.make_tensor(new_ptr, tensor.layout)

{quack_kernels-0.1.5 → quack_kernels-0.1.6/quack_kernels.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: quack-kernels
-Version: 0.1.5
+Version: 0.1.6
 Requires-Python: >=3.9
 License-File: LICENSE
 Requires-Dist: nvidia-cutlass-dsl==4.1.0.dev0

{quack_kernels-0.1.5 → quack_kernels-0.1.6}/quack_kernels.egg-info/SOURCES.txt

@@ -4,6 +4,7 @@ pyproject.toml
 setup.py
 quack/__init__.py
 quack/cross_entropy.py
+quack/layernorm.py
 quack/reduction_base.py
 quack/rmsnorm.py
 quack/softmax.py
@@ -14,5 +15,6 @@ quack_kernels.egg-info/dependency_links.txt
 quack_kernels.egg-info/requires.txt
 quack_kernels.egg-info/top_level.txt
 tests/test_cross_entropy.py
+tests/test_layernorm.py
 tests/test_rmsnorm.py
 tests/test_softmax.py

{quack_kernels-0.1.5 → quack_kernels-0.1.6}/quack_kernels.egg-info/top_level.txt

@@ -1,2 +1,3 @@
 dist
+media
 quack

{quack_kernels-0.1.5 → quack_kernels-0.1.6}/tests/test_cross_entropy.py

@@ -16,18 +16,19 @@ import cutlass
 )
 @pytest.mark.parametrize("M", [1, 77, 289])
 # @pytest.mark.parametrize("M", [1])
-def test_cross_entropy_forward(M, N, input_dtype):
+def test_cross_entropy(M, N, input_dtype):
     """Test Cross Entropy forward pass against reference implementation."""
     device = "cuda"
-    atol, rtol = 1e-5, 1e-5
+    atol, rtol = 5e-5, 1e-5
     torch.random.manual_seed(0)
+    cutlass.cuda.initialize_cuda_context()
     # Create input tensors (scale down to avoid overflow)
-    x = 0.1 * torch.randn(M, N, device=device, dtype=input_dtype, requires_grad=False)
+    x = (0.1 * torch.randn(M, N, device=device, dtype=input_dtype)).requires_grad_()
     target = torch.randint(0, N, (M,), device=device, dtype=torch.int64)
-    x_ref = x.detach().clone()
+    x_ref = x.detach().clone().requires_grad_()
     target_ref = target.detach().clone()
     # Forward pass
-    loss = _cross_entropy(x, target)
+    loss = cross_entropy(x, target)
     loss_ref = F.cross_entropy(x_ref.float(), target_ref, reduction='none')
     # Check output shape and dtype
     assert loss.shape == (M,)
@@ -40,6 +41,13 @@ def test_cross_entropy_forward(M, N, input_dtype):
     # Check that loss is reasonable (not inf or nan)
     assert not torch.isnan(loss).any()
     assert not torch.isinf(loss).any()
+    # Test backward pass
+    dloss = torch.randn_like(loss)
+    torch.cuda.synchronize()
+    dx_ref, = torch.autograd.grad(loss_ref, x_ref, grad_outputs=dloss)
+    dx, = torch.autograd.grad(loss, x, grad_outputs=dloss)
+    assert dx.shape == x.shape
+    torch.testing.assert_close(dx, dx_ref.to(input_dtype), atol=atol, rtol=rtol)
 
 
 @pytest.mark.parametrize("input_dtype", [torch.float16, torch.float32])
@@ -99,47 +107,3 @@ def test_cross_entropy_edge_targets():
     loss_last = _cross_entropy(x, target_last)
     loss_ref_last = F.cross_entropy(x, target_last, reduction='none')
     torch.testing.assert_close(loss_last, loss_ref_last, atol=1e-4, rtol=1e-4)
-
-
-
-
-
-@pytest.mark.parametrize("input_dtype", [torch.bfloat16, torch.float16, torch.float32])
-@pytest.mark.parametrize(
-    "N",
-    [192, 256, 512, 760, 1024, 1128, 2048, 4096, 8192, 16384, 32768, 65536, 128256, 131072, 256128, 262144]  # A representative subset to keep compile time reasonable
-)
-@pytest.mark.parametrize("M", [1, 37, 77])
-def test_cross_entropy_autograd_backward(M, N, input_dtype):
-    device = "cuda"
-    
-    if input_dtype == torch.bfloat16:
-        atol = 1e-3
-        rtol = 1e-3
-    else:
-        atol = 1e-5
-        rtol = 1e-5
-
-    torch.random.manual_seed(0)
-
-    x = 0.1 * torch.randn(M, N, device=device, dtype=input_dtype, requires_grad=True)
-    target = torch.randint(0, N, (M,), device=device, dtype=torch.int64)
-
-    x_ref = x.detach().clone().requires_grad_(True)
-    target_ref = target.detach().clone()
-
-    cutlass.cuda.initialize_cuda_context()
-
-    loss = cross_entropy(x, target)  # our autograd-enabled op
-    loss_ref = F.cross_entropy(x_ref.float(), target_ref, reduction='none')
-
-    torch.testing.assert_close(loss, loss_ref, atol=atol, rtol=rtol)
-
-    dloss = torch.randn_like(loss)
-
-    dx_ref, = torch.autograd.grad(loss_ref, x_ref, grad_outputs=dloss)
-
-    dx, = torch.autograd.grad(loss, x, grad_outputs=dloss)
-
-    assert dx.shape == x.shape
-    torch.testing.assert_close(dx, dx_ref.to(input_dtype), atol=atol, rtol=rtol)

quack_kernels-0.1.6/tests/test_layernorm.py

@@ -0,0 +1,162 @@
+# tests/test_layernorm.py
+
+import pytest
+import torch
+
+from quack.layernorm import layernorm, layernorm_ref, rstd_ref, mean_ref
+
+
+@pytest.mark.parametrize("eps", [1e-5, 1e-6])
+@pytest.mark.parametrize("input_dtype", [torch.bfloat16, torch.float16, torch.float32])
+@pytest.mark.parametrize("M", [1, 37, 199])
+@pytest.mark.parametrize(
+    "N", [256, 512, 760, 1024, 1128, 2048, 4096, 8192, 16384, 32768, 65536, 131072, 262144]
+)  # , 32768])
+def test_layernorm_forward(M, N, input_dtype, eps):
+    """Test LayerNorm forward pass against reference implementation."""
+    device = "cuda"
+
+    # tolerance depends on precision
+    if input_dtype == torch.bfloat16:
+        atol = 1e-2
+        rtol = 1e-2
+    elif input_dtype == torch.float16:
+        atol = 1e-3
+        rtol = 1e-3
+    else:
+        atol = 1e-4
+        rtol = 1e-4
+
+    torch.random.manual_seed(0)
+    x = torch.randn(M, N, device=device, dtype=input_dtype, requires_grad=True)
+    weight = torch.randn(N, device=device, dtype=torch.float32, requires_grad=True)
+
+    # pure‐PyTorch refs
+    x_ref = x.detach().clone().requires_grad_()
+    weight_ref = weight.detach().clone().requires_grad_()
+
+    out, rstd, mean = layernorm(x, weight, eps=eps, return_rstd=True, return_mean=True)
+    out_ref = layernorm_ref(x_ref, weight_ref, eps=eps)
+    rstd_ref_val = rstd_ref(x_ref, eps=eps)
+    mean_ref_val = mean_ref(x_ref)
+
+    # shapes & dtypes
+    assert out.shape == x.shape
+    assert out.dtype == input_dtype
+    assert rstd.shape == (M,) and rstd.dtype == torch.float32
+    assert mean.shape == (M,) and mean.dtype == torch.float32
+
+    # numeric check
+    torch.testing.assert_close(out, out_ref, atol=atol, rtol=rtol)
+    torch.testing.assert_close(rstd, rstd_ref_val, atol=6e-4, rtol=6e-4)
+    torch.testing.assert_close(mean, mean_ref_val, atol=6e-4, rtol=6e-4)
+
+
+@pytest.mark.parametrize("return_rstd", [True, False])
+@pytest.mark.parametrize("return_mean", [True, False])
+def test_layernormnorm_return_rstd_option(return_rstd, return_mean):
+    """Test that return_rstd option works correctly."""
+    device = "cuda"
+    M, N = 32, 1024
+    eps = 1e-6
+
+    x = torch.randn(M, N, device=device, dtype=torch.float16)
+    weight = torch.randn(N, device=device, dtype=torch.float32)
+
+    if return_rstd and return_mean:
+        out, rstd, mean = layernorm(x, weight, eps=eps, return_rstd=True, return_mean=True)
+        assert out.shape == (M, N)
+        assert rstd.shape == (M,)
+        assert rstd.dtype == torch.float32
+        assert mean.shape == (M,)
+        assert mean.dtype == torch.float32
+    elif return_rstd and not return_mean:
+        out, rstd = layernorm(x, weight, eps=eps, return_rstd=True, return_mean=False)
+        assert out.shape == (M, N)
+        assert rstd.shape == (M,)
+        assert rstd.dtype == torch.float32
+    elif not return_rstd and return_mean:
+        out, mean = layernorm(x, weight, eps=eps, return_rstd=False, return_mean=True)
+        assert out.shape == (M, N)
+        assert mean.shape == (M,)
+        assert mean.dtype == torch.float32
+    else:
+        out = layernorm(x, weight, eps=eps, return_rstd=False, return_mean=False)
+        assert out.shape == (M, N)
+        assert isinstance(out, torch.Tensor)
+
+
+def test_layernorm_input_validation():
+    """Test input validation and error handling."""
+    device = "cuda"
+
+    # Test 3D input (should fail)
+    x_3d = torch.randn(2, 32, 1024, device=device, dtype=torch.float16)
+    weight = torch.randn(1024, device=device, dtype=torch.float32)
+
+    with pytest.raises(AssertionError, match="Input must be 2D"):
+        layernorm(x_3d, weight)
+
+    # Test weight dimension mismatch
+    x = torch.randn(32, 1024, device=device, dtype=torch.float16)
+    weight_wrong = torch.randn(512, device=device, dtype=torch.float32)
+
+    with pytest.raises(AssertionError, match="Last dimension of input must match weight dimension"):
+        layernorm(x, weight_wrong)
+
+    # Test CPU tensors (should fail)
+    x_cpu = torch.randn(32, 1024, dtype=torch.float16)
+    weight_cpu = torch.randn(1024, dtype=torch.float32)
+
+    with pytest.raises(AssertionError, match="Tensors must be on CUDA device"):
+        layernorm(x_cpu, weight_cpu)
+
+    # Test unsupported dtype
+    x = torch.randn(32, 1024, device=device, dtype=torch.float64)
+    weight = torch.randn(1024, device=device, dtype=torch.float32)
+
+    with pytest.raises(AssertionError, match="Unsupported dtype"):
+        layernorm(x, weight)
+
+    # Test wrong weight dtype
+    x = torch.randn(32, 1024, device=device, dtype=torch.float16)
+    weight_wrong_dtype = torch.randn(1024, device=device, dtype=torch.float16)
+
+    with pytest.raises(AssertionError, match="Weight must be float32"):
+        layernorm(x, weight_wrong_dtype)
+
+
+def test_layernorm_compile_cache():
+    """Test that compile cache works correctly for repeated calls."""
+    device = "cuda"
+    M, N = 32, 1024
+    eps = 1e-6
+
+    # Clear cache
+    layernorm.compile_cache.clear()
+    assert len(layernorm.compile_cache) == 0
+
+    x1 = torch.randn(M, N, device=device, dtype=torch.float16)
+    weight1 = torch.randn(N, device=device, dtype=torch.float32)
+
+    # First call should compile
+    out1 = layernorm(x1, weight1, eps=eps)
+    assert len(layernorm.compile_cache) == 1
+
+    # Same shape should reuse cache
+    x2 = torch.randn(M, N, device=device, dtype=torch.float16)
+    weight2 = torch.randn(N, device=device, dtype=torch.float32)
+    out2 = layernorm(x2, weight2, eps=eps)
+    assert len(layernorm.compile_cache) == 1
+
+    # Different shape should create new cache entry
+    x3 = torch.randn(M, N * 2, device=device, dtype=torch.float16)
+    weight3 = torch.randn(N * 2, device=device, dtype=torch.float32)
+    out3 = layernorm(x3, weight3, eps=eps)
+    assert len(layernorm.compile_cache) == 2
+
+    # Different dtype should create new cache entry
+    x4 = torch.randn(M, N, device=device, dtype=torch.float32)
+    weight4 = torch.randn(N, device=device, dtype=torch.float32)
+    out4 = layernorm(x4, weight4, eps=eps)
+    assert len(layernorm.compile_cache) == 3

{quack_kernels-0.1.5 → quack_kernels-0.1.6}/tests/test_rmsnorm.py

@@ -32,19 +32,17 @@ def test_rmsnorm_forward(M, N, input_dtype, eps):
     weight = torch.randn(N, device=device, dtype=torch.float32, requires_grad=True)
     x_ref = x.detach().clone().requires_grad_()
     weight_ref = weight.detach().clone().requires_grad_()
-    out, rstd = rmsnorm(x, weight, eps=eps, return_rstd=True)
+    out = rmsnorm(x, weight, eps=eps)
     out_ref = rmsnorm_ref(x_ref, weight_ref, eps=eps)
-    rstd_ref_val = rstd_ref(x_ref, eps=eps)
+    # rstd_ref_val = rstd_ref(x_ref, eps=eps)
 
     # Check output shape and dtype
     assert out.shape == x.shape
     assert out.dtype == input_dtype
-    assert rstd.shape == (M,)
-    assert rstd.dtype == torch.float32
 
     # Check accuracy
     torch.testing.assert_close(out, out_ref, atol=atol, rtol=1e-3)
-    torch.testing.assert_close(rstd, rstd_ref_val, atol=atol, rtol=1e-3)
+    # torch.testing.assert_close(rstd, rstd_ref_val, atol=atol, rtol=1e-3)
 
 
 # @pytest.mark.parametrize("eps", [1e-5, 1e-6])
@@ -88,6 +86,39 @@ def test_rmsnorm_forward(M, N, input_dtype, eps):
 #     torch.testing.assert_close(weight.grad, weight_ref.grad, atol=atol, rtol=1e-3)
 
 
+@pytest.mark.parametrize("eps", [1e-5])
+@pytest.mark.parametrize("input_dtype", [torch.bfloat16])
+@pytest.mark.parametrize(
+    "N",
+    [131072, 262144]
+    # [262144]
+)
+@pytest.mark.parametrize("M", [32 * 1024])
+def test_rmsnorm_large_tensor(M, N, input_dtype, eps):
+    """Test RMSNorm forward pass against reference implementation."""
+    device = "cuda"
+    # Set tolerance based on dtype
+    if input_dtype == torch.bfloat16:
+        atol = 1e-1
+    elif input_dtype == torch.float16:
+        atol = 1e-2
+    else:
+        atol = 1e-4
+    torch.random.manual_seed(0)
+    torch.cuda.empty_cache()
+    x = torch.randn(M, N, device=device, dtype=input_dtype, requires_grad=False)
+    weight = torch.randn(N, device=device, dtype=torch.float32, requires_grad=False)
+    out = rmsnorm(x, weight, eps=eps)
+    # Need to compile, otherwise it OOMs
+    rmsnorm_compiled = torch.compile(rmsnorm_ref)
+    # Run once with smaller input to avoid OOMs
+    rmsnorm_compiled(x[:32], weight, eps=eps)
+    out_ref = rmsnorm_compiled(x, weight, eps=eps)
+    # Need to chunk, otherwise it OOMs
+    assert all((out_c - out_ref_c).abs().max() < atol
+               for out_c, out_ref_c in zip(out.chunk(16), out_ref.chunk(16)))
+
+
 @pytest.mark.parametrize("return_rstd", [True, False])
 def test_rmsnorm_return_rstd_option(return_rstd):
     """Test that return_rstd option works correctly."""

{quack_kernels-0.1.5 → quack_kernels-0.1.6}/tests/test_softmax.py

@@ -34,7 +34,7 @@ def test_softmax(M, N, input_dtype):
 
     torch.random.manual_seed(0)
     # Create input tensors (scale down to avoid overflow in softmax)
-    x = 0.1 * torch.randn(M, N, device=device, dtype=input_dtype, requires_grad=True)
+    x = (0.1 * torch.randn(M, N, device=device, dtype=input_dtype)).requires_grad_()
     x_ref = x.detach().clone().requires_grad_(True)
 
     # Forward pass
@@ -58,13 +58,12 @@ def test_softmax(M, N, input_dtype):
 
     # Test backward pass
     dy = torch.randn_like(out)
+    torch.cuda.synchronize()  # without sync, torch.autograd gets wrong results
     dx_ref, = torch.autograd.grad(out_ref, x_ref, grad_outputs=dy)
     # Call our implementation later, otherwise getting CUDA_ERROR_INVALID_CONTEXT
     dx, = torch.autograd.grad(out, x, grad_outputs=dy)
-    # Check output shape and dtype
     assert dx.shape == dy.shape
     assert dx.dtype == input_dtype
-    # Check accuracy against reference
     torch.testing.assert_close(dx, dx_ref, atol=atol, rtol=rtol)