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

quack/rmsnorm.py

@@ -161,30 +161,33 @@ class RMSNorm(ReductionBase):
         copy_atom_load_X_async = cute.make_copy_atom(
             cute.nvgpu.cpasync.CopyG2SOp(), mX.element_type, num_bits_per_copy=128
         )
+        num_bits_per_copy_W = cutlass.const_expr(
+            min(128, 128 // mX.element_type.width * mW.element_type.width)
+        )
         copy_atom_load_W = cute.make_copy_atom(
-            cute.nvgpu.CopyUniversalOp(), mW.element_type, num_bits_per_copy=128
+            cute.nvgpu.CopyUniversalOp(), mW.element_type, num_bits_per_copy=num_bits_per_copy_W
+        )
+        num_bits_per_copy_O = cutlass.const_expr(
+            min(128, 128 // mX.element_type.width * mO.element_type.width)
         )
         copy_atom_store_O = cute.make_copy_atom(
-            cute.nvgpu.CopyUniversalOp(), mO.element_type, num_bits_per_copy=128
+            cute.nvgpu.CopyUniversalOp(), mO.element_type, num_bits_per_copy=num_bits_per_copy_O
         )
 
         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)
+        tXgW = thr_copy_X.partition_S(gW)
         tXgX = thr_copy_X.partition_S(gX)
         tXsX = thr_copy_X.partition_D(sX)
-        tXgO = thr_copy_O.partition_D(gO)
-        tXrRstd = thr_copy_O.partition_D(gRstd) if cutlass.const_expr(mRstd is not None) else None
+        tXgO = thr_copy_X.partition_D(gO)
+        tXrRstd = thr_copy_X.partition_D(gRstd) if cutlass.const_expr(mRstd is not None) else None
         tXcX = thr_copy_X.partition_S(cX)[(0, None), None, None]
 
         # allocate fragments for gmem->rmem
-        tWrW = cute.make_fragment_like(tWgW)
-        tWrW.fill(0.0)
-        tXrW = thr_copy_X.retile(tWrW)
+        tXrW = cute.make_fragment_like(tXgW)
+        tXrW.fill(0.0)
         tXrX, tXrO = [cute.make_fragment_like(thr) for thr in (tXgX, tXgO)]
 
         num_warps = cute.size(tv_layout, mode=[0]) // cute.arch.WARP_SIZE
@@ -196,9 +199,9 @@ class RMSNorm(ReductionBase):
             cute.copy(copy_atom_load_X_async, tXgX, tXsX, pred=tXpX)
         cute.arch.cp_async_commit_group()
 
-        tWpW = utils.predicate_k(thr_copy_W.partition_S(cX), limit=shape[1])
+        tXpW = utils.predicate_k(thr_copy_X.partition_S(cX), limit=shape[1])
         if cutlass.const_expr(not delay_w_load):
-            cute.copy(copy_atom_load_W, tWgW, tWrW, pred=tWpW)
+            cute.copy(copy_atom_load_W, tXgW, tXrW, pred=tXpW)
 
         cute.arch.cp_async_wait_group(0)
         cute.autovec_copy(tXsX, tXrX)
@@ -223,7 +226,7 @@ class RMSNorm(ReductionBase):
             ):
                 tXrRstd[0] = rstd
         if cutlass.const_expr(delay_w_load):
-            cute.copy(copy_atom_load_W, tWgW, tWrW, pred=tWpW)
+            cute.copy(copy_atom_load_W, tXgW, tXrW, pred=tXpW)
         if cutlass.const_expr(reload_from == "smem"):
             cute.autovec_copy(tXsX, tXrX)
             x = tXrX.load().to(cute.Float32)
@@ -234,9 +237,9 @@ class RMSNorm(ReductionBase):
         w = tXrW.load().to(cute.Float32)
         y = x_hat * w
         tXrO.store(y.to(tXrO.element_type))
-        tOpO = utils.predicate_k(thr_copy_O.partition_S(cX), limit=shape[1])
+        tXpO = utils.predicate_k(thr_copy_X.partition_S(cX), limit=shape[1])
         if row < shape[0]:
-            cute.copy(copy_atom_store_O, tXrO, tXgO, pred=tOpO)
+            cute.copy(copy_atom_store_O, tXrO, tXgO, pred=tXpO)
 
 
 def _rmsnorm_fwd(
@@ -460,39 +463,41 @@ class RMSNormBackward(ReductionBase):
         copy_atom_load_X_async = cute.make_copy_atom(
             cute.nvgpu.cpasync.CopyG2SOp(), mX.element_type, num_bits_per_copy=128
         )
+        num_bits_per_copy_W = cutlass.const_expr(
+            min(128, 128 // mX.element_type.width * mW.element_type.width)
+        )
         copy_atom_load_W = cute.make_copy_atom(
-            cute.nvgpu.CopyUniversalOp(), mW.element_type, num_bits_per_copy=128
+            cute.nvgpu.CopyUniversalOp(), mW.element_type, num_bits_per_copy=num_bits_per_copy_W
+        )
+        num_bits_per_copy_dX = cutlass.const_expr(
+            min(128, 128 // mX.element_type.width * mdX.element_type.width)
         )
         copy_atom_store_dX = cute.make_copy_atom(
-            cute.nvgpu.CopyUniversalOp(), mdX.element_type, num_bits_per_copy=128
+            cute.nvgpu.CopyUniversalOp(), mdX.element_type, num_bits_per_copy=num_bits_per_copy_dX
+        )
+        num_bits_per_copy_dW = cutlass.const_expr(
+            min(128, 128 // mX.element_type.width * mdW.element_type.width)
         )
         copy_atom_store_dW = cute.make_copy_atom(
-            cute.nvgpu.CopyUniversalOp(), mdW.element_type, num_bits_per_copy=128
+            cute.nvgpu.CopyUniversalOp(), mdW.element_type, num_bits_per_copy=num_bits_per_copy_dW
         )
 
         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_dW = cute.make_tiled_copy(copy_atom_store_dW, tv_layout, tiler_mn).get_slice(tidx)
-        thr_store_dX = cute.make_tiled_copy(copy_atom_store_dX, tv_layout, tiler_mn).get_slice(tidx)
 
         gW = cute.local_tile(mW, tiler_mn, (0, cluster_y))
-        tWgW = thr_copy_W.partition_S(gW)
-        tWrW = cute.make_fragment_like(tWgW)
+        tXgW = thr_copy_X.partition_S(gW)
+        tXrW = cute.make_fragment_like(tXgW)
         # Need this, otherwise rW can have arbitrary values that changes the reduction
         if not is_even_N:
-            tWrW.fill(0.0)
-        tXrW = thr_copy_X.retile(tWrW)
+            tXrW.fill(0.0)
 
         gW_coord = cute.local_tile(idX, tiler_mn, (0, cluster_y))
-        tWpW = (
-            utils.predicate_k(thr_copy_W.partition_S(gW_coord), limit=shape[1])
+        tXpW = (
+            utils.predicate_k(thr_copy_X.partition_S(gW_coord), limit=shape[1])
             if not is_even_N
             else None
         )
-        cute.copy(copy_atom_load_W, tWgW, tWrW, pred=tWpW)
+        cute.copy(copy_atom_load_W, tXgW, tXrW, pred=tXpW)
         weight = tXrW.load().to(cute.Float32)
 
         num_warps = cute.size(tv_layout, mode=[0]) // cute.arch.WARP_SIZE
@@ -500,17 +505,16 @@ class RMSNormBackward(ReductionBase):
         self._initialize_cluster(tidx, mbar_ptr, num_warps, is_persistent=True)
 
         dw_coord = cute.local_tile(idX, tiler_mn, (0, cluster_y))
-        tdWpdW = (
-            utils.predicate_k(thr_copy_dW.partition_S(dw_coord), limit=shape[1])
+        tXpdW = (
+            utils.predicate_k(thr_copy_X.partition_S(dw_coord), limit=shape[1])
             if not is_even_N
             else None
         )
 
         gdW = cute.local_tile(mdW, (1, tiler_mn[1]), (bidx_start, cluster_y))
-        tdWgdW = thr_copy_dW.partition_D(gdW)
+        tXgdW = thr_copy_X.partition_S(gdW)
         # Always compute partial weight gradients in fp32
-        tdWrdW = cute.make_fragment_like(tdWgdW, Float32)
-        tXrdW = thr_copy_X.retile(tdWrdW)
+        tXrdW = cute.make_fragment_like(tXgdW, Float32)
 
         gX = cute.local_tile(mX, tiler_mn, (None, cluster_y))
         gdOut = cute.local_tile(mdOut, tiler_mn, (None, cluster_y))
@@ -520,7 +524,7 @@ class RMSNormBackward(ReductionBase):
         tXsX = thr_copy_X.partition_D(sX)
         tXgdOut = thr_copy_X.partition_S(gdOut)
         tXsdOut = thr_copy_X.partition_D(sdOut)
-        tXgdX = thr_store_dX.partition_D(gdX)
+        tXgdX = thr_copy_X.partition_D(gdX)
         tXcX = thr_copy_X.partition_S(cX)[(0, None), None, None, None]
         # This doesn't change across iterations
         tXpX = (
@@ -670,11 +674,10 @@ class RMSNormBackward(ReductionBase):
                     tXsdW_other = cute.make_tensor(tXsdW.iterator + i * sdW.stride[0], tXsdW.layout)
                     cute.autovec_copy(tXsdW_other, tXrdW_other)
                     tXrdW.store(tXrdW.load() + tXrdW_other.load())
-                cute.copy(copy_atom_store_dW, tdWrdW, tdWgdW, pred=tdWpdW)
-
+                cute.copy(copy_atom_store_dW, tXrdW, tXgdW, pred=tXpdW)
         else:
             # dw is already in fp32, so we can directly copy to global memory
-            cute.copy(copy_atom_store_dW, tdWrdW, tdWgdW, pred=tdWpdW)
+            cute.copy(copy_atom_store_dW, tXrdW, tXgdW, pred=tXpdW)
 
 
 def _rmsnorm_backward(

quack/utils.py

@@ -315,7 +315,7 @@ def exp2f(x: cute.TensorSSA | Float32) -> cute.TensorSSA | Float32:
     if cutlass.const_expr(isinstance(x, cute.TensorSSA)):
         res = cute.make_fragment(x.shape, Float32)
         res.store(x)
-        for i in cutlass.range_constexpr(cute.size(x.shape)):
+        for i in cutlass.range(cute.size(x.shape), unroll_full=True):
             res[i] = cute.arch.exp2(res[i])
         return res.load()
     else:

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

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

quack_kernels-0.1.10.dist-info/RECORD

@@ -0,0 +1,13 @@
+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,,

quack_kernels-0.1.8.dist-info/RECORD

@@ -1,12 +0,0 @@
-quack/__init__.py,sha256=tDgX5MF1ttfEyDVFWi47DA8tDooYcBQlkuzvabGUoQI,203
-quack/cross_entropy.py,sha256=VYSAd28GmtnMoKQwLrorvySDtJfRhoqVd-aeM52FmsI,20866
-quack/layernorm.py,sha256=1WUspbr6ktPZ25O00kKs-FK_lm_Fejat72BMV8tBSfw,13504
-quack/reduction_base.py,sha256=4nAzkZR1yoQVA4Lc-GpU0XMjS5ARAmvYdeE0Doy7UCU,3789
-quack/rmsnorm.py,sha256=-qrKqPKk0fUuq0a5-vJmZZ7nQsHgyaqTg0EKhWT44r0,32738
-quack/softmax.py,sha256=3-5P_ORBrfQ6JYTIzgDs9jwmV7Za73SogaX7q9M7GCM,16698
-quack/utils.py,sha256=aiyzBc9BEwq8s965elfiR331hAaLLBKL9kDHjuls86Q,17791
-quack_kernels-0.1.8.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-quack_kernels-0.1.8.dist-info/METADATA,sha256=b_2PxFEoVqWJbT2FtuP9FJyF-jpL2Z3q9OHoOEipqo4,289
-quack_kernels-0.1.8.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-quack_kernels-0.1.8.dist-info/top_level.txt,sha256=6e4Jr_vNJbZTYwlO_Ahf_sDeHDE0zcqcf7Le11FKxxo,6
-quack_kernels-0.1.8.dist-info/RECORD,,