quack-kernels 0.1.2__py3-none-any.whl → 0.1.3__py3-none-any.whl

quack/__init__.py

@@ -1,5 +1,11 @@
-__version__ = "0.1.2"
+__version__ = "0.1.3"
 
 from quack.rmsnorm import rmsnorm
 from quack.softmax import softmax
 from quack.cross_entropy import cross_entropy
+
+__all__ = [
+    "rmsnorm",
+    "softmax",
+    "cross_entropy",
+]

quack/cross_entropy.py

@@ -1,7 +1,6 @@
 import math
 import torch
-import operator
-from typing import Callable, Union, Optional
+from typing import Optional, Type
 
 import cuda.bindings.driver as cuda
 
@@ -10,169 +9,192 @@ 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
 
 
-@cute.kernel
-def cross_entropy_kernel(
-    mX: cute.Tensor,  # (M, N)
-    mTarget: cute.Tensor, # (M,)
-    mLoss: cute.Tensor, # (M,)
-    mLSE: Optional[cute.Tensor], # (M,)
-    tv_layout: cute.Layout,
-    tiler_mn: cute.Shape,
-    cluster_n: cutlass.Constexpr = 1,
-):
-    tidx, _, _ = cute.arch.thread_idx()
-    bidx, cluster_y, _ = cute.arch.block_idx()
-
-    shape: cute.Shape = mX.shape
-    idX = cute.make_identity_tensor(shape)
-    # slice for CTAs
-    gX, cX = [
-        cute.local_tile(mT, tiler_mn, (bidx, 0 if cluster_n == 1 else cluster_y))
-        for mT in (mX, idX)
-    ]
-
-    smem = cutlass.utils.SmemAllocator()
-    sX = smem.allocate_tensor(mX.element_type, cute.make_ordered_layout(tiler_mn, order=(1, 0)), byte_alignment=16)
-    num_warps = cute.size(tv_layout, mode=[0]) // cute.arch.WARP_SIZE
-    warps_per_row = utils.max_constexpr(tv_layout.shape[0][0] // cute.arch.WARP_SIZE, 1)
-    reduction_buffer_layout = cute.make_ordered_layout(
+class CrossEntropy(ReductionBase):
+    def __init__(self, dtype: Type[cutlass.Numeric], N: int, online_softmax: bool = True):
         # 2 stages: 1 for max, 1 for sum
-        (num_warps // warps_per_row, (warps_per_row, cluster_n), 2),
-        order=(1, 0, 2)
-    )
-    reduction_buffer = smem.allocate_tensor(cutlass.Float32, reduction_buffer_layout, byte_alignment=4)
-    if cutlass.const_expr(cluster_n > 1):
-        # 1 mbar for max reduction, 1 mbar for sum reduction
-        mbar_ptr = smem.allocate_array(cutlass.Int64, num_elems=2)
-    else:
-        mbar_ptr = None
-
-    # declare the atoms which will be used later for memory copy
-    copy_atom_load_X = cute.make_copy_atom(cute.nvgpu.cpasync.CopyG2SOp(), gX.element_type, num_bits_per_copy=128)
-    thr_copy_X = cute.make_tiled_copy(copy_atom_load_X, tv_layout, tiler_mn).get_slice(tidx)
-
-    #### Thread View
-    tXgX = thr_copy_X.partition_S(gX)
-    tXsX = thr_copy_X.partition_D(sX)
-    tXcX = thr_copy_X.partition_S(cX)[(0, None), None, None]
-    tXrX = cute.make_fragment_like(tXgX)
-
-    if cluster_n > 1:
-        if tidx < 2:
-            cute.arch.mbarrier_init_arrive_cnt(mbar_ptr + tidx, 1)
-        cute.arch.mbarrier_init_fence()
-        if tidx < 2:
-            cute.arch.mbarrier_init_tx_bytes(mbar_ptr + tidx, num_warps * cluster_n * cutlass.Float32.width // 8)
-        # Cluster arrive after barrier init
-        cute.arch.cluster_arrive_relaxed()
-
-    row = tXcX[0][0]
-    target = cute.Int32.zero
-    if row < shape[0] and tXcX[0][1] == 0:
-        target = cute.Int32(mTarget[row])
-
-    is_even_N = cutlass.const_expr(shape[1] == tiler_mn[1] * cluster_n)
-    tXpX = utils.predicate_k(thr_copy_X.partition_S(cX), limit=shape[1]) if not is_even_N else None
-    if row < shape[0]:
-        cute.copy(copy_atom_load_X, tXgX, tXsX, pred=tXpX)
-    cute.arch.cp_async_commit_group()
-    cute.arch.cp_async_wait_group(0)
-    cute.autovec_copy(tXsX, tXrX)
-    x = tXrX.load().to(cute.Float32)
-    # Fill OOB values with -inf
-    if cutlass.const_expr(not is_even_N):
-        tXrX_fp32 = cute.make_fragment_like(tXrX, cutlass.Float32)
-        tXrX_fp32.store(x)
-        for rest_v in range(tXpX.shape[0]):
-            for rest_k in range(tXpX.shape[2]):
-                if not tXpX[rest_v, 0, rest_k]:
-                    tXrX_fp32[(None, rest_v), None, rest_k].fill(-cutlass.Float32.inf)
-        x = tXrX_fp32.load()
-
-    target_logit = cute.Float32.zero
-    if row < shape[0] and tXcX[0][1] == 0:
-        target_logit = cute.Float32(mX[row, target])
-
-    threads_per_row = tv_layout.shape[0][0]
-    max_x = utils.row_reduce(
-        x,
-        cute.ReductionOp.MAX,
-        threads_per_row,
-        reduction_buffer[None, None, 0],
-        mbar_ptr + 0 if cluster_n > 1 else None,
-        init_val=-cutlass.Float32.inf,
-        hook_fn=cute.arch.cluster_wait if cutlass.const_expr(cluster_n > 1) else None
-    )
-    log2_e = math.log2(math.e)
-    # exp_x = cute.math.exp2((x - max_x) * log2_e, fastmath=True)
-    exp_x = utils.exp2f((x - max_x) * log2_e)  # a bit faster, idk why
-    denom = utils.row_reduce(
-        exp_x,
-        cute.ReductionOp.ADD,
-        threads_per_row,
-        reduction_buffer[None, None, 1],
-        mbar_ptr + 1 if cluster_n > 1 else None,
-        init_val=0.0,
-    )
+        super().__init__(
+            dtype,
+            N,
+            stage=2 if not online_softmax else 1,
+            reduction_dtype=cutlass.Float32 if not online_softmax else cutlass.Int64,
+        )
+        self.online_softmax = online_softmax
+        self.reload_from = None if N <= 16384 or online_softmax else "smem"
+
+    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)))
+            )
+        )
 
-    if tXcX[0][1] == 0 and row < shape[0] and (cluster_n == 1 or cute.arch.block_idx_in_cluster() == 0):
-        ln_2 = math.log(2.0)
-        lse = max_x + utils.log2f(denom) * ln_2
-        loss_val = lse - target_logit
-        mLoss[row] = loss_val.to(mLoss.element_type)
-        if cutlass.const_expr(mLSE is not None):
-            mLSE[row] = lse
-
-
-@cute.jit
-def cross_entropy_interface(
-    mX: cute.Tensor,
-    mTarget: cute.Tensor,
-    mLoss: cute.Tensor,
-    mLSE: Optional[cute.Tensor],
-    stream: cuda.CUstream,
-    N: cutlass.Constexpr,
-    copy_bits: cutlass.Constexpr = 128
-):
-    vecsize = copy_bits // mX.element_type.width
-    assert N % vecsize == 0, f"Input N {N} is not divisible by vector size {vecsize}"
-    num_threads = 128 if N <= 16384 else 256
-
-    num_warps = num_threads // cute.arch.WARP_SIZE
-    assert num_threads % cute.arch.WARP_SIZE == 0
-    threads_per_row = 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))))
-
-    if cutlass.const_expr(mX.element_type.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 <= 16 * 1024 else (2 if N <= 64 * 1024 else (4 if N <= 128 * 1024 else 8))
-
-    num_blocks_N = cute.ceil_div(N // vecsize, threads_per_row * cluster_n)
-    cols_per_block = num_threads // threads_per_row
-    tiler_mn = (cols_per_block, vecsize * num_blocks_N * threads_per_row)  # This rounds up N
-    tv_layout = cute.make_layout(
-        ((threads_per_row, cols_per_block), (vecsize, num_blocks_N)),
-        stride=((vecsize * cols_per_block, 1), (cols_per_block, cols_per_block * vecsize * threads_per_row))
-    )
+    def _set_cluster_n(self):
+        N = self.N
+        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 <= 16 * 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,
+        mTarget: cute.Tensor,
+        mLoss: cute.Tensor,
+        mLSE: Optional[cute.Tensor],
+        stream: cuda.CUstream,
+    ):
+        assert mX.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
+        self.kernel(mX, mTarget, mLoss, mLSE, tv_layout, tiler_mn).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 self.cluster_n > 1 else None,
+            smem=self._smem_size_in_bytes(tiler_mn, num_warps),
+            stream=stream,
+        )
 
-    smem_allocated = cute.size_in_bytes(mX.element_type, cute.make_layout(tiler_mn)) + 2 * num_warps * cluster_n * (cutlass.Float32.width // 8) + 2 * (cutlass.Int64.width // 8)
-    cross_entropy_kernel(mX, mTarget, mLoss, mLSE, tv_layout, tiler_mn, cluster_n).launch(
-        grid=[cute.ceil_div(mX.shape[0], tiler_mn[0]), cluster_n, 1],
-        block=[cute.size(tv_layout, mode=[0]), 1, 1],
-        # Launching with cluster=[1, 1, 1] instead of None slows down the kernel by ~8us
-        cluster=[1, cluster_n, 1] if cluster_n > 1 else None,
-        smem=smem_allocated,
-        stream=stream,
-    )
+    @cute.kernel
+    def kernel(
+        self,
+        mX: cute.Tensor,  # (M, N)
+        mTarget: cute.Tensor,  # (M,)
+        mLoss: cute.Tensor,  # (M,)
+        mLSE: Optional[cute.Tensor],  # (M,)
+        tv_layout: cute.Layout,
+        tiler_mn: cute.Shape,
+    ):
+        tidx, _, _ = cute.arch.thread_idx()
+        bidx, cluster_y, _ = cute.arch.block_idx()
+
+        shape: cute.Shape = mX.shape
+        idX = cute.make_identity_tensor(shape)
+        # slice for CTAs
+        gX, cX = [
+            cute.local_tile(mT, tiler_mn, (bidx, 0 if self.cluster_n == 1 else cluster_y))
+            for mT in (mX, idX)
+        ]
+
+        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)
 
+        # declare the atoms which will be used later for memory copy
+        copy_atom_load_X = cute.make_copy_atom(
+            cute.nvgpu.cpasync.CopyG2SOp(), gX.element_type, num_bits_per_copy=128
+        )
+        thr_copy_X = cute.make_tiled_copy(copy_atom_load_X, tv_layout, tiler_mn).get_slice(tidx)
+
+        #### Thread View
+        tXgX = thr_copy_X.partition_S(gX)
+        tXsX = thr_copy_X.partition_D(sX)
+        tXcX = thr_copy_X.partition_S(cX)[(0, None), None, None]
+        tXrX = cute.make_fragment_like(tXgX)
+
+        num_warps = cute.size(tv_layout, mode=[0]) // cute.arch.WARP_SIZE
+        self._initialize_cluster(tidx, mbar_ptr, num_warps)
 
-torch2cute_dtype_map = {
-    torch.float16: cutlass.Float16,
-    torch.bfloat16: cutlass.BFloat16,
-    torch.float32: cutlass.Float32,
-}
+        row = tXcX[0][0]
+        target = cute.Int32.zero
+        if row < shape[0] and tXcX[0][1] == 0:
+            target = cute.Int32(mTarget[row])
+
+        is_even_N = cutlass.const_expr(shape[1] == tiler_mn[1] * self.cluster_n)
+        tXpX = (
+            utils.predicate_k(thr_copy_X.partition_S(cX), limit=shape[1]) if not is_even_N else None
+        )
+        if row < shape[0]:
+            cute.copy(copy_atom_load_X, tXgX, tXsX, pred=tXpX)
+        cute.arch.cp_async_commit_group()
+        cute.arch.cp_async_wait_group(0)
+        # Fill OOB values with -inf
+        if cutlass.const_expr(not is_even_N):
+            utils.fill_oob(tXsX, tXpX, -tXsX.element_type.inf)
+        cute.autovec_copy(tXsX, tXrX)
+        x = tXrX.load().to(cute.Float32)
+
+        target_logit = cute.Float32.zero
+        if row < shape[0] and tXcX[0][1] == 0:
+            target_logit = cute.Float32(mX[row, target])
+
+        threads_per_row = tv_layout.shape[0][0]
+        if cutlass.const_expr(not self.online_softmax):
+            max_x = utils.row_reduce(
+                x,
+                cute.ReductionOp.MAX,
+                threads_per_row,
+                reduction_buffer[None, None, 0],
+                mbar_ptr + 0 if self.cluster_n > 1 else None,
+                init_val=-cutlass.Float32.inf,
+                hook_fn=cute.arch.cluster_wait if cutlass.const_expr(self.cluster_n > 1) else None,
+            )
+            if cutlass.const_expr(self.reload_from == "smem"):
+                cute.autovec_copy(tXsX, tXrX)
+                x = tXrX.load().to(cute.Float32)
+            log2_e = math.log2(math.e)
+            # exp_x = cute.math.exp2((x - max_x) * log2_e, fastmath=True)
+            # a bit faster, probably because it's calling ex2.approx.ftz instead of ex2.approx?
+            # exp_x = utils.exp2f((x - max_x) * log2_e)
+            # This would use ffma instead of fadd then fmul
+            exp_x = utils.exp2f(x * log2_e - (max_x * log2_e))
+            denom = utils.row_reduce(
+                exp_x,
+                cute.ReductionOp.ADD,
+                threads_per_row,
+                reduction_buffer[None, None, 1],
+                mbar_ptr + 1 if self.cluster_n > 1 else None,
+                init_val=0.0,
+            )
+        else:
+            max_x, denom, _ = utils.online_softmax_reduce(
+                x,
+                threads_per_row,
+                reduction_buffer[None, None, 0],
+                mbar_ptr,
+                hook_fn=cute.arch.cluster_wait if cutlass.const_expr(self.cluster_n > 1) else None,
+            )
+
+        if (
+            tXcX[0][1] == 0
+            and row < shape[0]
+            and (self.cluster_n == 1 or cute.arch.block_idx_in_cluster() == 0)
+        ):
+            ln_2 = math.log(2.0)
+            lse = max_x + utils.log2f(denom) * ln_2
+            loss_val = lse - target_logit
+            mLoss[row] = loss_val.to(mLoss.element_type)
+            if cutlass.const_expr(mLSE is not None):
+                mLSE[row] = lse
 
 
 def cross_entropy(
@@ -194,27 +216,36 @@ def cross_entropy(
     assert x.shape[0] == target.shape[0], "Batch dimensions must match"
     assert x.is_cuda and target.is_cuda, "Tensors must be on CUDA device"
     assert x.dtype in [torch.float16, torch.bfloat16, torch.float32], "Unsupported input dtype"
-    assert target.dtype == torch.int64, "Target must be int64"
+    assert target.dtype in [torch.int32, torch.int64], "Target must be int32 or int64"
     M, N = x.shape
     device = x.device
     loss = torch.empty(M, device=device, dtype=torch.float32)
     lse = torch.empty(M, device=device, dtype=torch.float32) if return_lse else None
     dtype = torch2cute_dtype_map[x.dtype]
     convert_from_dlpack = lambda tensor: (
-        from_dlpack(tensor.detach(), assumed_align=16)
-        .mark_compact_shape_dynamic(mode=0, stride_order=(0, 1))
+        from_dlpack(tensor.detach(), assumed_align=16).mark_compact_shape_dynamic(
+            mode=0, stride_order=(0, 1)
+        )
     )
-    x_tensor, = [convert_from_dlpack(tensor) for tensor in (x,)]
+    x_tensor = convert_from_dlpack(x)
     loss_tensor = from_dlpack(loss.detach(), assumed_align=4).mark_compact_shape_dynamic(mode=0)
-    lse_tensor = from_dlpack(loss.detach(), assumed_align=4).mark_compact_shape_dynamic(mode=0) if lse is not None else None
+    lse_tensor = (
+        from_dlpack(lse.detach(), assumed_align=4).mark_compact_shape_dynamic(mode=0)
+        if lse is not None
+        else None
+    )
     target_tensor = from_dlpack(target.detach(), assumed_align=8).mark_compact_shape_dynamic(mode=0)
     stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
-    compile_key = (dtype, N, lse_tensor is not None)
+
+    compile_key = (dtype, N, lse is not None)
     if compile_key not in cross_entropy.compile_cache:
+        cross_entropy_op = CrossEntropy(dtype, N)
         cross_entropy.compile_cache[compile_key] = cute.compile(
-            cross_entropy_interface, x_tensor, target_tensor, loss_tensor, lse_tensor, stream, N
+            cross_entropy_op, x_tensor, target_tensor, loss_tensor, lse_tensor, stream
         )
-    cross_entropy.compile_cache[compile_key](x_tensor, target_tensor, loss_tensor, lse_tensor, stream)
+    cross_entropy.compile_cache[compile_key](
+        x_tensor, target_tensor, loss_tensor, lse_tensor, stream
+    )
     return loss if not return_lse else (loss, lse)
 
 

quack/reduction_base.py

@@ -0,0 +1,98 @@
+# Copyright (c) 2025, Wentao Guo, Ted Zadouri, Tri Dao.
+
+import torch
+from typing import Type, Tuple, Optional
+
+import cutlass
+import cutlass.cute as cute
+
+import quack.utils as utils
+
+
+torch2cute_dtype_map = {
+    torch.float16: cutlass.Float16,
+    torch.bfloat16: cutlass.BFloat16,
+    torch.float32: cutlass.Float32,
+}
+
+
+class ReductionBase:
+    def __init__(
+        self, dtype: Type[cutlass.Numeric], N: int, stage: int, reduction_dtype=cutlass.Float32
+    ):
+        self.dtype = dtype
+        self.N = N
+        self.stage = stage
+        self.reduction_dtype = reduction_dtype
+
+    def _calculate_threads_per_row(self):
+        raise NotImplementedError()
+
+    def _set_cluster_n(self):
+        self.cluster_n = 1
+
+    def _get_num_threads(self):
+        return 128 if self.N <= 16384 else 256
+
+    def _get_tv_layout(self):
+        copy_bits = 128
+        vecsize = copy_bits // self.dtype.width
+        assert self.N % vecsize == 0, f"Input N {self.N} is not divisible by vector size {vecsize}"
+        num_threads = self._get_num_threads()
+        num_warps = num_threads // cute.arch.WARP_SIZE
+        assert num_threads % cute.arch.WARP_SIZE == 0
+
+        threads_per_row = self._calculate_threads_per_row()
+        num_blocks_N = cute.ceil_div(self.N // vecsize, threads_per_row * self.cluster_n)
+        cols_per_block = num_threads // threads_per_row
+        tiler_mn = (cols_per_block, vecsize * num_blocks_N * threads_per_row)
+        tv_layout = cute.make_layout(
+            ((threads_per_row, cols_per_block), (vecsize, num_blocks_N)),
+            stride=(
+                (vecsize * cols_per_block, 1),
+                (cols_per_block, cols_per_block * vecsize * threads_per_row),
+            ),
+        )
+        return tiler_mn, tv_layout
+
+    def _smem_size_in_bytes(self, tiler_mn, num_warps):
+        return (
+            cute.size_in_bytes(self.dtype, cute.make_layout(tiler_mn))
+            + self.stage * num_warps * self.cluster_n * (self.reduction_dtype.width // 8)
+            + self.stage * (cutlass.Int64.width // 8)
+        )
+
+    def _get_reduction_buffer_layout(self, tv_layout: cute.Layout, cluster_n: int):
+        num_warps = cute.size(tv_layout, mode=[0]) // cute.arch.WARP_SIZE
+        warps_per_row = utils.max_constexpr(tv_layout.shape[0][0] // cute.arch.WARP_SIZE, 1)
+        return cute.make_ordered_layout(
+            (num_warps // warps_per_row, (warps_per_row, cluster_n), self.stage),
+            order=(1, 0, 2),
+        )
+
+    def _allocate_reduction_buffer_and_mbar(
+        self, smem: cutlass.utils.SmemAllocator, tv_layout: cute.Layout
+    ) -> Tuple[cute.Tensor, Optional[cute.Pointer]]:
+        reduction_buffer = smem.allocate_tensor(
+            self.reduction_dtype,
+            self._get_reduction_buffer_layout(tv_layout, self.cluster_n),
+            byte_alignment=4,
+        )
+        if cutlass.const_expr(self.cluster_n > 1):
+            mbar_ptr = smem.allocate_array(cutlass.Int64, num_elems=self.stage)
+        else:
+            mbar_ptr = None
+        return reduction_buffer, mbar_ptr
+
+    @cute.jit
+    def _initialize_cluster(self, tidx: cutlass.Int32, mbar_ptr: cute.Pointer, num_warps: int):
+        if cutlass.const_expr(self.cluster_n > 1):
+            if tidx < self.stage:
+                cute.arch.mbarrier_init_arrive_cnt(mbar_ptr + tidx, 1)
+            cute.arch.mbarrier_init_fence()
+            if tidx < self.stage:
+                cute.arch.mbarrier_init_tx_bytes(
+                    mbar_ptr + tidx, num_warps * self.cluster_n * self.reduction_dtype.width // 8
+                )
+            # Cluster arrive after barrier init
+            cute.arch.cluster_arrive_relaxed()