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

quack/utils.py

@@ -1,8 +1,7 @@
 # Copyright (c) 2025, Wentao Guo, Ted Zadouri, Tri Dao.
 
-import operator
 import math
-from typing import Callable, Optional, Tuple, Type, Union
+from typing import Optional, Tuple, Type, Union
 
 import cutlass
 import cutlass.cute as cute
@@ -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
@@ -114,197 +87,6 @@ def store_shared_remote(
     )
 
 
-@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,
-        )
-    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)
-
-
 @dsl_user_op
 def fmin(a: Union[float, Float32], b: Union[float, Float32], *, loc=None, ip=None) -> Float32:
     return Float32(
@@ -381,21 +163,6 @@ def rsqrt(a: float | Float32, *, loc=None, ip=None) -> Float32:
     )
 
 
-@dsl_user_op
-def tanh(a: float | Float32, *, loc=None, ip=None) -> Float32:
-    return Float32(
-        llvm.inline_asm(
-            T.f32(),
-            [Float32(a).ir_value(loc=loc, ip=ip)],
-            "tanh.approx.f32 $0, $1;",
-            "=f,f",
-            has_side_effects=False,
-            is_align_stack=False,
-            asm_dialect=llvm.AsmDialect.AD_ATT,
-        )
-    )
-
-
 @dsl_user_op
 def ceil(a: float | Float32, *, loc=None, ip=None) -> Int32:
     return Int32(
@@ -411,16 +178,6 @@ def ceil(a: float | Float32, *, loc=None, ip=None) -> Int32:
     )
 
 
-@dsl_user_op
-def silu(a: float | Float32, *, loc=None, ip=None) -> Float32:
-    """
-    silu(a) = a * sigmoid(a) = a * (1 + tanh(a / 2)) / 2 = (0.5 * a) * tanh(0.5 * a) + (0.5 * a)
-    This compiles down to 3 SASS instructions: FMUL to get 0.5 * a, MUFU.TANH, and FFMA.
-    """
-    a_half = 0.5 * a
-    return a_half * tanh(a_half) + a_half
-
-
 @dsl_user_op
 def prmt(a: int | Int32, b: int | Int32, c: int | Int32, *, loc=None, ip=None) -> Int32:
     return Int32(
@@ -498,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]):
@@ -538,9 +295,9 @@ def i64_to_f32x2(c: cutlass.Int64, *, loc=None, ip=None) -> Tuple[Float32, Float
 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)
-    assert len(flat_coord_i64) == len(flat_stride), (
-        "Coordinate and stride must have the same length"
-    )
+    assert len(flat_coord_i64) == len(
+        flat_stride
+    ), "Coordinate and stride must have the same length"
     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
@@ -662,5 +419,3 @@ def atomic_inc_i32(a: int | Int32, gmem_ptr: cute.Pointer, *, loc=None, ip=None)
     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.11.dist-info → quack_kernels-0.2.0.dist-info}/METADATA

@@ -1,9 +1,9 @@
 Metadata-Version: 2.4
 Name: quack-kernels
-Version: 0.1.11
+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/lse.py

@@ -1,62 +0,0 @@
-# Copyright (c) 2025, Tri Dao.
-# TODO: we probably dont' need this kernel, just use torch.logsumexp
-import torch
-
-import triton
-import triton.language as tl
-
-
-@triton.jit
-def _lse_kernel(
-    lse_ptr,
-    logits_ptr,
-    n_rows,
-    n_cols,
-    logits_row_stride,
-    logits_col_stride,
-    BLOCK_SIZE_M: tl.constexpr,
-    BLOCK_SIZE_N: tl.constexpr,
-):
-    row_start = tl.program_id(0) * BLOCK_SIZE_M
-    rows = row_start + tl.arange(0, BLOCK_SIZE_M)
-    cols = tl.arange(0, BLOCK_SIZE_N)
-    logits = tl.load(
-        logits_ptr + rows[:, None] * logits_row_stride + cols[None, :] * logits_col_stride,
-        mask=(rows[:, None] < n_rows) & (cols[None, :] < n_cols),
-        other=-float("inf"),
-    ).to(tl.float32)
-    m = tl.max(logits, 1)
-    lse = tl.log(tl.sum(tl.exp(logits - m[:, None]), 1)) + m
-    tl.store(lse_ptr + rows, lse, mask=rows < n_rows)
-
-
-def logsumexp(logits):
-    n_rows, n_cols = logits.shape
-    BLOCK_SIZE_M = 32 if logits.stride(1) != 1 else 1
-    MAX_BLOCK_SIZE = 64 * 1024
-    # BLOCK_SIZE_N = min(triton.next_power_of_2(n_cols), MAX_BLOCK_SIZE // BLOCK_SIZE_M)
-    BLOCK_SIZE_N = triton.next_power_of_2(n_cols)
-    assert (
-        BLOCK_SIZE_M * BLOCK_SIZE_N <= MAX_BLOCK_SIZE
-    ), f"Only support max dimension {MAX_BLOCK_SIZE // BLOCK_SIZE_M}"
-    num_warps = (
-        4
-        if BLOCK_SIZE_N < 2048
-        else (8 if BLOCK_SIZE_N < 8192 else (16 if BLOCK_SIZE_N < 128 * 1024 else 32))
-    )
-    lse = torch.empty(n_rows, dtype=torch.float, device=logits.device)
-    # Need this, otherwise Triton tries to launch from cuda:0 and we get
-    # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?)
-    with torch.cuda.device(logits.device.index):
-        _lse_kernel[(triton.cdiv(n_rows, BLOCK_SIZE_M),)](
-            lse,
-            logits,
-            n_rows,
-            n_cols,  # shapes
-            logits.stride(0),  # strides
-            logits.stride(1),
-            BLOCK_SIZE_M=BLOCK_SIZE_M,  # constants
-            BLOCK_SIZE_N=BLOCK_SIZE_N,  # constants
-            num_warps=num_warps,
-        )
-    return lse

quack_kernels-0.1.11.dist-info/RECORD

@@ -1,31 +0,0 @@
-quack/__init__.py,sha256=AD0T-rBhSfKXpwZ6E4JIPiugvlFaAePjl-3pUhWOlPE,292
-quack/autotuner.py,sha256=aF9-Cw47gaX7_LZvyVbLsj6Z2AWi4UZ-0Qwjy06Xd5I,10733
-quack/cross_entropy.py,sha256=xsg2bXZ4wNvusBARhN4PwAzm5PbejEcfwj71nR7bzuE,20852
-quack/cute_dsl_utils.py,sha256=LkNyFEKwYrgp-tLt_775EZWuBR3v7G80El3UAObHY2U,1292
-quack/dense_gemm_sm100.py,sha256=W_j8BO-ilb1YUYFuclo7_itfPIRTkjPV_ittWgQy8t4,109937
-quack/dense_gemm_sm90.py,sha256=Dff0GbIv92uTjrtsUE1GjVKCtwSf6_5KZbrqYZm-ZMY,110418
-quack/fast_math.py,sha256=XqXVvKLSxXC3c9tIGLvKVRWdPsmjAa_O4C0plmsfZ0w,3106
-quack/gemm_config.py,sha256=Gz4dkHH1Uwg9IdW-x5W_5tjdaFHBfxq4bn7hJx_xu5s,1789
-quack/gemm_interface.py,sha256=XHgxo08d8LIu6dTlQKBOBJtjCegUB5uLh4k9hC-5mvY,9525
-quack/layernorm.py,sha256=1WUspbr6ktPZ25O00kKs-FK_lm_Fejat72BMV8tBSfw,13504
-quack/linear.py,sha256=Wd0KeXWvWjbkKrgW4Av1ud2v_mbhzf1RvubF7BYhcw4,6425
-quack/lse.py,sha256=aANOleIYREyrkUQM9cfJ9Gt63eawMb2KVd7YAGWNoZU,2092
-quack/mlp.py,sha256=D9V7aIfvoBMzhKwN8ZE6GlSOmwFJe_JGqgOvQprU0OQ,8224
-quack/pipeline.py,sha256=SwvRZAR4RqYH60wAFC3OTu5DisN1XDMv5umQF4czJW4,5867
-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/symmetric_dense_gemm_sm90.py,sha256=t-6eLasZwyu1NW4HpnvVBBPOvfqUzOg8VHe9sJQYdmg,88637
-quack/tensormap_manager.py,sha256=pzBNwLCB8kV_yp8X8_BoDdtbwWeht2jrgRhyyfVIcMI,5261
-quack/tile_scheduler.py,sha256=mImjD2LuIVchM6USJoJY4-CSG54jGuwyLIvFG6LTP9Y,42205
-quack/topk.py,sha256=1pObblNJnxKLaE_T3qGvaMnUua0dqG2en9OU5PSp71s,9020
-quack/utils.py,sha256=4ViEFgHecaX5wcYpO6XzTCzdnuZv2rniUJAJH5Ta0bA,24981
-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.1.11.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-quack_kernels-0.1.11.dist-info/METADATA,sha256=WTYlk9lmhr4Jkin71stp3h-NrBdme-8OrBc7lAf4vSw,286
-quack_kernels-0.1.11.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-quack_kernels-0.1.11.dist-info/top_level.txt,sha256=6e4Jr_vNJbZTYwlO_Ahf_sDeHDE0zcqcf7Le11FKxxo,6
-quack_kernels-0.1.11.dist-info/RECORD,,