quack-kernels 0.1.10__py3-none-any.whl → 0.2.0__py3-none-any.whl

quack/utils.py

@@ -1,15 +1,14 @@
 # Copyright (c) 2025, Wentao Guo, Ted Zadouri, Tri Dao.
 
-import operator
 import math
-from typing import Callable, Optional, Tuple
+from typing import Optional, Tuple, Type, Union
 
 import cutlass
 import cutlass.cute as cute
 
-from cutlass import Float32
+from cutlass import Float32, Int32
 from cutlass.cutlass_dsl import T, dsl_user_op
-from cutlass._mlir.dialects import llvm, vector
+from cutlass._mlir.dialects import llvm, nvvm, vector
 from cutlass.cute.runtime import from_dlpack
 
 
@@ -23,46 +22,20 @@ def convert_from_dlpack(x, leading_dim, alignment=16, divisibility=1) -> cute.Te
     )
 
 
-@cute.jit
-def warp_reduce(
-    val: cute.TensorSSA | cute.Numeric,
-    op: Callable,
-    width: cutlass.Constexpr[int] = cute.arch.WARP_SIZE,
-) -> cute.TensorSSA | cute.Numeric:
-    if cutlass.const_expr(isinstance(val, cute.TensorSSA)):
-        res = cute.make_fragment(val.shape, val.dtype)
-        res.store(val)
-        for i in cutlass.range_constexpr(cute.size(val.shape)):
-            res[i] = warp_reduce(res[i], op, width)
-        return res.load()
-    else:
-        for i in cutlass.range_constexpr(int(math.log2(width))):
-            val = op(val, cute.arch.shuffle_sync_bfly(val, offset=1 << i))
-    return val
-
-
-@cute.jit
-def block_reduce(
-    val: cute.Numeric, op: Callable, reduction_buffer: cute.Tensor, init_val: cute.Numeric = 0.0
-) -> cute.Numeric:
-    """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 = 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
-    cute.arch.barrier()
-    block_reduce_val = init_val
-    if lane_idx < warps_per_row:
-        block_reduce_val = reduction_buffer[row_idx, lane_idx]
-    return warp_reduce(block_reduce_val, op)
-
-
 @dsl_user_op
 def elem_pointer(x: cute.Tensor, coord: cute.Coord, *, loc=None, ip=None) -> cute.Pointer:
     return x.iterator + cute.crd2idx(coord, x.layout, loc=loc, ip=ip)
 
 
+@cute.jit
+def load_scalar_or_pointer(x: Float32 | cute.Pointer) -> Float32:
+    if cutlass.const_expr(isinstance(x, cute.Pointer)):
+        return Float32(cute.make_tensor(x, cute.make_layout(1))[0])
+    else:
+        assert isinstance(x, Float32)
+        return x
+
+
 @dsl_user_op
 def set_block_rank(
     smem_ptr: cute.Pointer, peer_cta_rank_in_cluster: cute.Int32, *, loc=None, ip=None
@@ -100,208 +73,31 @@ def store_shared_remote(
     ).ir_value()
     if cutlass.const_expr(isinstance(val, float)):
         val = Float32(val)
-    assert isinstance(val, (Float32, cutlass.Int64)), "val must be Float32 or Int64"
-    suffix = "f32" if cutlass.const_expr(isinstance(val, Float32)) else "s64"
+    assert isinstance(val, (Float32, Int32, cutlass.Int64)), "val must be Float32, Int32, or Int64"
+    suffix = {Float32: "f32", Int32: "s32", cutlass.Int64: "s64"}[type(val)]
+    constraint = {Float32: "f", Int32: "r", cutlass.Int64: "l"}[type(val)]
     llvm.inline_asm(
         None,
         [remote_smem_ptr_i32, val.ir_value(loc=loc, ip=ip), remote_mbar_ptr_i32],
         f"st.async.shared::cluster.mbarrier::complete_tx::bytes.{suffix} [$0], $1, [$2];",
-        f"r,{'f' if cutlass.const_expr(isinstance(val, Float32)) else 'l'},r",
+        f"r,{constraint},r",
         has_side_effects=True,
         is_align_stack=False,
         asm_dialect=llvm.AsmDialect.AD_ATT,
     )
 
 
-@cute.jit
-def cluster_reduce(
-    val: cute.Numeric,
-    op: Callable,
-    reduction_buffer: cute.Tensor,
-    mbar_ptr: cute.Pointer,
-    init_val: cute.Numeric = 0.0,
-    phase: Optional[cutlass.Int32] = None,
-) -> cute.Numeric:
-    """reduction_buffer has shape (num_warps / warps_per_row, (warps_per_row, cluster_n))"""
-    cta_rank_in_cluster = cute.arch.block_idx_in_cluster()
-    lane_idx, warp_idx = cute.arch.lane_idx(), cute.arch.warp_idx()
-    rows_per_block, (warps_per_row, cluster_n) = reduction_buffer.shape
-    row_idx, col_idx = warp_idx // warps_per_row, warp_idx % warps_per_row
-    if warp_idx == 0:
-        with cute.arch.elect_one():
-            num_warps = rows_per_block * warps_per_row
-            cute.arch.mbarrier_arrive_and_expect_tx(
-                mbar_ptr,
-                num_warps * cluster_n * reduction_buffer.element_type.width // 8,
-            )
-    if lane_idx < cluster_n:
-        store_shared_remote(
-            val,
-            elem_pointer(reduction_buffer, (row_idx, (col_idx, cta_rank_in_cluster))),
-            mbar_ptr,
-            peer_cta_rank_in_cluster=lane_idx,
+@dsl_user_op
+def fmin(a: Union[float, Float32], b: Union[float, Float32], *, loc=None, ip=None) -> Float32:
+    return Float32(
+        nvvm.fmin(
+            T.f32(),
+            Float32(a).ir_value(loc=loc, ip=ip),
+            Float32(b).ir_value(loc=loc, ip=ip),
+            loc=loc,
+            ip=ip,
         )
-    cute.arch.mbarrier_wait(mbar_ptr, phase=phase if phase is not None else 0)
-    block_reduce_val = init_val
-    num_iter = cute.ceil_div(warps_per_row * cluster_n, cute.arch.WARP_SIZE)
-    for i in cutlass.range_constexpr(num_iter):
-        idx = lane_idx + i * cute.arch.WARP_SIZE
-        if idx < cute.size(reduction_buffer, mode=[1]):
-            block_reduce_val = op(block_reduce_val, reduction_buffer[row_idx, idx])
-    return warp_reduce(block_reduce_val, op)
-
-
-@cute.jit
-def block_or_cluster_reduce(
-    val: cute.Numeric,
-    op: Callable,
-    reduction_buffer: cute.Tensor,
-    mbar_ptr: Optional[cute.Pointer],
-    phase: Optional[cutlass.Int32] = None,
-    init_val: cute.Numeric = 0.0,
-) -> cute.Numeric:
-    """Perform either block or cluster reduction based on whether mbar_ptr is provided."""
-    if cutlass.const_expr(mbar_ptr is None):
-        return block_reduce(val, op, reduction_buffer, init_val=init_val)
-    else:
-        return cluster_reduce(val, op, reduction_buffer, mbar_ptr, phase=phase, 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,
-    phase: Optional[cutlass.Int32] = 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 == Float32) else max,
-        cute.ReductionOp.MIN: min,
-        cute.ReductionOp.MUL: operator.mul,
-    }[op]
-    val = warp_reduce(
-        val,
-        warp_op,
-        width=min(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, phase=phase, init_val=init_val
-            )
-    return val
-
-
-@cute.jit
-def online_softmax_reduce(
-    x: cute.TensorSSA,
-    threads_per_row: cutlass.Constexpr[int],
-    reduction_buffer: Optional[cute.Tensor] = None,
-    mbar_ptr: Optional[cute.Pointer] = None,
-    hook_fn: Optional[Callable] = None,
-    phase: Optional[cutlass.Int32] = None,
-    return_exp_x: bool = False,
-) -> [Float32, Float32, Optional[cute.TensorSSA]]:
-    assert x.dtype == Float32, "x must be of type Float32"
-    """reduction_buffer must have shape (num_warps / warps_per_row, (warps_per_row, cluster_n), 2)"""
-    max_x = warp_reduce(
-        x.reduce(cute.ReductionOp.MAX, init_val=-Float32.inf, reduction_profile=0),
-        cute.arch.fmax,
-        width=min(threads_per_row, cute.arch.WARP_SIZE),
-    )
-    log2_e = math.log2(math.e)
-    exp_x = exp2f(x * log2_e - (max_x * log2_e))
-    # exp_x = exp2f((x - max_x) * log2_e)
-    sum_exp_x = warp_reduce(
-        exp_x.reduce(cute.ReductionOp.ADD, init_val=0.0, reduction_profile=0),
-        operator.add,
-        width=min(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):
-        rows_per_block, (warps_per_row, cluster_n) = reduction_buffer.shape
-        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):
-            assert (
-                reduction_buffer.element_type == cutlass.Int64
-            ), "reduction_buffer must be of type cute.Int64"
-            lane_idx, warp_idx = cute.arch.lane_idx(), cute.arch.warp_idx()
-            row_idx, col_idx = warp_idx // warps_per_row, warp_idx % warps_per_row
-            if cutlass.const_expr(mbar_ptr is None):
-                if lane_idx == 0:
-                    reduction_buffer[row_idx, col_idx] = f32x2_to_i64(max_x, sum_exp_x)
-                cute.arch.barrier()
-                max_x_single_warp = -Float32.inf
-                sum_exp_x = 0.0
-                if lane_idx < warps_per_row:
-                    max_x_single_warp, sum_exp_x = i64_to_f32x2(reduction_buffer[row_idx, lane_idx])
-                max_x_final = warp_reduce(max_x_single_warp, cute.arch.fmax)
-                sum_exp_x *= exp2f((max_x_single_warp - max_x_final) * log2_e)
-                sum_exp_x = warp_reduce(sum_exp_x, operator.add)
-                if cutlass.const_expr(return_exp_x):
-                    exp_x *= exp2f((max_x - max_x_final) * log2_e)
-                max_x = max_x_final
-            else:
-                cta_rank_in_cluster = cute.arch.block_idx_in_cluster()
-                if warp_idx == 0:
-                    with cute.arch.elect_one():
-                        num_warps = rows_per_block * warps_per_row
-                        cute.arch.mbarrier_arrive_and_expect_tx(
-                            mbar_ptr,
-                            num_warps * cluster_n * reduction_buffer.element_type.width // 8,
-                        )
-                if lane_idx < cluster_n:
-                    store_shared_remote(
-                        f32x2_to_i64(max_x, sum_exp_x),
-                        elem_pointer(reduction_buffer, (row_idx, (col_idx, cta_rank_in_cluster))),
-                        mbar_ptr,
-                        peer_cta_rank_in_cluster=lane_idx,
-                    )
-                cute.arch.mbarrier_wait(mbar_ptr, phase=phase if phase is not None else 0)
-                num_iter = cute.ceil_div(warps_per_row * cluster_n, cute.arch.WARP_SIZE)
-                max_x_single_warp = cute.make_fragment(num_iter, Float32)
-                max_x_single_warp.fill(-Float32.inf)
-                sum_exp_x_single_warp = cute.make_fragment(num_iter, Float32)
-                sum_exp_x_single_warp.fill(0.0)
-                for i in cutlass.range_constexpr(num_iter):
-                    idx = lane_idx + i * cute.arch.WARP_SIZE
-                    if idx < cute.size(reduction_buffer, mode=[1]):
-                        max_x_single_warp[i], sum_exp_x_single_warp[i] = i64_to_f32x2(
-                            reduction_buffer[row_idx, idx]
-                        )
-                max_x_final = max_x_single_warp.load().reduce(
-                    cute.ReductionOp.MAX, init_val=-Float32.inf, reduction_profile=0
-                )
-                max_x_final = warp_reduce(max_x_final, cute.arch.fmax)
-                sum_exp_x = 0.0
-                for i in cutlass.range_constexpr(num_iter):
-                    sum_exp_x += sum_exp_x_single_warp[i] * exp2f(
-                        (max_x_single_warp[i] - max_x_final) * log2_e
-                    )
-                sum_exp_x = warp_reduce(sum_exp_x, operator.add)
-                if cutlass.const_expr(return_exp_x):
-                    exp_x *= exp2f((max_x - max_x_final) * log2_e)
-                max_x = max_x_final
-    return max_x, sum_exp_x, (exp_x if cutlass.const_expr(return_exp_x) else None)
 
 
 @cute.jit
@@ -337,6 +133,21 @@ def log2f(a: float | Float32, *, loc=None, ip=None) -> Float32:
     )
 
 
+@dsl_user_op
+def sqrt(a: float | Float32, *, loc=None, ip=None) -> Float32:
+    return Float32(
+        llvm.inline_asm(
+            T.f32(),
+            [Float32(a).ir_value(loc=loc, ip=ip)],
+            "sqrt.approx.ftz.f32 $0, $1;",
+            "=f,f",
+            has_side_effects=False,
+            is_align_stack=False,
+            asm_dialect=llvm.AsmDialect.AD_ATT,
+        )
+    )
+
+
 @dsl_user_op
 def rsqrt(a: float | Float32, *, loc=None, ip=None) -> Float32:
     return Float32(
@@ -352,6 +163,73 @@ def rsqrt(a: float | Float32, *, loc=None, ip=None) -> Float32:
     )
 
 
+@dsl_user_op
+def ceil(a: float | Float32, *, loc=None, ip=None) -> Int32:
+    return Int32(
+        llvm.inline_asm(
+            T.i32(),
+            [Float32(a).ir_value(loc=loc, ip=ip)],
+            "cvt.rpi.ftz.s32.f32 $0, $1;",
+            "=r,f",
+            has_side_effects=False,
+            is_align_stack=False,
+            asm_dialect=llvm.AsmDialect.AD_ATT,
+        )
+    )
+
+
+@dsl_user_op
+def prmt(a: int | Int32, b: int | Int32, c: int | Int32, *, loc=None, ip=None) -> Int32:
+    return Int32(
+        llvm.inline_asm(
+            T.i32(),
+            [
+                Int32(a).ir_value(loc=loc, ip=ip),
+                Int32(b).ir_value(loc=loc, ip=ip),
+                Int32(c).ir_value(loc=loc, ip=ip),
+            ],
+            "prmt.b32 $0, $1, $2, $3;",
+            "=r,r,r,r",
+            has_side_effects=False,
+            is_align_stack=False,
+            asm_dialect=llvm.AsmDialect.AD_ATT,
+        )
+    )
+
+
+@cute.jit
+def permute_gated_Cregs_b16(t: cute.Tensor) -> None:
+    assert t.element_type.width == 16
+    assert cute.size(t.shape) % 4 == 0, "Tensor size must be a multiple of 4 for b16 permutation"
+    t_u32 = cute.recast_tensor(t, Int32)
+
+    quad_idx = cute.arch.lane_idx() % 4
+    lane_03 = quad_idx == 0 or quad_idx == 3
+    selector_upper = Int32(0x5410) if lane_03 else Int32(0x1054)
+    selector_lower = Int32(0x7632) if lane_03 else Int32(0x3276)
+    # upper_map = [0, 3, 1, 2]
+    # lower_map = [1, 2, 0, 3]
+    # upper_idx = upper_map[quad_idx]
+    # indexing isn't supported so we have to do arithmetic
+    upper_idx = quad_idx // 2 if quad_idx % 2 == 0 else 3 - quad_idx // 2
+    lower_idx = upper_idx ^ 1
+
+    # 1 -> 0b11111, 2 -> 0b11110, 4 -> 0b11100, 8 -> 0b11000, 16 -> 0b10000, 32 -> 0b00000
+    width = 4
+    mask = cute.arch.WARP_SIZE - width
+    clamp = cute.arch.WARP_SIZE - 1
+    mask_and_clamp = mask << 8 | clamp
+
+    for i in cutlass.range(cute.size(t_u32.shape) // 2, unroll_full=True):
+        upper, lower = t_u32[i * 2 + 0], t_u32[i * 2 + 1]
+        upper0 = upper if lane_03 else lower
+        lower0 = lower if lane_03 else upper
+        upper0 = cute.arch.shuffle_sync(upper0, offset=upper_idx, mask_and_clamp=mask_and_clamp)
+        lower0 = cute.arch.shuffle_sync(lower0, offset=lower_idx, mask_and_clamp=mask_and_clamp)
+        t_u32[i * 2 + 0] = prmt(upper0, lower0, selector_upper)
+        t_u32[i * 2 + 1] = prmt(upper0, lower0, selector_lower)
+
+
 @cute.jit
 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"
@@ -377,7 +255,7 @@ def fill_oob(tXsX: cute.Tensor, tXpX: Optional[cute.Tensor], fill_value: cute.Nu
         tXpX: Predicate tensor indicating valid elements
         fill_value: Value to fill OOB locations with
     """
-    tXrX_fill = cute.make_fragment_like(tXsX[(None, 0), 0, 0])
+    tXrX_fill = cute.make_fragment_like(tXsX[(None, 0), None, 0])
     tXrX_fill.fill(fill_value)
     for rest_v in cutlass.range_constexpr(tXsX.shape[0][1]):
         for rest_k in cutlass.range_constexpr(tXsX.shape[2]):
@@ -446,3 +324,98 @@ def coord_offset_i64(
         assumed_align=tensor.iterator.max_alignment,
     )
     return cute.make_tensor(new_ptr, tensor.layout)
+
+
+@cute.jit
+def warp_prefix_sum(val: cutlass.Int32, lane: Optional[cutlass.Int32] = None) -> cutlass.Int32:
+    if cutlass.const_expr(lane is None):
+        lane = cute.arch.lane_idx()
+    for i in cutlass.range_constexpr(int(math.log2(cute.arch.WARP_SIZE))):
+        offset = 1 << i
+        # Very important that we set mask_and_clamp to 0
+        partial_sum = cute.arch.shuffle_sync_up(val, offset=offset, mask_and_clamp=0)
+        if lane >= offset:
+            val += partial_sum
+    return val
+
+
+def convert_layout_acc_mn(acc_layout: cute.Layout) -> cute.Layout:
+    """
+    For Sm80, convert ((2, 2), MMA_M, MMA_N, ...) to ((2, MMA_M), (2, MMA_N), ...).
+    For Sm90, convert ((2, 2, V), MMA_M, MMA_N, ...) to ((2, MMA_M), (2, V, MMA_N), ...).
+    """
+    acc_layout_col_major = cute.make_layout(acc_layout.shape)
+    acc_layout_mn = cute.make_layout(
+        (
+            (acc_layout_col_major.shape[0][1], acc_layout_col_major.shape[1]),  # MMA_M
+            (
+                acc_layout_col_major.shape[0][0],
+                *acc_layout_col_major.shape[0][2:],
+                acc_layout_col_major.shape[2],
+            ),  # MMA_N
+            *acc_layout_col_major.shape[3:],
+        ),
+        stride=(
+            (acc_layout_col_major.stride[0][1], acc_layout_col_major.stride[1]),  # MMA_M
+            (
+                acc_layout_col_major.stride[0][0],
+                *acc_layout_col_major.stride[0][2:],
+                acc_layout_col_major.stride[2],
+            ),  # MMA_N
+            *acc_layout_col_major.stride[3:],
+        ),
+    )
+    return cute.composition(acc_layout, acc_layout_mn)
+
+
+def make_acc_tensor_mn_view(acc: cute.Tensor) -> cute.Tensor:
+    return cute.make_tensor(acc.iterator, convert_layout_acc_mn(acc.layout))
+
+
+@dsl_user_op
+def sm90_get_smem_load_op(
+    layout_c: cutlass.utils.LayoutEnum,
+    elem_ty_c: Type[cutlass.Numeric],
+    *,
+    loc=None,
+    ip=None,
+) -> cute.CopyAtom:
+    """
+    Selects the largest vectorized smem load atom available subject to constraint of gmem layout.
+
+    Parameters:
+    -----------
+    layout_c : LayoutEnum
+        The layout enum of the output tensor D.
+
+    elem_ty_c : Type[Numeric]
+        The element type for output tensor D.
+
+    Returns:
+    --------
+    Either SmemLoadMatrix or SimtSyncCopy, based on the input parameters.
+    """
+
+    if not isinstance(elem_ty_c, cutlass.cutlass_dsl.NumericMeta):
+        raise TypeError(f"elem_ty_c must be a Numeric, but got {elem_ty_c}")
+    is_m_major = layout_c.is_m_major_c()
+    if elem_ty_c.width == 16:
+        return cute.make_copy_atom(
+            cute.nvgpu.warp.LdMatrix8x8x16bOp(is_m_major, 4), elem_ty_c, loc=loc, ip=ip
+        )
+    else:
+        return cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), elem_ty_c, loc=loc, ip=ip)
+
+
+@dsl_user_op
+def atomic_add_i32(a: int | Int32, gmem_ptr: cute.Pointer, *, loc=None, ip=None) -> Int32:
+    return nvvm.atomicrmw(
+        res=T.i32(), op=nvvm.AtomicOpKind.ADD, ptr=gmem_ptr.llvm_ptr, a=Int32(a).ir_value()
+    )
+
+
+@dsl_user_op
+def atomic_inc_i32(a: int | Int32, gmem_ptr: cute.Pointer, *, loc=None, ip=None) -> Int32:
+    return nvvm.atomicrmw(
+        res=T.i32(), op=nvvm.AtomicOpKind.INC, ptr=gmem_ptr.llvm_ptr, a=Int32(a).ir_value()
+    )

quack/varlen_utils.py

@@ -0,0 +1,22 @@
+# Copyright (c) 2025, Tri Dao.
+
+from typing import Optional
+from dataclasses import dataclass
+
+import cutlass.cute as cute
+
+from quack.cute_dsl_utils import ArgumentsBase
+
+
+# Grouping arguments together that should be passed to __call__
+@dataclass
+class VarlenArguments(ArgumentsBase):
+    mCuSeqlensM: Optional[cute.Tensor] = None
+    mCuSeqlensK: Optional[cute.Tensor] = None
+    mTensormaps: Optional[cute.Tensor] = None
+
+    def __post_init__(self):
+        if self.mCuSeqlensM is not None or self.mCuSeqlensK is not None:
+            assert (
+                self.mTensormaps is not None
+            ), "mTensormaps must be provided if mCuSeqlensM or mCuSeqlensK is provided"

{quack_kernels-0.1.10.dist-info → quack_kernels-0.2.0.dist-info}/METADATA

@@ -1,9 +1,9 @@
 Metadata-Version: 2.4
 Name: quack-kernels
-Version: 0.1.10
+Version: 0.2.0
 Requires-Python: >=3.12
 License-File: LICENSE
-Requires-Dist: nvidia-cutlass-dsl==4.1.0
+Requires-Dist: nvidia-cutlass-dsl==4.2.0
 Requires-Dist: torch
 Provides-Extra: dev
 Requires-Dist: pre-commit; extra == "dev"

quack_kernels-0.2.0.dist-info/RECORD

@@ -0,0 +1,37 @@
+quack/__init__.py,sha256=fGBYbb9JlaNT7HdtUTbUnuAkL5G2Dg8XZAA5Ir1R-ow,364
+quack/activation.py,sha256=ysXaVUXX2yGQC5o4ZVeRXw_fDIHOrqnzpHJaIsc0kHc,10271
+quack/autotuner.py,sha256=czO6JrYL0EJpOeJOYDSsVdrJaFuwfL3vTdG8QfL1F34,10792
+quack/cross_entropy.py,sha256=Kc3P83Vsu1nGaCu7llsO3vct3J_t3frRYPxij7JfHMA,28619
+quack/cute_dsl_utils.py,sha256=D2Pw7rzX9jY8u8wikIPvPvinmFLCDeZg95HPBLqGej4,4635
+quack/dense_gemm_sm100.py,sha256=hKBNC34UxdctrTKVP68nvANZl4Dq2rnUjRcweESEq3g,109965
+quack/dense_gemm_sm90.py,sha256=TjnjHnjhAwWH5YQWsFlADq07xSxtsprkw_p2Cy0yw7I,100407
+quack/fast_math.py,sha256=E1XUqfUt0_n9BPZNggF-UDzZ6anso9bYUrwqafemWvQ,2297
+quack/gemm_act_sm90.py,sha256=N5UAFWZvw1na22Vh5JSGgcdqZ2zI6kQMBVOLxYbCAUU,14332
+quack/gemm_config.py,sha256=gbYjPFeyT5wAhVwFQroRHlHoMKEJqAWX9P8wWy04l8Q,2258
+quack/gemm_dact_sm90.py,sha256=KCXgjOzdamSDexwrwf_pX2r-ippPRirbClrlU6BP7b8,4990
+quack/gemm_interface.py,sha256=_JTpE7zQw6NUw-v65Wql_XUOZBfW0oSEgiMnharTJU4,20501
+quack/gemm_wrapper_utils.py,sha256=aMMtu-Ojhtjay_5xJH4AjP-JRVks1AB8jmtNme_DIqU,5960
+quack/layernorm.py,sha256=JkK0sVdUfZ-SmoBmNqLF3wCiszDbdorvcBH2julv0Vg,13560
+quack/linear.py,sha256=SrhRiAFjC7ONIMVmiNu-kSPLHNUyaCXt59a1f_5nNXo,9383
+quack/linear_cross_entropy.py,sha256=Zhy_gdMsKHOie-jntBaqIuiDJtkiq6qEBwnyuWwIRw4,10092
+quack/mlp.py,sha256=YjdwQRwEePA9KyidFXp5H1-lxiJc8dZ41vl8Fv8pgss,2259
+quack/pipeline.py,sha256=DyCwZX8WvoUBFcMBz7CeYm9VUM31haEGgBhAzmxu8cE,5519
+quack/reduce.py,sha256=hsYByu6haCZjLTLB-qpYmKDjqS2UqlwPgfWTup38GNA,10341
+quack/reduction_base.py,sha256=CT-t_j7z8H1ByD9FkQYDRik_-THMDFv9QoXHmr9Xx9E,3636
+quack/rmsnorm.py,sha256=93qlTPjY9JBm3R5M-HeHse1PbAfD9931G3OFs71yo_g,48998
+quack/softmax.py,sha256=Mq3_2Ul8H64zeGUI9wOKEpIISJnrCcHQpZvk2sb10Tg,17101
+quack/symmetric_dense_gemm_sm90.py,sha256=2UXooIpClT2izdyGis1XaIgYYlLj-7MrcOMg2yR7YCk,88694
+quack/tensormap_manager.py,sha256=Ts3Mxp0_es2RNA0ffvUjWMXN79lsfWEBZ0DQYhtbcnw,5338
+quack/tile_scheduler.py,sha256=8qqYmx6GpQzt8XiidcrdLIaWf0TGbJVdwKFfeb1X_us,42265
+quack/topk.py,sha256=RQl-23lIicQ9ry9Njur8i0JGem_WbO_Gchr6jy8EtVM,9185
+quack/utils.py,sha256=tiqeJZiPPFl5irQWCUd7dTPA_OAv4SjHUW5S-u9wO8Y,14526
+quack/varlen_utils.py,sha256=vkduMEpo5bJJvZRNnIcKPb6pp1wD34vaIpMIB0ZGIZA,681
+quack/sort/bitonic_sort.py,sha256=8t0SG1a6iEpYIlY8YM_AWvm4aN-4AA4vEzdBuJMJm9g,4768
+quack/sort/generate_sorting_networks.py,sha256=vkJBOjTVEinQkWT4OtFqOWxFVdTIPoNAQocneKc9-rM,14477
+quack/sort/sorting_networks.py,sha256=l_26zi3gXD_z-tnm2eAczRrmE-mbaz00KmqH6ONivL8,9686
+quack/sort/utils.py,sha256=Mkr-l97RMAV-ZoNrwuzA1U3KO0Wjr38CV9Jm7ScyZoI,1090
+quack_kernels-0.2.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+quack_kernels-0.2.0.dist-info/METADATA,sha256=DAeQymRUqp7lSfSTNyS7TZF3oWcFzCKriGJ2p8JLu6A,285
+quack_kernels-0.2.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+quack_kernels-0.2.0.dist-info/top_level.txt,sha256=6e4Jr_vNJbZTYwlO_Ahf_sDeHDE0zcqcf7Le11FKxxo,6
+quack_kernels-0.2.0.dist-info/RECORD,,

quack_kernels-0.1.10.dist-info/RECORD

@@ -1,13 +0,0 @@
-quack/__init__.py,sha256=4tLchTx7d0d1ZVg6psRjjoXAWKHqzIWRF5mUk8ZdgkQ,204
-quack/cross_entropy.py,sha256=xsg2bXZ4wNvusBARhN4PwAzm5PbejEcfwj71nR7bzuE,20852
-quack/dense_gemm_sm90.py,sha256=jULXfAQkRh1SUAOpesx8wouY-GLDCm05Fb5LynozSl8,59932
-quack/layernorm.py,sha256=1WUspbr6ktPZ25O00kKs-FK_lm_Fejat72BMV8tBSfw,13504
-quack/reduction_base.py,sha256=4nAzkZR1yoQVA4Lc-GpU0XMjS5ARAmvYdeE0Doy7UCU,3789
-quack/rmsnorm.py,sha256=bJEHqc8ila-LTGco-tNNCUyFBjJ2UdXeoMplYNJPXFI,32740
-quack/softmax.py,sha256=3-5P_ORBrfQ6JYTIzgDs9jwmV7Za73SogaX7q9M7GCM,16698
-quack/utils.py,sha256=RZq-7YA8UMUizHpVyZM1we4zGm9NaC178M2g2HXdjmE,17799
-quack_kernels-0.1.10.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-quack_kernels-0.1.10.dist-info/METADATA,sha256=baMTwibt6u0IQb8YJFFhCY0RD3Aervf5sl6EpYF6IQ8,286
-quack_kernels-0.1.10.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-quack_kernels-0.1.10.dist-info/top_level.txt,sha256=6e4Jr_vNJbZTYwlO_Ahf_sDeHDE0zcqcf7Le11FKxxo,6
-quack_kernels-0.1.10.dist-info/RECORD,,