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

mslk/attention/fmha/flash_mtia.py

@@ -0,0 +1,245 @@
+# 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 typing import Any, Iterable, Mapping, Optional, Set, Tuple
+
+import torch
+
+from .attn_bias import (
+    BlockDiagonalCausalFromBottomRightMask,
+    BlockDiagonalCausalLocalAttentionFromBottomRightMask,
+    BlockDiagonalCausalLocalAttentionMask,
+    BlockDiagonalCausalLocalAttentionPaddedKeysMask,
+    BlockDiagonalCausalMask,
+    BlockDiagonalCausalWithOffsetGappyKeysMask,
+    BlockDiagonalCausalWithOffsetPaddedKeysMask,
+    BlockDiagonalGappyKeysMask,
+    BlockDiagonalLocalAttentionPaddedKeysMask,
+    BlockDiagonalMask,
+    BlockDiagonalPaddedKeysMask,
+    LocalAttentionFromBottomRightMask,
+    LowerTriangularFromBottomRightLocalAttentionMask,
+    LowerTriangularFromBottomRightMask,
+    LowerTriangularMask,
+)
+from .flash import BwOp as BwOpCUDA, FwOp as FwOpCUDA
+from .utils.op_common import get_operator, register_operator
+
+try:
+    import mtia.host_runtime.torch_mtia.dynamic_library  # noqa  # type: ignore
+
+    @torch.library.custom_op(
+        "mslk_flash_mtia::flash_fwd",
+        mutates_args=(),
+        device_types=["mtia"],
+    )
+    def _flash_fwd(
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        cu_seqlens_q: Optional[torch.Tensor],
+        cu_seqlens_k: Optional[torch.Tensor],
+        seqused_k: Optional[torch.Tensor],
+        max_seqlen_q: int,
+        max_seqlen_k: int,
+        p: float,
+        softmax_scale: float,
+        is_causal: bool,
+        window_left: int,
+        window_right: int,
+        return_softmax: bool,
+        block_tables: Optional[torch.Tensor],
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        ret = torch.ops.aten._flash_attention_forward(
+            query,
+            key,
+            value,
+            cu_seqlens_q,  # cum_seq_q
+            cu_seqlens_k,  # cum_seq_k
+            max_seqlen_q,  # max_q
+            max_seqlen_k,  # max_k
+            p,  # dropout_p
+            is_causal,
+            return_debug_mask=False,
+            scale=softmax_scale,
+            window_size_left=window_left,
+            window_size_right=window_right,
+            seqused_k=seqused_k,
+            alibi_slopes=None,  # alibi_slopes
+        )
+        attention, logsumexp, rng_state, _, _ = ret
+        return attention, logsumexp, rng_state
+
+    @torch.library.register_fake("mslk_flash_mtia::flash_fwd")
+    def _flash_fwd_abstract(
+        query,
+        key,
+        value,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        seqused_k,
+        max_seqlen_q,
+        max_seqlen_k,
+        p,
+        softmax_scale,
+        is_causal,
+        window_left,
+        window_right,
+        return_softmax,
+        block_tables,
+    ):
+        out = torch.empty_like(query)
+        if cu_seqlens_q is None:
+            B, M, H, K = query.shape
+            lse_shape = [B, H, M]  # XXXX ?
+        else:
+            M, H, K = query.shape
+            B = cu_seqlens_q.shape[0] - 1
+            lse_shape = [H, M]
+        softmax_lse = torch.empty(lse_shape, device=query.device, dtype=torch.float32)
+        rng_state = torch.empty([2], device=query.device, dtype=torch.int64)
+        return out, softmax_lse, rng_state
+
+    @torch.library.custom_op(
+        "mslk_flash_mtia::flash_bwd",
+        mutates_args=(),
+        device_types=["mtia"],
+    )
+    def _flash_bwd(
+        grads_share_storage: bool,
+        grad: torch.Tensor,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        out: torch.Tensor,
+        lse: torch.Tensor,
+        cu_seqlens_q: torch.Tensor,
+        cu_seqlens_k: torch.Tensor,
+        max_seqlen_q: int,
+        max_seqlen_k: int,
+        p: float,
+        softmax_scale: float,
+        is_causal: bool,
+        window_left: int,
+        window_right: int,
+        rng_state: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        rng_state0 = rng_state1 = rng_state
+        dq, dk, dv = torch.ops.aten._flash_attention_backward(
+            grad,
+            query,
+            key,
+            value,
+            out,
+            lse,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            max_seqlen_q,
+            max_seqlen_k,
+            p,
+            is_causal,
+            rng_state0,
+            rng_state1,
+            scale=softmax_scale,
+            window_size_left=window_left,
+            window_size_right=window_right,
+        )
+        return dq, dk, dv
+
+    @torch.library.register_fake("mslk_flash_mtia::flash_bwd")
+    def _flash_bwd_abstract(
+        grads_share_storage,
+        grad,
+        query,
+        key,
+        value,
+        *args,
+        **kwargs,
+    ):
+        return _create_dq_dk_dv(grads_share_storage, query, key, value)
+
+except (ImportError, OSError):
+    pass
+
+
+def _create_dq_dk_dv(
+    grads_share_storage: bool, query, key, value
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    # Create dq,dk,dv
+    # If Q/K/V come from a single QKV tensor, let's put the gradient in the
+    # right strides, so we can avoid a `cat`
+    if grads_share_storage:
+        chunk = torch.empty(
+            (*query.shape[0:-2], 3, query.shape[-2], query.shape[-1]),
+            dtype=query.dtype,
+            device=query.device,
+        )
+        return chunk.select(-3, 0), chunk.select(-3, 1), chunk.select(-3, 2)
+    return torch.empty_like(query), torch.empty_like(key), torch.empty_like(value)
+
+
+FLASH_VERSION = torch.nn.attention._get_flash_version()  # noqa  # type: ignore
+
+
+@register_operator
+class FwOp(FwOpCUDA):
+    """Operator that computes memory-efficient attention using MTIA devicesa"""
+
+    OPERATOR = get_operator("mslk_flash_mtia", "flash_fwd")
+    SUPPORTED_DEVICES: Set[str] = {"mtia"}
+    SUPPORTED_ATTN_BIAS_TYPES: Iterable[Any] = (
+        type(None),
+        LowerTriangularMask,
+        LowerTriangularFromBottomRightMask,
+        LowerTriangularFromBottomRightLocalAttentionMask,
+        BlockDiagonalMask,
+        BlockDiagonalCausalMask,
+        BlockDiagonalCausalLocalAttentionMask,
+        BlockDiagonalCausalLocalAttentionFromBottomRightMask,
+        BlockDiagonalLocalAttentionPaddedKeysMask,
+        BlockDiagonalCausalLocalAttentionPaddedKeysMask,
+        BlockDiagonalCausalFromBottomRightMask,
+        BlockDiagonalCausalWithOffsetGappyKeysMask,
+        BlockDiagonalCausalWithOffsetPaddedKeysMask,
+        BlockDiagonalGappyKeysMask,
+        BlockDiagonalPaddedKeysMask,
+        LocalAttentionFromBottomRightMask,
+    )
+    NAME = f"fa2F@{FLASH_VERSION}-mtia"
+
+    ERROR_ATOL: Mapping[torch.dtype, float] = {
+        torch.float: 3e-4,
+        torch.half: 7e-3,
+        torch.bfloat16: 2e-2,
+    }
+    ERROR_RTOL: Mapping[torch.dtype, float] = {
+        torch.float: 2e-5,
+        torch.half: 4e-4,
+        torch.bfloat16: 5e-3,
+    }
+
+
+@register_operator
+class BwOp(BwOpCUDA):
+    """Operator that computes memory-efficient attention using MTIA devicesa"""
+
+    OPERATOR = get_operator("mslk_flash_mtia", "flash_bwd")
+    SUPPORTED_DEVICES: Set[str] = {"mtia"}
+    SUPPORTED_ATTN_BIAS_TYPES: Iterable[Any] = (
+        type(None),
+        LowerTriangularMask,
+        LowerTriangularFromBottomRightMask,
+        LowerTriangularFromBottomRightLocalAttentionMask,
+        BlockDiagonalMask,
+        BlockDiagonalCausalMask,
+        BlockDiagonalCausalLocalAttentionMask,
+        BlockDiagonalCausalLocalAttentionFromBottomRightMask,
+        BlockDiagonalCausalFromBottomRightMask,
+        LocalAttentionFromBottomRightMask,
+    )
+    NAME = f"fa2B@{FLASH_VERSION}-mtia"

mslk/attention/fmha/merge_training.py

@@ -0,0 +1,192 @@
+# 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 typing import Callable, Optional, Tuple, Type, Union
+
+import torch
+
+from .. import fmha
+
+
+"""
+Friendly wrapper around merge_attentions which works with autograd.
+
+Use as follows
+
+```
+partial1 = memory_efficient_attention_partial_autograd(q, k1, v1, ...)
+partial2 = memory_efficient_attention_partial_autograd(q, k2, v2, ...)
+attn_out = merge_attentions_autograd(partial1, partial2)
+```
+
+merge_attentions_autograd() can take any number of inputs. Note that
+partial1 and partial2 are not tensors, but rather objects of type
+`Partial`.
+
+If you have partial1 and you changed your mind and don't
+want to merge it with anything, you can do
+```
+attn_out = merge_attentions_autograd(partial1)
+```
+
+"""
+
+
+class _PartialFunc(torch.autograd.Function):
+    @staticmethod
+    def forward(  # type: ignore[override]
+        ctx: torch.autograd.function.FunctionCtx,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_bias: Optional[Union[torch.Tensor, fmha.AttentionBias]],
+        p: float = 0.0,
+        scale: Optional[float] = None,
+        op: Optional[Union[fmha.AttentionOp, Type[fmha.AttentionFwOpBase]]] = None,
+        output_dtype: Optional[torch.dtype] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        ctx.bias = attn_bias  # type: ignore
+        ctx.save_for_backward(query, key, value)
+        ctx.p = p  # type: ignore
+        ctx.scale = scale  # type: ignore
+        ctx.op = op[1] if isinstance(op, tuple) else None  # type: ignore
+        attn, lse = fmha.memory_efficient_attention_partial(
+            query, key, value, attn_bias, p, scale, op=op, output_dtype=output_dtype
+        )
+        placeholder = torch.empty_like(attn)
+        return attn, lse, placeholder
+
+    @staticmethod
+    def backward(  # type: ignore[override]
+        ctx: torch.autograd.function.FunctionCtx,
+        grad_attn: torch.Tensor,
+        lse: torch.Tensor,
+        out: torch.Tensor,
+    ) -> Tuple[Optional[torch.Tensor], ...]:
+        query, key, value = ctx.saved_tensors  # type: ignore
+        grad_q, grad_k, grad_v = fmha.memory_efficient_attention_backward(
+            grad_attn,
+            out,
+            lse.contiguous(),
+            query,
+            key,
+            value,
+            ctx.bias,  # type: ignore
+            ctx.p,  # type: ignore
+            ctx.scale,  # type: ignore
+            op=ctx.op,  # type: ignore
+        )
+        return grad_q, grad_k, grad_v, None, None, None, None, None
+
+
+class _MergeFunc(torch.autograd.Function):
+    @staticmethod
+    def forward(  # type: ignore[override]
+        ctx: torch.autograd.function.FunctionCtx,
+        *inputs: torch.Tensor,
+    ) -> torch.Tensor:
+        ctx.length = len(inputs) // 3  # type: ignore
+        if ctx.length > 1:  # type: ignore
+            attns = inputs[::3]
+            lses = inputs[1::3]
+            out, lse = fmha.merge_attentions(attns, lses)
+            assert lse is not None
+        else:
+            out, lse = inputs[:2]
+        ctx.save_for_backward(out, lse)
+        return out
+
+    @staticmethod
+    def backward(  # type: ignore[override]
+        ctx: torch.autograd.function.FunctionCtx, grad_out: torch.Tensor
+    ) -> Tuple[Optional[torch.Tensor], ...]:
+        out, lse = ctx.saved_tensors  # type: ignore
+        return (grad_out, lse, out) * ctx.length  # type: ignore
+
+
+class Partial:
+    """
+    This class is used to represent a partial attention output, which is
+    returned by `memory_efficient_attention_partial_autograd`.
+
+    Attributes: (Do not access them directly, use the methods instead.)
+
+        _attn: torch.Tensor
+        _lse: torch.Tensor . (Its grad is the full LSE to be used by
+            the individual backward passes.)
+        _placeholder: torch.Tensor, whose grad is used for passing the full attention
+            output to the individual backward passes.
+    """
+
+    def __init__(
+        self,
+        attn: torch.Tensor,
+        lse: torch.Tensor,
+        placeholder: torch.Tensor,
+    ) -> None:
+        """
+        Internal use only
+        """
+        self._attn = attn
+        self._lse = lse
+        self._placeholder = placeholder
+
+    def is_bmghk(self) -> bool:
+        return self._attn.ndim == 5
+
+    def apply(self, fn: Callable[[torch.Tensor], torch.Tensor]) -> "Partial":
+        """
+        Applies fn to the attention output, as if we were a tensor.
+        fn must expect tensors of shape BMGHK or BMHK, but cannot actually
+        manipulate the K dimension (because the LSE doesn't have one).
+
+        For example, to slice on the sequence dimension, you might apply
+        `lambda x: x[:, start:end]`.
+        """
+        attn = fn(self._attn)
+        if self.is_bmghk():
+            rearranged = self._lse.permute(0, 3, 1, 2).unsqueeze(-1)
+            lse = fn(rearranged).squeeze(-1).permute(0, 2, 3, 1)
+        else:
+            rearranged = self._lse.permute(0, 2, 1).unsqueeze(-1)
+            lse = fn(rearranged).squeeze(-1).permute(0, 2, 1)
+        placeholder = fn(self._placeholder)
+        return self.__class__(attn, lse, placeholder)
+
+    def _tuple(self) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        return self._attn, self._lse, self._placeholder
+
+
+def memory_efficient_attention_partial_autograd(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attn_bias: Optional[Union[torch.Tensor, fmha.AttentionBias]] = None,
+    p: float = 0.0,
+    scale: Optional[float] = None,
+    *,
+    op: Optional[Union[fmha.AttentionOp, Type[fmha.AttentionFwOpBase]]] = None,
+    output_dtype: Optional[torch.dtype] = None,
+) -> Partial:
+    """
+    Wrapper around `memory_efficient_attention_partial` which works with autograd.
+    Arguments are the same as for `memory_efficient_attention_partial`.
+    """
+    return Partial(
+        *_PartialFunc.apply(query, key, value, attn_bias, p, scale, op, output_dtype)
+    )
+
+
+def merge_attentions_autograd(
+    *partials: Partial,
+) -> torch.Tensor:
+    """
+    Wrapper around merge_attentions which works with autograd.
+    """
+    args = [i for part in partials for i in part._tuple()]
+    if len(args) == 0:
+        raise ValueError("No partials to merge")
+    return _MergeFunc.apply(*args)