Initial import: grid-bot — grid trading bot for BTC-USDT on Cifra Markets
This commit is contained in:
Kolp
@@ -0,0 +1,197 @@
|
||||
# mypy: allow-untyped-defs
|
||||
"""This module contains functions and classes that alter the behavior of torch.nn.functional.scaled_dot_product_attention"""
|
||||
|
||||
import contextlib
|
||||
from collections.abc import Iterable
|
||||
from typing import Union
|
||||
from warnings import warn
|
||||
|
||||
import torch.backends.cuda
|
||||
from torch._C import _SDPBackend as SDPBackend
|
||||
from torch.backends.cuda import (
|
||||
can_use_efficient_attention,
|
||||
can_use_flash_attention,
|
||||
SDPAParams,
|
||||
)
|
||||
|
||||
|
||||
__all__: list[str] = [
|
||||
"SDPBackend",
|
||||
"sdpa_kernel",
|
||||
"WARN_FOR_UNFUSED_KERNELS",
|
||||
"register_flash_attention_impl",
|
||||
"activate_flash_attention_impl",
|
||||
"list_flash_attention_impls",
|
||||
"current_flash_attention_impl",
|
||||
"restore_flash_attention_impl",
|
||||
]
|
||||
|
||||
|
||||
# Note: [SDPA warnings]
|
||||
# TODO: Consider using this for sdpa regardless of subclasses
|
||||
# This only effects users of bias subclasses
|
||||
# If this is set to True, we will warn the user if they are not using the fused kernels
|
||||
# As well, it will raise warnings for all the reasons why the fused kernels can't be run.
|
||||
# To set this to True, run
|
||||
# torch.nn.attention.WARN_FOR_UNFUSED_KERNELS = True
|
||||
WARN_FOR_UNFUSED_KERNELS = False
|
||||
|
||||
|
||||
r"""An enum-like class that contains the different backends for scaled dot product attention.
|
||||
This backend class is designed to be used with the sdpa_kernel context manager.
|
||||
|
||||
The following Enums are available:
|
||||
- ERROR: An error occurred when trying to determine the backend.
|
||||
- MATH: The math backend for scaled dot product attention.
|
||||
- FLASH_ATTENTION: The flash attention backend for scaled dot product attention.
|
||||
- EFFICIENT_ATTENTION: The efficient attention backend for scaled dot product attention.
|
||||
- CUDNN_ATTENTION: The cuDNN backend for scaled dot product attention.
|
||||
- OVERRIDEABLE: The overridable backend for extension.
|
||||
|
||||
See :func:`torch.nn.attention.sdpa_kernel` for more details.
|
||||
|
||||
.. warning:: This class is in beta and subject to change.
|
||||
"""
|
||||
SDPBackend.__module__ = __name__
|
||||
SDPBackend.__name__ = "SDPBackend"
|
||||
|
||||
|
||||
def _raise_kernel_warnings(params: SDPAParams) -> None:
|
||||
"""
|
||||
If WARN_FOR_UNFUSED_KERNELS is set to True, this will raise warnings
|
||||
for all the reasons why the fused kernels can't be run. If using subclasses
|
||||
"""
|
||||
if WARN_FOR_UNFUSED_KERNELS:
|
||||
if not can_use_efficient_attention(params):
|
||||
warn("Efficient attention can't be used because:", stacklevel=2)
|
||||
can_use_efficient_attention(params, True)
|
||||
if not can_use_flash_attention(params):
|
||||
warn("Flash attention can't be used because:", stacklevel=2)
|
||||
can_use_flash_attention(params, True)
|
||||
|
||||
|
||||
_backend_names = {
|
||||
"cudnn": "CUDNN_ATTENTION",
|
||||
"flash": "FLASH_ATTENTION",
|
||||
"mem_efficient": "EFFICIENT_ATTENTION",
|
||||
"math": "MATH",
|
||||
"overrideable": "OVERRIDEABLE",
|
||||
}
|
||||
|
||||
|
||||
def _backend_from_string(name: str):
|
||||
return getattr(SDPBackend, name)
|
||||
|
||||
|
||||
def _cur_sdpa_kernel_backends(with_priority: bool = False):
|
||||
backends = []
|
||||
for name, val in _backend_names.items():
|
||||
if getattr(torch._C, f"_get_{name}_sdp_enabled")():
|
||||
backends.append(getattr(SDPBackend, val))
|
||||
if with_priority:
|
||||
curr_priority = torch._C._get_sdp_priority_order()
|
||||
backends = sorted(
|
||||
backends, key=lambda backend: curr_priority.index(int(backend))
|
||||
)
|
||||
return backends
|
||||
|
||||
|
||||
def _sdpa_kernel(backends: Iterable, set_priority: bool = False) -> None:
|
||||
for name, val in _backend_names.items():
|
||||
enabled = getattr(SDPBackend, val) in backends
|
||||
getattr(torch._C, f"_set_sdp_use_{name}")(enabled)
|
||||
if set_priority:
|
||||
# backends should be a unique list
|
||||
user_priority = [int(backend) for backend in backends]
|
||||
previous_priority = torch._C._get_sdp_priority_order()
|
||||
for backend in previous_priority:
|
||||
if backend not in user_priority:
|
||||
user_priority.append(int(backend))
|
||||
torch._C._set_sdp_priority_order(user_priority)
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def sdpa_kernel(backends: list[SDPBackend] | SDPBackend, set_priority: bool = False):
|
||||
r"""
|
||||
Context manager to select which backend to use for scaled dot product attention.
|
||||
|
||||
.. warning:: This function is beta and subject to change.
|
||||
|
||||
Args:
|
||||
backends (Union[List[SDPBackend], SDPBackend]): A backend or list of backends for scaled dot product attention.
|
||||
set_priority (bool=False): Whether the ordering of the backends is interpreted as their priority order.
|
||||
|
||||
Example:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from torch.nn.functional import scaled_dot_product_attention
|
||||
from torch.nn.attention import SDPBackend, sdpa_kernel
|
||||
|
||||
# Only enable flash attention backend
|
||||
with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
|
||||
scaled_dot_product_attention(...)
|
||||
|
||||
# Enable the Math or Efficient attention backends
|
||||
with sdpa_kernel([SDPBackend.MATH, SDPBackend.EFFICIENT_ATTENTION]):
|
||||
scaled_dot_product_attention(...)
|
||||
|
||||
# Enable the cuDNN or flash attention backends, and in that order
|
||||
with sdpa_kernel(
|
||||
[SDPBackend.CUDNN_ATTENTION, SDPBackend.FLASH_ATTENTION], set_priority=True
|
||||
):
|
||||
scaled_dot_product_attention(...)
|
||||
|
||||
This context manager can be used to select which backend to use for scaled dot product attention.
|
||||
Upon exiting the context manager, the previous state of the flags will be restored, enabling all backends.
|
||||
"""
|
||||
if not isinstance(backends, (list, SDPBackend)):
|
||||
raise AssertionError(
|
||||
f"Backend must be an instance of SDPBackend or a list of SDPBackend instances, got {type(backends).__name__}"
|
||||
)
|
||||
|
||||
if isinstance(backends, SDPBackend):
|
||||
backends = [backends]
|
||||
|
||||
backends = list(dict.fromkeys(backends))
|
||||
|
||||
previous_backends = _cur_sdpa_kernel_backends(with_priority=set_priority)
|
||||
try:
|
||||
_sdpa_kernel(backends, set_priority)
|
||||
yield {}
|
||||
finally:
|
||||
_sdpa_kernel(previous_backends, set_priority)
|
||||
|
||||
|
||||
# variadic version of sdpa_kernel for dynamo to use while reconstructing
|
||||
@contextlib.contextmanager
|
||||
def _sdpa_kernel_variadic(*backends: SDPBackend):
|
||||
with sdpa_kernel(list(backends)):
|
||||
yield
|
||||
|
||||
|
||||
def _get_flash_version() -> str:
|
||||
"""This returns the closest matching tag for the flash attention backend"""
|
||||
return "2.5.7"
|
||||
|
||||
|
||||
from . import _registry
|
||||
|
||||
|
||||
# Re-export registry types and functions for public API
|
||||
_FlashAttentionImpl = _registry._FlashAttentionImpl
|
||||
_RegisterFn = _registry._RegisterFn
|
||||
register_flash_attention_impl = _registry.register_flash_attention_impl
|
||||
activate_flash_attention_impl = _registry.activate_flash_attention_impl
|
||||
list_flash_attention_impls = _registry.list_flash_attention_impls
|
||||
current_flash_attention_impl = _registry.current_flash_attention_impl
|
||||
restore_flash_attention_impl = _registry.restore_flash_attention_impl
|
||||
|
||||
register_flash_attention_impl.__module__ = __name__
|
||||
activate_flash_attention_impl.__module__ = __name__
|
||||
list_flash_attention_impls.__module__ = __name__
|
||||
current_flash_attention_impl.__module__ = __name__
|
||||
restore_flash_attention_impl.__module__ = __name__
|
||||
|
||||
# Import built-in implementations to trigger self-registration
|
||||
from . import _fa3, _fa4 # noqa: F401 # noqa: F401
|
||||
@@ -0,0 +1,763 @@
|
||||
"""
|
||||
PROTOTYPE!
|
||||
Flash Attention 3 implementation.
|
||||
For fp8: only supports forward pass right now.
|
||||
For fp16/bf16: supports forward and backward pass.
|
||||
"""
|
||||
# mypy: allow-untyped-defs
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib
|
||||
import warnings
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable
|
||||
|
||||
from dataclasses import dataclass
|
||||
from functools import cache
|
||||
from typing_extensions import TypeVarTuple, Unpack
|
||||
|
||||
import torch
|
||||
from torch.library import Library
|
||||
|
||||
from . import _registry
|
||||
|
||||
|
||||
__all__ = [
|
||||
"register_flash_attention_fa3",
|
||||
]
|
||||
|
||||
|
||||
_FA3_CUDA_FWD: Callable | None = None # Cache for torch.ops.flash_attn_3.fwd
|
||||
_FA3_CUDA_BWD: Callable | None = None # Cache for torch.ops.flash_attn_3.bwd
|
||||
|
||||
|
||||
@dataclass
|
||||
class _FA3Handle:
|
||||
library: Library | None
|
||||
|
||||
def remove(self) -> None:
|
||||
self.library = None
|
||||
# Clear the C++ flag
|
||||
torch._C._set_sdp_use_fa3(False)
|
||||
|
||||
|
||||
@cache
|
||||
def _get_device_major(device: torch.device) -> int:
|
||||
major, _ = torch.cuda.get_device_capability(device)
|
||||
return major
|
||||
|
||||
|
||||
def register_flash_attention_fa3(
|
||||
module_path: str = "flash_attn_interface",
|
||||
) -> _FA3Handle:
|
||||
"""
|
||||
Register FA3 flash attention kernels with the PyTorch dispatcher.
|
||||
|
||||
Args:
|
||||
module_path: Python module path to the FA3 implementation.
|
||||
"""
|
||||
_fa3_import_module(module_path)
|
||||
|
||||
# Expose FA3 registration status to C++
|
||||
torch._C._set_sdp_use_fa3(True)
|
||||
|
||||
return _FA3Handle(_fa3_register_kernels())
|
||||
|
||||
|
||||
def _fa3_import_module(module_path: str) -> None:
|
||||
importlib.import_module(module_path)
|
||||
if not hasattr(torch.ops, "flash_attn_3"):
|
||||
raise RuntimeError(f"Module '{module_path}' does not expose FA3 kernels")
|
||||
if not hasattr(torch.ops.flash_attn_3, "fwd"):
|
||||
raise RuntimeError(
|
||||
f"Module '{module_path}' does not expose FA3 forward kernels"
|
||||
)
|
||||
if not hasattr(torch.ops.flash_attn_3, "bwd"):
|
||||
raise RuntimeError(
|
||||
f"Module '{module_path}' does not expose FA3 backward kernels"
|
||||
)
|
||||
global _FA3_CUDA_FWD, _FA3_CUDA_BWD
|
||||
_FA3_CUDA_FWD = torch.ops.flash_attn_3.fwd
|
||||
_FA3_CUDA_BWD = torch.ops.flash_attn_3.bwd
|
||||
|
||||
|
||||
def _fa3_register_kernels() -> Library:
|
||||
lib = Library("aten", "IMPL", "CUDA") # noqa: TOR901
|
||||
lib.impl(
|
||||
"_flash_attention_forward.quantized", _fa3_flash_attention_forward_impl, "CUDA"
|
||||
)
|
||||
lib.impl(
|
||||
"_scaled_dot_product_flash_attention.quantized",
|
||||
_fa3_scaled_dot_product_flash_attention_forward_impl,
|
||||
"CUDA",
|
||||
)
|
||||
lib.impl(
|
||||
"_flash_attention_forward", _fa3_flash_attention_forward_impl_default, "CUDA"
|
||||
)
|
||||
lib.impl(
|
||||
"_flash_attention_forward_no_dropout_inplace",
|
||||
_fa3_flash_attention_forward_no_dropout_inplace_impl,
|
||||
"CUDA",
|
||||
)
|
||||
lib.impl(
|
||||
"_scaled_dot_product_flash_attention",
|
||||
_fa3_scaled_dot_product_flash_attention_forward_impl_default,
|
||||
"CUDA",
|
||||
)
|
||||
|
||||
lib.impl("_flash_attention_backward", _fa3_flash_attention_backward_impl, "CUDA")
|
||||
lib.impl(
|
||||
"_scaled_dot_product_flash_attention_backward",
|
||||
_fa3_scaled_dot_product_flash_attention_backward_impl,
|
||||
"CUDA",
|
||||
)
|
||||
return lib
|
||||
|
||||
|
||||
def _fa3_common_support_error(
|
||||
query: torch.Tensor,
|
||||
tensors: tuple[torch.Tensor, ...],
|
||||
dropout_p: float,
|
||||
cum_seq_q: torch.Tensor | None,
|
||||
q_descale: torch.Tensor | None,
|
||||
k_descale: torch.Tensor | None,
|
||||
v_descale: torch.Tensor | None,
|
||||
) -> str | None:
|
||||
if dropout_p != 0.0:
|
||||
return "dropout_p must be 0"
|
||||
|
||||
if not all(t.is_cuda for t in tensors):
|
||||
return "inputs must be CUDA tensors"
|
||||
if len({t.device for t in tensors}) != 1:
|
||||
return "inputs must share device"
|
||||
if query.dtype == torch.float8_e4m3fn and (
|
||||
q_descale is None or k_descale is None or v_descale is None
|
||||
):
|
||||
warnings.warn(
|
||||
"When using SDPA with fp8, descale tensor should always be used"
|
||||
" for accurate dequantization. Please use "
|
||||
"_scaled_dot_product_attention_quantized and "
|
||||
"provide the descale tensors.",
|
||||
UserWarning,
|
||||
)
|
||||
if cum_seq_q is None and query.dim() != 4:
|
||||
return "dense query must be 4D"
|
||||
if cum_seq_q is not None and query.dim() != 3:
|
||||
return "ragged query must be 3D"
|
||||
if not torch.cuda.is_available():
|
||||
return "CUDA not available"
|
||||
if _get_device_major(query.device) != 9:
|
||||
return "FA3 requires compute capability 9.0"
|
||||
return None
|
||||
|
||||
|
||||
def _fa3_forward_support_error(
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
dropout_p: float,
|
||||
return_debug_mask: bool,
|
||||
alibi_slopes: torch.Tensor | None,
|
||||
seqused_k: torch.Tensor | None,
|
||||
cum_seq_q: torch.Tensor | None,
|
||||
q_descale: torch.Tensor | None,
|
||||
k_descale: torch.Tensor | None,
|
||||
v_descale: torch.Tensor | None,
|
||||
) -> str | None:
|
||||
if return_debug_mask:
|
||||
return "return_debug_mask must be False"
|
||||
if alibi_slopes is not None:
|
||||
return "alibi_slopes not supported"
|
||||
if seqused_k is not None:
|
||||
if seqused_k.dtype != torch.int32:
|
||||
return "seqused_k must be int32"
|
||||
if not seqused_k.is_cuda:
|
||||
return "seqused_k must be CUDA"
|
||||
supported_dtypes = (torch.float8_e4m3fn, torch.float16, torch.bfloat16)
|
||||
if not all(t.dtype in supported_dtypes for t in {query, key, value}):
|
||||
return f"inputs must be one of {supported_dtypes}"
|
||||
if len({t.dtype for t in {query, key, value}}) != 1:
|
||||
return "all inputs must have the same dtype"
|
||||
error = _fa3_common_support_error(
|
||||
query,
|
||||
(query, key, value),
|
||||
dropout_p,
|
||||
cum_seq_q,
|
||||
q_descale,
|
||||
k_descale,
|
||||
v_descale,
|
||||
)
|
||||
if error is not None:
|
||||
if error == "inputs must share device":
|
||||
return "query, key, value must be on same device"
|
||||
return error
|
||||
return None
|
||||
|
||||
|
||||
def _fa3_backward_support_error(
|
||||
grad_out: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
out: torch.Tensor,
|
||||
logsumexp: torch.Tensor,
|
||||
dropout_p: float,
|
||||
cum_seq_q: torch.Tensor | None,
|
||||
window_size_left: int | None,
|
||||
window_size_right: int | None,
|
||||
) -> str | None:
|
||||
# FA3 backward ONLY supports fp16/bf16, NOT fp8
|
||||
if query.dtype == torch.float8_e4m3fn:
|
||||
return (
|
||||
"FA3 backward does not support fp8 - use inference only (torch.no_grad())"
|
||||
)
|
||||
if logsumexp.dtype != torch.float32:
|
||||
return "logsumexp dtype must be float32"
|
||||
supported_dtypes = (torch.float16, torch.bfloat16)
|
||||
if not all(t.dtype in supported_dtypes for t in {grad_out, query, key, value, out}):
|
||||
return f"inputs must be one of {supported_dtypes}"
|
||||
if len({t.dtype for t in {grad_out, query, key, value, out}}) != 1:
|
||||
return "all inputs must have the same dtype"
|
||||
error = _fa3_common_support_error(
|
||||
query,
|
||||
(grad_out, query, key, value, out, logsumexp),
|
||||
dropout_p,
|
||||
cum_seq_q,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
if error is not None:
|
||||
return error
|
||||
return None
|
||||
|
||||
|
||||
Ts = TypeVarTuple("Ts")
|
||||
|
||||
|
||||
def _transpose_dense(*tensors: Unpack[Ts]) -> tuple[Unpack[Ts]]:
|
||||
return tuple(t.transpose(1, 2) for t in tensors) # type: ignore[attr-defined]
|
||||
|
||||
|
||||
def _maybe_contiguous(x: torch.Tensor | None) -> torch.Tensor | None:
|
||||
"""Ensure tensor is contiguous in the last dimension."""
|
||||
return x.contiguous() if x is not None and x.stride(-1) != 1 else x
|
||||
|
||||
|
||||
def _fa3_run_forward(
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
cu_seq_q: torch.Tensor | None,
|
||||
cu_seq_k: torch.Tensor | None,
|
||||
max_q: int,
|
||||
max_k: int,
|
||||
scale: float | None,
|
||||
is_causal: bool,
|
||||
window_size_left: int | None,
|
||||
window_size_right: int | None,
|
||||
seqused_k: torch.Tensor | None,
|
||||
out: torch.Tensor | None = None,
|
||||
q_descale: torch.Tensor | None = None,
|
||||
k_descale: torch.Tensor | None = None,
|
||||
v_descale: torch.Tensor | None = None,
|
||||
block_table: torch.Tensor | None = None,
|
||||
num_splits: int | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Run the FA3 forward pass by calling the C++ kernel directly.
|
||||
"""
|
||||
if _FA3_CUDA_FWD is None:
|
||||
raise RuntimeError("FA3 not registered")
|
||||
# Ensure contiguous in the last dimension
|
||||
q = _maybe_contiguous(query)
|
||||
k = _maybe_contiguous(key)
|
||||
v = (
|
||||
value.contiguous()
|
||||
if value.dtype == torch.float8_e4m3fn
|
||||
and value.stride(-1) != 1
|
||||
and value.stride(-3) != 1
|
||||
else _maybe_contiguous(value)
|
||||
)
|
||||
|
||||
cu_seqlens_q = _maybe_contiguous(cu_seq_q)
|
||||
cu_seqlens_k = _maybe_contiguous(cu_seq_k)
|
||||
seqused_k = _maybe_contiguous(seqused_k)
|
||||
block_table = _maybe_contiguous(block_table)
|
||||
|
||||
out, softmax_lse, out_accum, softmax_lse_accum = _FA3_CUDA_FWD(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
None, # k_new
|
||||
None, # v_new
|
||||
None, # qv
|
||||
out, # out_ (pre-allocated output)
|
||||
cu_seqlens_q, # cu_seqlens_q
|
||||
cu_seqlens_k, # cu_seqlens_k
|
||||
None, # cu_seqlens_k_new
|
||||
None, # seqused_q
|
||||
seqused_k, # seqused_k
|
||||
max_q, # max_seqlen_q
|
||||
max_k, # max_seqlen_k
|
||||
block_table, # block_table,
|
||||
None, # kv_batch_idx,
|
||||
None, # leftpad_k,
|
||||
None, # rotary_cos,
|
||||
None, # rotary_sin,
|
||||
None, # seqlens_rotary,
|
||||
q_descale, # q_descale,
|
||||
k_descale, # k_descale,
|
||||
v_descale, # v_descale,
|
||||
scale, # softmax_scale,
|
||||
is_causal, # causal,
|
||||
window_size_left if window_size_left is not None else -1, # window_size_left
|
||||
window_size_right if window_size_right is not None else -1, # window_size_right
|
||||
0, # attention_chunk,
|
||||
0.0, # softcap,
|
||||
True, # rotary_interleaved,
|
||||
None, # scheduler_metadata,
|
||||
num_splits
|
||||
or (1 if torch.are_deterministic_algorithms_enabled() else 0), # num_splits,
|
||||
None, # pack_gqa,
|
||||
torch._C._get_sm_carveout_experimental() or 0, # sm_margin,
|
||||
)
|
||||
return out, softmax_lse.contiguous()
|
||||
|
||||
|
||||
def _fa3_run_backward(
|
||||
grad_out: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
out: torch.Tensor,
|
||||
logsumexp: torch.Tensor,
|
||||
cu_seq_q: torch.Tensor | None,
|
||||
cu_seq_k: torch.Tensor | None,
|
||||
max_seqlen_q: int | None,
|
||||
max_seqlen_k: int | None,
|
||||
scale: float | None,
|
||||
is_causal: bool,
|
||||
window_size_left: int,
|
||||
window_size_right: int,
|
||||
deterministic: bool = False,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
if _FA3_CUDA_BWD is None:
|
||||
raise RuntimeError("FA3 not registered")
|
||||
|
||||
# Ensure contiguous
|
||||
dout = _maybe_contiguous(grad_out)
|
||||
q = query.contiguous() if query.stride(-1) != 1 else query
|
||||
k = key.contiguous() if key.stride(-1) != 1 else key
|
||||
v = value.contiguous() if value.stride(-1) != 1 else value
|
||||
o = _maybe_contiguous(out)
|
||||
lse = _maybe_contiguous(logsumexp)
|
||||
|
||||
# Pre-allocate gradient tensors
|
||||
dq = torch.empty_like(q)
|
||||
dk = torch.empty_like(k)
|
||||
dv = torch.empty_like(v)
|
||||
_FA3_CUDA_BWD(
|
||||
dout,
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
o,
|
||||
lse,
|
||||
dq,
|
||||
dk,
|
||||
dv,
|
||||
cu_seq_q,
|
||||
cu_seq_k,
|
||||
None,
|
||||
None,
|
||||
max_seqlen_q,
|
||||
max_seqlen_k,
|
||||
scale,
|
||||
is_causal,
|
||||
window_size_left,
|
||||
window_size_right,
|
||||
0.0,
|
||||
deterministic,
|
||||
torch._C._get_sm_carveout_experimental() or 0,
|
||||
)
|
||||
return dq, dk, dv
|
||||
|
||||
|
||||
def _fa3_flash_attention_forward_impl(
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
cum_seq_q: torch.Tensor | None,
|
||||
cum_seq_k: torch.Tensor | None,
|
||||
max_q: int,
|
||||
max_k: int,
|
||||
dropout_p: float,
|
||||
is_causal: bool,
|
||||
return_debug_mask: bool,
|
||||
q_descale: torch.Tensor | None = None,
|
||||
k_descale: torch.Tensor | None = None,
|
||||
v_descale: torch.Tensor | None = None,
|
||||
*,
|
||||
scale: float | None = None,
|
||||
window_size_left: int = -1,
|
||||
window_size_right: int = -1,
|
||||
seqused_k: torch.Tensor | None = None,
|
||||
alibi_slopes: torch.Tensor | None = None,
|
||||
out: torch.Tensor | None = None,
|
||||
block_table: torch.Tensor | None = None,
|
||||
compute_auxiliary: bool = True,
|
||||
num_splits: int | None = None,
|
||||
):
|
||||
error = _fa3_forward_support_error(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
dropout_p,
|
||||
return_debug_mask,
|
||||
alibi_slopes,
|
||||
seqused_k,
|
||||
cum_seq_q,
|
||||
q_descale,
|
||||
k_descale,
|
||||
v_descale,
|
||||
)
|
||||
if error is not None:
|
||||
raise RuntimeError(f"FA3 flash_attention forward unsupported: {error}")
|
||||
out, lse = _fa3_run_forward(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
cum_seq_q,
|
||||
cum_seq_k,
|
||||
max_q,
|
||||
max_k,
|
||||
scale,
|
||||
is_causal,
|
||||
window_size_left,
|
||||
window_size_right,
|
||||
seqused_k,
|
||||
out,
|
||||
q_descale,
|
||||
k_descale,
|
||||
v_descale,
|
||||
block_table,
|
||||
num_splits,
|
||||
)
|
||||
if compute_auxiliary:
|
||||
rng_state = torch.zeros((2,), dtype=torch.uint64, device=query.device)
|
||||
philox_offset = torch.zeros((), dtype=torch.uint64, device=query.device)
|
||||
debug_mask = torch.empty(0, dtype=query.dtype, device=query.device)
|
||||
else:
|
||||
rng_state = None
|
||||
philox_offset = None
|
||||
debug_mask = None
|
||||
return out, lse, rng_state, philox_offset, debug_mask
|
||||
|
||||
|
||||
def _fa3_flash_attention_forward_no_dropout_inplace_impl(
|
||||
out: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
cum_seq_q: torch.Tensor | None,
|
||||
cum_seq_k: torch.Tensor | None,
|
||||
max_q: int,
|
||||
max_k: int,
|
||||
dropout_p: float,
|
||||
is_causal: bool,
|
||||
return_debug_mask: bool,
|
||||
*,
|
||||
scale: float | None = None,
|
||||
window_size_left: int = -1,
|
||||
window_size_right: int = -1,
|
||||
seqused_k: torch.Tensor | None = None,
|
||||
alibi_slopes: torch.Tensor | None = None,
|
||||
block_table: torch.Tensor | None = None,
|
||||
num_splits: int | None = None,
|
||||
):
|
||||
_, lse, _, _, _ = _fa3_flash_attention_forward_impl(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
cum_seq_q,
|
||||
cum_seq_k,
|
||||
max_q,
|
||||
max_k,
|
||||
dropout_p,
|
||||
is_causal,
|
||||
return_debug_mask,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
scale=scale,
|
||||
window_size_left=window_size_left,
|
||||
window_size_right=window_size_right,
|
||||
seqused_k=seqused_k,
|
||||
alibi_slopes=alibi_slopes,
|
||||
out=out,
|
||||
block_table=block_table,
|
||||
compute_auxiliary=False,
|
||||
num_splits=num_splits,
|
||||
)
|
||||
return lse
|
||||
|
||||
|
||||
def _fa3_flash_attention_forward_impl_default(
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
cum_seq_q: torch.Tensor | None,
|
||||
cum_seq_k: torch.Tensor | None,
|
||||
max_q: int,
|
||||
max_k: int,
|
||||
dropout_p: float,
|
||||
is_causal: bool,
|
||||
return_debug_mask: bool,
|
||||
*,
|
||||
scale: float | None = None,
|
||||
window_size_left: int = -1,
|
||||
window_size_right: int = -1,
|
||||
seqused_k: torch.Tensor | None = None,
|
||||
alibi_slopes: torch.Tensor | None = None,
|
||||
block_table: torch.Tensor | None = None,
|
||||
out: torch.Tensor | None = None,
|
||||
num_splits: int | None = None,
|
||||
):
|
||||
return _fa3_flash_attention_forward_impl(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
cum_seq_q,
|
||||
cum_seq_k,
|
||||
max_q,
|
||||
max_k,
|
||||
dropout_p,
|
||||
is_causal,
|
||||
return_debug_mask,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
scale=scale,
|
||||
window_size_left=window_size_left,
|
||||
window_size_right=window_size_right,
|
||||
seqused_k=seqused_k,
|
||||
alibi_slopes=alibi_slopes,
|
||||
out=out,
|
||||
block_table=block_table,
|
||||
num_splits=num_splits,
|
||||
)
|
||||
|
||||
|
||||
def _fa3_flash_attention_backward_impl(
|
||||
grad_out: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
out: torch.Tensor,
|
||||
logsumexp: torch.Tensor,
|
||||
cum_seq_q: torch.Tensor | None,
|
||||
cum_seq_k: torch.Tensor | None,
|
||||
max_q: int,
|
||||
max_k: int,
|
||||
dropout_p: float,
|
||||
is_causal: bool,
|
||||
rng_state: torch.Tensor,
|
||||
unused: torch.Tensor,
|
||||
*,
|
||||
scale: float | None = None,
|
||||
window_size_left: int | None = None,
|
||||
window_size_right: int | None = None,
|
||||
):
|
||||
"""FA3 implementation of _flash_attention_backward."""
|
||||
error = _fa3_backward_support_error(
|
||||
grad_out,
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
out,
|
||||
logsumexp,
|
||||
dropout_p,
|
||||
cum_seq_q,
|
||||
window_size_left,
|
||||
window_size_right,
|
||||
)
|
||||
|
||||
if error is not None:
|
||||
raise RuntimeError(f"FA3 flash_attention backward unsupported: {error}")
|
||||
|
||||
deterministic = torch.are_deterministic_algorithms_enabled()
|
||||
|
||||
dq, dk, dv = _fa3_run_backward(
|
||||
grad_out,
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
out,
|
||||
logsumexp,
|
||||
cum_seq_q,
|
||||
cum_seq_k,
|
||||
max_q,
|
||||
max_k,
|
||||
scale,
|
||||
is_causal,
|
||||
window_size_left if window_size_left is not None else -1,
|
||||
window_size_right if window_size_right is not None else -1,
|
||||
deterministic,
|
||||
)
|
||||
return dq, dk, dv
|
||||
|
||||
|
||||
def _fa3_scaled_dot_product_flash_attention_forward_impl(
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
q_descale: torch.Tensor | None = None,
|
||||
k_descale: torch.Tensor | None = None,
|
||||
v_descale: torch.Tensor | None = None,
|
||||
dropout_p: float = 0.0,
|
||||
is_causal: bool = False,
|
||||
return_debug_mask: bool = False,
|
||||
*,
|
||||
scale: float | None = None,
|
||||
):
|
||||
error = _fa3_forward_support_error(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
dropout_p,
|
||||
return_debug_mask,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
q_descale,
|
||||
k_descale,
|
||||
v_descale,
|
||||
)
|
||||
if error is not None:
|
||||
raise RuntimeError(f"FA3 SDPA forward unsupported: {error}")
|
||||
q, k, v = _transpose_dense(query, key, value)
|
||||
|
||||
# Pre-allocate output with query's strides (BHSD layout), then create
|
||||
# a BSHD view for the kernel. This ensures the returned output has
|
||||
# the same memory layout as the input query.
|
||||
out_dtype = torch.bfloat16 if query.dtype == torch.float8_e4m3fn else query.dtype
|
||||
out_bhsd = torch.empty_like(query, dtype=out_dtype)
|
||||
out_bshd = out_bhsd.transpose(1, 2)
|
||||
|
||||
max_q_flash = q.size(1)
|
||||
max_k_flash = k.size(1)
|
||||
_, lse, rng_state, philox_offset, debug_mask = _fa3_flash_attention_forward_impl(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
None,
|
||||
None,
|
||||
max_q_flash,
|
||||
max_k_flash,
|
||||
dropout_p,
|
||||
is_causal,
|
||||
return_debug_mask,
|
||||
scale=scale,
|
||||
out=out_bshd,
|
||||
q_descale=q_descale,
|
||||
k_descale=k_descale,
|
||||
v_descale=v_descale,
|
||||
)
|
||||
max_q = query.size(2)
|
||||
max_k = key.size(2)
|
||||
return (
|
||||
out_bhsd,
|
||||
lse,
|
||||
None,
|
||||
None,
|
||||
max_q,
|
||||
max_k,
|
||||
rng_state,
|
||||
philox_offset,
|
||||
debug_mask,
|
||||
)
|
||||
|
||||
|
||||
def _fa3_scaled_dot_product_flash_attention_forward_impl_default(
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
dropout_p: float = 0.0,
|
||||
is_causal: bool = False,
|
||||
return_debug_mask: bool = False,
|
||||
*,
|
||||
scale: float | None = None,
|
||||
):
|
||||
return _fa3_scaled_dot_product_flash_attention_forward_impl(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
dropout_p,
|
||||
is_causal,
|
||||
return_debug_mask,
|
||||
scale=scale,
|
||||
)
|
||||
|
||||
|
||||
def _fa3_scaled_dot_product_flash_attention_backward_impl(
|
||||
grad_out: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
out: torch.Tensor,
|
||||
logsumexp: torch.Tensor,
|
||||
cum_seq_q: torch.Tensor | None,
|
||||
cum_seq_k: torch.Tensor | None,
|
||||
max_q: int,
|
||||
max_k: int,
|
||||
dropout_p: float,
|
||||
is_causal: bool,
|
||||
philox_seed: torch.Tensor,
|
||||
philox_offset: torch.Tensor,
|
||||
*,
|
||||
scale: float | None = None,
|
||||
):
|
||||
"""FA3 implementation of _scaled_dot_product_flash_attention_backward."""
|
||||
error = _fa3_backward_support_error(
|
||||
grad_out, query, key, value, out, logsumexp, dropout_p, None, None, None
|
||||
)
|
||||
if error is not None:
|
||||
raise RuntimeError(f"FA3 SDPA backward unsupported: {error}")
|
||||
|
||||
# SDPA uses BHSD layout, FA3 uses BSHD - transpose
|
||||
grad_out_t, q_t, k_t, v_t, out_t = _transpose_dense(
|
||||
grad_out, query, key, value, out
|
||||
)
|
||||
|
||||
dq, dk, dv = _fa3_flash_attention_backward_impl(
|
||||
grad_out_t,
|
||||
q_t,
|
||||
k_t,
|
||||
v_t,
|
||||
out_t,
|
||||
logsumexp,
|
||||
None, # cum_seq_q (dense attention)
|
||||
None, # cum_seq_k
|
||||
max_q, # max_seqlen_q
|
||||
max_k, # max_seqlen_k
|
||||
dropout_p,
|
||||
is_causal,
|
||||
philox_seed,
|
||||
philox_offset,
|
||||
scale=scale,
|
||||
)
|
||||
|
||||
# Transpose gradients back to BHSD layout
|
||||
dq_out, dk_out, dv_out = _transpose_dense(dq, dk, dv)
|
||||
return dq_out, dk_out, dv_out
|
||||
|
||||
|
||||
_registry.register_flash_attention_impl("FA3", register_fn=register_flash_attention_fa3)
|
||||
@@ -0,0 +1,541 @@
|
||||
"""UBER PROTOTYPE!!!"""
|
||||
# mypy: allow-untyped-defs
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib
|
||||
from dataclasses import dataclass
|
||||
from functools import cache
|
||||
from typing import Any, TYPE_CHECKING
|
||||
from typing_extensions import TypeVarTuple, Unpack
|
||||
|
||||
from . import _registry
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from types import ModuleType
|
||||
|
||||
import torch
|
||||
from torch.library import Library
|
||||
|
||||
|
||||
__all__ = [
|
||||
"register_flash_attention_fa4",
|
||||
]
|
||||
|
||||
|
||||
_FA4_MODULE_PATH: str | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class _FA4Handle:
|
||||
library: Library | None
|
||||
|
||||
def remove(self) -> None:
|
||||
self.library = None
|
||||
|
||||
|
||||
@cache
|
||||
def _get_device_major(device: torch.device) -> int:
|
||||
major, _ = torch.cuda.get_device_capability(device)
|
||||
return major
|
||||
|
||||
|
||||
def register_flash_attention_fa4(
|
||||
module_path: str = "flash_attn.cute.interface",
|
||||
) -> _FA4Handle:
|
||||
"""
|
||||
Register FA4 flash attention kernels with the PyTorch dispatcher.
|
||||
|
||||
Args:
|
||||
module_path: Python module path to the FA4 implementation.
|
||||
"""
|
||||
global _FA4_MODULE_PATH
|
||||
_ = _fa4_import_module(module_path)
|
||||
_FA4_MODULE_PATH = module_path
|
||||
return _FA4Handle(_fa4_register_kernels())
|
||||
|
||||
|
||||
@cache
|
||||
def _fa4_import_module(module_path: str) -> ModuleType:
|
||||
module = importlib.import_module(module_path)
|
||||
if not hasattr(module, "_flash_attn_fwd") or not hasattr(module, "_flash_attn_bwd"):
|
||||
raise RuntimeError(f"Module '{module_path}' does not expose FA4 kernels")
|
||||
return module
|
||||
|
||||
|
||||
def _fa4_register_kernels() -> Library:
|
||||
lib = Library("aten", "IMPL", "CUDA") # noqa: TOR901
|
||||
lib.impl("_flash_attention_forward", _fa4_flash_attention_forward_impl, "CUDA")
|
||||
lib.impl(
|
||||
"_flash_attention_forward_no_dropout_inplace",
|
||||
_fa4_flash_attention_forward_no_dropout_inplace_impl,
|
||||
"CUDA",
|
||||
)
|
||||
lib.impl("_flash_attention_backward", _fa4_flash_attention_backward_impl, "CUDA")
|
||||
lib.impl(
|
||||
"_scaled_dot_product_flash_attention",
|
||||
_fa4_scaled_dot_product_flash_attention_forward_impl,
|
||||
"CUDA",
|
||||
)
|
||||
lib.impl(
|
||||
"_scaled_dot_product_flash_attention_backward",
|
||||
_fa4_scaled_dot_product_flash_attention_backward_impl,
|
||||
"CUDA",
|
||||
)
|
||||
return lib
|
||||
|
||||
|
||||
def _fa4_common_support_error(
|
||||
query: torch.Tensor,
|
||||
tensors: tuple[torch.Tensor, ...],
|
||||
cum_seq_q: torch.Tensor | None,
|
||||
require_fp32: tuple[tuple[str, torch.Tensor], ...] = (),
|
||||
) -> str | None:
|
||||
if not all(t.is_cuda for t in tensors):
|
||||
return "inputs must be CUDA tensors"
|
||||
if len({t.device for t in tensors}) != 1:
|
||||
return "inputs must share device"
|
||||
if query.dtype not in (torch.float16, torch.bfloat16):
|
||||
return "query dtype must be float16 or bfloat16"
|
||||
for name, tensor in require_fp32:
|
||||
if tensor.dtype != torch.float32:
|
||||
return f"{name} dtype must be float32"
|
||||
if cum_seq_q is None and query.dim() != 4:
|
||||
return "dense query must be 4D"
|
||||
if cum_seq_q is not None and query.dim() != 3:
|
||||
return "ragged query must be 3D"
|
||||
if not torch.cuda.is_available():
|
||||
return "CUDA not available"
|
||||
if _get_device_major(query.device) not in (9, 10):
|
||||
return "FA4 requires compute capability 9.0 or 10.0"
|
||||
return None
|
||||
|
||||
|
||||
def _fa4_forward_support_error(
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
dropout_p: float,
|
||||
return_debug_mask: bool,
|
||||
alibi_slopes: torch.Tensor | None,
|
||||
seqused_k: torch.Tensor | None,
|
||||
cum_seq_q: torch.Tensor | None,
|
||||
block_table: torch.Tensor | None = None,
|
||||
num_splits: int | None = None,
|
||||
) -> str | None:
|
||||
if dropout_p != 0.0:
|
||||
return "dropout_p must be 0"
|
||||
if return_debug_mask:
|
||||
return "return_debug_mask must be False"
|
||||
if alibi_slopes is not None:
|
||||
return "alibi_slopes not supported"
|
||||
if seqused_k is not None:
|
||||
if seqused_k.dtype != torch.int32:
|
||||
return "seqused_k must be int32"
|
||||
if not seqused_k.is_cuda:
|
||||
return "seqused_k must be CUDA"
|
||||
major = _get_device_major(query.device)
|
||||
if block_table is not None and major != 10:
|
||||
return f"paged KV (block_table) not supported on SM {major}0"
|
||||
if num_splits is not None and num_splits > 1 and major != 10:
|
||||
return f"SplitKV (num_splits > 1) not supported on SM {major}0"
|
||||
error = _fa4_common_support_error(
|
||||
query,
|
||||
(query, key, value),
|
||||
cum_seq_q,
|
||||
)
|
||||
if error is not None:
|
||||
if error == "inputs must share device":
|
||||
return "query, key, value must be on same device"
|
||||
return error
|
||||
return None
|
||||
|
||||
|
||||
def _fa4_backward_support_error(
|
||||
grad_out: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
out: torch.Tensor,
|
||||
logsumexp: torch.Tensor,
|
||||
dropout_p: float,
|
||||
cum_seq_q: torch.Tensor | None,
|
||||
) -> str | None:
|
||||
if dropout_p != 0.0:
|
||||
return "dropout_p must be 0"
|
||||
error = _fa4_common_support_error(
|
||||
query,
|
||||
(grad_out, query, key, value, out, logsumexp),
|
||||
cum_seq_q,
|
||||
require_fp32=(("logsumexp", logsumexp),),
|
||||
)
|
||||
if error is not None:
|
||||
return error
|
||||
return None
|
||||
|
||||
|
||||
def _aten_to_fa4_window_size(val: int | None) -> int | None:
|
||||
"""need to convert -1 to None for FA4"""
|
||||
return None if val == -1 else val
|
||||
|
||||
|
||||
Ts = TypeVarTuple("Ts")
|
||||
|
||||
|
||||
def _transpose_dense(*tensors: Unpack[Ts]) -> tuple[Unpack[Ts]]:
|
||||
return tuple(t.transpose(1, 2) for t in tensors) # type: ignore[attr-defined]
|
||||
|
||||
|
||||
def _fa4_run_forward(
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
cu_seq_q: torch.Tensor | None,
|
||||
cu_seq_k: torch.Tensor | None,
|
||||
max_q: int | None,
|
||||
max_k: int | None,
|
||||
scale: float | None,
|
||||
is_causal: bool,
|
||||
window_size_left: int | None,
|
||||
window_size_right: int | None,
|
||||
seqused_k: torch.Tensor | None,
|
||||
out: torch.Tensor | None = None,
|
||||
block_table: torch.Tensor | None = None,
|
||||
num_splits: int | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
if _FA4_MODULE_PATH is None:
|
||||
raise RuntimeError("FA4 not registered")
|
||||
module = _fa4_import_module(_FA4_MODULE_PATH)
|
||||
|
||||
kwargs: dict[str, Any] = {
|
||||
"softmax_scale": scale,
|
||||
"causal": is_causal,
|
||||
"window_size_left": _aten_to_fa4_window_size(window_size_left),
|
||||
"window_size_right": _aten_to_fa4_window_size(window_size_right),
|
||||
"return_lse": True,
|
||||
"cu_seqlens_q": cu_seq_q,
|
||||
"cu_seqlens_k": cu_seq_k,
|
||||
"max_seqlen_q": max_q,
|
||||
"max_seqlen_k": max_k,
|
||||
"seqused_k": seqused_k.contiguous() if seqused_k is not None else None,
|
||||
"page_table": block_table,
|
||||
"num_splits": num_splits or 1,
|
||||
"out": out,
|
||||
}
|
||||
out, lse = module._flash_attn_fwd(query, key, value, **kwargs)
|
||||
return out, lse.contiguous()
|
||||
|
||||
|
||||
def _fa4_run_backward(
|
||||
grad_out: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
out: torch.Tensor,
|
||||
logsumexp: torch.Tensor,
|
||||
cu_seq_q: torch.Tensor | None,
|
||||
cu_seq_k: torch.Tensor | None,
|
||||
scale: float | None,
|
||||
is_causal: bool,
|
||||
window_size_left: int | None,
|
||||
window_size_right: int | None,
|
||||
deterministic: bool = False,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
if _FA4_MODULE_PATH is None:
|
||||
raise RuntimeError("FA4 not registered")
|
||||
module = _fa4_import_module(_FA4_MODULE_PATH)
|
||||
dq, dk, dv = module._flash_attn_bwd(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
out,
|
||||
grad_out,
|
||||
logsumexp.contiguous(),
|
||||
softmax_scale=scale,
|
||||
causal=is_causal,
|
||||
window_size_left=_aten_to_fa4_window_size(window_size_left),
|
||||
window_size_right=_aten_to_fa4_window_size(window_size_right),
|
||||
cu_seqlens_q=cu_seq_q,
|
||||
cu_seqlens_k=cu_seq_k,
|
||||
deterministic=deterministic,
|
||||
)
|
||||
return dq, dk, dv
|
||||
|
||||
|
||||
def _fa4_flash_attention_forward_impl(
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
cum_seq_q: torch.Tensor | None,
|
||||
cum_seq_k: torch.Tensor | None,
|
||||
max_q: int,
|
||||
max_k: int,
|
||||
dropout_p: float,
|
||||
is_causal: bool,
|
||||
return_debug_mask: bool,
|
||||
*,
|
||||
scale: float | None = None,
|
||||
window_size_left: int | None = None,
|
||||
window_size_right: int | None = None,
|
||||
seqused_k: torch.Tensor | None = None,
|
||||
alibi_slopes: torch.Tensor | None = None,
|
||||
out: torch.Tensor | None = None,
|
||||
block_table: torch.Tensor | None = None,
|
||||
compute_auxiliary: bool = True,
|
||||
num_splits: int | None = None,
|
||||
):
|
||||
error = _fa4_forward_support_error(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
dropout_p,
|
||||
return_debug_mask,
|
||||
alibi_slopes,
|
||||
seqused_k,
|
||||
cum_seq_q,
|
||||
block_table,
|
||||
num_splits,
|
||||
)
|
||||
if error is not None:
|
||||
raise RuntimeError(f"FA4 flash_attention forward unsupported: {error}")
|
||||
out, lse = _fa4_run_forward(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
cum_seq_q,
|
||||
cum_seq_k,
|
||||
max_q,
|
||||
max_k,
|
||||
scale,
|
||||
is_causal,
|
||||
window_size_left,
|
||||
window_size_right,
|
||||
seqused_k,
|
||||
out,
|
||||
block_table,
|
||||
num_splits,
|
||||
)
|
||||
if compute_auxiliary:
|
||||
rng_state = torch.zeros((2,), dtype=torch.uint64, device=query.device)
|
||||
philox_offset = torch.zeros((), dtype=torch.uint64, device=query.device)
|
||||
debug_mask = torch.empty(0, dtype=query.dtype, device=query.device)
|
||||
else:
|
||||
rng_state = None
|
||||
philox_offset = None
|
||||
debug_mask = None
|
||||
return out, lse, rng_state, philox_offset, debug_mask
|
||||
|
||||
|
||||
def _fa4_flash_attention_forward_no_dropout_inplace_impl(
|
||||
out: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
cum_seq_q: torch.Tensor | None,
|
||||
cum_seq_k: torch.Tensor | None,
|
||||
max_q: int,
|
||||
max_k: int,
|
||||
dropout_p: float,
|
||||
is_causal: bool,
|
||||
return_debug_mask: bool,
|
||||
*,
|
||||
scale: float | None = None,
|
||||
window_size_left: int | None = None,
|
||||
window_size_right: int | None = None,
|
||||
seqused_k: torch.Tensor | None = None,
|
||||
alibi_slopes: torch.Tensor | None = None,
|
||||
block_table: torch.Tensor | None = None,
|
||||
num_splits: int | None = None,
|
||||
):
|
||||
_, lse, _, _, _ = _fa4_flash_attention_forward_impl(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
cum_seq_q,
|
||||
cum_seq_k,
|
||||
max_q,
|
||||
max_k,
|
||||
dropout_p,
|
||||
is_causal,
|
||||
return_debug_mask,
|
||||
scale=scale,
|
||||
window_size_left=window_size_left,
|
||||
window_size_right=window_size_right,
|
||||
seqused_k=seqused_k,
|
||||
alibi_slopes=alibi_slopes,
|
||||
out=out,
|
||||
block_table=block_table,
|
||||
compute_auxiliary=False,
|
||||
num_splits=num_splits,
|
||||
)
|
||||
return lse
|
||||
|
||||
|
||||
def _fa4_flash_attention_backward_impl(
|
||||
grad_out: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
out: torch.Tensor,
|
||||
logsumexp: torch.Tensor,
|
||||
cum_seq_q: torch.Tensor | None,
|
||||
cum_seq_k: torch.Tensor | None,
|
||||
max_q: int,
|
||||
max_k: int,
|
||||
dropout_p: float,
|
||||
is_causal: bool,
|
||||
rng_state: torch.Tensor,
|
||||
unused: torch.Tensor,
|
||||
*,
|
||||
scale: float | None = None,
|
||||
window_size_left: int | None = None,
|
||||
window_size_right: int | None = None,
|
||||
):
|
||||
error = _fa4_backward_support_error(
|
||||
grad_out,
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
out,
|
||||
logsumexp,
|
||||
dropout_p,
|
||||
cum_seq_q,
|
||||
)
|
||||
if error is not None:
|
||||
raise RuntimeError(f"FA4 flash_attention backward unsupported: {error}")
|
||||
deterministic = torch.are_deterministic_algorithms_enabled()
|
||||
dq, dk, dv = _fa4_run_backward(
|
||||
grad_out,
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
out,
|
||||
logsumexp,
|
||||
cum_seq_q,
|
||||
cum_seq_k,
|
||||
scale,
|
||||
is_causal,
|
||||
window_size_left,
|
||||
window_size_right,
|
||||
deterministic,
|
||||
)
|
||||
return dq, dk, dv
|
||||
|
||||
|
||||
def _fa4_scaled_dot_product_flash_attention_forward_impl(
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
dropout_p: float = 0.0,
|
||||
is_causal: bool = False,
|
||||
return_debug_mask: bool = False,
|
||||
*,
|
||||
scale: float | None = None,
|
||||
):
|
||||
error = _fa4_forward_support_error(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
dropout_p,
|
||||
return_debug_mask,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
if error is not None:
|
||||
raise RuntimeError(f"FA4 SDPA forward unsupported: {error}")
|
||||
q, k, v = _transpose_dense(query, key, value)
|
||||
|
||||
# Pre-allocate output with query's strides (BHSD layout), then create
|
||||
# a BSHD view for the kernel. This ensures the returned output has
|
||||
# the same memory layout as the input query.
|
||||
out_bhsd = torch.empty_like(query)
|
||||
out_bshd = out_bhsd.transpose(1, 2)
|
||||
|
||||
max_q_flash = q.size(1)
|
||||
max_k_flash = k.size(1)
|
||||
_, lse, rng_state, philox_offset, debug_mask = _fa4_flash_attention_forward_impl(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
None,
|
||||
None,
|
||||
max_q_flash,
|
||||
max_k_flash,
|
||||
dropout_p,
|
||||
is_causal,
|
||||
return_debug_mask,
|
||||
scale=scale,
|
||||
out=out_bshd,
|
||||
)
|
||||
max_q = query.size(2)
|
||||
max_k = key.size(2)
|
||||
return (
|
||||
out_bhsd,
|
||||
lse,
|
||||
None,
|
||||
None,
|
||||
max_q,
|
||||
max_k,
|
||||
rng_state,
|
||||
philox_offset,
|
||||
debug_mask,
|
||||
)
|
||||
|
||||
|
||||
def _fa4_scaled_dot_product_flash_attention_backward_impl(
|
||||
grad_out: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
out: torch.Tensor,
|
||||
logsumexp: torch.Tensor,
|
||||
cum_seq_q: torch.Tensor | None,
|
||||
cum_seq_k: torch.Tensor | None,
|
||||
max_q: int,
|
||||
max_k: int,
|
||||
dropout_p: float,
|
||||
is_causal: bool,
|
||||
philox_seed: torch.Tensor,
|
||||
philox_offset: torch.Tensor,
|
||||
*,
|
||||
scale: float | None = None,
|
||||
):
|
||||
error = _fa4_backward_support_error(
|
||||
grad_out,
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
out,
|
||||
logsumexp,
|
||||
dropout_p,
|
||||
None,
|
||||
)
|
||||
if error is not None:
|
||||
raise RuntimeError(f"FA4 SDPA backward unsupported: {error}")
|
||||
q, k, v, o, go = _transpose_dense(query, key, value, out, grad_out)
|
||||
max_q = query.size(2)
|
||||
max_k = key.size(2)
|
||||
dq, dk, dv = _fa4_flash_attention_backward_impl(
|
||||
go,
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
o,
|
||||
logsumexp,
|
||||
None,
|
||||
None,
|
||||
max_q,
|
||||
max_k,
|
||||
dropout_p,
|
||||
is_causal,
|
||||
philox_seed,
|
||||
philox_offset,
|
||||
scale=scale,
|
||||
)
|
||||
dq, dk, dv = _transpose_dense(dq, dk, dv)
|
||||
return dq, dk, dv
|
||||
|
||||
|
||||
_registry.register_flash_attention_impl("FA4", register_fn=register_flash_attention_fa4)
|
||||
@@ -0,0 +1,137 @@
|
||||
# mypy: allow-untyped-defs
|
||||
"""Registry for flash attention implementations.
|
||||
|
||||
This module contains the registration system for flash attention implementations.
|
||||
It has no torch dependencies to avoid circular imports during initialization.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from collections.abc import Callable
|
||||
from typing import Literal, Protocol
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class FlashAttentionHandle(Protocol):
|
||||
def remove(self) -> None: ...
|
||||
|
||||
|
||||
_RegisterFn = Callable[..., FlashAttentionHandle | None]
|
||||
_FlashAttentionImpl = Literal["FA3", "FA4"]
|
||||
|
||||
_FLASH_ATTENTION_IMPLS: dict[str, _RegisterFn] = {}
|
||||
|
||||
_FLASH_ATTENTION_ACTIVE: tuple[str, FlashAttentionHandle] | None = None
|
||||
|
||||
|
||||
def register_flash_attention_impl(
|
||||
impl: str | _FlashAttentionImpl,
|
||||
*,
|
||||
register_fn: _RegisterFn,
|
||||
) -> None:
|
||||
"""
|
||||
Register the callable that activates a flash attention impl.
|
||||
|
||||
.. note::
|
||||
This function is intended for SDPA backend providers to register their
|
||||
implementations. End users should use :func:`activate_flash_attention_impl`
|
||||
to activate a registered implementation.
|
||||
|
||||
Args:
|
||||
impl: Implementation identifier (e.g., ``"FA4"``).
|
||||
register_fn: Callable that performs the actual dispatcher registration.
|
||||
This function will be invoked by :func:`activate_flash_attention_impl`
|
||||
and should register custom kernels with the PyTorch dispatcher.
|
||||
It may optionally return a handle implementing
|
||||
:class:`FlashAttentionHandle` to keep any necessary state alive.
|
||||
|
||||
Example:
|
||||
>>> def my_impl_register(module_path: str = "my_flash_impl"):
|
||||
... # Register custom kernels with torch dispatcher
|
||||
... pass # doctest: +SKIP
|
||||
>>> register_flash_attention_impl(
|
||||
... "MyImpl", register_fn=my_impl_register
|
||||
... ) # doctest: +SKIP
|
||||
"""
|
||||
global _FLASH_ATTENTION_IMPLS
|
||||
_FLASH_ATTENTION_IMPLS[impl] = register_fn
|
||||
|
||||
|
||||
def activate_flash_attention_impl(
|
||||
impl: str | _FlashAttentionImpl,
|
||||
) -> None:
|
||||
"""
|
||||
Activate into the dispatcher a previously registered flash attention impl.
|
||||
|
||||
.. note::
|
||||
Backend providers should NOT automatically activate their implementation
|
||||
on import. Users should explicitly opt-in by calling this function or via
|
||||
environment variables to ensure multiple provider libraries can coexist.
|
||||
|
||||
Args:
|
||||
impl: Implementation identifier to activate. See
|
||||
:func:`~torch.nn.attention.list_flash_attention_impls` for available
|
||||
implementations.
|
||||
If the backend's :func:`register_flash_attention_impl` callable
|
||||
returns a :class:`FlashAttentionHandle`, the registry keeps that
|
||||
handle alive for the lifetime of the process (until explicit
|
||||
uninstall support exists).
|
||||
|
||||
Example:
|
||||
>>> activate_flash_attention_impl("FA4") # doctest: +SKIP
|
||||
"""
|
||||
global _FLASH_ATTENTION_ACTIVE, _FLASH_ATTENTION_IMPLS
|
||||
|
||||
restore_flash_attention_impl(
|
||||
_raise_warn=False
|
||||
) # first restore any prev overrides (if any) to default
|
||||
|
||||
register_fn = _FLASH_ATTENTION_IMPLS.get(impl)
|
||||
if register_fn is None:
|
||||
raise ValueError(
|
||||
f"Unknown flash attention impl '{impl}'. "
|
||||
f"Available implementations: {list_flash_attention_impls()}"
|
||||
)
|
||||
|
||||
handle = register_fn()
|
||||
if handle is not None:
|
||||
_FLASH_ATTENTION_ACTIVE = (impl, handle)
|
||||
|
||||
|
||||
def list_flash_attention_impls() -> list[str]:
|
||||
"""Return the names of all available flash attention implementations."""
|
||||
return sorted(_FLASH_ATTENTION_IMPLS.keys())
|
||||
|
||||
|
||||
def current_flash_attention_impl() -> str | None:
|
||||
"""
|
||||
Return the currently activated flash attention impl name, if any.
|
||||
|
||||
``None`` indicates that no custom impl has been activated.
|
||||
"""
|
||||
return (
|
||||
_FLASH_ATTENTION_ACTIVE[0]
|
||||
if _FLASH_ATTENTION_ACTIVE is not None
|
||||
else _FLASH_ATTENTION_ACTIVE
|
||||
)
|
||||
|
||||
|
||||
def restore_flash_attention_impl(_raise_warn: bool = True) -> None:
|
||||
"""
|
||||
Restore the default FA2 implementation
|
||||
"""
|
||||
global _FLASH_ATTENTION_ACTIVE
|
||||
|
||||
handle = None
|
||||
if _FLASH_ATTENTION_ACTIVE is not None:
|
||||
handle = _FLASH_ATTENTION_ACTIVE[1]
|
||||
|
||||
if handle is not None:
|
||||
handle.remove()
|
||||
elif _raise_warn:
|
||||
logger.warning(
|
||||
"Trying to restore default FA2 impl when no custom impl was activated"
|
||||
)
|
||||
|
||||
_FLASH_ATTENTION_ACTIVE = None # default
|
||||
@@ -0,0 +1,61 @@
|
||||
# mypy: allow-untyped-defs
|
||||
"""Defines utilities for interacting with scaled_dot_product_attention"""
|
||||
|
||||
import math
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
__all__: list[str] = []
|
||||
|
||||
|
||||
def _input_requires_grad(*tensors: torch.Tensor) -> bool:
|
||||
"""Returns True if any of the tensors requires grad"""
|
||||
return any(t.requires_grad for t in tensors)
|
||||
|
||||
|
||||
def _postprocess_flash_output(inpt_tensor: torch.Tensor, og_size: int) -> torch.Tensor:
|
||||
"""Handles the unpad of the last dimension"""
|
||||
if inpt_tensor.size(-1) != og_size:
|
||||
return inpt_tensor[..., :og_size]
|
||||
return inpt_tensor
|
||||
|
||||
|
||||
def _calculate_scale(head_dim_size: int, scale: float | None) -> float:
|
||||
"""
|
||||
For FlashAttention we pad the head dimension to be a multiple of 8 so we need to scale the output
|
||||
by the original head size and not the padded.
|
||||
"""
|
||||
if scale is not None:
|
||||
return scale
|
||||
return 1.0 / math.sqrt(head_dim_size)
|
||||
|
||||
|
||||
def _validate_sdpa_input(
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
attn_mask: torch.Tensor | None = None,
|
||||
dropout_p=0.0,
|
||||
is_causal=False,
|
||||
scale=None,
|
||||
allow_lowp_kv=False,
|
||||
) -> None:
|
||||
if not allow_lowp_kv:
|
||||
if query.dtype != key.dtype or query.dtype != value.dtype:
|
||||
raise ValueError(
|
||||
f"Expected query, key, and value to have the same dtype, "
|
||||
f"but got query.dtype: {query.dtype}, key.dtype: {key.dtype}, "
|
||||
f"and value.dtype: {value.dtype} instead."
|
||||
)
|
||||
if query.device != key.device or query.device != value.device:
|
||||
raise ValueError(
|
||||
f"Expected query, key, and value to have the same device type, "
|
||||
f"but got query.device: {query.device}, key.device: {key.device}, "
|
||||
f"and value.device: {value.device} instead."
|
||||
)
|
||||
if query.dim() < 2 or key.dim() < 2 or value.dim() < 2:
|
||||
raise ValueError(
|
||||
f"Expected query, key, and value to all be at least 2 dimensional, but got query.dim: "
|
||||
f"{query.dim()}, key.dim: {key.dim()} and value.dim: {value.dim()} instead."
|
||||
)
|
||||
@@ -0,0 +1,376 @@
|
||||
# mypy: allow-untyped-defs
|
||||
"""Defines bias subclasses that work with scaled_dot_product_attention"""
|
||||
|
||||
from enum import auto, IntEnum
|
||||
from warnings import warn
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch.backends.cuda import (
|
||||
can_use_efficient_attention,
|
||||
can_use_flash_attention,
|
||||
is_flash_attention_available,
|
||||
SDPAParams,
|
||||
)
|
||||
from torch.nn.attention import _raise_kernel_warnings
|
||||
from torch.nn.attention._utils import (
|
||||
_calculate_scale,
|
||||
_input_requires_grad,
|
||||
_postprocess_flash_output,
|
||||
_validate_sdpa_input,
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["causal_upper_left", "causal_lower_right", "CausalVariant", "CausalBias"]
|
||||
|
||||
|
||||
torch._dynamo.allow_in_graph(is_flash_attention_available)
|
||||
torch._dynamo.allow_in_graph(can_use_flash_attention)
|
||||
torch._dynamo.allow_in_graph(can_use_efficient_attention)
|
||||
torch._dynamo.allow_in_graph(SDPAParams)
|
||||
|
||||
|
||||
class CausalVariant(IntEnum):
|
||||
r"""
|
||||
Enum for causal variants used in attention mechanisms.
|
||||
|
||||
Defines two types of causal biases:
|
||||
|
||||
``UPPER_LEFT``: Represents upper-left triangular bias for standard causal attention.
|
||||
The equivalent pytorch code for constructing this bias is:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
torch.tril(torch.ones(size, dtype=torch.bool))
|
||||
|
||||
For instance, with ``shape=(3,4)``, the materialized bias tensor will be:
|
||||
|
||||
.. code-block:: text
|
||||
|
||||
[[1, 0, 0, 0],
|
||||
[1, 1, 0, 0],
|
||||
[1, 1, 1, 0]]
|
||||
|
||||
|
||||
``LOWER_RIGHT``: Represents lower-right triangular bias, the include values are aligned to the lower
|
||||
right corner of the matrix.
|
||||
|
||||
The equivalent pytorch code for constructing this bias is:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
diagonal_offset = size[1] - size[0]
|
||||
torch.tril(
|
||||
torch.ones(size, dtype=torch.bool),
|
||||
diagonal=diagonal_offset,
|
||||
)
|
||||
|
||||
For instance, with ``shape=(3,4)``, the materialized bias tensor will be:
|
||||
|
||||
.. code-block:: text
|
||||
|
||||
[[1, 1, 0, 0],
|
||||
[1, 1, 1, 0],
|
||||
[1, 1, 1, 1]]
|
||||
|
||||
Note that these variants are equivalent to each other when the sequence lengths of the query and key/value
|
||||
tensors are equal since the triangular matrix is square.
|
||||
|
||||
.. warning:: This enum is a prototype and subject to change.
|
||||
"""
|
||||
|
||||
UPPER_LEFT = auto()
|
||||
LOWER_RIGHT = auto()
|
||||
|
||||
|
||||
class CausalBias(torch.Tensor):
|
||||
"""
|
||||
A bias representing causal attention patterns. For an overview of the bias structure, see the :class:`CausalVariant` enum.
|
||||
|
||||
This class is used for defining causal (triangular) attention biases. For construing the bias, there exist
|
||||
two factory functions: :func:`causal_upper_left` and :func:`causal_lower_right`.
|
||||
|
||||
Example:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from torch.nn.attention.bias import causal_lower_right
|
||||
|
||||
bsz, num_heads, seqlen_q, seqlen_kv, head_dim = 32, 8, 4, 12, 8
|
||||
|
||||
# Create a lower-right causal bias
|
||||
attn_bias = causal_lower_right(seqlen_q, seqlen_kv)
|
||||
|
||||
q = torch.randn(
|
||||
bsz, num_heads, seqlen_q, head_dim, device="cuda", dtype=torch.float16
|
||||
)
|
||||
k = torch.randn(
|
||||
bsz, num_heads, seqlen_kv, head_dim, device="cuda", dtype=torch.float16
|
||||
)
|
||||
v = torch.randn(
|
||||
bsz, num_heads, seqlen_kv, head_dim, device="cuda", dtype=torch.float16
|
||||
)
|
||||
|
||||
out = F.scaled_dot_product_attention(q, k, v, attn_bias)
|
||||
|
||||
.. warning:: This class is a prototype and subject to change.
|
||||
"""
|
||||
|
||||
def __init__(self, variant: CausalVariant, seq_len_q: int, seq_len_kv: int) -> None:
|
||||
"""
|
||||
Initializes the CausalBias instance with a specified variant and sequence lengths.
|
||||
|
||||
Args:
|
||||
variant (CausalVariant): The type of causal bias to use (either UPPER_LEFT or LOWER_RIGHT).
|
||||
seq_len_q (int): The sequence length of the query tensor.
|
||||
seq_len_kv (int): The sequence length of the key/value tensor.
|
||||
|
||||
Raises a warning if the LOWER_RIGHT variant is used with seq_len_q > seq_len_kv, as it may produce NaNs.
|
||||
"""
|
||||
if not isinstance(variant, CausalVariant):
|
||||
raise AssertionError(
|
||||
f"variant must be a CausalVariant, got {type(variant).__name__}"
|
||||
)
|
||||
super().__init__()
|
||||
self.variant = variant
|
||||
self.seq_len_q = seq_len_q
|
||||
self.seq_len_kv = seq_len_kv
|
||||
if seq_len_q > seq_len_kv and variant == CausalVariant.LOWER_RIGHT:
|
||||
warn(
|
||||
"Lower right causal bias will produce NaNs in the output when seq_len_q > seq_len_kv!",
|
||||
stacklevel=2,
|
||||
)
|
||||
|
||||
def _upper_left(self, device: torch.device) -> torch.Tensor:
|
||||
"""Upper left causal bias"""
|
||||
return torch.tril(
|
||||
torch.ones(self.seq_len_q, self.seq_len_kv, device=device, dtype=torch.bool)
|
||||
)
|
||||
|
||||
def _lower_right(self, device: torch.device) -> torch.Tensor:
|
||||
"""Lower right causal bias"""
|
||||
diagonal_offset = self.seq_len_kv - self.seq_len_q
|
||||
return torch.tril(
|
||||
torch.ones(
|
||||
self.seq_len_q, self.seq_len_kv, device=device, dtype=torch.bool
|
||||
),
|
||||
diagonal=diagonal_offset,
|
||||
)
|
||||
|
||||
# pyrefly: ignore [bad-return]
|
||||
def _materialize(self, device: torch.device | None = None) -> torch.Tensor:
|
||||
"""
|
||||
Materializes the causal bias into a tensor form.
|
||||
|
||||
Depending on the variant, this method generates either an upper-left or lower-right
|
||||
triangular matrix to represent the causal bias.
|
||||
|
||||
Args:
|
||||
device (Optional[torch.device]): The device on which to create the tensor. Defaults to CPU.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: The materialized bias tensor.
|
||||
"""
|
||||
if device is None:
|
||||
device = torch.device("cpu")
|
||||
if self.variant == CausalVariant.UPPER_LEFT:
|
||||
return self._upper_left(device)
|
||||
elif self.variant == CausalVariant.LOWER_RIGHT:
|
||||
return self._lower_right(device)
|
||||
|
||||
@staticmethod
|
||||
def _dispatch(
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
attn_mask: "CausalBias",
|
||||
dropout_p: float = 0.0,
|
||||
is_causal: bool = False,
|
||||
scale: float | None = None,
|
||||
enable_gqa: bool = False,
|
||||
) -> torch.Tensor:
|
||||
r"""
|
||||
Handles the logic for computing attention with the specified causal bias.
|
||||
|
||||
Args:
|
||||
query (Tensor): Query tensor; shape :math:`(N, ..., L, E)`.
|
||||
key (Tensor): Key tensor; shape :math:`(N, ..., S, E)`.
|
||||
value (Tensor): Value tensor; shape :math:`(N, ..., S, Ev)`.
|
||||
attn_mask (CausalBias): The type of causal attention to apply.
|
||||
A boolean mask where a value of True indicates that the element *should* take part in attention.
|
||||
A float mask of the same type as query, key, value that is added to the attention score.
|
||||
dropout_p (float): Dropout probability; if greater than 0.0, dropout is applied
|
||||
is_causal (bool): If true, assumes upper left causal attention masking and errors if both attn_mask and is_causal
|
||||
are set.
|
||||
scale (optional float): Scaling factor applied prior to softmax. If None, the default value is set
|
||||
to :math:`\frac{1}{\sqrt{E}}`.
|
||||
enable_gqa (optional bool): If set to True, Grouped Query Attention (GQA) is enabled, by default it is set to False.
|
||||
|
||||
Returns:
|
||||
output (Tensor): Attention output; shape :math:`(N, ..., L, Ev)`.
|
||||
|
||||
Raises:
|
||||
ValueError: If the causal bias variant is not a CausalVariant type.
|
||||
|
||||
"""
|
||||
if is_causal:
|
||||
raise ValueError("CausalBias should not be used with causal=True")
|
||||
|
||||
if (
|
||||
attn_mask.seq_len_q == attn_mask.seq_len_kv
|
||||
or attn_mask.variant == CausalVariant.UPPER_LEFT
|
||||
):
|
||||
return F.scaled_dot_product_attention(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
attn_mask=None,
|
||||
dropout_p=dropout_p,
|
||||
is_causal=True,
|
||||
scale=scale,
|
||||
enable_gqa=enable_gqa,
|
||||
)
|
||||
elif attn_mask.variant == CausalVariant.LOWER_RIGHT:
|
||||
_validate_sdpa_input(query, key, value, None, dropout_p, is_causal, scale)
|
||||
sdpa_params = SDPAParams(
|
||||
query, key, value, None, dropout_p, is_causal, enable_gqa
|
||||
)
|
||||
if can_use_flash_attention(sdpa_params):
|
||||
alignment = 64 if query.device.type == "xpu" else 8
|
||||
og_head_size = query.size(-1)
|
||||
og_scale = _calculate_scale(og_head_size, scale)
|
||||
needs_padding = og_head_size % alignment != 0
|
||||
if needs_padding:
|
||||
pad_len = alignment - (og_head_size % alignment)
|
||||
query = torch.nn.functional.pad(query, (0, pad_len))
|
||||
key = torch.nn.functional.pad(key, (0, pad_len))
|
||||
value = torch.nn.functional.pad(value, (0, pad_len))
|
||||
out = torch.ops.aten._scaled_dot_product_flash_attention(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
dropout_p,
|
||||
is_causal=True, # TODO: Flash accepts causal = True and for this particular op it means lower right
|
||||
return_debug_mask=False,
|
||||
scale=og_scale,
|
||||
)[0]
|
||||
return _postprocess_flash_output(out, og_head_size)
|
||||
if can_use_efficient_attention(sdpa_params):
|
||||
compute_log_sumexp = False
|
||||
if _input_requires_grad(query, key, value):
|
||||
compute_log_sumexp = True
|
||||
return torch.ops.aten._efficient_attention_forward(
|
||||
query.transpose(1, 2),
|
||||
key.transpose(1, 2),
|
||||
value.transpose(1, 2),
|
||||
bias=None,
|
||||
cu_seqlens_q=None,
|
||||
cu_seqlens_k=None,
|
||||
max_seqlen_q=None,
|
||||
max_seqlen_k=None,
|
||||
dropout_p=dropout_p,
|
||||
custom_mask_type=int(attn_mask.variant),
|
||||
compute_log_sumexp=compute_log_sumexp,
|
||||
scale=scale,
|
||||
seqlen_k=None,
|
||||
)[0].transpose(1, 2)
|
||||
else:
|
||||
_raise_kernel_warnings(sdpa_params)
|
||||
# We can't use efficient attention the only support for lower right is via materialization
|
||||
return F.scaled_dot_product_attention(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
attn_mask=attn_mask._materialize(query.device),
|
||||
dropout_p=dropout_p,
|
||||
is_causal=False,
|
||||
scale=scale,
|
||||
enable_gqa=enable_gqa,
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"CausalBias.variant must be a CausalVariant type, but found: {attn_mask.variant}"
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def __torch_function__(cls, func, types, args=(), kwargs=None):
|
||||
"""Defines the behavior of torch.nn.functional.scaled_dot_product_attention when the attn_bias is an AttnBias"""
|
||||
if kwargs is None:
|
||||
kwargs = {}
|
||||
if func is torch.nn.functional.scaled_dot_product_attention:
|
||||
return cls._dispatch(*args, **kwargs)
|
||||
return super().__torch_function__(func, types, args, kwargs)
|
||||
|
||||
def __repr__(self) -> str: # type:ignore[override]
|
||||
return self._materialize().__repr__()
|
||||
|
||||
|
||||
def causal_upper_left(*size) -> CausalBias:
|
||||
"""
|
||||
Creates an upper-left triangular causal bias.
|
||||
|
||||
This function generates a upper-left triangular matrix to represent causal attention bias with a
|
||||
diagonal offset set so that the inclusive values are aligned to the upper left corner of the matrix.
|
||||
This equivalent to the `is_causal=True` argument in `scaled_dot_product_attention`.
|
||||
|
||||
The equivalent pytorch code for constructing this bias is:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
torch.tril(torch.ones(size, dtype=torch.bool))
|
||||
|
||||
For instance, with `shape=(3,4)`, the materialized bias tensor will be:
|
||||
|
||||
.. code-block:: text
|
||||
|
||||
[[1, 0, 0, 0],
|
||||
[1, 1, 0, 0],
|
||||
[1, 1, 1, 0]]
|
||||
|
||||
Args:
|
||||
size: The size of the bias matrix.
|
||||
|
||||
Returns:
|
||||
CausalBias: The UPPER_LEFT triangular causal bias variant.
|
||||
"""
|
||||
if len(size) != 2:
|
||||
raise AssertionError("causal_upper_left only supports 2D tensors")
|
||||
seq_len_q, seq_len_kv = size
|
||||
return CausalBias(CausalVariant.UPPER_LEFT, seq_len_q, seq_len_kv)
|
||||
|
||||
|
||||
def causal_lower_right(*size) -> CausalBias:
|
||||
"""
|
||||
Creates a lower-right triangular causal bias.
|
||||
|
||||
This function generates a lower-right triangular matrix to represent causal attention bias with a
|
||||
diagonal offset set so that the inclusive values are aligned to the lower right corner of the matrix.
|
||||
|
||||
The equivalent pytorch code for constructing this bias is:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
diagonal_offset = size[1] - size[0]
|
||||
torch.tril(
|
||||
torch.ones(size, dtype=torch.bool),
|
||||
diagonal=diagonal_offset,
|
||||
)
|
||||
|
||||
For instance, with `shape=(3,4)`, the materialized bias tensor will be:
|
||||
|
||||
.. code-block:: text
|
||||
|
||||
[[1, 1, 0, 0],
|
||||
[1, 1, 1, 0],
|
||||
[1, 1, 1, 1]]
|
||||
|
||||
Args:
|
||||
size: The size of the bias matrix.
|
||||
|
||||
Returns:
|
||||
CausalBias: The LOWER_RIGHT triangular causal bias variant.
|
||||
"""
|
||||
if len(size) != 2:
|
||||
raise AssertionError("causal_lower_right only supports 2D tensors")
|
||||
seq_len_q, seq_len_kv = size
|
||||
return CausalBias(CausalVariant.LOWER_RIGHT, seq_len_q, seq_len_kv)
|
||||
@@ -0,0 +1,2 @@
|
||||
# Experimental features are not mature yet and are subject to change.
|
||||
# We do not provide any BC/FC guarantees
|
||||
+358
@@ -0,0 +1,358 @@
|
||||
# mypy: allow-untyped-defs
|
||||
"""
|
||||
This module implements Paged Attention on top of flex_attention.
|
||||
This module is experimental and subject to change.
|
||||
"""
|
||||
|
||||
import torch
|
||||
from torch.nn.attention.flex_attention import (
|
||||
_identity,
|
||||
_mask_mod_signature,
|
||||
_score_mod_signature,
|
||||
BlockMask,
|
||||
noop_mask,
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["PagedAttention"]
|
||||
|
||||
|
||||
def _cdiv(x: int | float | torch.Tensor, multiple: int | float | torch.Tensor):
|
||||
return (x + multiple - 1) // multiple
|
||||
|
||||
|
||||
class PagedAttention:
|
||||
"""
|
||||
PagedAttention supports flex attention inference with a large batch size.
|
||||
With PagedAttention, a batch of key/value tensors with varying kv length
|
||||
is split into tensor blocks of fixed length and cached in a compact way.
|
||||
Thus we can avoid redundant memory consumption due to varying kv length and
|
||||
support a larger batch size.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
n_pages: int,
|
||||
page_size: int,
|
||||
max_batch_size: int,
|
||||
device: str = "cuda",
|
||||
) -> None:
|
||||
# number of pages
|
||||
self.n_pages = n_pages
|
||||
|
||||
# number of tokens per page
|
||||
self.page_size = page_size
|
||||
|
||||
# page table: [batch, logical_block_idx] -> physical_page_idx
|
||||
self.page_table = -torch.ones(
|
||||
(max_batch_size, self.n_pages), dtype=torch.int64, device=device
|
||||
)
|
||||
|
||||
# capacity: batch_idx -> allocated sequence length
|
||||
self.capacity = torch.zeros(max_batch_size, dtype=torch.int64, device=device)
|
||||
|
||||
# index of empty pages that is available for allocation
|
||||
self.empty_pages = list(range(n_pages - 1, -1, -1))
|
||||
|
||||
# mapping from physical page index to logical page index
|
||||
self.physical_to_logical = -torch.ones(
|
||||
(max_batch_size, n_pages), dtype=torch.int64, device=device
|
||||
)
|
||||
|
||||
def reserve(self, batch_idx: torch.Tensor, seq_len: torch.Tensor) -> None:
|
||||
"""
|
||||
Requests the capacity of a given batch to be at least enough to
|
||||
hold `seq_len` elements.
|
||||
|
||||
Args:
|
||||
batch_idx (Tensor): batch index to be reserved; shape :math:`(1)`.
|
||||
seq_len (Tensor): minimum capacity for the given batch; shape :math:`(1)`.
|
||||
"""
|
||||
|
||||
if seq_len <= self.capacity[batch_idx]:
|
||||
return
|
||||
|
||||
num_pages_to_allocate = _cdiv(
|
||||
seq_len - self.capacity[batch_idx], self.page_size
|
||||
)
|
||||
|
||||
if len(self.empty_pages) < num_pages_to_allocate:
|
||||
raise AssertionError(
|
||||
f"requested {num_pages_to_allocate.item()} pages "
|
||||
f"but there are only {len(self.empty_pages)} empty pages"
|
||||
)
|
||||
|
||||
start_page_idx = self.capacity[batch_idx] // self.page_size
|
||||
end_page_idx = start_page_idx + num_pages_to_allocate
|
||||
|
||||
# find empty physical pages
|
||||
allocated_pages = torch.tensor(
|
||||
self.empty_pages[-num_pages_to_allocate:],
|
||||
device=num_pages_to_allocate.device,
|
||||
)
|
||||
self.empty_pages = self.empty_pages[:-num_pages_to_allocate]
|
||||
|
||||
# update page table
|
||||
self.page_table[
|
||||
batch_idx,
|
||||
start_page_idx:end_page_idx,
|
||||
] = allocated_pages
|
||||
|
||||
# update metadata
|
||||
self.physical_to_logical[batch_idx, allocated_pages] = torch.arange(
|
||||
start_page_idx.item(),
|
||||
end_page_idx.item(),
|
||||
device=num_pages_to_allocate.device,
|
||||
)
|
||||
self.capacity[batch_idx] += num_pages_to_allocate * self.page_size
|
||||
|
||||
def erase(self, batch_idx: torch.Tensor) -> None:
|
||||
"""
|
||||
Removes a single batch from paged attention.
|
||||
|
||||
Args:
|
||||
batch_idx (Tensor): batch index to be removed; shape :math:`(1)`.
|
||||
"""
|
||||
|
||||
# find allocated pages
|
||||
allocated_page_idx = self.page_table[batch_idx] != -1
|
||||
allocated_pages = self.page_table[batch_idx][allocated_page_idx]
|
||||
|
||||
# clean metadata
|
||||
self.capacity[batch_idx] = 0
|
||||
self.empty_pages += allocated_pages.tolist()
|
||||
self.physical_to_logical[batch_idx][:, allocated_pages] = -1
|
||||
self.page_table[batch_idx] = -1
|
||||
|
||||
def assign(
|
||||
self,
|
||||
batch_idx: torch.Tensor,
|
||||
input_pos: torch.Tensor,
|
||||
k_val: torch.Tensor,
|
||||
v_val: torch.Tensor,
|
||||
k_cache: torch.Tensor,
|
||||
v_cache: torch.Tensor,
|
||||
) -> None:
|
||||
"""
|
||||
Assigns new contents `val` to the storage `cache` at the location
|
||||
`batch_idx` and `input_pos`.
|
||||
|
||||
Args:
|
||||
batch_idx (Tensor): batch index; shape :math:`(B)`.
|
||||
input_pos (Tensor): input positions to be assigned for the given batch; shape :math:`(B, S)`.
|
||||
val (Tensor): value to be assigned; shape :math:`(B, H, S, D)`
|
||||
cache (Tensor): the cache to store the values; shape:`(1, H, MAX_S, D)`
|
||||
"""
|
||||
if k_val.requires_grad:
|
||||
raise RuntimeError("val must not require gradient")
|
||||
|
||||
B, H, S, K_D = k_val.shape
|
||||
V_D = v_val.shape[3]
|
||||
if B != batch_idx.shape[0]:
|
||||
raise RuntimeError(
|
||||
f"Expect val and batch_idx have the same batch size "
|
||||
f"but got B={B} and B={batch_idx.shape[0]}."
|
||||
)
|
||||
if H != k_cache.shape[1]:
|
||||
raise RuntimeError(
|
||||
f"Expect val and cache has the same number of heads "
|
||||
f"but got H={H} and H={k_cache.shape[1]}."
|
||||
)
|
||||
if S != input_pos.shape[1]:
|
||||
raise RuntimeError(
|
||||
f"Expect val and input_pos has the same length "
|
||||
f"but got S={S} and S={input_pos.shape[0]}."
|
||||
)
|
||||
if K_D != k_cache.shape[3]:
|
||||
raise RuntimeError(
|
||||
f"Expect k_val and k_cache has the same hidden dim "
|
||||
f"but got D={K_D} and D={k_cache.shape[3]}."
|
||||
)
|
||||
if V_D != v_cache.shape[3]:
|
||||
raise RuntimeError(
|
||||
f"Expect v_val and v_cache has the same hidden dim "
|
||||
f"but got D={V_D} and D={v_cache.shape[3]}."
|
||||
)
|
||||
|
||||
# find address
|
||||
logical_block_idx = input_pos // self.page_size # [B, S]
|
||||
logical_block_offset = input_pos % self.page_size # [B, S]
|
||||
physical_block_idx = torch.gather(
|
||||
self.page_table[batch_idx], 1, logical_block_idx.to(torch.int64)
|
||||
).to(torch.int32) # [B, S]
|
||||
|
||||
addr = (physical_block_idx * self.page_size + logical_block_offset).view(
|
||||
-1
|
||||
) # [B*S]
|
||||
|
||||
k_val = k_val.permute(1, 0, 2, 3).contiguous().view(1, H, B * S, K_D)
|
||||
v_val = v_val.permute(1, 0, 2, 3).contiguous().view(1, H, B * S, V_D)
|
||||
|
||||
k_cache[:, :, addr, :] = k_val
|
||||
v_cache[:, :, addr, :] = v_val
|
||||
|
||||
def convert_logical_block_mask(
|
||||
self,
|
||||
block_mask: BlockMask,
|
||||
batch_idx: torch.Tensor | None = None,
|
||||
kv_len: torch.Tensor | None = None,
|
||||
) -> BlockMask:
|
||||
"""
|
||||
Converts a logical block mask by mapping its logical kv indices to the corresponding
|
||||
physical kv indices.
|
||||
|
||||
Args:
|
||||
block_mask (BlockMask): logical block mask;
|
||||
kv_indices shape :math:`(B, H, ROWS, MAX_BLOCKS_IN_COL)`.
|
||||
batch_idx (Tensor): batch index corresponding to the block_mask
|
||||
batch dimension. This provides flexibility to convert a
|
||||
block mask with smaller batch size than the page table;
|
||||
shape :math:`(B)`.
|
||||
kv_len (Optional[Tensor]): actual KV sequence length for upper bound check;
|
||||
shape :math:`(B,)` to handle multiple batches.
|
||||
"""
|
||||
B, H, ROWS, MAX_BLOCKS_IN_COL = block_mask.kv_indices.shape
|
||||
|
||||
if block_mask.BLOCK_SIZE[1] != self.page_size:
|
||||
raise RuntimeError(
|
||||
f"Expect block_mask has the same column block size as page_size"
|
||||
f"but got size={block_mask.BLOCK_SIZE[1]} and size={self.page_size}"
|
||||
)
|
||||
|
||||
# Increase the num columns of converted block mask from logical block mask's
|
||||
# num columns to n_pages, since a) the converted block mask
|
||||
# may have larger indices values; and b) `_ordered_to_dense` realizes
|
||||
# a dense tensor with these converted indices. There would be an IndexError
|
||||
# if using the logical block mask's num columns.
|
||||
|
||||
device = block_mask.kv_num_blocks.device
|
||||
|
||||
if batch_idx is None:
|
||||
batch_idx = torch.arange(B, device=device)
|
||||
page_table = self.page_table[batch_idx]
|
||||
|
||||
new_kv_num_blocks = block_mask.kv_num_blocks.clone()
|
||||
|
||||
new_kv_indices = torch.zeros(
|
||||
(B, H, ROWS, self.n_pages), dtype=torch.int32, device=device
|
||||
)
|
||||
new_kv_indices[:, :, :, :MAX_BLOCKS_IN_COL] = (
|
||||
torch.gather(
|
||||
page_table, 1, block_mask.kv_indices.view(B, -1).to(torch.int64)
|
||||
)
|
||||
.view(block_mask.kv_indices.shape)
|
||||
.to(torch.int32)
|
||||
)
|
||||
|
||||
new_full_kv_indices, new_full_kv_num_blocks = None, None
|
||||
if block_mask.full_kv_num_blocks is not None:
|
||||
if block_mask.full_kv_indices is None:
|
||||
raise AssertionError(
|
||||
"block_mask.full_kv_indices must not be None when full_kv_num_blocks is not None"
|
||||
)
|
||||
new_full_kv_num_blocks = block_mask.full_kv_num_blocks.clone()
|
||||
new_full_kv_indices = torch.zeros(
|
||||
(B, H, ROWS, self.n_pages), dtype=torch.int32, device=device
|
||||
)
|
||||
new_full_kv_indices[:, :, :, :MAX_BLOCKS_IN_COL] = (
|
||||
torch.gather(
|
||||
page_table,
|
||||
1,
|
||||
block_mask.full_kv_indices.view(B, -1).to(torch.int64),
|
||||
)
|
||||
.view(block_mask.full_kv_indices.shape)
|
||||
.to(torch.int32)
|
||||
)
|
||||
|
||||
new_mask_mod = self.get_mask_mod(block_mask.mask_mod, kv_len)
|
||||
|
||||
seq_lengths = (block_mask.seq_lengths[0], self.n_pages * self.page_size)
|
||||
return BlockMask.from_kv_blocks(
|
||||
new_kv_num_blocks,
|
||||
new_kv_indices,
|
||||
new_full_kv_num_blocks,
|
||||
new_full_kv_indices,
|
||||
block_mask.BLOCK_SIZE,
|
||||
new_mask_mod,
|
||||
seq_lengths=seq_lengths,
|
||||
)
|
||||
|
||||
def get_mask_mod(
|
||||
self,
|
||||
mask_mod: _mask_mod_signature | None,
|
||||
kv_len: torch.Tensor | None = None,
|
||||
) -> _mask_mod_signature:
|
||||
"""
|
||||
Converts a mask_mod based on mapping from the physical block index to the logical
|
||||
block index.
|
||||
|
||||
Args:
|
||||
mask_mod (_mask_mod_signature): mask_mod based on the logical block index.
|
||||
kv_len (Optional[torch.Tensor]): actual KV sequence length for upper bound check.
|
||||
"""
|
||||
if mask_mod is None:
|
||||
mask_mod = noop_mask
|
||||
|
||||
def new_mask_mod(
|
||||
b: torch.Tensor,
|
||||
h: torch.Tensor,
|
||||
q_idx: torch.Tensor,
|
||||
physical_kv_idx: torch.Tensor,
|
||||
):
|
||||
physical_kv_block = physical_kv_idx // self.page_size
|
||||
physical_kv_offset = physical_kv_idx % self.page_size
|
||||
logical_block_idx = self.physical_to_logical[b, physical_kv_block]
|
||||
logical_kv_idx = logical_block_idx * self.page_size + physical_kv_offset
|
||||
live_block = logical_block_idx >= 0
|
||||
within_upper_bound = (
|
||||
logical_kv_idx < kv_len[b] if kv_len is not None else True
|
||||
)
|
||||
within_lower_bound = logical_kv_idx >= 0
|
||||
is_valid = live_block & within_upper_bound & within_lower_bound
|
||||
|
||||
return torch.where(is_valid, mask_mod(b, h, q_idx, logical_kv_idx), False)
|
||||
|
||||
return new_mask_mod
|
||||
|
||||
def get_score_mod(
|
||||
self,
|
||||
score_mod: _score_mod_signature | None,
|
||||
kv_len: torch.Tensor | None = None,
|
||||
) -> _score_mod_signature:
|
||||
"""
|
||||
Converts a score_mod based on mapping from the physical block index to the logical
|
||||
block index.
|
||||
|
||||
Args:
|
||||
score_mod (_score_mod_signature): score_mod based on the logical block index.
|
||||
`kv_len (Optional[torch.Tensor]): actual KV sequence length for upper bound check.
|
||||
|
||||
"""
|
||||
if score_mod is None:
|
||||
score_mod = _identity
|
||||
|
||||
def new_score_mod(
|
||||
score: torch.Tensor,
|
||||
b: torch.Tensor,
|
||||
h: torch.Tensor,
|
||||
q_idx: torch.Tensor,
|
||||
physical_kv_idx: torch.Tensor,
|
||||
):
|
||||
physical_kv_block = physical_kv_idx // self.page_size
|
||||
physical_kv_offset = physical_kv_idx % self.page_size
|
||||
logical_block_idx = self.physical_to_logical[b, physical_kv_block]
|
||||
logical_kv_idx = logical_block_idx * self.page_size + physical_kv_offset
|
||||
live_block = logical_block_idx >= 0
|
||||
within_upper_bound = (
|
||||
logical_kv_idx < kv_len[b] if kv_len is not None else True
|
||||
)
|
||||
within_lower_bound = logical_kv_idx >= 0
|
||||
is_valid = live_block & within_upper_bound & within_lower_bound
|
||||
|
||||
return torch.where(
|
||||
is_valid,
|
||||
score_mod(score, b, h, q_idx, logical_kv_idx),
|
||||
float("-inf"),
|
||||
)
|
||||
|
||||
return new_score_mod
|
||||
+154
@@ -0,0 +1,154 @@
|
||||
# mypy: allow-untyped-defs
|
||||
"""
|
||||
This operator implements FP8 scaled dot product attention using Flash Attention 3.
|
||||
This operator is experimental and subject to change.
|
||||
"""
|
||||
|
||||
import warnings
|
||||
from enum import IntEnum
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
|
||||
|
||||
class DescaleType(IntEnum):
|
||||
"""Describes the scaling granularity for FP8 descale tensors.
|
||||
|
||||
Used with _scaled_dot_product_attention_quantized to explicitly specify
|
||||
how the descale factors are applied to the quantized inputs.
|
||||
|
||||
.. warning::
|
||||
This enum is experimental and subject to change.
|
||||
"""
|
||||
|
||||
PER_HEAD = 0
|
||||
"""Per-head descaling. Descale tensor shape: (batch_size, num_kv_heads)."""
|
||||
|
||||
|
||||
def _validate_descale(
|
||||
descale: Tensor | None,
|
||||
name: str,
|
||||
query: Tensor,
|
||||
key: Tensor,
|
||||
descale_type: DescaleType,
|
||||
) -> None:
|
||||
"""Validate descale tensor for the specified scaling type.
|
||||
|
||||
Args:
|
||||
descale: The descale tensor to validate (may be None)
|
||||
name: Name of the descale tensor ("q", "k", or "v") for error messages
|
||||
query: Query tensor to get batch size
|
||||
key: Key tensor to get num_kv_heads
|
||||
descale_type: The scaling granularity being used
|
||||
|
||||
Raises:
|
||||
ValueError: If the descale tensor has invalid dtype, device, or shape
|
||||
|
||||
Note:
|
||||
All descale tensors (q, k, v) use num_kv_heads for the head dimension.
|
||||
For GQA/MQA where num_query_heads > num_kv_heads, q_descale is broadcast
|
||||
from (B, H_kv) to match the query heads internally.
|
||||
"""
|
||||
if descale is None:
|
||||
return
|
||||
|
||||
# Check dtype
|
||||
if descale.dtype != torch.float32:
|
||||
raise ValueError(f"{name}_descale must have dtype float32, got {descale.dtype}")
|
||||
|
||||
# Check device
|
||||
if not descale.is_cuda:
|
||||
raise ValueError(f"{name}_descale must be a CUDA tensor")
|
||||
|
||||
# Check shape based on descale type
|
||||
if descale_type == DescaleType.PER_HEAD:
|
||||
batch_size = query.size(0)
|
||||
# All descale tensors use num_kv_heads, even q_descale (broadcast internally)
|
||||
# For BHSD layout, num_kv_heads is at dim 1 of key
|
||||
num_kv_heads = key.size(1)
|
||||
|
||||
if descale.dim() != 2:
|
||||
raise ValueError(
|
||||
f"{name}_descale must be a 2D tensor with shape (batch_size, num_kv_heads) "
|
||||
f"for PER_HEAD descaling, got {descale.dim()}D tensor"
|
||||
)
|
||||
|
||||
if descale.size(0) != batch_size:
|
||||
raise ValueError(
|
||||
f"{name}_descale batch dimension must match query batch size, "
|
||||
f"expected {batch_size}, got {descale.size(0)}"
|
||||
)
|
||||
|
||||
if descale.size(1) != num_kv_heads:
|
||||
raise ValueError(
|
||||
f"{name}_descale head dimension must match num_kv_heads, "
|
||||
f"expected {num_kv_heads}, got {descale.size(1)}"
|
||||
)
|
||||
|
||||
|
||||
def _scaled_dot_product_attention_quantized(
|
||||
query: Tensor,
|
||||
key: Tensor,
|
||||
value: Tensor,
|
||||
is_causal: bool = False,
|
||||
scale: float | None = None,
|
||||
q_descale: Tensor | None = None,
|
||||
k_descale: Tensor | None = None,
|
||||
v_descale: Tensor | None = None,
|
||||
q_descale_type: DescaleType = DescaleType.PER_HEAD,
|
||||
k_descale_type: DescaleType = DescaleType.PER_HEAD,
|
||||
v_descale_type: DescaleType = DescaleType.PER_HEAD,
|
||||
) -> Tensor:
|
||||
r"""Scaled dot product attention for FP8 inputs.
|
||||
|
||||
This is a specialized version of scaled_dot_product_attention that supports
|
||||
FP8 quantized inputs (float8_e4m3fn) with per-head descaling. Requires the
|
||||
Flash Attention 3 backend to be activated.
|
||||
|
||||
.. warning::
|
||||
This function is experimental and only supports forward pass.
|
||||
|
||||
Args:
|
||||
query (Tensor): Query tensor; shape :math:`(N, H_q, L, E)` dtype float8_e4m3fn
|
||||
key (Tensor): Key tensor; shape :math:`(N, H, S, E)` dtype float8_e4m3fn
|
||||
value (Tensor): Value tensor; shape :math:`(N, H, S, E_v)` dtype float8_e4m3fn
|
||||
is_causal (bool): Apply causal attention mask
|
||||
scale (float, optional): Scaling factor for attention weights
|
||||
q_descale (Tensor, optional): Query descale tensor; shape :math:`(N, H)` for PER_HEAD
|
||||
k_descale (Tensor, optional): Key descale tensor; shape :math:`(N, H)` for PER_HEAD
|
||||
v_descale (Tensor, optional): Value descale tensor; shape :math:`(N, H)` for PER_HEAD
|
||||
q_descale_type (DescaleType): Specifies the descaling granularity for query. Default: PER_HEAD
|
||||
k_descale_type (DescaleType): Specifies the descaling granularity for key. Default: PER_HEAD
|
||||
v_descale_type (DescaleType): Specifies the descaling granularity for value. Default: PER_HEAD
|
||||
|
||||
Returns:
|
||||
Tensor: Attention output; shape :math:`(N, H_q, L, E_v)` dtype bfloat16
|
||||
"""
|
||||
# Validate descale tensors
|
||||
_validate_descale(q_descale, "q", query, key, q_descale_type)
|
||||
_validate_descale(k_descale, "k", query, key, k_descale_type)
|
||||
_validate_descale(v_descale, "v", query, key, v_descale_type)
|
||||
|
||||
if torch.is_grad_enabled() and (
|
||||
query.requires_grad or key.requires_grad or value.requires_grad
|
||||
):
|
||||
warnings.warn(
|
||||
"_scaled_dot_product_attention_quantized does not support backward pass. "
|
||||
"Gradients will not be computed for query, key, or value.",
|
||||
UserWarning,
|
||||
)
|
||||
# Directly call the internal flash attention operator which has descale support
|
||||
# NOTE: This should be torch._scaled_dot_product_flash_attention, but it does not work with torch.compile
|
||||
result = torch.ops.aten._scaled_dot_product_flash_attention.quantized(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
q_descale,
|
||||
k_descale,
|
||||
v_descale,
|
||||
0.0,
|
||||
is_causal,
|
||||
False,
|
||||
scale=scale,
|
||||
)
|
||||
return result[0] # Return the output tensor, mirroring scaled_dot_product_attention
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,671 @@
|
||||
"""
|
||||
Variable-length attention implementation using Flash Attention.
|
||||
|
||||
This module provides a high-level Python interface for variable-length attention
|
||||
that calls into the optimized Flash Attention kernels.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from functools import lru_cache
|
||||
from typing import Any, NamedTuple
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
__all__ = ["varlen_attn", "varlen_attn_out", "AuxRequest"]
|
||||
|
||||
|
||||
def _normalize_window_size(window_size: list[int] | None) -> list[int]:
|
||||
if window_size is None:
|
||||
window_size = [-1, -1]
|
||||
|
||||
if len(window_size) != 2:
|
||||
raise ValueError(f"window_size must have length 2, got {len(window_size)}")
|
||||
return window_size
|
||||
|
||||
|
||||
@lru_cache(maxsize=8)
|
||||
def _should_use_cudnn(device_index: int) -> bool:
|
||||
"""Cache device capability check to avoid repeated CUDA calls."""
|
||||
return False
|
||||
|
||||
|
||||
class AuxRequest(NamedTuple):
|
||||
"""
|
||||
Request which auxiliary outputs to compute from varlen_attn.
|
||||
|
||||
Each field is a boolean indicating whether that auxiliary output should be computed.
|
||||
"""
|
||||
|
||||
lse: bool = False
|
||||
|
||||
|
||||
@torch.library.custom_op("torch_attn::_varlen_attn", mutates_args={})
|
||||
def _varlen_attn(
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
cu_seq_q: torch.Tensor,
|
||||
cu_seq_k: torch.Tensor | None,
|
||||
max_q: int,
|
||||
max_k: int,
|
||||
is_causal: bool = False,
|
||||
scale: float | None = None,
|
||||
window_size: list[int] | None = None,
|
||||
enable_gqa: bool = False,
|
||||
seqused_k: torch.Tensor | None = None,
|
||||
block_table: torch.Tensor | None = None,
|
||||
num_splits: int | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Private custom op for variable-length attention.
|
||||
|
||||
This is the internal implementation. Users should use the public varlen_attn function instead.
|
||||
"""
|
||||
window_size = _normalize_window_size(window_size)
|
||||
|
||||
use_cudnn = query.is_cuda and _should_use_cudnn(query.device.index)
|
||||
|
||||
if use_cudnn:
|
||||
log.info("Using cuDNN backend for varlen_attn")
|
||||
|
||||
if enable_gqa:
|
||||
# TODO: check this
|
||||
raise RuntimeError("GQA is not supported with the cuDNN backend.")
|
||||
if num_splits is not None:
|
||||
# TODO: check this
|
||||
raise RuntimeError("num_splits is not supported with the cuDNN backend.")
|
||||
if window_size[0] != -1 or window_size[1] != -1:
|
||||
raise RuntimeError(
|
||||
"cuDNN backend does not support window attention. Please use Flash Attention backend."
|
||||
)
|
||||
if seqused_k is not None or block_table is not None:
|
||||
# TODO: cuDNN supports per-sequence KV lengths via SEQ_LEN_KV + padding_mask,
|
||||
# but _cudnn_attention_forward doesn't expose it yet.
|
||||
raise RuntimeError(
|
||||
"seqused_k/block_table is not yet supported with the cuDNN backend."
|
||||
)
|
||||
|
||||
result = torch.ops.aten._cudnn_attention_forward(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
None, # attn_bias
|
||||
cu_seq_q,
|
||||
cu_seq_k,
|
||||
max_q,
|
||||
max_k,
|
||||
True, # compute_log_sumexp
|
||||
0.0, # dropout_p hardcoded to 0.0
|
||||
is_causal,
|
||||
False, # return_debug_mask
|
||||
scale=scale,
|
||||
)
|
||||
# cuDNN returns: (output, logsumexp, cum_seq_q, cum_seq_k, max_q, max_k, philox_seed, philox_offset, debug_attn_mask)
|
||||
output, softmax_lse, rng_state = result[0], result[1], result[6]
|
||||
else:
|
||||
log.info("Using Flash Attention backend for varlen_attn")
|
||||
output, softmax_lse, rng_state, _, _ = torch.ops.aten._flash_attention_forward(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
cu_seq_q,
|
||||
cu_seq_k,
|
||||
max_q,
|
||||
max_k,
|
||||
0.0, # dropout_p hardcoded to 0.0
|
||||
is_causal,
|
||||
return_debug_mask=False,
|
||||
scale=scale,
|
||||
window_size_left=window_size[0],
|
||||
window_size_right=window_size[1],
|
||||
seqused_k=seqused_k,
|
||||
block_table=block_table,
|
||||
num_splits=num_splits,
|
||||
)
|
||||
|
||||
rng_state_ = torch.zeros(
|
||||
(2,), dtype=torch.uint64, device=query.device
|
||||
) # hardcoded since dropout is hardcoded to 0
|
||||
return output, softmax_lse, rng_state_
|
||||
|
||||
|
||||
@_varlen_attn.register_fake
|
||||
def _varlen_attn_fake(
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
cu_seq_q: torch.Tensor,
|
||||
cu_seq_k: torch.Tensor | None,
|
||||
max_q: int,
|
||||
max_k: int,
|
||||
is_causal: bool = False,
|
||||
scale: float | None = None,
|
||||
window_size: list[int] | None = None,
|
||||
enable_gqa: bool = False,
|
||||
seqused_k: torch.Tensor | None = None,
|
||||
block_table: torch.Tensor | None = None,
|
||||
num_splits: int | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Fake implementation for meta tensor computation and tracing.
|
||||
|
||||
Based on the 3D varlen path from meta__flash_attention_forward:
|
||||
- query shape: (total, num_heads, head_dim)
|
||||
- logsumexp shape: (num_heads, total_q)
|
||||
"""
|
||||
window_size = _normalize_window_size(window_size)
|
||||
|
||||
# Output has same shape as query
|
||||
output = torch.empty_like(query)
|
||||
|
||||
# For varlen path: logsumexp shape is (num_heads, total_q)
|
||||
total_q = query.size(0)
|
||||
num_heads = query.size(1)
|
||||
logsumexp = torch.empty(
|
||||
(num_heads, total_q), dtype=torch.float, device=query.device
|
||||
)
|
||||
|
||||
if torch.version.hip:
|
||||
preferred = torch._C._get_rocm_fa_preferred_backend()
|
||||
if preferred == torch._C._ROCmFABackend.AOTriton:
|
||||
# AOTriton ROCm path uses batched 3D
|
||||
batch_size = cu_seq_q.size(0) - 1
|
||||
logsumexp = torch.empty(
|
||||
(batch_size, num_heads, max_q), dtype=torch.float, device=query.device
|
||||
)
|
||||
|
||||
rng_state = torch.empty((2,), dtype=torch.uint64, device=query.device)
|
||||
|
||||
return output, logsumexp, rng_state
|
||||
|
||||
|
||||
def varlen_attn(
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
cu_seq_q: torch.Tensor,
|
||||
cu_seq_k: torch.Tensor | None,
|
||||
max_q: int,
|
||||
max_k: int,
|
||||
*,
|
||||
return_aux: AuxRequest | None = None,
|
||||
scale: float | None = None,
|
||||
window_size: tuple[int, int] = (-1, -1),
|
||||
enable_gqa: bool = False,
|
||||
seqused_k: torch.Tensor | None = None,
|
||||
block_table: torch.Tensor | None = None,
|
||||
num_splits: int | None = None,
|
||||
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
|
||||
r"""Compute variable-length attention using Flash Attention.
|
||||
|
||||
This function is similar to scaled_dot_product_attention but optimized for
|
||||
variable-length sequences using cumulative sequence position tensors.
|
||||
|
||||
Args:
|
||||
query (Tensor): Query tensor; shape :math:`(T_q, H_q, D)`
|
||||
key (Tensor): Key tensor; shape :math:`(T_k, H_{kv}, D)`, or
|
||||
:math:`(\text{total\_pages}, \text{page\_size}, H_{kv}, D)` when ``block_table`` is provided.
|
||||
value (Tensor): Value tensor; shape :math:`(T_k, H_{kv}, D)`, or
|
||||
:math:`(\text{total\_pages}, \text{page\_size}, H_{kv}, D)` when ``block_table`` is provided.
|
||||
cu_seq_q (Tensor): Cumulative sequence positions for queries; shape :math:`(N+1,)`
|
||||
cu_seq_k (Tensor): Cumulative sequence positions for keys/values; shape :math:`(N+1,)`
|
||||
max_q (int): Maximum query sequence length in the batch.
|
||||
max_k (int): Maximum key/value sequence length in the batch.
|
||||
return_aux (Optional[AuxRequest]): If not None and ``return_aux.lse`` is True, also returns the logsumexp tensor.
|
||||
scale (float, optional): Scaling factor for attention scores
|
||||
window_size (tuple[int, int], optional): Window size for sliding window attention as (left, right).
|
||||
Use (-1, -1) for full attention (default), (-1, 0) for causal attention,
|
||||
or (W, 0) for causal attention with sliding window of size W.
|
||||
enable_gqa (bool): If set to True, enables Grouped Query Attention (GQA)
|
||||
and allows key/value to have fewer heads than query.
|
||||
Each KV head is shared by a group of :math:`H_q / H_{kv}` query heads,
|
||||
so :math:`H_q` must be divisible by :math:`H_{kv}`.
|
||||
Default is False.
|
||||
seqused_k (Tensor, optional): Number of valid KV tokens per batch element; shape :math:`(N,)`.
|
||||
When set, only the first ``seqused_k[i]`` tokens in the key/value sequence for batch
|
||||
element *i* participate in attention. Useful for KV-cache decoding where the cache slot
|
||||
is larger than the actual sequence. Inference-only (not supported in backward).
|
||||
block_table (Tensor, optional): Block table for paged KV cache; shape
|
||||
:math:`(N, \text{max\_pages\_per\_seq})`, dtype ``int32``.
|
||||
Requires ``seqused_k``. Inference-only (not supported in backward).
|
||||
|
||||
When ``block_table`` is provided, ``key`` and ``value`` are a "pool" of
|
||||
pages of tokens of KV data and the pages belong to any sequence/order.
|
||||
The ``block_table`` is what maps each sequence's logical chunks
|
||||
back to physical pages in this pool.
|
||||
|
||||
``seqused_k[i]`` tells the kernel how many tokens in sequence *i* are
|
||||
actually valid, since the last page is typically only partially filled.
|
||||
num_splits (int, optional): Number of splits for split-KV. Set to ``1``
|
||||
to disable split-KV which enables batch invariance. Split-KV
|
||||
parallelizes the key/value sequence dimension across multiple thread
|
||||
blocks and combines partial results. The split decision depends
|
||||
on ``max_k`` (the longest sequence in the batch), so different batch
|
||||
compositions can change the reduction order and produce different
|
||||
floating-point results for the same sequence. When this is disabled,
|
||||
bitwise identical outputs are guaranteed for a given sequence
|
||||
regardless of what other sequences are in the batch, at the
|
||||
cost of lower GPU utilization when there are few queries. When
|
||||
``None`` (default), the kernel chooses automatically.
|
||||
|
||||
Returns:
|
||||
output (Tensor): Output tensor from attention computation; shape :math:`(T_q, H_q, D)`.
|
||||
|
||||
If ``return_aux`` is not None and ``return_aux.lse`` is True:
|
||||
lse (Tensor): Log-sum-exp of attention scores; shape :math:`(T_q, H_q)`.
|
||||
|
||||
Shape legend:
|
||||
- :math:`N`: Batch size
|
||||
- :math:`T_q`: Total number of query tokens in the batch (sum of all query sequence lengths)
|
||||
- :math:`T_k`: Total number of key/value tokens in the batch (sum of all key/value sequence lengths)
|
||||
- :math:`H_q`: Number of query attention heads
|
||||
- :math:`H_{kv}`: Number of key/value attention heads (equal to :math:`H_q` unless GQA is enabled)
|
||||
- :math:`D`: Head dimension
|
||||
|
||||
Example::
|
||||
|
||||
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA)
|
||||
>>> batch_size, max_seq_len, embed_dim, num_heads = 2, 512, 1024, 16
|
||||
>>> head_dim = embed_dim // num_heads
|
||||
>>> seq_lengths = []
|
||||
>>> for _ in range(batch_size):
|
||||
... length = torch.randint(1, max_seq_len // 64 + 1, (1,)).item() * 64
|
||||
... seq_lengths.append(min(length, max_seq_len))
|
||||
>>> seq_lengths = torch.tensor(seq_lengths, device="cuda")
|
||||
>>> total_tokens = seq_lengths.sum().item()
|
||||
>>>
|
||||
>>> # Create packed query, key, value tensors
|
||||
>>> query = torch.randn(
|
||||
... total_tokens, num_heads, head_dim, dtype=torch.float16, device="cuda"
|
||||
... )
|
||||
>>> key = torch.randn(
|
||||
... total_tokens, num_heads, head_dim, dtype=torch.float16, device="cuda"
|
||||
... )
|
||||
>>> value = torch.randn(
|
||||
... total_tokens, num_heads, head_dim, dtype=torch.float16, device="cuda"
|
||||
... )
|
||||
>>>
|
||||
>>> # Build cumulative sequence tensor
|
||||
>>> cu_seq = torch.zeros(batch_size + 1, device="cuda", dtype=torch.int32)
|
||||
>>> cu_seq[1:] = seq_lengths.cumsum(0)
|
||||
>>> max_len = seq_lengths.max().item()
|
||||
>>>
|
||||
>>> # Call varlen_attn
|
||||
>>> output = varlen_attn(
|
||||
... query, key, value, cu_seq, cu_seq, max_len, max_len
|
||||
... )
|
||||
"""
|
||||
|
||||
num_heads_q = query.size(1)
|
||||
num_heads_k = key.size(2) if block_table is not None else key.size(1)
|
||||
if not enable_gqa and num_heads_q != num_heads_k:
|
||||
raise ValueError(
|
||||
f"Expect query and key/value to have the same number of heads "
|
||||
f"but got Hq={num_heads_q} and Hkv={num_heads_k}. "
|
||||
f"Try setting enable_gqa=True for GQA."
|
||||
)
|
||||
if enable_gqa and num_heads_q % num_heads_k != 0:
|
||||
raise ValueError(
|
||||
f"Expect number of query heads to be a multiple of kv heads for GQA "
|
||||
f"but got Hq={num_heads_q} and Hkv={num_heads_k}."
|
||||
)
|
||||
|
||||
is_causal = window_size == (-1, 0)
|
||||
out, lse, _ = torch.ops.torch_attn._varlen_attn(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
cu_seq_q,
|
||||
cu_seq_k,
|
||||
max_q,
|
||||
max_k,
|
||||
is_causal,
|
||||
scale,
|
||||
list(window_size),
|
||||
enable_gqa,
|
||||
seqused_k,
|
||||
block_table,
|
||||
num_splits,
|
||||
)
|
||||
if return_aux is not None and return_aux.lse:
|
||||
return out, lse
|
||||
return out
|
||||
|
||||
|
||||
@torch.library.custom_op("torch_attn::_varlen_attn_out", mutates_args={"out"})
|
||||
def _varlen_attn_out(
|
||||
out: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
cu_seq_q: torch.Tensor,
|
||||
cu_seq_k: torch.Tensor | None,
|
||||
max_q: int,
|
||||
max_k: int,
|
||||
is_causal: bool = False,
|
||||
scale: float | None = None,
|
||||
window_size: list[int] | None = None,
|
||||
enable_gqa: bool = False,
|
||||
seqused_k: torch.Tensor | None = None,
|
||||
block_table: torch.Tensor | None = None,
|
||||
num_splits: int | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Private custom op for variable-length attention with pre-allocated output.
|
||||
Same as _varlen_attn but writes the attention output into the provided out tensor.
|
||||
"""
|
||||
window_size = _normalize_window_size(window_size)
|
||||
|
||||
use_cudnn = query.is_cuda and _should_use_cudnn(query.device.index)
|
||||
|
||||
if use_cudnn:
|
||||
# TODO: look into this
|
||||
raise RuntimeError("cuDNN backend does not support out variant.")
|
||||
|
||||
log.info("Using Flash Attention backend for varlen_attn_out")
|
||||
softmax_lse = torch.ops.aten._flash_attention_forward_no_dropout_inplace(
|
||||
out,
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
cu_seq_q,
|
||||
cu_seq_k,
|
||||
max_q,
|
||||
max_k,
|
||||
0.0, # dropout_p hardcoded to 0.0
|
||||
is_causal,
|
||||
False, # return_debug_mask
|
||||
scale=scale,
|
||||
window_size_left=window_size[0],
|
||||
window_size_right=window_size[1],
|
||||
seqused_k=seqused_k,
|
||||
block_table=block_table,
|
||||
num_splits=num_splits,
|
||||
)
|
||||
|
||||
return softmax_lse
|
||||
|
||||
|
||||
@_varlen_attn_out.register_fake
|
||||
def _varlen_attn_out_fake(
|
||||
out: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
cu_seq_q: torch.Tensor,
|
||||
cu_seq_k: torch.Tensor | None,
|
||||
max_q: int,
|
||||
max_k: int,
|
||||
is_causal: bool = False,
|
||||
scale: float | None = None,
|
||||
window_size: list[int] | None = None,
|
||||
enable_gqa: bool = False,
|
||||
seqused_k: torch.Tensor | None = None,
|
||||
block_table: torch.Tensor | None = None,
|
||||
num_splits: int | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Fake implementation for meta tensor computation and tracing.
|
||||
"""
|
||||
total_q = query.size(0)
|
||||
num_heads = query.size(1)
|
||||
logsumexp = torch.empty(
|
||||
(num_heads, total_q), dtype=torch.float, device=query.device
|
||||
)
|
||||
|
||||
if torch.version.hip:
|
||||
preferred = torch._C._get_rocm_fa_preferred_backend()
|
||||
if preferred == torch._C._ROCmFABackend.AOTriton:
|
||||
batch_size = cu_seq_q.size(0) - 1
|
||||
logsumexp = torch.empty(
|
||||
(batch_size, num_heads, max_q), dtype=torch.float, device=query.device
|
||||
)
|
||||
|
||||
return logsumexp
|
||||
|
||||
|
||||
def varlen_attn_out(
|
||||
out: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
cu_seq_q: torch.Tensor,
|
||||
cu_seq_k: torch.Tensor | None,
|
||||
max_q: int,
|
||||
max_k: int,
|
||||
*,
|
||||
return_aux: AuxRequest | None = None,
|
||||
scale: float | None = None,
|
||||
window_size: tuple[int, int] = (-1, -1),
|
||||
enable_gqa: bool = False,
|
||||
seqused_k: torch.Tensor | None = None,
|
||||
block_table: torch.Tensor | None = None,
|
||||
num_splits: int | None = None,
|
||||
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
|
||||
r"""Compute variable-length attention using Flash Attention with a pre-allocated output tensor.
|
||||
|
||||
Same as :func:`varlen_attn` but writes the attention output into the provided ``out`` tensor
|
||||
instead of allocating a new one.
|
||||
|
||||
"""
|
||||
num_heads_q = query.size(1)
|
||||
num_heads_k = key.size(2) if block_table is not None else key.size(1)
|
||||
if not enable_gqa and num_heads_q != num_heads_k:
|
||||
raise ValueError(
|
||||
f"Expect query and key/value to have the same number of heads "
|
||||
f"but got Hq={num_heads_q} and Hkv={num_heads_k}. "
|
||||
f"Try setting enable_gqa=True for GQA."
|
||||
)
|
||||
if enable_gqa and num_heads_q % num_heads_k != 0:
|
||||
raise ValueError(
|
||||
f"Expect number of query heads to be a multiple of kv heads for GQA "
|
||||
f"but got Hq={num_heads_q} and Hkv={num_heads_k}."
|
||||
)
|
||||
|
||||
is_causal = window_size == (-1, 0)
|
||||
lse = torch.ops.torch_attn._varlen_attn_out(
|
||||
out,
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
cu_seq_q,
|
||||
cu_seq_k,
|
||||
max_q,
|
||||
max_k,
|
||||
is_causal,
|
||||
scale,
|
||||
list(window_size),
|
||||
enable_gqa,
|
||||
seqused_k,
|
||||
block_table,
|
||||
num_splits,
|
||||
)
|
||||
if return_aux is not None and return_aux.lse:
|
||||
return out, lse
|
||||
return out
|
||||
|
||||
|
||||
def _setup_context(ctx: Any, inputs: tuple[Any, ...], output: Any) -> None:
|
||||
(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
cu_seq_q,
|
||||
cu_seq_k,
|
||||
max_q,
|
||||
max_k,
|
||||
is_causal,
|
||||
scale,
|
||||
window_size,
|
||||
enable_gqa,
|
||||
seqused_k,
|
||||
block_table,
|
||||
num_splits,
|
||||
) = inputs
|
||||
out, lse, rng_state = output
|
||||
|
||||
if seqused_k is not None:
|
||||
raise RuntimeError("seqused_k is an inference-only parameter.")
|
||||
if block_table is not None:
|
||||
raise RuntimeError("block_table is an inference-only parameter.")
|
||||
|
||||
ctx.save_for_backward(query, key, value, cu_seq_q, cu_seq_k, out, lse, rng_state)
|
||||
|
||||
ctx.max_q = max_q
|
||||
ctx.max_k = max_k
|
||||
ctx.is_causal = is_causal
|
||||
ctx.scale = scale
|
||||
ctx.window_size = window_size
|
||||
|
||||
|
||||
@torch.library.custom_op("torch_attn::_varlen_attn_backward", mutates_args={})
|
||||
def _varlen_attn_backward(
|
||||
grad_out: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
out: torch.Tensor,
|
||||
lse: torch.Tensor,
|
||||
cu_seq_q: torch.Tensor,
|
||||
cu_seq_k: torch.Tensor,
|
||||
max_q: int,
|
||||
max_k: int,
|
||||
is_causal: bool,
|
||||
rng_state: torch.Tensor,
|
||||
scale: float | None = None,
|
||||
window_size: list[int] | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
window_size = _normalize_window_size(window_size)
|
||||
|
||||
unused = torch.empty(0, device=query.device)
|
||||
|
||||
use_cudnn = query.is_cuda and _should_use_cudnn(query.device.index)
|
||||
if use_cudnn:
|
||||
log.info("Using cuDNN backend for varlen_attn")
|
||||
if window_size[0] != -1 or window_size[1] != -1:
|
||||
raise RuntimeError(
|
||||
"cuDNN backend does not support window attention. Please use Flash Attention backend."
|
||||
)
|
||||
dq, dk, dv = torch.ops.aten._cudnn_attention_backward(
|
||||
grad_out,
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
out,
|
||||
lse,
|
||||
cu_seq_q,
|
||||
cu_seq_k,
|
||||
max_q,
|
||||
max_k,
|
||||
0.0,
|
||||
is_causal,
|
||||
rng_state,
|
||||
unused,
|
||||
scale=scale,
|
||||
)
|
||||
else:
|
||||
log.info("Using Flash Attention backend for varlen_attn")
|
||||
dq, dk, dv = torch.ops.aten._flash_attention_backward(
|
||||
grad_out,
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
out,
|
||||
lse,
|
||||
cu_seq_q,
|
||||
cu_seq_k,
|
||||
max_q,
|
||||
max_k,
|
||||
0.0,
|
||||
is_causal,
|
||||
rng_state,
|
||||
unused,
|
||||
scale=scale,
|
||||
window_size_left=window_size[0],
|
||||
window_size_right=window_size[1],
|
||||
)
|
||||
return dq, dk, dv
|
||||
|
||||
|
||||
@_varlen_attn_backward.register_fake
|
||||
def _varlen_attn_backward_fake(
|
||||
grad_out: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
out: torch.Tensor,
|
||||
lse: torch.Tensor,
|
||||
cu_seq_q: torch.Tensor,
|
||||
cu_seq_k: torch.Tensor,
|
||||
max_q: int,
|
||||
max_k: int,
|
||||
is_causal: bool,
|
||||
rng_state: torch.Tensor,
|
||||
scale: float | None = None,
|
||||
window_size: list[int] | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Fake implementation for meta tensor computation and tracing.
|
||||
"""
|
||||
window_size = _normalize_window_size(window_size)
|
||||
|
||||
grad_query = torch.empty_like(query)
|
||||
grad_key = torch.empty_like(key)
|
||||
grad_value = torch.empty_like(value)
|
||||
|
||||
return grad_query, grad_key, grad_value
|
||||
|
||||
|
||||
def _backward(
|
||||
ctx: Any, grad_out: torch.Tensor, grad_lse: torch.Tensor, grad_rng: torch.Tensor
|
||||
) -> tuple[torch.Tensor | None, ...]:
|
||||
query, key, value, cu_seq_q, cu_seq_k, out, lse, rng_state = ctx.saved_tensors
|
||||
|
||||
max_q = ctx.max_q
|
||||
max_k = ctx.max_k
|
||||
is_causal = ctx.is_causal
|
||||
scale = ctx.scale
|
||||
window_size = ctx.window_size
|
||||
|
||||
dq, dk, dv = torch.ops.torch_attn._varlen_attn_backward(
|
||||
grad_out,
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
out,
|
||||
lse,
|
||||
cu_seq_q,
|
||||
cu_seq_k,
|
||||
max_q,
|
||||
max_k,
|
||||
is_causal,
|
||||
rng_state,
|
||||
scale,
|
||||
window_size,
|
||||
)
|
||||
# cu_seq_q, cu_seq_k, max_q, max_k, is_causal, scale, window_size, \
|
||||
# enable_gqa, seqused_k, block_table, num_splits
|
||||
num_params = 11
|
||||
return (dq, dk, dv, *((None,) * num_params))
|
||||
|
||||
|
||||
_varlen_attn.register_autograd(_backward, setup_context=_setup_context)
|
||||
|
||||
torch._dynamo.disallow_in_graph(
|
||||
torch.ops.aten._flash_attention_forward_no_dropout_inplace
|
||||
)
|
||||
|
||||
from torch.utils.flop_counter import (
|
||||
_varlen_attn_backward_flop,
|
||||
_varlen_attn_forward_flop,
|
||||
_varlen_attn_out_flop,
|
||||
flop_registry,
|
||||
)
|
||||
|
||||
|
||||
flop_registry[torch.ops.torch_attn._varlen_attn] = _varlen_attn_forward_flop
|
||||
flop_registry[torch.ops.torch_attn._varlen_attn_out] = _varlen_attn_out_flop
|
||||
flop_registry[torch.ops.torch_attn._varlen_attn_backward] = _varlen_attn_backward_flop
|
||||
Reference in New Issue
Block a user