mslk-cuda-nightly 2026.1.19__cp310-cp310-manylinux_2_28_x86_64.whl

mslk/testing/__init__.py

@@ -0,0 +1,7 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-unsafe

mslk/testing/attributes.py

@@ -0,0 +1,60 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-strict
+
+import os
+import subprocess
+from typing import Tuple
+
+import mslk
+import torch
+
+################################################################################
+# Unit test skip attributes for environments
+################################################################################
+
+gpu_unavailable: Tuple[bool, str] = (
+    not torch.cuda.is_available() or torch.cuda.device_count() == 0,
+    "GPU is not available or no GPUs detected",
+)
+
+running_in_github: Tuple[bool, str] = (
+    os.getenv("GITHUB_ENV") is not None,
+    "Test fails or hangs when run in the GitHub runners",
+)
+
+running_in_oss: Tuple[bool, str] = (
+    # pyre-ignore [16]
+    getattr(mslk, "open_source", False),
+    "Test is currently known to fail in OSS mode",
+)
+
+################################################################################
+# Unit test skip attributes for platforms
+################################################################################
+
+running_on_arm: Tuple[bool, str] = (
+    subprocess.run(["uname", "-m"], stdout=subprocess.PIPE)
+    .stdout.decode("utf-8")
+    .strip()
+    == "aarch64",
+    "Test is currently known to fail when running on ARM platform",
+)
+
+running_on_cuda: Tuple[bool, str] = (
+    torch.cuda.is_available()
+    and torch.cuda.device_count() > 0
+    and torch.version.hip is not None,
+    "Test currently doesn't work on the ROCm stack",
+)
+
+running_on_rocm: Tuple[bool, str] = (
+    torch.cuda.is_available()
+    and torch.cuda.device_count() > 0
+    and torch.version.hip is not None,
+    "Test currently doesn't work on the ROCm stack",
+)

mslk/testing/rocm.py

@@ -0,0 +1,91 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-strict
+
+import os
+import unittest
+from functools import wraps
+from typing import Any, Callable
+
+import torch
+
+running_on_rocm: bool = (
+    torch.cuda.is_available()
+    and torch.cuda.device_count() > 0
+    and torch.version.hip is not None
+)
+
+
+def skipIfRocm(
+    reason: str = "Test does not work on ROCm",
+) -> Any:
+    # pyre-fixme[24]: Generic type `Callable` expects 2 type parameters.
+    def decorator(fn: Callable) -> Any:
+        @wraps(fn)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            if running_on_rocm:
+                raise unittest.SkipTest(reason)
+            else:
+                fn(*args, **kwargs)
+
+        return wrapper
+
+    return decorator
+
+
+def skipIfNotRocm(
+    reason: str = "Test only works on ROCm",
+) -> Any:
+    # pyre-fixme[24]: Generic type `Callable` expects 2 type parameters.
+    def decorator(fn: Callable) -> Any:
+        @wraps(fn)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            if running_on_rocm:
+                fn(*args, **kwargs)
+            else:
+                raise unittest.SkipTest(reason)
+
+        return wrapper
+
+    return decorator
+
+
+def skipIfRocmLessThan(min_version: int) -> Any:
+    # pyre-fixme[24]: Generic type `Callable` expects 2 type parameters.
+    def decorator(testfn: Callable) -> Any:
+        @wraps(testfn)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            ROCM_VERSION_FILEPATH = "/opt/rocm/.info/version"
+            if running_on_rocm:
+                # Fail if ROCm version file is missing.
+                if not os.path.isfile(ROCM_VERSION_FILEPATH):
+                    raise AssertionError(
+                        f"ROCm version file {ROCM_VERSION_FILEPATH} is missing!"
+                    )
+
+                # Parse the version number from the file.
+                with open(ROCM_VERSION_FILEPATH, "r") as file:
+                    version = file.read().strip()
+                version = version.replace("-", "").split(".")
+                version = (
+                    int(version[0]) * 10000 + int(version[1]) * 100 + int(version[2])
+                )
+
+                # Fail if ROCm version is less than the minimum version.
+                if version < min_version:
+                    raise unittest.SkipTest(
+                        f"Skip the test since the ROCm version is less than {min_version}"
+                    )
+                else:
+                    testfn(*args, **kwargs)
+
+            else:
+                testfn(*args, **kwargs)
+
+        return wrapper
+
+    return decorator

mslk/utils/__init__.py

@@ -0,0 +1,7 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-strict

mslk/utils/torch/__init__.py

@@ -0,0 +1,7 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-strict

mslk/utils/torch/library.py

@@ -0,0 +1,150 @@
+#!/usr/bin/env python3
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-strict
+
+import logging
+import re
+from typing import Callable
+
+import torch
+
+
+def load_library_buck(buck_target: str) -> None:
+    import mslk  # noqa: F401
+
+    # pyre-ignore [16]
+    open_source: bool = getattr(mslk, "open_source", False)
+
+    try:
+        torch.ops.load_library(buck_target)
+    except OSError as e:
+        if open_source:
+            pass
+        else:
+            logging.error(
+                f"Failed to load buck target {buck_target}, ops will not be available via torch.ops! Error: {e}"
+            )
+
+
+class TorchLibraryFragment:
+    """
+    A wrapper class around PyTorch library fragments, which are used to define
+    and register PyTorch operators.  Handles duplicate operator definitions and
+    registrations under the hood.
+    """
+
+    def __init__(self, namespace: str) -> None:
+        """
+        Constructs the TorchLibraryFragment class.
+
+        Args:
+            namespace: The namespace for the operators.
+
+        Returns:
+            None
+
+        Example:
+            lib = TorchLibrary("mslk")
+        """
+        self.namespace = namespace
+        self.lib = torch.library.Library(namespace, "FRAGMENT")
+
+    def define(self, schema: str) -> None:
+        """
+        Defines an operator schema.  This function handles the case where the
+        opeator name has already been defined.
+
+        Args:
+            schema: The schema of the operator to be defined.  The operator name
+            should NOT be prefixed with the operator namespace.
+
+        Returns:
+            None
+
+        Example:
+            lib = TorchLibrary("mslk")
+            lib.define("sll_jagged_jagged_bmm(Tensor x, Tensor y, bool flag=True) -> Tensor")
+        """
+        pattern = re.compile(
+            r"""
+            (\w+)               # Match the function name (capturing group)
+            \s*\(               # Match the opening parenthesis with optional whitespace
+            ([^)]*)             # Match params list (capturing group)
+            \s*\)               # Match the closing parenthesis with optional whitespace
+            \s*->\s*.+          # Match '-> <Return Type>'
+            """,
+            re.VERBOSE,
+        )
+
+        match = pattern.search(schema.strip())
+        if match:
+            name = match.group(1)
+            if f"{self.namespace}::{name}" not in torch.library._defs:
+                self.lib.define(schema)
+        else:
+            raise ValueError(
+                f"PyTorch operator schema appears to be ill-defined: '''{schema}'''"
+            )
+
+    # pyre-ignore[24]
+    def register_dispatch(self, op_name: str, dispatch_key: str, fn: Callable) -> None:
+        """
+        Registers a single dispatch for an operator with the given name and dispatch key.
+
+        Args:
+            op_name: operator name
+            dispatch_key: dispatch key that the function should be registered for (e.g., "CUDA")
+            fn: a function that is the operator implementation for the input dispatch key
+
+        Returns:
+            None
+
+        Example:
+            lib = TorchLibrary("mslk")
+            lib.define(...)
+            lib.register_dispatch(lib, "jagged_dense_bmm", jagged_dense_bmm, "CUDA")
+        """
+
+        valid_backends = [
+            "CUDA",
+            "AutogradCUDA",
+            "CPU",
+            "AutogradCPU",
+            "AutogradMeta",
+            "Meta",
+            "CompositeImplicitAutograd",
+        ]
+        assert dispatch_key in valid_backends
+
+        if not torch._C._dispatch_has_kernel_for_dispatch_key(
+            f"{self.namespace}::{op_name}", dispatch_key
+        ):
+            if dispatch_key == "Meta":
+                self.lib._register_fake(op_name, fn)
+            else:
+                self.lib.impl(op_name, fn, dispatch_key)
+
+    # pyre-ignore[24]
+    def register(self, op_name: str, functors: dict[str, Callable]) -> None:
+        """
+        Registers a set of dispatches for a defined operator.
+
+        Args:
+            op_name: operator name
+            functors: A dictionary of dispatch keys to dispatch implementations
+
+        Returns:
+            None
+
+        Example:
+            lib = TorchLibrary("mslk")
+            lib.define(...)
+            lib.register(lib, "jagged_dense_bmm", {"CUDA": jagged_dense_bmm, "Meta": jagged_dense_bmm_meta })
+        """
+        for dispatch, func in functors.items():
+            self.register_dispatch(op_name, dispatch, func)

mslk/utils/triton/__init__.py

@@ -0,0 +1,7 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-strict

mslk/utils/triton/fp8_utils.py

@@ -0,0 +1,72 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-unsafe
+import functools
+import logging
+import os
+from typing import Tuple
+
+import torch
+import triton.language as tl  # @manual
+from triton.runtime.jit import reinterpret as tl_reinterpret, TensorWrapper  # @manual
+
+
+running_on_github: bool = os.getenv("GITHUB_ENV") is not None
+
+
+@functools.lru_cache
+def supports_float8_fnuz(throw_on_hip_incompatibility: bool = True) -> bool:
+    if torch.version.hip:
+        device_capability = torch.cuda.get_device_capability()
+
+        if device_capability < (9, 4):
+            gpu_arch = torch.cuda.get_device_properties("cuda").gcnArchName
+            msg = f"Unsupported GPU arch: {gpu_arch} for FP8"
+            if throw_on_hip_incompatibility:
+                raise RuntimeError(msg)
+            else:
+                logging.error(msg)
+                return False
+
+        elif device_capability == (9, 4):
+            return True
+
+    return False
+
+
+def get_fp8_constants() -> Tuple[torch.dtype, tl.dtype, float, float]:
+    """
+    Helper function to get constant values for the current platform.
+
+    Returns:
+        pt_dtype (torch.dtype): The correct torch fp8 datatype.
+        tl_dtype (tl.dtype): The correct triton fp8 datatype.
+        max_fp8 (float): The maximum reprsentable value for the fp8 datatype.
+        eps (float): Minimum clip value to prevent divide by zero.
+    """
+    if supports_float8_fnuz(throw_on_hip_incompatibility=(not running_on_github)):
+        pt_fp8_dtype = torch.float8_e4m3fnuz
+        tl_fp8_dtype = tl.float8e4b8
+    else:
+        pt_fp8_dtype = torch.float8_e4m3fn
+        tl_fp8_dtype = tl.float8e4nv
+
+    return pt_fp8_dtype, tl_fp8_dtype, torch.finfo(pt_fp8_dtype).max, 1e-12
+
+
+def reinterpret_fp8_type(tensor: torch.Tensor, dtype: tl.dtype) -> TensorWrapper:
+    """
+    Converts tensor to triton fp8 type.
+
+    Args:
+        tensor (torch.Tensor): input tensor.
+        dtype (tl.dtype): target triton dtype.
+
+    Returns:
+        triton.TensorWrapper: fp8 tensor.
+    """
+    return tl_reinterpret(tensor, dtype=dtype)

mslk/utils/triton/utils.py

@@ -0,0 +1,128 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-unsafe
+import sys
+
+import torch
+import triton  # @manual
+import triton.language as tl  # @manual
+
+
+def map_dtype_to_triton(dtype: torch.dtype) -> tl.dtype:
+    """
+    Maps torch dtype to triton dtype.
+
+    Args:
+        dtype (torch.dtype): input dtype.
+
+    Returns:
+        tl.dtype: triton dtype.
+    """
+    if dtype == torch.float16:
+        return tl.float16
+    elif dtype == torch.bfloat16:
+        return tl.bfloat16
+    elif dtype == torch.float32:
+        return tl.float32
+    elif dtype == torch.int32:
+        return tl.int32
+    elif dtype == torch.float8_e4m3fn and torch.version.hip is None:
+        return tl.float8e4nv
+    else:
+        raise ValueError(f"Unsupported dtype {dtype}")
+
+
+# check if we have the TMA version in Triton PR #4498 (https://github.com/triton-lang/triton/pull/4498).
+HAS_TMA_DESC = "nv_tma_desc_type" in dir(tl)
+
+if HAS_TMA_DESC:
+    print(
+        "TMA benchmarks will be running with experimental grid constant TMA descriptor.",
+        file=sys.stderr,
+    )
+else:
+    print(
+        "TMA benchmarks will be running without grid constant TMA descriptor.",
+        file=sys.stderr,
+    )
+
+
+class TmaAutoTuneHelper:
+    # duck typing wrapper to implement the same interface as TmaDescKernelParam in Triton PR #4498
+    class KernelParamWrapper:
+        def __init__(self, desc):
+            self.desc = desc
+
+        def tma_desc_cpu_ptr(self):
+            return self.desc.data_ptr()
+
+    TMA_SIZE = 128
+
+    def __init__(self):
+        self.fill_1d_tma_descriptor_inner = (
+            triton.runtime.driver.active.utils.fill_1d_tma_descriptor
+        )
+        self.fill_2d_tma_descriptor_inner = (
+            triton.runtime.driver.active.utils.fill_2d_tma_descriptor
+        )
+        if HAS_TMA_DESC:
+            self.descriptors = {}
+        else:
+            self.cuda_descriptors = {}
+
+    # Call this method outside of the lambda function for grid size
+    def init_tma_descriptor(self, name):
+        if HAS_TMA_DESC:
+            self.descriptors[name] = torch.empty(
+                TmaAutoTuneHelper.TMA_SIZE, device="cpu", dtype=torch.int8
+            )
+        else:
+            self.cuda_descriptors[name] = torch.empty(
+                TmaAutoTuneHelper.TMA_SIZE, device="cuda", dtype=torch.int8
+            )
+
+    # Call this method inside the lambda function for grid size
+    def fill_1d_tma_descriptor(self, name, ptr, dim, block_dim, element_size):
+        if HAS_TMA_DESC:
+            desc_x = self.descriptors[name]
+            assert desc_x.data_ptr() % 64 == 0
+            self.fill_1d_tma_descriptor_inner(
+                ptr, dim, block_dim, element_size, desc_x.data_ptr()
+            )
+        else:
+            desc_x = self.cuda_descriptors[name]
+            buf_x = torch.empty_like(desc_x, device="cpu", pin_memory=True)
+            self.fill_1d_tma_descriptor_inner(
+                ptr, dim, block_dim, element_size, buf_x.data_ptr()
+            )
+            desc_x.copy_(buf_x, non_blocking=True)
+
+    # Call this method inside the lambda function for grid size
+    def fill_2d_tma_descriptor(
+        self, name, ptr, dim1, dim0, block_dim1, block_dim0, element_size
+    ):
+        if HAS_TMA_DESC:
+            desc_x = self.descriptors[name]
+            assert desc_x.data_ptr() % 64 == 0
+            self.fill_2d_tma_descriptor_inner(
+                ptr, dim1, dim0, block_dim1, block_dim0, element_size, desc_x.data_ptr()
+            )
+        else:
+            desc_x = self.cuda_descriptors[name]
+            buf_x = torch.empty_like(desc_x, device="cpu", pin_memory=True)
+            self.fill_2d_tma_descriptor_inner(
+                ptr, dim1, dim0, block_dim1, block_dim0, element_size, buf_x.data_ptr()
+            )
+            desc_x.copy_(buf_x, non_blocking=True)
+
+    def get_tma_descriptor_kernel_param(self, name):
+        if HAS_TMA_DESC:
+            assert self.descriptors[name] is not None
+            return self.KernelParamWrapper(self.descriptors[name])
+        else:
+            assert self.cuda_descriptors[name] is not None
+            return self.cuda_descriptors[name]

mslk/version.py

@@ -0,0 +1,11 @@
+
+#!/usr/bin/env python3
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+__version__: str = "2026.1.19"
+__target__: str = "default"
+__variant__: str = "cuda"

mslk_cuda_nightly-2026.1.19.dist-info/METADATA

@@ -0,0 +1,102 @@
+Metadata-Version: 2.4
+Name: mslk-cuda-nightly
+Version: 2026.1.19
+Home-page: https://github.com/pytorch/MSLK
+Author: MSLK Team
+Author-email: packages@pytorch.org
+License: BSD-3
+Keywords: PyTorch,Generative AI,High Performance Computing,GPU,CUDA,ROCm
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: BSD License
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: keywords
+Dynamic: license
+Dynamic: license-file
+Dynamic: requires-dist
+
+# MSLK Library
+
+
+MSLK (Meta Superintelligence Labs Kernels, formerly known as **[FBGEMM GenAI](https://github.com/pytorch/FBGEMM/tree/main/fbgemm_gpu/experimental/gen_ai)**)
+is a collection of high-performance kernels and optimizations built on top of PyTorch
+primitives for GenAI training and inference.
+
+## **Installation**
+
+```bash
+# Install MSLK for CUDA
+pip install mslk-cuda==1.0.0
+# Install MSLK for ROCm
+pip install mslk-rocm==1.0.0
+# Install a nightly version
+pip3 install --pre mslk --index-url https://download.pytorch.org/whl/nightly/cu128
+```
+
+## Release Compatibility Table
+
+MSLK is released in accordance to the PyTorch release schedule, and each
+release has no guarantee to work in conjunction with PyTorch releases that are
+older than the one that the MSLK release corresponds to.
+
+| MSLK Release | Corresponding PyTorch Release | Supported Python Versions | Supported CUDA Versions | Supported CUDA Architectures | Supported ROCm Versions | Supported ROCm Architectures |
+|---------|---------|---------|---------|----------|-------------|-------------|
+| 1.0.0 | 2.10.x | 3.10, 3.11, 3.12, 3.13, 3.14 | 12.6, 12.8, 12.9, 13.0 | 8.0, 9.0a, 10.0a, 12.0a | 7.0, 7.1 | gfx908, gfx90a, gfx942, gfx950 |
+
+## **Running Benchmarks**
+```bash
+python bench/gemm/gemm_bench.py
+python bench/quantize/quantize_bench.py
+```
+
+## **Running Tests**
+```bash
+python test/gemm/gemm_test.py
+python test/gemm/quantize_test.py
+```
+
+## **Build From Source**
+We only support building on Linux. See the release compatibility table above for supported versions of Python, CUDA, ROCm.
+```bash
+# Clone repo
+git clone https://github.com/meta-pytorch/MSLK
+cd MSLK
+git submodule sync
+git submodule update --init --recursive
+# Build and install
+# The script will create a conda environment and install the required dependencies.
+# The conda environment will look something like: build-py3.14-torchnightly-cuda12.9.1
+./ci/integration/mslk_oss_build.bash
+# After the initial environment setup, you can activate the environment and iterate faster:
+conda activate build-py3.14-torchnightly-cuda12.9.1
+python setup.py install
+```
+
+## Join the MSLK community
+
+For questions, support, news updates, or feature requests, please feel free to:
+
+* File a ticket in [GitHub Issues](https://github.com/meta-pytorch/MSLK/issues)
+
+For contributions, please see the [`CONTRIBUTING`](./CONTRIBUTING.md) file for
+ways to help out.
+
+## License
+
+MSLK is BSD licensed, as found in the [`LICENSE`](./LICENSE) file.