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

quack/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.1.1"
+__version__ = "0.1.2"
 
 from quack.rmsnorm import rmsnorm
 from quack.softmax import softmax

quack/cross_entropy.py

@@ -1,7 +1,7 @@
 import math
 import torch
 import operator
-from typing import Callable, Union
+from typing import Callable, Union, Optional
 
 import cuda.bindings.driver as cuda
 
@@ -17,37 +17,29 @@ 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()
-    gdim, _, _ = cute.arch.grid_dim()
 
     shape: cute.Shape = mX.shape
-    idX = cute.make_identity_tensor(mX.shape)
-    gX, cX = [cute.zipped_divide(mT, tiler_mn) for mT in (mX, idX)]
-    blkX, blkCrd = [gT[(None, None), bidx if cluster_n == 1 else (bidx, cluster_y)] for gT in (gX, cX)]
-
-    # declare the atoms which will be used later for memory copy
-    copy_atom_load_X = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), gX.element_type, num_bits_per_copy=128)
-    copy_atom_load_X_async = cute.make_copy_atom(cute.nvgpu.cpasync.CopyG2SOp(), gX.element_type, num_bits_per_copy=128)
-    copy_atom_scalar = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), gX.element_type, num_bits_per_copy=gX.element_type.width)
-
-    thr_copy_X = cute.make_tiled_copy(copy_atom_load_X, tv_layout, tiler_mn).get_slice(tidx)
-    thr_copy_X_async = cute.make_tiled_copy(copy_atom_load_X_async, tv_layout, tiler_mn).get_slice(tidx)
+    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()
-
-    # Don't use blkX.layout here, because the stride is N, not N_rounded
-    sX = smem.allocate_tensor(gX.element_type, cute.make_ordered_layout(blkX.shape, order=(1, 0)), byte_alignment=16)
+    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(
         # 2 stages: 1 for max, 1 for sum
-        (num_warps // warps_per_row, warps_per_row if cluster_n == 1 else (warps_per_row, cluster_n), 2),
+        (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)
@@ -57,14 +49,15 @@ def cross_entropy_kernel(
     else:
         mbar_ptr = None
 
-    #### Thread View
-    tXgX = thr_copy_X_async.partition_S(blkX)
-    tXsX = thr_copy_X_async.partition_S(sX)
-
-    tXcX = thr_copy_X.partition_S(blkCrd)[(0, None), None, 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)
 
-    # allocate fragments for gmem->rmem
-    tXrX = cute.make_fragment_like(tXgX)  # only logits fragment needed
+    #### 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:
@@ -80,54 +73,57 @@ def cross_entropy_kernel(
     if row < shape[0] and tXcX[0][1] == 0:
         target = cute.Int32(mTarget[row])
 
-    tXpX = cute.make_fragment_like(tXgX[(0, None), None, None], cutlass.Boolean)
-    for i in range(cute.size(tXpX)):
-        tXpX[i] = cute.elem_less(tXcX[i][1], shape[1])
+    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_async, tXgX, tXsX, pred=tXpX)
+        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])
 
-    max_x = utils.warp_reduce(
-        x.reduce(cute.ReductionOp.MAX, init_val=float('-inf'), reduction_profile=0),
-        cute.arch.fmax,
-        width=utils.min_constexpr(tv_layout.shape[0][0], cute.arch.WARP_SIZE),
+    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
     )
-    if cutlass.const_expr(cluster_n > 1):
-        cute.arch.cluster_wait()
-    if cutlass.const_expr(warps_per_row > 1 or cluster_n > 1):
-        max_mbar_ptr = mbar_ptr + 0 if cluster_n > 1 else None
-        max_x = utils.block_or_cluster_reduce(
-            max_x, cute.arch.fmax, reduction_buffer[None, None, 0], max_mbar_ptr, init_val=-cutlass.Float32.inf
-        )
     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.warp_reduce(
-        exp_x.reduce(cute.ReductionOp.ADD, init_val=0.0, reduction_profile=0),
-        operator.add,
-        width=utils.min_constexpr(tv_layout.shape[0][0], cute.arch.WARP_SIZE),
+    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,
     )
-    if cutlass.const_expr(warps_per_row > 1 or cluster_n > 1):
-        sum_mbar_ptr = mbar_ptr + 1 if cluster_n > 1 else None
-        denom = utils.block_or_cluster_reduce(
-            denom, operator.add, reduction_buffer[None, None, 1], sum_mbar_ptr, init_val=0.0
-        )
 
-    if tXcX[0][1] == 0 and row < shape[0]:
+    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)
-        loss_val = -target_logit + max_x + utils.log2f(denom) * ln_2
-        if cutlass.const_expr(cluster_n == 1):
-            mLoss[row] = loss_val.to(mLoss.element_type)
-        else:
-            if cute.arch.block_idx_in_cluster() == 0:
-                mLoss[row] = loss_val.to(mLoss.element_type)
+        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
@@ -135,6 +131,7 @@ 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
@@ -161,7 +158,7 @@ def cross_entropy_interface(
     )
 
     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, tv_layout, tiler_mn, cluster_n).launch(
+    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
@@ -181,6 +178,7 @@ torch2cute_dtype_map = {
 def cross_entropy(
     x: torch.Tensor,
     target: torch.Tensor,
+    return_lse: bool = False,
 ) -> torch.Tensor:
     """Cross entropy forward pass.
 
@@ -199,7 +197,8 @@ def cross_entropy(
     assert target.dtype == torch.int64, "Target must be int64"
     M, N = x.shape
     device = x.device
-    loss = torch.empty(M, device=device, dtype=x.dtype)
+    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)
@@ -207,15 +206,16 @@ def cross_entropy(
     )
     x_tensor, = [convert_from_dlpack(tensor) for tensor in (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
     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)
+    compile_key = (dtype, N, lse_tensor is not None)
     if compile_key not in cross_entropy.compile_cache:
         cross_entropy.compile_cache[compile_key] = cute.compile(
-            cross_entropy_interface, x_tensor, target_tensor, loss_tensor, stream, N
+            cross_entropy_interface, x_tensor, target_tensor, loss_tensor, lse_tensor, stream, N
         )
-    cross_entropy.compile_cache[compile_key](x_tensor, target_tensor, loss_tensor, stream)
-    return loss
+    cross_entropy.compile_cache[compile_key](x_tensor, target_tensor, loss_tensor, lse_tensor, stream)
+    return loss if not return_lse else (loss, lse)
 
 
 cross_entropy.compile_cache = {}

quack/rmsnorm.py

@@ -16,13 +16,11 @@ import quack.utils as utils
 
 @cute.kernel
 def rmsnorm_kernel(
-    gX: cute.Tensor,
-    gW: cute.Tensor,
-    gO: cute.Tensor,
-    gRstd: cute.Tensor,
-    cX: cute.Tensor,  # coordinate tensor
+    mX: cute.Tensor,
+    mW: cute.Tensor,
+    mO: cute.Tensor,
+    mRstd: cute.Tensor,
     eps: cute.Float32,
-    shape: cute.Shape,
     tv_layout: cute.Layout,
     tiler_mn: cute.Shape,
     cluster_n: cutlass.Constexpr = 1,
@@ -31,42 +29,45 @@ def rmsnorm_kernel(
 ):
     tidx, _, _ = cute.arch.thread_idx()
     bidx, cluster_y, _ = cute.arch.block_idx()
-    gdim, _, _ = cute.arch.grid_dim()
-
-    # slice for CTAs
-    # logical id -> address
-    blkX, blkOut, blkRstd, blkCrd = [gT[(None, None), bidx if cluster_n == 1 else (bidx, cluster_y)] for gT in (gX, gO, gRstd, cX)]
-    blkW = gW[(None, None), 0 if cluster_n == 1 else (0, cluster_y)]
-
-    # declare the atoms which will be used later for memory copy
-    copy_atom_load_X = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), gX.element_type, num_bits_per_copy=128)
-    copy_atom_load_X_async = cute.make_copy_atom(cute.nvgpu.cpasync.CopyG2SOp(), gX.element_type, num_bits_per_copy=128)
-    copy_atom_load_W = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), gW.element_type, num_bits_per_copy=128)
-    copy_atom_store_O = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), gO.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)
-    thr_copy_X_async = 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)
 
     smem = cutlass.utils.SmemAllocator()
-    # Don't use blkX.layout here, because the stride is N, not N_rounded
-    sX = smem.allocate_tensor(gX.element_type, cute.make_ordered_layout(blkX.shape, order=(1, 0)), byte_alignment=16)
+    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((num_warps // warps_per_row, warps_per_row), order=(1, 0))
-    reduction_buffer_layout = cute.make_ordered_layout((num_warps // warps_per_row, warps_per_row if cluster_n == 1 else (warps_per_row, cluster_n)), order=(1, 0))
+    reduction_buffer_layout = cute.make_ordered_layout(
+        (num_warps // warps_per_row, (warps_per_row, cluster_n)),
+        order=(1, 0)
+    )
     reduction_buffer = smem.allocate_tensor(cutlass.Float32, reduction_buffer_layout, byte_alignment=4)
     if cutlass.const_expr(cluster_n > 1):
         mbar_ptr = smem.allocate(cutlass.Int64.width // 8, byte_alignment=8)
     else:
         mbar_ptr = None
 
-    tWgW = thr_copy_W.partition_S(blkW)
-    tXgX = thr_copy_X_async.partition_S(blkX)
-    tXsX = thr_copy_X_async.partition_S(sX)
-    tXgO, tXrRstd = [thr_copy_O.partition_D(blk) for blk in (blkOut, blkRstd)]
-    tXcX = thr_copy_X.partition_S(blkCrd)[(0, None), None, None]
+    shape = mX.shape
+    idX = cute.make_identity_tensor(shape)
+    # slice for CTAs
+    gX, gO, gRstd, cX = [
+        cute.local_tile(mT, tiler_mn, (bidx, 0 if cluster_n == 1 else cluster_y))
+        for mT in (mX, mO, mRstd, idX)
+    ]
+    gW = cute.local_tile(mW, tiler_mn, (0, 0 if cluster_n == 1 else cluster_y))
+
+    # 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, tXrRstd = [thr_copy_O.partition_D(gT) for gT in (gO, gRstd)]
+    tXcX = thr_copy_X.partition_S(cX)[(0, None), None, None]
 
     # allocate fragments for gmem->rmem
     tWrW = cute.make_fragment_like(tWgW)
@@ -82,44 +83,33 @@ def rmsnorm_kernel(
         # Cluster arrive after barrier init
         cute.arch.cluster_arrive_relaxed()
 
-    tXpX = cute.make_fragment_like(tXgX[(0, None), None, None], cutlass.Boolean)
-    for i in range(cute.size(tXpX)):
-        tXpX[i] = cute.elem_less(tXcX[i][1], shape[1])
-    # tXrX.fill(0.0)
-    if tXcX[0][0] < shape[0]:
-        # cute.copy(copy_atom_load_X, tXgX, tXrX, pred=tXpX)
+    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 = cute.make_fragment_like(tWgW[(0, None), None, None], cutlass.Boolean)
-    tWcX = thr_copy_W.partition_S(blkCrd)[(0, None), None, None]
-    for i in range(cute.size(tWpW)):
-        tWpW[i] = cute.elem_less(tWcX[i][1], shape[1])
+    tWpW = utils.predicate_k(thr_copy_W.partition_S(cX), limit=shape[1])
     if 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)
-    sum_sq_x = utils.warp_reduce(
-        (x * x).reduce(cute.ReductionOp.ADD, init_val=0.0, reduction_profile=0),
-        operator.add,
-        width=utils.min_constexpr(tv_layout.shape[0][0], cute.arch.WARP_SIZE),
+    threads_per_row = tv_layout.shape[0][0]
+    sum_sq_x = utils.row_reduce(
+        x * x,
+        cute.ReductionOp.ADD,
+        threads_per_row,
+        reduction_buffer,
+        mbar_ptr,
+        init_val=0.0,
+        hook_fn=cute.arch.cluster_wait if cutlass.const_expr(cluster_n > 1) else None
     )
-    if cutlass.const_expr(cluster_n > 1):
-        cute.arch.cluster_wait()
-    if cutlass.const_expr(warps_per_row > 1 or cluster_n > 1):
-        sum_sq_x = utils.block_or_cluster_reduce(
-            sum_sq_x, operator.add, reduction_buffer, mbar_ptr, init_val=0.0
-        )
     rstd = utils.rsqrt(sum_sq_x / shape[1] + eps)
     # Only the thread corresponding to column 0 writes out the rstd to gmem
-    if tXcX[0][1] == 0 and tXcX[0][0] < shape[0]:
-        if cutlass.const_expr(cluster_n == 1):
-            tXrRstd[0] = rstd
-        else:
-            if cute.arch.block_idx_in_cluster() == 0:
-                tXrRstd[0] = rstd
+    if tXcX[0][1] == 0 and row < shape[0] and (cluster_n == 1 or cute.arch.block_idx_in_cluster() == 0):
+        tXrRstd[0] = rstd
     if delay_w_load:
         cute.copy(copy_atom_load_W, tWgW, tWrW, pred=tWpW)
     if reload_from == "smem":
@@ -132,20 +122,16 @@ def rmsnorm_kernel(
     w = tXrW.load().to(cute.Float32)
     y = x_hat * w
     tXrO.store(y.to(tXrO.element_type))
-    tOcX = thr_copy_O.partition_S(blkCrd)[(0, None), None, None]
-    tOpO = cute.make_fragment_like(tXgO[(0, None), None, None], cutlass.Boolean)
-    for i in range(cute.size(tOpO)):
-        tOpO[i] = cute.elem_less(tOcX[i][1], shape[1])
-    if tXcX[0][0] < shape[0]:
+    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)
 
 
 @cute.jit
 def rmsnorm_interface(
-    # mX_: cute.Tensor,
     mX: cute.Tensor,
     mW: cute.Tensor,
-    mOut: cute.Tensor,
+    mO: cute.Tensor,
     mRstd: cute.Tensor,
     stream: cuda.CUstream,
     N: cutlass.Constexpr,
@@ -180,21 +166,18 @@ def rmsnorm_interface(
     mW_expanded = cute.make_tensor(mW.iterator, mW_expanded_layout)
     mRstd_expanded_layout = cute.append(mRstd.layout, cute.make_layout((N,), stride=(0,)))
     mRstd_expanded = cute.make_tensor(mRstd.iterator, mRstd_expanded_layout)
-    idX = cute.make_identity_tensor(mX.shape)
-    gX, gW, gO, gRstd, cX = [cute.zipped_divide(mT, tiler_mn) for mT in (mX, mW_expanded, mOut, mRstd_expanded, idX)]  # ((TileM,TileN),(RestM,RestN))
 
     # reload_from = None if N <= 16384 else ("smem" if N <= 32768 else "gmem")
     reload_from = None if N <= 16384 else "smem"
     # delay_w_load = N > 64 * 1024
     delay_w_load = False
     N_rounded = tiler_mn[1]
-    rmsnorm_kernel(gX, gW, gO, gRstd, cX, eps, mX.shape, tv_layout, tiler_mn, cluster_n, reload_from).launch(
-        grid=[cute.size(gX, mode=[1, 0]), cluster_n, 1],
+    rmsnorm_kernel(mX, mW_expanded, mO, mRstd_expanded, eps, tv_layout, tiler_mn, cluster_n, reload_from).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,
-        # We don't want to use gX.layout[0] here since that has stride in N, not N_rounded, leading IMA on smem
-        smem=cute.size_in_bytes(mX.element_type, cute.make_layout(gX.shape[0])) + num_warps * cluster_n * (cutlass.Float32.width // 8) + (cutlass.Int64.width // 8),
+        smem=cute.size_in_bytes(mX.element_type, cute.make_layout(tiler_mn)) + num_warps * cluster_n * (cutlass.Float32.width // 8) + (cutlass.Int64.width // 8),
         stream=stream,
     )
 

quack/softmax.py

@@ -15,40 +15,30 @@ import quack.utils as utils
 
 @cute.kernel
 def softmax_kernel(
-    gX: cute.Tensor,
-    gO: cute.Tensor,
-    cX: cute.Tensor,  # coordinate tensor
-    shape: cute.Shape,
+    mX: cute.Tensor,
+    mO: cute.Tensor,
     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()
-    gdim, _, _ = cute.arch.grid_dim()
 
+    shape = mX.shape
+    idX = cute.make_identity_tensor(shape)
     # slice for CTAs
-    # logical id -> address
-    blkX, blkOut, blkCrd = [gT[(None, None), bidx if cluster_n == 1 else (bidx, cluster_y)] for gT in (gX, gO, cX)]
-
-    # declare the atoms which will be used later for memory copy
-    copy_atom_load_X = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), gX.element_type, num_bits_per_copy=128)
-    copy_atom_load_X_async = cute.make_copy_atom(cute.nvgpu.cpasync.CopyG2SOp(), gX.element_type, num_bits_per_copy=128)
-    copy_atom_store_O = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), gO.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)
-    thr_copy_X_async = cute.make_tiled_copy(copy_atom_load_X_async, 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)
+    gX, gO, cX = [
+        cute.local_tile(mT, tiler_mn, (bidx, 0 if cluster_n == 1 else cluster_y))
+        for mT in (mX, mO, idX)
+    ]
 
     smem = cutlass.utils.SmemAllocator()
-    # Don't use blkX.layout here, because the stride is N, not N_rounded
-    sX = smem.allocate_tensor(gX.element_type, cute.make_ordered_layout(blkX.shape, order=(1, 0)), byte_alignment=16)
+    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(
         # 2 stages: 1 for max, 1 for sum
-        (num_warps // warps_per_row, warps_per_row if cluster_n == 1 else (warps_per_row, cluster_n), 2),
+        (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)
@@ -58,10 +48,17 @@ def softmax_kernel(
     else:
         mbar_ptr = None
 
-    tXgX = thr_copy_X_async.partition_S(blkX)
-    tXsX = thr_copy_X_async.partition_S(sX)
-    tXgO = thr_copy_O.partition_D(blkOut)
-    tXcX = thr_copy_X.partition_S(blkCrd)[(0, None), None, None]
+    # declare the atoms which will be used later for memory copy
+    copy_atom_load_X = cute.make_copy_atom(cute.nvgpu.cpasync.CopyG2SOp(), mX.element_type, num_bits_per_copy=128)
+    copy_atom_store_O = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), gO.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)
+    thr_copy_O = cute.make_tiled_copy(copy_atom_store_O, tv_layout, tiler_mn).get_slice(tidx)
+
+    tXgX = thr_copy_X.partition_S(gX)
+    tXsX = thr_copy_X.partition_D(sX)
+    tXgO = thr_copy_O.partition_D(gO)
+    tXcX = thr_copy_X.partition_S(cX)[(0, None), None, None]
 
     # allocate fragments for gmem->rmem
     tXrX, tXrO = [cute.make_fragment_like(thr) for thr in (tXgX, tXgO)]
@@ -75,49 +72,48 @@ def softmax_kernel(
         # Cluster arrive after barrier init
         cute.arch.cluster_arrive_relaxed()
 
-    tXpX = cute.make_fragment_like(tXgX[(0, None), None, None], cutlass.Boolean)
-    for i in range(cute.size(tXpX)):
-        tXpX[i] = cute.elem_less(tXcX[i][1], shape[1])
-
+    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 tXcX[0][0] < shape[0]:
-        cute.copy(copy_atom_load_X_async, tXgX, tXsX, pred=tXpX)
+        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)
-    max_x = utils.warp_reduce(
-        x.reduce(cute.ReductionOp.MAX, init_val=float('-inf'), reduction_profile=0),
-        cute.arch.fmax,
-        width=utils.min_constexpr(tv_layout.shape[0][0], cute.arch.WARP_SIZE),
+    # 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()
+    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
     )
-    if cutlass.const_expr(cluster_n > 1):
-        cute.arch.cluster_wait()
-    if cutlass.const_expr(warps_per_row > 1 or cluster_n > 1):
-        max_mbar_ptr = mbar_ptr + 0 if cluster_n > 1 else None
-        max_x = utils.block_or_cluster_reduce(
-            max_x, cute.arch.fmax, reduction_buffer[None, None, 0], max_mbar_ptr, init_val=-cutlass.Float32.inf
-        )
     log2_e = math.log2(math.e)
     exp_x = cute.math.exp2((x - max_x) * log2_e, fastmath=True)
-    denom = utils.warp_reduce(
-        exp_x.reduce(cute.ReductionOp.ADD, init_val=0.0, reduction_profile=0),
-        operator.add,
-        width=utils.min_constexpr(tv_layout.shape[0][0], cute.arch.WARP_SIZE),
+    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,
     )
-    if cutlass.const_expr(warps_per_row > 1 or cluster_n > 1):
-        sum_mbar_ptr = mbar_ptr + 1 if cluster_n > 1 else None
-        denom = utils.block_or_cluster_reduce(
-            denom, operator.add, reduction_buffer[None, None, 1], sum_mbar_ptr, init_val=0.0
-        )
     inv = 1.0 / denom
     y = exp_x * inv
-
     tXrO.store(y.to(tXrO.element_type))
-    tOcX = thr_copy_O.partition_S(blkCrd)[(0, None), None, None]
-    tOpO = cute.make_fragment_like(tXgO[(0, None), None, None], cutlass.Boolean)
-    for i in range(cute.size(tOpO)):
-        tOpO[i] = cute.elem_less(tOcX[i][1], shape[1])
+    tOpO = utils.predicate_k(thr_copy_O.partition_S(cX), limit=shape[1]) if not is_even_N else None
     if tXcX[0][0] < shape[0]:
         cute.copy(copy_atom_store_O, tXrO, tXgO, pred=tOpO)
 
@@ -125,7 +121,7 @@ def softmax_kernel(
 @cute.jit
 def softmax_interface(
     mX: cute.Tensor,
-    mOut: cute.Tensor,
+    mO: cute.Tensor,
     stream: cuda.CUstream,
     N: cutlass.Constexpr,
     copy_bits: cutlass.Constexpr = 128
@@ -149,12 +145,9 @@ def softmax_interface(
         stride=((vecsize * cols_per_block, 1), (cols_per_block, cols_per_block * vecsize * threads_per_row))
     )
 
-    idX = cute.make_identity_tensor(mX.shape)
-    gX, gO, cX = [cute.zipped_divide(mT, tiler_mn) for mT in (mX, mOut, idX)]  # ((TileM,TileN),(RestM,RestN))
-
-    smem_allocated = cute.size_in_bytes(mX.element_type, cute.make_layout(gX.shape[0])) + 2 * num_warps * cluster_n * (cutlass.Float32.width // 8) + 2 * (cutlass.Int64.width // 8)
-    softmax_kernel(gX, gO, cX, mX.shape, tv_layout, tiler_mn, cluster_n).launch(
-        grid=[cute.size(gX, mode=[1, 0]), cluster_n, 1],
+    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)
+    softmax_kernel(mX, mO, 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,

quack/utils.py

@@ -1,5 +1,6 @@
 # Copyright (c) 2025, Wentao Guo, Ted Zadouri, Tri Dao.
 
+import operator
 import math
 from typing import Type, Callable, Optional
 
@@ -57,7 +58,7 @@ def block_reduce(val: cute.Numeric, op: Callable, reduction_buffer: cute.Tensor,
     """reduction_buffer has shape (num_warps / warp_per_row, warps_per_row)
     """
     lane_idx, warp_idx = cute.arch.lane_idx(), cute.arch.warp_idx()
-    warps_per_row = reduction_buffer.shape[1]
+    warps_per_row = cute.size(reduction_buffer.shape[1])
     row_idx, col_idx = warp_idx // warps_per_row, warp_idx % warps_per_row
     if lane_idx == 0:
         reduction_buffer[row_idx, col_idx] = val
@@ -142,6 +143,46 @@ def block_or_cluster_reduce(val: cute.Numeric, op: Callable, reduction_buffer: c
         return cluster_reduce(val, op, reduction_buffer, mbar_ptr, init_val=init_val)
 
 
+@cute.jit
+def row_reduce(
+    x: cute.TensorSSA | cute.Numeric,
+    op: cute.ReductionOp,
+    threads_per_row: cutlass.Constexpr[int],
+    reduction_buffer: Optional[cute.Tensor] = None,
+    mbar_ptr: Optional[cute.Pointer] = None,
+    init_val: cute.Numeric = 0.0,
+    hook_fn: Optional[Callable] = None,
+) -> cute.Numeric:
+    """reduction_buffer must have shape (num_warps / warps_per_row, (warps_per_row, cluster_n))
+    """
+    if cutlass.const_expr(isinstance(x, cute.TensorSSA)):
+        val = x.reduce(op, init_val=init_val, reduction_profile=0)
+    else:
+        val = x
+    warp_op = {
+        cute.ReductionOp.ADD: operator.add,
+        cute.ReductionOp.MAX: cute.arch.fmax if cutlass.const_expr(x.dtype == cute.Float32) else max,
+        cute.ReductionOp.MIN: min,
+        cute.ReductionOp.MUL: operator.mul,
+    }[op]
+    val = warp_reduce(
+        val,
+        warp_op,
+        width=min_constexpr(threads_per_row, cute.arch.WARP_SIZE),
+    )
+    if cutlass.const_expr(hook_fn is not None):
+        hook_fn()
+    if cutlass.const_expr(reduction_buffer is not None):
+        warps_per_row, cluster_n = reduction_buffer.shape[1]
+        assert cluster_n == 1 or mbar_ptr is not None, "mbar_ptr must be provided for cluster reduction"
+        if cutlass.const_expr(warps_per_row > 1 or cluster_n > 1):
+            val = block_or_cluster_reduce(
+                val, warp_op, reduction_buffer, mbar_ptr, init_val=init_val
+            )
+    return val
+
+
+
 def exp2f(x: cute.TensorSSA | cutlass.Float32) -> cute.TensorSSA | cutlass.Float32:
     """exp2f calculation for both vector and scalar.
 
@@ -188,3 +229,18 @@ def rsqrt(a: float | cute.Float32, *, loc=None, ip=None) -> cute.Float32:
             asm_dialect=llvm.AsmDialect.AD_ATT,
         )
     )
+
+
+def predicate_k(tAcA: cute.Tensor, limit: cutlass.Int32) -> cute.Tensor:
+    # Only compute predicates for the "k" dimension. For the mn dimension, we will use "if"
+    tApA = cute.make_fragment(
+        cute.make_layout(
+            (cute.size(tAcA, mode=[0, 1]), cute.size(tAcA, mode=[1]), cute.size(tAcA, mode=[2])),
+            stride=(cute.size(tAcA, mode=[2]), 0, 1),
+        ),
+        cutlass.Boolean,
+    )
+    for rest_v in range(tApA.shape[0]):
+        for rest_k in range(tApA.shape[2]):
+            tApA[rest_v, 0, rest_k] = cute.elem_less(tAcA[(0, rest_v), 0, rest_k][1], limit)
+    return tApA

{quack_kernels-0.1.1.dist-info → quack_kernels-0.1.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: quack-kernels
-Version: 0.1.1
+Version: 0.1.2
 Requires-Python: >=3.9
 License-File: LICENSE
 Requires-Dist: nvidia-cutlass-dsl==4.0.0

quack_kernels-0.1.2.dist-info/RECORD

@@ -0,0 +1,10 @@
+quack/__init__.py,sha256=Nf01m1CGrOjSkqGJom6P65hSLkckljRMhlkSoqqlO9k,137
+quack/cross_entropy.py,sha256=gdo8sR9KT5TsrShbgAmy-bwRZLu0gTs_ykXBF2RMbFI,8900
+quack/rmsnorm.py,sha256=JhwJSAPDDpB_hV90xU9ymiLU-zu4WScrSHc5JX2JarY,10470
+quack/softmax.py,sha256=C8e8ZNaF5ePJ1NlrWZN1goCcvsx1C60FWlRyuFCcYoM,7737
+quack/utils.py,sha256=PRdu-P7azA_PeHUNdtoy1zyxZwg_QyVrSiVwE1iXaWo,8961
+quack_kernels-0.1.2.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+quack_kernels-0.1.2.dist-info/METADATA,sha256=3WjugLu1IhLlgsg2qUcLBZq1HI4-BIyyJIuQc5Hk-rU,186
+quack_kernels-0.1.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+quack_kernels-0.1.2.dist-info/top_level.txt,sha256=6e4Jr_vNJbZTYwlO_Ahf_sDeHDE0zcqcf7Le11FKxxo,6
+quack_kernels-0.1.2.dist-info/RECORD,,

quack_kernels-0.1.1.dist-info/RECORD

@@ -1,10 +0,0 @@
-quack/__init__.py,sha256=y3Oa4OVPqaGU_P1miI435DzfpMgIwKVmU8-Eogv58jQ,137
-quack/cross_entropy.py,sha256=V0kG8DCNh2735sPIDwe68NB50rAqDF3XQApnGyo-sKg,9220
-quack/rmsnorm.py,sha256=RNqcT-q4uvMbF6ejpzuqQH8l8VVuTRlnueXf28V47sc,11954
-quack/softmax.py,sha256=QABgOESH5JjDm3yuUkyZZKXXpzn7CTuMSs0NEBnFD80,8536
-quack/utils.py,sha256=ofV7QLDuq80h3nEA3TwZW-ti8CnYwMgnz1dpxpvhHpk,6859
-quack_kernels-0.1.1.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-quack_kernels-0.1.1.dist-info/METADATA,sha256=XG3zS0_q48TzkoR7CemzaJGVYHS731yVOrzH49_uRK8,186
-quack_kernels-0.1.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-quack_kernels-0.1.1.dist-info/top_level.txt,sha256=6e4Jr_vNJbZTYwlO_Ahf_sDeHDE0zcqcf7Le11FKxxo,6
-quack_kernels-0.1.1.dist-info/RECORD,,