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

mslk/__init__.py

@@ -0,0 +1,56 @@
+# 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.
+
+import logging
+import os
+
+import torch
+
+open_source: bool = True
+
+
+def _load_library(filename: str, version: str, no_throw: bool = False) -> None:
+    """Load a shared library from the given filename."""
+    try:
+        library_path = os.path.join(os.path.dirname(__file__), filename)
+        torch.ops.load_library(library_path)
+        torch.classes.load_library(library_path)
+        logging.info(f"Successfully loaded: '{filename}'")
+
+    except Exception as error:
+        logging.error(f"Could not load the library '{filename}'!\n\n\n{error}\n\n\n")
+        if not no_throw:
+            raise error
+
+
+try:
+    # Export the version string from the version file auto-generated by setup.py
+    from .version import __target__, __variant__, __version__  # noqa: F401, E402
+except Exception:
+    __variant__: str = "INTERNAL"
+    __version__: str = "INTERNAL"
+    __target__: str = "default"
+
+_default_libraries = [
+    "mslk",
+]
+
+libraries_to_load = {
+    "default": _default_libraries,
+}
+
+for library in libraries_to_load.get(__target__, []):
+    # NOTE: In all cases, we want to throw an error if we cannot load the
+    # library.  However, this appears to break the OSS documentation build,
+    # where the Python documentation doesn't show up in the generated docs.
+    #
+    # To work around this problem, we introduce a fake build variant called
+    # `docs` and we only throw a library load error when the variant is not
+    # `docs`.  For more information, see:
+    #
+    #   https://github.com/pytorch/FBGEMM/pull/3477
+    #   https://github.com/pytorch/FBGEMM/pull/3717
+    _load_library(f"{library}.so", __version__, __variant__ == "docs")

mslk/attention/__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/attention/cutlass_blackwell_fmha/__init__.py

@@ -0,0 +1,30 @@
+# 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
+
+from mslk.utils.torch.library import load_library_buck
+
+from . import cutlass_blackwell_fmha_custom_op  # noqa: F401
+from .cutlass_blackwell_fmha_interface import (  # noqa: F401
+    _cutlass_blackwell_fmha_forward,
+    cutlass_blackwell_fmha_decode_forward,
+    cutlass_blackwell_fmha_func,
+)
+
+load_library_buck(
+    "//mslk/csrc/attention/cuda/cutlass_blackwell_fmha:blackwell_attention_ops_gpu"
+)
+
+# Note: _cutlass_blackwell_fmha_forward is an internal function (indicated by leading underscore)
+# that is exported here specifically for testing purposes. It allows tests to access the LSE
+# (log-sum-exp) values returned by the forward pass without modifying the public API.
+# Production code should use cutlass_blackwell_fmha_func instead.
+__all__ = [
+    "_cutlass_blackwell_fmha_forward",
+    "cutlass_blackwell_fmha_decode_forward",
+    "cutlass_blackwell_fmha_func",
+]

mslk/attention/cutlass_blackwell_fmha/cutlass_blackwell_fmha_custom_op.py

@@ -0,0 +1,332 @@
+# 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.
+
+from typing import Optional
+
+import torch
+from torch.library import register_fake
+
+
+torch.library.define(
+    "mslk::cutlass_blackwell_fmha_fwd",
+    "(Tensor q, Tensor k, Tensor v, Tensor? cu_seqlens_q, Tensor? cu_seqlens_k, int? max_seq_len_q, int? max_seq_len_k, float? softmax_scale, bool? causal, Tensor? seqlen_kv, Tensor? page_table, int seqlen_k=-1, int window_size_left=-1, int window_size_right=-1, bool bottom_right=True) -> (Tensor, Tensor)",
+    tags=torch.Tag.pt2_compliant_tag,
+)
+
+torch.library.define(
+    "mslk::cutlass_blackwell_fmha_bwd",
+    "(Tensor dout, Tensor q, Tensor k, Tensor v, Tensor out, Tensor softmax_lse, Tensor? cu_seqlens_q, Tensor? cu_seqlens_k, int? max_seq_len_q, int? max_seq_len_k, float? softmax_scale, bool? causal, int window_size_left=-1, int window_size_right=-1, bool bottom_right=True, bool deterministic=False) -> (Tensor, Tensor, Tensor)",
+    tags=torch.Tag.pt2_compliant_tag,
+)
+
+
+@torch.library.impl("mslk::cutlass_blackwell_fmha_fwd", "cuda")
+def custom_op_fmha(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: Optional[torch.Tensor] = None,
+    cu_seqlens_k: Optional[torch.Tensor] = None,
+    max_seq_len_q: Optional[int] = None,
+    max_seq_len_k: Optional[int] = None,
+    softmax_scale: Optional[float] = None,
+    causal: bool = False,
+    seqlen_kv: Optional[torch.Tensor] = None,
+    page_table: Optional[torch.Tensor] = None,
+    seqlen_k: Optional[int] = None,
+    window_size_left: int = -1,
+    window_size_right: int = -1,
+    bottom_right: bool = True,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    assert q.is_contiguous(), "q is not contiguous"
+    assert k.is_contiguous(), "k is not contiguous"
+    assert v.is_contiguous(), "v is not contiguous"
+    assert q.is_cuda, "q must be on GPU"
+    assert k.is_cuda, "k must be on GPU"
+    assert v.is_cuda, "v must be on GPU"
+
+    return torch.ops.mslk.fmha_fwd(
+        q,
+        k,
+        v,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k=cu_seqlens_k,
+        max_seq_len_q=max_seq_len_q,
+        max_seq_len_k=max_seq_len_k,
+        softmax_scale=softmax_scale,
+        causal=causal,
+        seqlen_kv=seqlen_kv,
+        page_table=page_table,
+        seqlen_k=seqlen_k,
+        window_size_left=window_size_left,
+        window_size_right=window_size_right,
+        bottom_right=bottom_right,
+    )
+
+
+@register_fake("mslk::cutlass_blackwell_fmha_fwd")
+def fmha_fwd_meta(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: Optional[torch.Tensor] = None,
+    cu_seqlens_k: Optional[torch.Tensor] = None,
+    max_seq_len_q: Optional[int] = None,
+    max_seq_len_k: Optional[int] = None,
+    softmax_scale: Optional[float] = None,
+    causal: bool = False,
+    seqlen_kv: Optional[torch.Tensor] = None,
+    page_table: Optional[torch.Tensor] = None,
+    seqlen_k: Optional[int] = None,
+    window_size_left: int = -1,
+    window_size_right: int = -1,
+    bottom_right: bool = True,
+):
+    if q.dtype == torch.float16:
+        out_dtype = torch.float16
+    elif q.dtype == torch.bfloat16:
+        out_dtype = torch.bfloat16
+    elif q.dtype == torch.float8_e4m3fn:
+        # Output is BF16 when input is FP8
+        out_dtype = torch.bfloat16
+    else:
+        raise RuntimeError(f"Unsupported dtype for q: {q.dtype}")
+
+    kIsVarlen = max_seq_len_q is not None
+    if kIsVarlen:
+        assert cu_seqlens_q is not None
+        SQ = q.shape[0]
+        H_Q = q.shape[1]
+        B = cu_seqlens_q.shape[0] - 1
+    else:
+        SQ = q.shape[1]
+        H_Q = q.shape[2]
+        B = q.shape[0]
+    device = q.device
+    options2 = {"dtype": torch.float32, "device": device}
+    if kIsVarlen:
+        assert max_seq_len_q is not None
+        out = torch.empty_like(q, dtype=out_dtype)
+        size = out.size()
+        stride = out.stride()
+        storage_offset = q.shape[-1] * max_seq_len_q * H_Q  # example scalar offset
+        out1 = torch.as_strided(
+            out, size=size, stride=stride, storage_offset=storage_offset
+        )
+    else:
+        out1 = torch.empty_like(q, dtype=out_dtype)
+
+    if kIsVarlen:
+        out2 = torch.empty((1, H_Q, SQ), **options2)  # type: ignore
+    else:
+        out2 = torch.empty((B, H_Q, SQ), **options2)  # type: ignore
+    return out1, out2
+
+
+@torch.library.impl("mslk::cutlass_blackwell_fmha_bwd", "cuda")
+def custom_op_fmha_bwd(
+    dOutput: torch.Tensor,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    output: torch.Tensor,
+    softmax_lse: torch.Tensor,
+    cu_seqlens_q: Optional[torch.Tensor] = None,
+    cu_seqlens_k: Optional[torch.Tensor] = None,
+    max_seq_len_q: Optional[int] = None,
+    max_seq_len_k: Optional[int] = None,
+    softmax_scale: Optional[float] = None,
+    causal: bool = False,
+    window_size_left: int = -1,
+    window_size_right: int = -1,
+    bottom_right: bool = True,
+    deterministic: bool = False,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    return torch.ops.mslk.fmha_bwd(
+        dOutput,
+        query,
+        key,
+        value,
+        output,
+        softmax_lse,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k=cu_seqlens_k,
+        max_seq_len_q=max_seq_len_q,
+        max_seq_len_k=max_seq_len_k,
+        softmax_scale=softmax_scale,
+        causal=causal,
+        window_size_left=window_size_left,
+        window_size_right=window_size_right,
+        bottom_right=bottom_right,
+        deterministic=deterministic,
+    )
+
+
+@register_fake("mslk::cutlass_blackwell_fmha_bwd")
+def fmha_bwd_meta(
+    dOutput: torch.Tensor,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    output: torch.Tensor,
+    softmax_lse: torch.Tensor,
+    cu_seqlens_q: Optional[torch.Tensor] = None,
+    cu_seqlens_k: Optional[torch.Tensor] = None,
+    max_seq_len_q: Optional[int] = None,
+    max_seq_len_k: Optional[int] = None,
+    softmax_scale: Optional[float] = None,
+    causal: bool = False,
+    window_size_left: int = -1,
+    window_size_right: int = -1,
+    bottom_right: bool = True,
+    deterministic: bool = False,
+):
+    return (
+        torch.empty_like(query),
+        torch.empty_like(key),
+        torch.empty_like(value),
+    )
+
+
+def _backward(ctx, *grad):
+    if ctx.is_gen:
+        # For gen case, no backward pass is needed (generation is inference only)
+        raise RuntimeError("Backward pass is not supported for generation phase (sq=1)")
+    q, k, v, out, softmax_lse = ctx.saved_tensors
+    if not grad[0].is_contiguous():
+        grad0 = grad[0].contiguous()
+    else:
+        grad0 = grad[0]
+    if not softmax_lse.is_contiguous:
+        softmax_lse = softmax_lse.contiguous()
+    if not out.is_contiguous:
+        out = out.contiguous()
+    if not q.is_contiguous:
+        q = q.contiguous()
+    if not k.is_contiguous:
+        k = k.contiguous()
+
+    if not softmax_lse.is_contiguous:
+        softmax_lse = softmax_lse.contiguous()
+    if not out.is_contiguous:
+        out = out.contiguous()
+    if not q.is_contiguous:
+        q = q.contiguous()
+    if not k.is_contiguous:
+        k = k.contiguous()
+
+    dq, dk, dv = torch.ops.mslk.cutlass_blackwell_fmha_bwd(
+        grad0,
+        q,
+        k,
+        v,
+        out,
+        softmax_lse,
+        ctx.cu_seqlens_q,
+        ctx.cu_seqlens_k,
+        ctx.max_seq_len_q,
+        ctx.max_seq_len_k,
+        ctx.softmax_scale,
+        ctx.causal,
+        ctx.window_size_left,
+        ctx.window_size_right,
+        ctx.bottom_right,
+        ctx.deterministic,
+    )
+    return (
+        dq,
+        dk,
+        dv,
+        None,
+        None,
+        None,
+        None,
+        None,
+        None,
+        None,
+        None,
+        None,
+        None,
+        None,
+        None,
+    )
+
+
+def _setup_context(ctx, inputs, output):
+    (
+        q,
+        k,
+        v,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        max_seq_len_q,
+        max_seq_len_k,
+        softmax_scale,
+        causal,
+        seqlen_kv,
+        page_table,
+        seqlen_k,
+        window_size_left,
+        window_size_right,
+        bottom_right,
+    ) = inputs
+    (out, softmax_lse) = output
+    ctx.save_for_backward(q, k, v, out, softmax_lse)
+    ctx.softmax_scale = softmax_scale
+    ctx.causal = causal
+    ctx.max_seq_len_q = max_seq_len_q
+    ctx.max_seq_len_k = max_seq_len_k
+    ctx.cu_seqlens_q = cu_seqlens_q
+    ctx.cu_seqlens_k = cu_seqlens_k
+    ctx.window_size_left = window_size_left
+    ctx.window_size_right = window_size_right
+    ctx.bottom_right = bottom_right
+    ctx.deterministic = False  # Set default value
+    ctx.is_gen = False
+
+
+# This code adds training support for the operator. You must provide us
+# the backward formula for the operator and a `setup_context` function
+# to save values to be used in the backward.
+torch.library.register_autograd(
+    "mslk::cutlass_blackwell_fmha_fwd", _backward, setup_context=_setup_context
+)
+
+
+def cutlass_blackwell_fmha_custom_op(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    softmax_scale: float | None = None,
+    causal: bool = False,
+    cu_seqlens_q: torch.Tensor | None = None,
+    cu_seqlens_k: torch.Tensor | None = None,
+    max_seq_len_q: int | None = None,
+    max_seq_len_k: int | None = None,
+    seqlen_kv: torch.Tensor | None = None,
+    page_table: torch.Tensor | None = None,
+    seqlen_k: int | None = -1,
+    window_size_left: int | None = -1,
+    window_size_right: int | None = -1,
+    bottom_right: bool | None = True,
+):
+    return torch.ops.mslk.cutlass_blackwell_fmha_fwd(
+        q=q,
+        k=k,
+        v=v,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k=cu_seqlens_k,
+        max_seq_len_q=max_seq_len_q,
+        max_seq_len_k=max_seq_len_k,
+        softmax_scale=softmax_scale,
+        causal=causal,
+        seqlen_kv=seqlen_kv,
+        page_table=page_table,
+        seqlen_k=seqlen_k,
+        window_size_left=window_size_left,
+        window_size_right=window_size_right,
+        bottom_right=bottom_right,
+    )[0]