cutez 0.1.0__py3-none-any.whl

cutex/__init__.py

@@ -0,0 +1,218 @@
+from typing import List, Type, Union, Tuple
+
+import cutlass.cute as cute
+
+from cutlass.cute.arch import numeric_conversion
+from cutlass.cute.nvgpu.common import CopyUniversalOp
+from cutlass.cute.nvgpu.warp import StMatrix8x8x16bOp, StMatrix16x8x8bOp
+from cutlass.cute.nvgpu.tcgen05 import (
+    MmaF16BF16Op,
+    MmaTF32Op,
+    MmaI8Op,
+    MmaFP8Op,
+    MmaMXF8Op,
+    MmaMXF4Op,
+    MmaMXF4NVF4Op,
+    OperandSource as Tcgen05OperandSource,
+    OperandMajorMode,
+    CtaGroup,
+    Ld16x64bOp,
+    Ld16x128bOp,
+    Ld16x256bOp,
+    Ld16x32bx2Op,
+    Ld32x32bOp,
+    Repetition,
+    Pack,
+    SmemLayoutAtomKind,
+    make_smem_layout_atom,
+    tile_to_mma_shape,
+    is_tmem_load,
+    get_tmem_copy_properties,
+)
+from cutlass.cute.nvgpu.cpasync import (
+    CopyBulkTensorTileG2SMulticastOp,
+    CopyBulkTensorTileG2SOp,
+)
+from cutlass.utils.layout import LayoutEnum
+from cutlass.cutlass_dsl import (
+    Float16,
+    BFloat16,
+    TFloat32,
+    Float32,
+    Uint8,
+    Int8,
+    Float8E4M3FN,
+    Float8E5M2,
+    Float4E2M1FN,
+    Numeric,
+    NumericMeta,
+    dsl_user_op,
+)
+from torch._prims_common import number_type
+
+@dsl_user_op
+def explain_get_smem_store_op(
+    layout_d: LayoutEnum,
+    elem_ty_d: Type[Numeric],
+    elem_ty_acc: Type[Numeric],
+    tiled_tmem_load: cute.TiledCopy,
+    *,
+    loc=None,
+    ip=None,
+) -> cute.CopyAtom:
+    """Selects the largest vectorized smem store atom available subject to
+    constraint of gmem layout and chosen TMEM_LOAD's thread-value ownership.
+
+    :param layout_d: The layout enum of the output tensor D.
+    :type layout_d: LayoutEnum
+    :param elem_ty_d: The element type for output tensor D.
+    :type elem_ty_d: Type[Numeric]
+    :param elem_ty_acc: The element type for accumulator.
+    :type elem_ty_acc: Type[Numeric]
+    :param tiled_tmem_load: An instance of TiledCopy that represents the tmem load operation.
+    :type tiled_tmem_load: cute.TiledCopy
+
+    :return: Either SmemStoreMatrix or SimtSyncCopy, based on the input parameters.
+    :rtype: cute.CopyAtom
+    """
+
+    def validate_type(ty, ty_name):
+        if not isinstance(ty, NumericMeta):
+            raise TypeError(f"{ty_name} must be a Numeric, but got {ty}")
+
+    validate_type(elem_ty_d, "elem_ty_d")
+    validate_type(elem_ty_acc, "elem_ty_acc")
+
+    is_m_major = layout_d.is_m_major_c()
+    is_n_major = layout_d.is_n_major_c()
+
+    if not is_tmem_load(tiled_tmem_load):
+        return cute.make_copy_atom(CopyUniversalOp(), elem_ty_d, loc=loc, ip=ip)
+
+    num_dp, num_bits, num_rep, pack = get_tmem_copy_properties(tiled_tmem_load)
+
+    use_stmatrix_m8n8_4x = (
+        all(
+            [
+                elem_ty_acc.width == 32,
+                elem_ty_d.width == 32,
+                is_n_major,
+                num_dp == 16,
+                num_bits == 128,
+                num_rep in (2, 4, 8, 16, 32, 64),
+                pack == Pack.NONE,
+            ]
+        )
+        or all(
+            [
+                elem_ty_acc.width == 32,
+                elem_ty_d.width == 16,
+                num_dp == 16,
+                num_bits == 256,
+                num_rep in (2, 4, 8, 16, 32),
+                pack == Pack.NONE,
+            ]
+        )
+        or all(
+            [
+                elem_ty_acc.width == 16,
+                elem_ty_d.width == 16,
+                num_dp == 16,
+                num_bits == 128,
+                num_rep in (2, 4, 8, 16, 32, 64),
+                pack == Pack.PACK_16b_IN_32b,
+            ]
+        )
+    )
+    use_stmatrix_m16n8_4x = all(
+        [
+            elem_ty_acc.width == 32,
+            elem_ty_d.width == 8,
+            is_m_major,
+            num_dp == 16,
+            num_bits == 256,
+            num_rep in (4, 8, 16, 32),
+            pack == Pack.NONE,
+        ]
+    )
+    use_stmatrix_m8n8_2x = (
+        all(
+            [
+                elem_ty_acc.width == 32,
+                elem_ty_d.width == 32,
+                is_n_major,
+                num_dp == 16,
+                num_bits == 128,
+                num_rep == 1,
+                pack == Pack.NONE,
+            ]
+        )
+        or all(
+            [
+                elem_ty_acc.width == 32,
+                elem_ty_d.width == 16,
+                num_dp == 16,
+                num_bits == 256,
+                num_rep == 1,
+                pack == Pack.NONE,
+            ]
+        )
+        or all(
+            [
+                elem_ty_acc.width == 16,
+                elem_ty_d.width == 16,
+                num_dp == 16,
+                num_bits == 128,
+                num_rep == 1,
+                pack == Pack.PACK_16b_IN_32b,
+            ]
+        )
+    )
+    use_stmatrix_m16n8_2x = all(
+        [
+            elem_ty_acc.width == 32,
+            elem_ty_d.width == 8,
+            is_m_major,
+            num_dp == 16,
+            num_bits == 256,
+            num_rep == 2,
+            pack == Pack.NONE,
+        ]
+    )
+    use_stmatrix_m16n8_1x = all(
+        [
+            elem_ty_acc.width == 32,
+            elem_ty_d.width == 8,
+            is_m_major,
+            num_dp == 16,
+            num_bits == 256,
+            num_rep == 1,
+            pack == Pack.NONE,
+        ]
+    )
+    print("===========================================")
+    print("explain_get_smem_store_op:")
+    if num_dp != 16 or num_bits < 128:
+        print("NOTE: to directly use stmatrix, tcgen05.ld must have 16 lanes and 128/256b instruction")
+    print(f"elem_ty_acc.width: {elem_ty_acc.width}")
+    print(f"elem_ty_d.width: {elem_ty_d.width}")
+    print(f"is_n_major: {is_n_major}")
+    print(f"num_dp: {num_dp}")
+    print(f"num_bits: {num_bits}")
+    print(f"num_rep: {num_rep}")
+    print(f"pack: {pack}")
+    print("===========================================")
+
+    if use_stmatrix_m8n8_4x:
+        op = StMatrix8x8x16bOp(is_m_major, 4)
+    elif use_stmatrix_m8n8_2x:
+        op = StMatrix8x8x16bOp(is_m_major, 2)
+    elif use_stmatrix_m16n8_4x:
+        op = StMatrix16x8x8bOp(transpose=True, num_matrices=4)
+    elif use_stmatrix_m16n8_2x:
+        op = StMatrix16x8x8bOp(transpose=True, num_matrices=2)
+    elif use_stmatrix_m16n8_1x:
+        op = StMatrix16x8x8bOp(transpose=True, num_matrices=1)
+    else:
+        op = CopyUniversalOp()
+    print(op)

cutex/__init__.pyi

@@ -0,0 +1,21 @@
+from cutlass.cute import *
+
+import cutlass.cute.nvgpu as nvgpu
+import cutlass.cute.runtime as runtime
+import cutlass.cute.testing as testing
+import cutlass.pipeline as pipeline
+import cutlass.utils as utils
+import cutlass.utils.blackwell_helpers as sm100
+
+from cutlass.cute.nvgpu import cpasync, tcgen05
+from cutlass.cute.runtime import *
+from cutlass.cute.runtime import from_dlpack, make_fake_stream
+from cutlass.cute.testing import *
+from cutlass.cute.nvgpu import *
+from cutlass.pipeline import *
+from cutlass.utils import *
+from cutlass.utils.blackwell_helpers import *
+
+def dummy() -> str: ...
+
+__all__: list[str]

cutex/pipeline.py

@@ -0,0 +1,7 @@
+import cutlass.pipeline as _pipeline
+
+for _name in getattr(_pipeline, "__all__", ()):
+    if hasattr(_pipeline, _name):
+        globals()[_name] = getattr(_pipeline, _name)
+
+__all__ = [name for name in getattr(_pipeline, "__all__", ()) if hasattr(_pipeline, name)]

cutex/pipeline.pyi

@@ -0,0 +1,3 @@
+from cutlass.pipeline import *
+
+__all__: list[str]

cutex/runtime.py

@@ -0,0 +1,14 @@
+import cutlass.cute.runtime as _runtime
+from cutlass.cute.runtime import from_dlpack
+
+_runtime_public_names = getattr(_runtime, "__all__", None)
+if _runtime_public_names is None:
+    _runtime_public_names = [name for name in dir(_runtime) if not name.startswith("_")]
+
+for _name in _runtime_public_names:
+    if hasattr(_runtime, _name):
+        globals()[_name] = getattr(_runtime, _name)
+
+__all__ = [name for name in _runtime_public_names if hasattr(_runtime, name)]
+if "from_dlpack" not in __all__:
+    __all__.append("from_dlpack")

cutex/runtime.pyi

@@ -0,0 +1,3 @@
+from cutlass.cute.runtime import *
+
+__all__: list[str]

cutex/testing.py

@@ -0,0 +1,11 @@
+import cutlass.cute.testing as _testing
+
+_testing_public_names = getattr(_testing, "__all__", None)
+if _testing_public_names is None:
+    _testing_public_names = [name for name in dir(_testing) if not name.startswith("_")]
+
+for _name in _testing_public_names:
+    if hasattr(_testing, _name):
+        globals()[_name] = getattr(_testing, _name)
+
+__all__ = [name for name in _testing_public_names if hasattr(_testing, name)]

cutex/testing.pyi

@@ -0,0 +1,3 @@
+from cutlass.cute.testing import *
+
+__all__: list[str]

cutex/utils/__init__.py

@@ -0,0 +1,7 @@
+import cutlass.utils as _utils
+
+for _name in getattr(_utils, "__all__", ()):
+    if hasattr(_utils, _name):
+        globals()[_name] = getattr(_utils, _name)
+
+__all__ = [name for name in getattr(_utils, "__all__", ()) if hasattr(_utils, name)]

cutex/utils/__init__.pyi

@@ -0,0 +1,3 @@
+from cutlass.utils import *
+
+__all__: list[str]

cutex/utils/sm100.py

@@ -0,0 +1,7 @@
+import cutlass.utils.blackwell_helpers as _sm100
+
+for _name in getattr(_sm100, "__all__", ()):
+    if hasattr(_sm100, _name):
+        globals()[_name] = getattr(_sm100, _name)
+
+__all__ = [name for name in getattr(_sm100, "__all__", ()) if hasattr(_sm100, name)]

cutex/utils/sm100.pyi

@@ -0,0 +1,3 @@
+from cutlass.utils.blackwell_helpers import *
+
+__all__: list[str]

cutez/__init__.py

@@ -0,0 +1,218 @@
+from typing import List, Type, Union, Tuple
+
+import cutlass.cute as cute
+
+from cutlass.cute.arch import numeric_conversion
+from cutlass.cute.nvgpu.common import CopyUniversalOp
+from cutlass.cute.nvgpu.warp import StMatrix8x8x16bOp, StMatrix16x8x8bOp
+from cutlass.cute.nvgpu.tcgen05 import (
+    MmaF16BF16Op,
+    MmaTF32Op,
+    MmaI8Op,
+    MmaFP8Op,
+    MmaMXF8Op,
+    MmaMXF4Op,
+    MmaMXF4NVF4Op,
+    OperandSource as Tcgen05OperandSource,
+    OperandMajorMode,
+    CtaGroup,
+    Ld16x64bOp,
+    Ld16x128bOp,
+    Ld16x256bOp,
+    Ld16x32bx2Op,
+    Ld32x32bOp,
+    Repetition,
+    Pack,
+    SmemLayoutAtomKind,
+    make_smem_layout_atom,
+    tile_to_mma_shape,
+    is_tmem_load,
+    get_tmem_copy_properties,
+)
+from cutlass.cute.nvgpu.cpasync import (
+    CopyBulkTensorTileG2SMulticastOp,
+    CopyBulkTensorTileG2SOp,
+)
+from cutlass.utils.layout import LayoutEnum
+from cutlass.cutlass_dsl import (
+    Float16,
+    BFloat16,
+    TFloat32,
+    Float32,
+    Uint8,
+    Int8,
+    Float8E4M3FN,
+    Float8E5M2,
+    Float4E2M1FN,
+    Numeric,
+    NumericMeta,
+    dsl_user_op,
+)
+from torch._prims_common import number_type
+
+@dsl_user_op
+def explain_get_smem_store_op(
+    layout_d: LayoutEnum,
+    elem_ty_d: Type[Numeric],
+    elem_ty_acc: Type[Numeric],
+    tiled_tmem_load: cute.TiledCopy,
+    *,
+    loc=None,
+    ip=None,
+) -> cute.CopyAtom:
+    """Selects the largest vectorized smem store atom available subject to
+    constraint of gmem layout and chosen TMEM_LOAD's thread-value ownership.
+
+    :param layout_d: The layout enum of the output tensor D.
+    :type layout_d: LayoutEnum
+    :param elem_ty_d: The element type for output tensor D.
+    :type elem_ty_d: Type[Numeric]
+    :param elem_ty_acc: The element type for accumulator.
+    :type elem_ty_acc: Type[Numeric]
+    :param tiled_tmem_load: An instance of TiledCopy that represents the tmem load operation.
+    :type tiled_tmem_load: cute.TiledCopy
+
+    :return: Either SmemStoreMatrix or SimtSyncCopy, based on the input parameters.
+    :rtype: cute.CopyAtom
+    """
+
+    def validate_type(ty, ty_name):
+        if not isinstance(ty, NumericMeta):
+            raise TypeError(f"{ty_name} must be a Numeric, but got {ty}")
+
+    validate_type(elem_ty_d, "elem_ty_d")
+    validate_type(elem_ty_acc, "elem_ty_acc")
+
+    is_m_major = layout_d.is_m_major_c()
+    is_n_major = layout_d.is_n_major_c()
+
+    if not is_tmem_load(tiled_tmem_load):
+        return cute.make_copy_atom(CopyUniversalOp(), elem_ty_d, loc=loc, ip=ip)
+
+    num_dp, num_bits, num_rep, pack = get_tmem_copy_properties(tiled_tmem_load)
+
+    use_stmatrix_m8n8_4x = (
+        all(
+            [
+                elem_ty_acc.width == 32,
+                elem_ty_d.width == 32,
+                is_n_major,
+                num_dp == 16,
+                num_bits == 128,
+                num_rep in (2, 4, 8, 16, 32, 64),
+                pack == Pack.NONE,
+            ]
+        )
+        or all(
+            [
+                elem_ty_acc.width == 32,
+                elem_ty_d.width == 16,
+                num_dp == 16,
+                num_bits == 256,
+                num_rep in (2, 4, 8, 16, 32),
+                pack == Pack.NONE,
+            ]
+        )
+        or all(
+            [
+                elem_ty_acc.width == 16,
+                elem_ty_d.width == 16,
+                num_dp == 16,
+                num_bits == 128,
+                num_rep in (2, 4, 8, 16, 32, 64),
+                pack == Pack.PACK_16b_IN_32b,
+            ]
+        )
+    )
+    use_stmatrix_m16n8_4x = all(
+        [
+            elem_ty_acc.width == 32,
+            elem_ty_d.width == 8,
+            is_m_major,
+            num_dp == 16,
+            num_bits == 256,
+            num_rep in (4, 8, 16, 32),
+            pack == Pack.NONE,
+        ]
+    )
+    use_stmatrix_m8n8_2x = (
+        all(
+            [
+                elem_ty_acc.width == 32,
+                elem_ty_d.width == 32,
+                is_n_major,
+                num_dp == 16,
+                num_bits == 128,
+                num_rep == 1,
+                pack == Pack.NONE,
+            ]
+        )
+        or all(
+            [
+                elem_ty_acc.width == 32,
+                elem_ty_d.width == 16,
+                num_dp == 16,
+                num_bits == 256,
+                num_rep == 1,
+                pack == Pack.NONE,
+            ]
+        )
+        or all(
+            [
+                elem_ty_acc.width == 16,
+                elem_ty_d.width == 16,
+                num_dp == 16,
+                num_bits == 128,
+                num_rep == 1,
+                pack == Pack.PACK_16b_IN_32b,
+            ]
+        )
+    )
+    use_stmatrix_m16n8_2x = all(
+        [
+            elem_ty_acc.width == 32,
+            elem_ty_d.width == 8,
+            is_m_major,
+            num_dp == 16,
+            num_bits == 256,
+            num_rep == 2,
+            pack == Pack.NONE,
+        ]
+    )
+    use_stmatrix_m16n8_1x = all(
+        [
+            elem_ty_acc.width == 32,
+            elem_ty_d.width == 8,
+            is_m_major,
+            num_dp == 16,
+            num_bits == 256,
+            num_rep == 1,
+            pack == Pack.NONE,
+        ]
+    )
+    print("===========================================")
+    print("explain_get_smem_store_op:")
+    if num_dp != 16 or num_bits < 128:
+        print("NOTE: to directly use stmatrix, tcgen05.ld must have 16 lanes and 128/256b instruction")
+    print(f"elem_ty_acc.width: {elem_ty_acc.width}")
+    print(f"elem_ty_d.width: {elem_ty_d.width}")
+    print(f"is_n_major: {is_n_major}")
+    print(f"num_dp: {num_dp}")
+    print(f"num_bits: {num_bits}")
+    print(f"num_rep: {num_rep}")
+    print(f"pack: {pack}")
+
+    if use_stmatrix_m8n8_4x:
+        op = StMatrix8x8x16bOp(is_m_major, 4)
+    elif use_stmatrix_m8n8_2x:
+        op = StMatrix8x8x16bOp(is_m_major, 2)
+    elif use_stmatrix_m16n8_4x:
+        op = StMatrix16x8x8bOp(transpose=True, num_matrices=4)
+    elif use_stmatrix_m16n8_2x:
+        op = StMatrix16x8x8bOp(transpose=True, num_matrices=2)
+    elif use_stmatrix_m16n8_1x:
+        op = StMatrix16x8x8bOp(transpose=True, num_matrices=1)
+    else:
+        op = CopyUniversalOp()
+    print(op)
+    print("===========================================")

cutez-0.1.0.dist-info/METADATA

@@ -0,0 +1,6 @@
+Metadata-Version: 2.4
+Name: cutez
+Version: 0.1.0
+License-File: LICENSE
+Requires-Dist: flash_attn_4==4.0.0b4
+Dynamic: license-file

cutez-0.1.0.dist-info/RECORD

@@ -0,0 +1,18 @@
+cutex/__init__.py,sha256=_zLi5I9jqk_WdAgiX0d--RVxScNy3Oo90iFQ5VV7WhY,6206
+cutex/__init__.pyi,sha256=Jy4j9ANKmE7SXBupKqH2IufsFm1xUO6WNPs7yZ_USwA,620
+cutex/pipeline.py,sha256=1wgcnhAex3mjVSadfG2r0fRLxsmiIlp4oKElYNq4OY0,265
+cutex/pipeline.pyi,sha256=LTgR1WXuv0H9KmSaM0CxJbJyGm5ft3fU3Ig4T1x5xcI,51
+cutex/runtime.py,sha256=DkMxTHUWwP6UUAh7r3q6Efn0XRgFD6hCMSdgL3KbpEw,536
+cutex/runtime.pyi,sha256=UQ_UbT1HfUuMtjdxP18zWzbqYClyjhtmLAEFV4dvBdg,55
+cutex/testing.py,sha256=Gnk6yME2Kmi4hBzhu2xc2K5ITpc1lix0a0M7Z2SwDAw,424
+cutex/testing.pyi,sha256=Cp2Ne0rhrUGaxpkg_9TVM8HSHFGGmovtVP-u6EG7XPc,55
+cutex/utils/__init__.py,sha256=G2ylaTDZAb7sIvHa5uxEVKjT4zeQP_WkDIx1K7Fh3us,244
+cutex/utils/__init__.pyi,sha256=98JcUMHBPpCbQ4GqiNgFVrzZmkGTzE0XlPw4jLa11vE,48
+cutex/utils/sm100.py,sha256=JjCmDJ9xyzI6zdTnngx0JRxpuKbds-TeKsp3uSV4VzM,262
+cutex/utils/sm100.pyi,sha256=yKzsc6H9fqPD95uJUnK38AIE1isoikNpYl4wE_xqS7Q,66
+cutez/__init__.py,sha256=NbPATzMmEYxzPWd5SKY1AGDQQGrl2oPJz0bARHOyM0g,6206
+cutez-0.1.0.dist-info/licenses/LICENSE,sha256=SQLFFQMkaYtZwAws2RPvCzAFCoGt4gctI1qmnGnh5lQ,1063
+cutez-0.1.0.dist-info/METADATA,sha256=X9wYU7jq3IkmVq7eH8e15Lm-jyCZ9VI5fasfjstAOos,130
+cutez-0.1.0.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
+cutez-0.1.0.dist-info/top_level.txt,sha256=RRDVDnvYVjgqHEhwgzWfzN0PFqZWRBbiToei-I0fhPk,6
+cutez-0.1.0.dist-info/RECORD,,

cutez-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (82.0.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

cutez-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 ZhangZ
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

cutez-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+cutez