quack-kernels 0.1.8__tar.gz → 0.1.10__tar.gz

{quack_kernels-0.1.8/quack_kernels.egg-info → quack_kernels-0.1.10}/PKG-INFO

@@ -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.8 → quack_kernels-0.1.10}/pyproject.toml

@@ -5,9 +5,9 @@ build-backend = "setuptools.build_meta"
 [project]
 name = "quack-kernels"
 dynamic = ["version"]
-requires-python = ">=3.9"
+requires-python = ">=3.12"
 dependencies = [
-    "nvidia-cutlass-dsl==4.1.0.dev0",
+    "nvidia-cutlass-dsl==4.1.0",
     "torch",
 ]
 
@@ -29,5 +29,6 @@ line-length = 100
 [tool.ruff.lint]
 ignore = [
     "E731",  # do not assign a lambda expression, use a def
+    "E741",  # Do not use variables named 'I', 'O', or 'l'
     "F841",  # local variable is assigned to but never used
 ]

{quack_kernels-0.1.8 → quack_kernels-0.1.10}/quack/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.1.8"
+__version__ = "0.1.10"
 
 from quack.rmsnorm import rmsnorm
 from quack.softmax import softmax

{quack_kernels-0.1.8 → quack_kernels-0.1.10}/quack/cross_entropy.py

@@ -446,13 +446,10 @@ class CrossEntropyBackward:
         log2_e = math.log2(math.e)
         probs = utils.exp2f((x - lse) * log2_e)
         prob_shifted = probs - 1.0
-
         mask = cute.make_fragment_like(tXrX, cutlass.Boolean)
-        for i in cutlass.range_constexpr(cute.size(tXcFull)):
+        for i in cutlass.range(cute.size(tXcFull), unroll_full=True):
             mask[i] = tXcFull[i][1] == label
-
-        mask = mask.load()
-        grad = cute.where(mask, prob_shifted, probs)
+        grad = cute.where(mask.load(), prob_shifted, probs)
         grad = grad * dloss
 
         tXrO.store(grad.to(tXrO.element_type))