diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000000000000000000000000000000000000..0cd58331b2a989b68be4ec5676383437fca8687b --- /dev/null +++ b/.gitattributes @@ -0,0 +1,36 @@ +*.7z filter=lfs diff=lfs merge=lfs -text +*.arrow filter=lfs diff=lfs merge=lfs -text +*.bin filter=lfs diff=lfs merge=lfs -text +*.bz2 filter=lfs diff=lfs merge=lfs -text +*.ckpt filter=lfs diff=lfs merge=lfs -text +*.ftz filter=lfs diff=lfs merge=lfs -text +*.gz filter=lfs diff=lfs merge=lfs -text +*.h5 filter=lfs diff=lfs merge=lfs -text +*.joblib filter=lfs diff=lfs merge=lfs -text +*.lfs.* filter=lfs diff=lfs merge=lfs -text +*.mlmodel filter=lfs diff=lfs merge=lfs -text +*.model filter=lfs diff=lfs merge=lfs -text +*.msgpack filter=lfs diff=lfs merge=lfs -text +*.npy filter=lfs diff=lfs merge=lfs -text +*.npz filter=lfs diff=lfs merge=lfs -text +*.onnx filter=lfs diff=lfs merge=lfs -text +*.ot filter=lfs diff=lfs merge=lfs -text +*.parquet filter=lfs diff=lfs merge=lfs -text +*.pb filter=lfs diff=lfs merge=lfs -text +*.pickle filter=lfs diff=lfs merge=lfs -text +*.pkl filter=lfs diff=lfs merge=lfs -text +*.pt filter=lfs diff=lfs merge=lfs -text +*.pth filter=lfs diff=lfs merge=lfs -text +*.rar filter=lfs diff=lfs merge=lfs -text +*.safetensors filter=lfs diff=lfs merge=lfs -text +saved_model/**/* filter=lfs diff=lfs merge=lfs -text +*.tar.* filter=lfs diff=lfs merge=lfs -text +*.tar filter=lfs diff=lfs merge=lfs -text +*.tflite filter=lfs diff=lfs merge=lfs -text +*.tgz filter=lfs diff=lfs merge=lfs -text +*.wasm filter=lfs diff=lfs merge=lfs -text +*.xz filter=lfs diff=lfs merge=lfs -text +*.zip filter=lfs diff=lfs merge=lfs -text +*.zst filter=lfs diff=lfs merge=lfs -text +*tfevents* filter=lfs diff=lfs merge=lfs -text +*.so filter=lfs diff=lfs merge=lfs -text diff --git a/README.md b/README.md new file mode 100644 index 0000000000000000000000000000000000000000..b06cb814449e36d8253bf34d9b4a5e69b3ebcb6b --- /dev/null +++ b/README.md @@ -0,0 +1,55 @@ +--- +library_name: kernels +license: apache-2.0 +--- + + + + +This is the repository card of {repo_id} that has been pushed on the Hub. It was built to be used with the [`kernels` library](https://github.com/huggingface/kernels). This card was automatically generated. + + +## How to use + +```python +# make sure `kernels` is installed: `pip install -U kernels` +from kernels import get_kernel + +kernel_module = get_kernel("kernels-community/flash-attn3") # <- change the ID if needed +flash_attn_combine = kernel_module.flash_attn_combine + +flash_attn_combine(...) +``` + +## Available functions + +- `flash_attn_combine` +- `flash_attn_func` +- `flash_attn_qkvpacked_func` +- `flash_attn_varlen_func` +- `flash_attn_with_kvcache` +- `get_scheduler_metadata` + +## Supported backends + +- cuda + +## CUDA Capabilities + +- 8.0 +- 9.0a + +## Benchmarks + +Benchmarking script is available for this kernel. Make sure to run `kernels benchmark org-id/repo-id` (replace "org-id" and "repo-id" with actual values). + +[TODO: provide benchmarks if available] + +## Source code + +[TODO: provide original source code and other relevant citations if available] + +## Notes + +[TODO: provide additional notes about this kernel if needed] \ No newline at end of file diff --git a/benchmark.py b/benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..d7b48ccbd63561f4f1e903fc5c7cb8fb864e003f --- /dev/null +++ b/benchmark.py @@ -0,0 +1,17 @@ +from kernels.benchmarks import ( + FlashAttentionBenchmark, + FlashAttentionCausalBenchmark, + FlashAttentionVarlenBenchmark, +) + + +class FlashAttn(FlashAttentionBenchmark): + pass + + +class FlashAttnCausal(FlashAttentionCausalBenchmark): + pass + + +class FlashAttnVarlen(FlashAttentionVarlenBenchmark): + pass diff --git a/benchmarks/benchmark.py b/benchmarks/benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..d7b48ccbd63561f4f1e903fc5c7cb8fb864e003f --- /dev/null +++ b/benchmarks/benchmark.py @@ -0,0 +1,17 @@ +from kernels.benchmarks import ( + FlashAttentionBenchmark, + FlashAttentionCausalBenchmark, + FlashAttentionVarlenBenchmark, +) + + +class FlashAttn(FlashAttentionBenchmark): + pass + + +class FlashAttnCausal(FlashAttentionCausalBenchmark): + pass + + +class FlashAttnVarlen(FlashAttentionVarlenBenchmark): + pass diff --git a/build/torch210-cxx11-cu128-x86_64-linux/__init__.py b/build/torch210-cxx11-cu128-x86_64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3d12edd447e84d6e65c4fa9ab91ae3c3192884c1 --- /dev/null +++ b/build/torch210-cxx11-cu128-x86_64-linux/__init__.py @@ -0,0 +1,17 @@ +from .flash_attn_interface import ( + flash_attn_combine, + flash_attn_func, + flash_attn_qkvpacked_func, + flash_attn_varlen_func, + flash_attn_with_kvcache, + get_scheduler_metadata, +) + +__all__ = [ + "flash_attn_combine", + "flash_attn_func", + "flash_attn_qkvpacked_func", + "flash_attn_varlen_func", + "flash_attn_with_kvcache", + "get_scheduler_metadata", +] diff --git a/build/torch210-cxx11-cu128-x86_64-linux/_flash_attn3_cuda_e1d5be2.abi3.so b/build/torch210-cxx11-cu128-x86_64-linux/_flash_attn3_cuda_e1d5be2.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..a60c24d41b1b793a99e813c2b8047c60eb3b1759 --- /dev/null +++ b/build/torch210-cxx11-cu128-x86_64-linux/_flash_attn3_cuda_e1d5be2.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c778c8063c340a4dfd13a5a4f372722dcef8f8b758af149eca664689d9c54847 +size 804191136 diff --git a/build/torch210-cxx11-cu128-x86_64-linux/_ops.py b/build/torch210-cxx11-cu128-x86_64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..b845f529a7b7c424bbe0db01a32c2cd6db45bb1a --- /dev/null +++ b/build/torch210-cxx11-cu128-x86_64-linux/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _flash_attn3_cuda_e1d5be2 +ops = torch.ops._flash_attn3_cuda_e1d5be2 + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_flash_attn3_cuda_e1d5be2::{op_name}" diff --git a/build/torch210-cxx11-cu128-x86_64-linux/flash_attn3/__init__.py b/build/torch210-cxx11-cu128-x86_64-linux/flash_attn3/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a9b2672c1cd85b74c1b3ded0fc0b2100e1aeac23 --- /dev/null +++ b/build/torch210-cxx11-cu128-x86_64-linux/flash_attn3/__init__.py @@ -0,0 +1,26 @@ +import ctypes +import importlib.util +import sys +from pathlib import Path +from types import ModuleType + + +def _import_from_path(file_path: Path) -> ModuleType: + # We cannot use the module name as-is, after adding it to `sys.modules`, + # it would also be used for other imports. So, we make a module name that + # depends on the path for it to be unique using the hex-encoded hash of + # the path. + path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value) + module_name = path_hash + spec = importlib.util.spec_from_file_location(module_name, file_path) + if spec is None: + raise ImportError(f"Cannot load spec for {module_name} from {file_path}") + module = importlib.util.module_from_spec(spec) + if module is None: + raise ImportError(f"Cannot load module {module_name} from spec") + sys.modules[module_name] = module + spec.loader.exec_module(module) # type: ignore + return module + + +globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py"))) diff --git a/build/torch210-cxx11-cu128-x86_64-linux/flash_attn_config.py b/build/torch210-cxx11-cu128-x86_64-linux/flash_attn_config.py new file mode 100644 index 0000000000000000000000000000000000000000..c4b7130324be2c4cabee7ada7000e170a2a81865 --- /dev/null +++ b/build/torch210-cxx11-cu128-x86_64-linux/flash_attn_config.py @@ -0,0 +1,7 @@ +# Auto-generated by flash attention 3 setup.py +CONFIG = {'build_flags': {'FLASHATTENTION_DISABLE_BACKWARD': False, 'FLASHATTENTION_DISABLE_SPLIT': False, 'FLASHATTENTION_DISABLE_PAGEDKV': False, 'FLASHATTENTION_DISABLE_APPENDKV': False, 'FLASHATTENTION_DISABLE_LOCAL': False, 'FLASHATTENTION_DISABLE_SOFTCAP': False, 'FLASHATTENTION_DISABLE_PACKGQA': False, 'FLASHATTENTION_DISABLE_FP16': False, 'FLASHATTENTION_DISABLE_FP8': False, 'FLASHATTENTION_DISABLE_VARLEN': False, 'FLASHATTENTION_DISABLE_CLUSTER': False, 'FLASHATTENTION_DISABLE_HDIM64': False, 'FLASHATTENTION_DISABLE_HDIM96': False, 'FLASHATTENTION_DISABLE_HDIM128': False, 'FLASHATTENTION_DISABLE_HDIM192': False, 'FLASHATTENTION_DISABLE_HDIM256': False, 'FLASHATTENTION_DISABLE_SM8x': False, 'FLASHATTENTION_ENABLE_VCOLMAJOR': False, 'FLASH_ATTENTION_DISABLE_HDIMDIFF64': False, 'FLASH_ATTENTION_DISABLE_HDIMDIFF192': False}} + +def show(): + from pprint import pprint + pprint(CONFIG) + diff --git a/build/torch210-cxx11-cu128-x86_64-linux/flash_attn_interface.py b/build/torch210-cxx11-cu128-x86_64-linux/flash_attn_interface.py new file mode 100644 index 0000000000000000000000000000000000000000..2e2a3cdd3e89cbc59bf6c956454e2776c305e0e5 --- /dev/null +++ b/build/torch210-cxx11-cu128-x86_64-linux/flash_attn_interface.py @@ -0,0 +1,1127 @@ +# Copyright (c) 2023, Tri Dao. + +from typing import Optional, Union, List, Tuple + +import torch +import torch.nn as nn + +from ._ops import ops as flash_attn_3_cuda +from ._ops import add_op_namespace_prefix + +def maybe_contiguous(x): + return x.contiguous() if x is not None and x.stride(-1) != 1 else x + + +def round_multiple(x, m): + return (x + m - 1) // m * m + + +def round_up_headdim(head_size: int) -> int: + from .flash_attn_config import CONFIG + + if not CONFIG["build_flags"]["FLASHATTENTION_DISABLE_HDIM64"]: + if head_size <= 64: + return 64 + if not CONFIG["build_flags"]["FLASHATTENTION_DISABLE_HDIM96"]: + if head_size <= 96: + return 96 + if not CONFIG["build_flags"]["FLASHATTENTION_DISABLE_HDIM128"]: + if head_size <= 128: + return 128 + if not CONFIG["build_flags"]["FLASHATTENTION_DISABLE_HDIM192"]: + if head_size <= 192: + return 192 + if not CONFIG["build_flags"]["FLASHATTENTION_DISABLE_HDIM256"]: + if head_size <= 256: + return 256 + return 256 + + +@torch.library.custom_op(add_op_namespace_prefix("_flash_attn_forward"), mutates_args=(), device_types="cuda") +def _flash_attn_forward( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_new: Optional[torch.Tensor] = None, + v_new: Optional[torch.Tensor] = None, + qv: Optional[torch.Tensor] = None, + out_: Optional[torch.Tensor] = None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + seqused_q: Optional[torch.Tensor] = None, + seqused_k: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_k: Optional[int] = None, + page_table: Optional[torch.Tensor] = None, + kv_batch_idx: Optional[torch.Tensor] = None, + leftpad_k: Optional[torch.Tensor] = None, + rotary_cos: Optional[torch.Tensor] = None, + rotary_sin: Optional[torch.Tensor] = None, + seqlens_rotary: Optional[torch.Tensor] = None, + q_descale: Optional[torch.Tensor] = None, + k_descale: Optional[torch.Tensor] = None, + v_descale: Optional[torch.Tensor] = None, + softmax_scale: Optional[float] = None, + causal: bool = False, + window_size_left: int = -1, + window_size_right: int = -1, + attention_chunk: int = 0, + softcap: float = 0.0, + rotary_interleaved: bool = True, + scheduler_metadata: Optional[torch.Tensor] = None, + num_splits: int = 1, + pack_gqa: Optional[bool] = None, + sm_margin: int = 0, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + q, k, k_new, v_new = [maybe_contiguous(x) for x in (q, k, k_new, v_new)] + v = v.contiguous() if v.stride(-1) != 1 and v.stride(-3) != 1 else v + cu_seqlens_q, cu_seqlens_k, cu_seqlens_k_new = [ + maybe_contiguous(x) for x in (cu_seqlens_q, cu_seqlens_k, cu_seqlens_k_new) + ] + seqused_q, seqused_k = [maybe_contiguous(x) for x in (seqused_q, seqused_k)] + page_table, kv_batch_idx, leftpad_k = [ + maybe_contiguous(x) for x in (page_table, kv_batch_idx, leftpad_k) + ] + rotary_cos, rotary_sin = [maybe_contiguous(x) for x in (rotary_cos, rotary_sin)] + seqlens_rotary = maybe_contiguous(seqlens_rotary) + out, softmax_lse, out_accum, softmax_lse_accum = flash_attn_3_cuda.fwd( + q, + k, + v, + k_new, + v_new, + qv, + out_, + cu_seqlens_q, + cu_seqlens_k, + cu_seqlens_k_new, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + page_table, + kv_batch_idx, + leftpad_k, + rotary_cos, + rotary_sin, + seqlens_rotary, + q_descale, + k_descale, + v_descale, + softmax_scale, + causal, + window_size_left, + window_size_right, + attention_chunk, + softcap, + rotary_interleaved, + scheduler_metadata, + num_splits, + pack_gqa, + sm_margin, + ) + + if out_accum is None: + out_accum = torch.tensor([], device=out.device) + + if softmax_lse_accum is None: + softmax_lse_accum = torch.tensor([], device=out.device) + + return out, softmax_lse, out_accum, softmax_lse_accum + + +@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_forward")) +def _flash_attn_forward_fake( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_new: Optional[torch.Tensor] = None, + v_new: Optional[torch.Tensor] = None, + qv: Optional[torch.Tensor] = None, + out_: Optional[torch.Tensor] = None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + seqused_q: Optional[torch.Tensor] = None, + seqused_k: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_k: Optional[int] = None, + page_table: Optional[torch.Tensor] = None, + kv_batch_idx: Optional[torch.Tensor] = None, + leftpad_k: Optional[torch.Tensor] = None, + rotary_cos: Optional[torch.Tensor] = None, + rotary_sin: Optional[torch.Tensor] = None, + seqlens_rotary: Optional[torch.Tensor] = None, + q_descale: Optional[torch.Tensor] = None, + k_descale: Optional[torch.Tensor] = None, + v_descale: Optional[torch.Tensor] = None, + softmax_scale: Optional[float] = None, + causal: bool = False, + window_size_left: int = -1, + window_size_right: int = -1, + attention_chunk: int = 0, + softcap: float = 0.0, + rotary_interleaved: bool = True, + scheduler_metadata: Optional[torch.Tensor] = None, + num_splits: int = 1, + pack_gqa: Optional[bool] = None, + sm_margin: int = 0, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Symbolic fake implementation of flash attention forward. + Returns tensors with the correct shapes and dtypes without actual computation. + """ + + # Determine if we're in varlen mode + is_varlen_q = cu_seqlens_q is not None + + # Get dimensions from query tensor + if is_varlen_q: + # varlen mode: q is (total_q, num_heads, head_size) + total_q, num_heads, head_size = q.shape + batch_size = cu_seqlens_q.shape[0] - 1 + + if max_seqlen_q is None: + raise ValueError("max_seqlen_q must be provided if cu_seqlens_q is provided") + seqlen_q = max_seqlen_q + else: + # batch mode: q is (batch_size, seqlen_q, num_heads, head_size) + batch_size, seqlen_q, num_heads, head_size = q.shape + total_q = batch_size * q.shape[1] + # Get value head dimension + head_size_v = v.shape[-1] + + # Determine output dtype (FP8 inputs produce BF16 outputs) + q_type = q.dtype + if q_type == torch.float8_e4m3fn: + out_dtype = torch.bfloat16 + else: + out_dtype = q_type + + # Create output tensor + if out_ is not None: + # If out_ is provided, _flash_attn_forward becomes non-functional + raise TypeError("Tracing (torch.compile/torch.export) with pre-allocated output tensor is not supported.") + + if is_varlen_q: + out = torch.empty((total_q, num_heads, head_size_v), dtype=out_dtype, device=q.device) + else: + out = torch.empty((batch_size, seqlen_q, num_heads, head_size_v), dtype=out_dtype, device=q.device) + + # Create softmax_lse tensor + if is_varlen_q: + softmax_lse = torch.empty((num_heads, total_q), dtype=torch.float32, device=q.device) + else: + softmax_lse = torch.empty((batch_size, num_heads, seqlen_q), dtype=torch.float32, device=q.device) + + # TODO(guilhermeleobas): Implement "get_num_splits" + # There's an heuristic to compute num_splits when "num_splits <= 0" + # assert that num_splits is > 0 for now + if num_splits <= 0: + raise ValueError(f"tracing (torch.compile/torch.export) with num_splits <= 0 not supported. Got {num_splits=}") + + if num_splits > 1: + if is_varlen_q: + out_accum = torch.empty((num_splits, num_heads, total_q, head_size_v), dtype=torch.float32, device=q.device) + softmax_lse_accum = torch.empty((num_splits, num_heads, total_q), dtype=torch.float32, device=q.device) + else: + out_accum = torch.empty((num_splits, batch_size, num_heads, seqlen_q, head_size_v), dtype=torch.float32, device=q.device) + softmax_lse_accum = torch.empty((num_splits, batch_size, num_heads, seqlen_q), dtype=torch.float32, device=q.device) + else: + # Tensors are not set when num_splits < 1 + out_accum = torch.tensor([], device=out.device) + softmax_lse_accum = torch.tensor([], device=out.device) + + return out, softmax_lse, out_accum, softmax_lse_accum + + +@torch.library.custom_op(add_op_namespace_prefix("_flash_attn_backward"), mutates_args=("dq", "dk", "dv"), device_types="cuda") +def _flash_attn_backward( + dout: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + out: torch.Tensor, + softmax_lse: torch.Tensor, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k: Optional[torch.Tensor] = None, + sequed_q: Optional[torch.Tensor] = None, + sequed_k: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_k: Optional[int] = None, + dq: Optional[torch.Tensor] = None, + dk: Optional[torch.Tensor] = None, + dv: Optional[torch.Tensor] = None, + softmax_scale: Optional[float] = None, + is_causal: bool = False, + window_size_left: int = -1, + window_size_right: int = -1, + softcap: float = 0.0, + deterministic: bool = False, + sm_margin: int = 0, +) -> torch.Tensor: + # dq, dk, dv are allocated by us so they should already be contiguous + dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] + softmax_d, *rest = flash_attn_3_cuda.bwd( + dout, + q, + k, + v, + out, + softmax_lse, + dq, + dk, + dv, + cu_seqlens_q, + cu_seqlens_k, + sequed_q, + sequed_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + is_causal, + window_size_left, + window_size_right, + softcap, + deterministic, + sm_margin, + ) + return softmax_d + + +@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_backward")) +def _flash_attn_backward_fake( + dout: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + out: torch.Tensor, + softmax_lse: torch.Tensor, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k: Optional[torch.Tensor] = None, + sequed_q: Optional[torch.Tensor] = None, + sequed_k: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_k: Optional[int] = None, + dq: Optional[torch.Tensor] = None, + dk: Optional[torch.Tensor] = None, + dv: Optional[torch.Tensor] = None, + softmax_scale: Optional[float] = None, + is_causal: bool = False, + window_size_left: int = -1, + window_size_right: int = -1, + softcap: float = 0.0, + deterministic: bool = False, + sm_margin: int = 0, +) -> torch.Tensor: + + is_varlen_q = cu_seqlens_q is not None + is_varlen_k = cu_seqlens_q is not None + is_varlen = is_varlen_q or is_varlen_k or sequed_q is not None or sequed_k is not None + + if not is_varlen_q: + batch_size = q.size(0) + seqlen_q = q.size(1) + seqlen_k = k.size(1) + total_q = batch_size * q.size(1) + else: + batch_size = cu_seqlens_q.size(0) - 1 + total_q = q.size(0) + seqlen_q = max_seqlen_q + seqlen_k = max_seqlen_k + + if window_size_left >= seqlen_k - 1: + window_size_left = -1 + + if window_size_right >= seqlen_q - 1: + window_size_right = -1 + + if is_causal: + window_size_right = 0 + + is_causal = window_size_left < 0 and window_size_right == 0 + + head_size = q.size(-1) + head_size_v = v.size(-1) + head_size_rounded = round_up_headdim(max(head_size, head_size_v)) + + # Hopper gpus uses cuda compute capabilities 9.0 + cap = torch.cuda.get_device_capability(q.device) + arch = cap[0] * 10 + cap[1] + + is_local = (window_size_left >= 0 or window_size_right >= 0) and not is_causal + + if head_size_rounded <= 64: + kBlockM_sm90 = 96 if (is_causal and softcap > 0.0) else 128 + elif head_size_rounded <= 96: + kBlockM_sm90 = 64 + elif head_size_rounded <= 128: + kBlockM_sm90 = 64 if (is_causal or is_local or softcap > 0.0) else 80 + else: + kBlockM_sm90 = 64 + + kBlockM_sm80 = 128 if head_size_rounded <= 64 else 64 + kBlockM_sm86 = 64 if head_size_rounded <= 192 else 32 + + if arch >= 90: + kBlockM = kBlockM_sm90 + elif arch == 86 or arch == 89: + kBlockM = kBlockM_sm86 + else: + kBlockM = kBlockM_sm80 + + num_heads = q.shape[-2] + seqlen_q_rounded = round_multiple(seqlen_q, kBlockM) + + total_q_padded_rounded = round_multiple(total_q + batch_size * kBlockM, kBlockM) + + dq = torch.empty_like(q) if dq is None else dq + dk = torch.empty_like(k) if dk is None else dk + dv = torch.empty_like(v) if dv is None else dv + + if not is_varlen: + softmax_d = torch.empty((batch_size, num_heads, seqlen_q_rounded), dtype=torch.float32, device=q.device) + else: + softmax_d = torch.empty((num_heads, total_q_padded_rounded), dtype=torch.float32, device=q.device) + + return softmax_d + + +def setup_context(ctx, inputs, output): + q, k, v = inputs[:3] + out, softmax_lse, _, _ = output + ctx.save_for_backward(q, k, v, out, softmax_lse) + ctx.softmax_scale = inputs[-11] + ctx.causal = inputs[-10] + ctx.window_size = [inputs[-9], inputs[-8]] + ctx.attention_chunk = inputs[-7] + ctx.softcap = inputs[-6] + ctx.sm_margin = inputs[-1] + + +def _backward(ctx, dout, *grads): + q, k, v, out, softmax_lse = ctx.saved_tensors + dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + None, None, # cu_seqlens_q, cu_seqlens_k, + None, None, # sequed_q, sequed_k, + None, None, # max_seqlen_q, max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size[0], + ctx.window_size[1], + ctx.softcap, + False, # deterministic + ctx.sm_margin, + ) + return dq, dk, dv, *((None,) * 21) + + +_flash_attn_forward.register_autograd(_backward, setup_context=setup_context) + + + +class FlashAttnQKVPackedFunc(torch.autograd.Function): + @staticmethod + def forward( + ctx, + qkv, + softmax_scale, + causal, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + deterministic=False, + num_heads_q=None, + sm_margin=0, + return_softmax=False, + ): + if softmax_scale is None: + softmax_scale = qkv.shape[-1] ** (-0.5) + if qkv.dim() == 5: + assert qkv.shape[-3] == 3 + q, k, v = qkv.unbind(dim=-3) + else: + assert qkv.dim() == 4 + assert num_heads_q is not None + num_heads_k = (qkv.shape[2] - num_heads_q) // 2 + assert num_heads_k * 2 + num_heads_q == qkv.shape[2] + q, k, v = qkv.split([num_heads_q, num_heads_k, num_heads_k], dim=-2) + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + None, # qv + None, # out + None, None, None, # cu_seqlens_q/k/k_new + None, None, # seqused_q/k + None, None, # max_seqlen_q/k + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size_left=window_size[0], + window_size_right=window_size[1], + attention_chunk=attention_chunk, + softcap=softcap, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse) + ctx.save_for_backward(q, k, v, out, softmax_lse) + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.ndim = qkv.dim() + ctx.sm_margin = sm_margin + return (out, softmax_lse) if return_softmax else out + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + if ctx.ndim == 5: + qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) + dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) + dq, dk, dv = dqkv.unbind(dim=-3) + else: + num_heads_q = q.shape[2] + num_heads_k = k.shape[2] + qkv_shape = q.shape[:-2] + (num_heads_q + num_heads_k * 2, *q.shape[-1:]) + dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) + dq, dk, dv = dqkv.split([num_heads_q, num_heads_k, num_heads_k], dim=-2) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + None, None, # cu_seqlens_q, cu_seqlens_k, + None, None, # sequed_q, sequed_k, + None, None, # max_seqlen_q, max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size[0], + ctx.window_size[1], + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dqkv = dqkv[..., : dout.shape[-1]] # We could have padded the head dimension + return dqkv, None, None, None, None, None, None, None, None, None, None, None, None + + +class FlashAttnFunc(torch.autograd.Function): + + @staticmethod + def forward( + ctx, + q, + k, + v, + softmax_scale, + causal, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + return_softmax=False, + ): + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + # out, q, k, v, out_padded, softmax_lse = _flash_attn_forward( + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + qv, # qv + None, # out + None, None, None, # cu_seqlens_q/k/k_new + None, None, # seqused_q/k + None, None, # max_seqlen_q/k + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size_left=window_size[0], + window_size_right=window_size[1], + attention_chunk=attention_chunk, + softcap=softcap, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse) + ctx.save_for_backward(q, k, v, out, softmax_lse) + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.sm_margin = sm_margin + return (out, softmax_lse) if return_softmax else out + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + None, None, # cu_seqlens_q, cu_seqlens_k, + None, None, # sequed_q, sequed_k, + None, None, # max_seqlen_q, max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size[0], + ctx.window_size[1], + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dq = dq[..., : q.shape[-1]] # We could have padded the head dimension + dk = dk[..., : k.shape[-1]] + dv = dv[..., : v.shape[-1]] + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None + + +class FlashAttnVarlenFunc(torch.autograd.Function): + + @staticmethod + def forward( + ctx, + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + return_softmax=False, + ): + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + # out, q, k, v, out_padded, softmax_lse = _flash_attn_varlen_forward( + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + qv, # qv + None, # out + cu_seqlens_q, + cu_seqlens_k, + None, # cu_seqlens_k_new + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size_left=window_size[0], + window_size_right=window_size[1], + attention_chunk=attention_chunk, + softcap=softcap, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k) + ctx.save_for_backward(q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k) + ctx.max_seqlen_q = max_seqlen_q + ctx.max_seqlen_k = max_seqlen_k + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.sm_margin = sm_margin + return (out, softmax_lse) if return_softmax else out + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + ctx.max_seqlen_q, + ctx.max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size[0], + ctx.window_size[1], + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dq = dq[..., : q.shape[-1]] # We could have padded the head dimension + dk = dk[..., : k.shape[-1]] + dv = dv[..., : v.shape[-1]] + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None + + +def flash_attn_qkvpacked_func( + qkv, + softmax_scale=None, + causal=False, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + deterministic=False, + num_heads_q=None, + sm_margin=0, + return_attn_probs=False, +): + """dropout_p should be set to 0.0 during evaluation + If Q, K, V are already stacked into 1 tensor, this function will be faster than + calling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation + of the gradients of Q, K, V. + For multi-query and grouped-query attention (MQA/GQA), please see + flash_attn_kvpacked_func and flash_attn_func. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between [i - window_size[0], i + window_size[1]] inclusive. + + Arguments: + qkv: (batch_size, seqlen, 3, nheads, headdim) + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + softcap: float. Anything > 0 activates softcapping attention. + alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of (-alibi_slope * |i - j|) is added to + the attention score of query i and key j. + deterministic: bool. Whether to use the deterministic implementation of the backward pass, + which is slightly slower and uses more memory. The forward pass is always deterministic. + return_attn_probs: bool. Whether to return the attention probabilities. This option is for + testing only. The returned probabilities are not guaranteed to be correct + (they might not have the right scaling). + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen). + The output of softmax (possibly with different scaling). It also encodes the dropout + pattern (negative means that location was dropped, nonnegative means it was kept). + """ + return FlashAttnQKVPackedFunc.apply( + qkv, + softmax_scale, + causal, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + deterministic, + num_heads_q, + sm_margin, + return_attn_probs, + ) + + +def flash_attn_func( + q, + k, + v, + softmax_scale=None, + causal=False, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + return_attn_probs=False, +): + """dropout_p should be set to 0.0 during evaluation + Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads + than Q. Note that the number of heads in Q must be divisible by the number of heads in KV. + For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head + 0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V. + + If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix. + For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is: + 1 1 1 1 0 + 1 1 1 1 1 + If seqlen_q = 5 and seqlen_k = 2, the causal mask is: + 0 0 + 0 0 + 0 0 + 1 0 + 1 1 + If the row of the mask is all zero, the output will be zero. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between + [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive. + + Arguments: + q: (batch_size, seqlen, nheads, headdim) + k: (batch_size, seqlen, nheads_k, headdim) + v: (batch_size, seqlen, nheads_k, headdim) + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of + (-alibi_slope * |i + seqlen_k - seqlen_q - j|) + is added to the attention score of query i and key j. + deterministic: bool. Whether to use the deterministic implementation of the backward pass, + which is slightly slower and uses more memory. The forward pass is always deterministic. + return_attn_probs: bool. Whether to return the attention probabilities. This option is for + testing only. The returned probabilities are not guaranteed to be correct + (they might not have the right scaling). + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + """ + return FlashAttnFunc.apply( + q, + k, + v, + softmax_scale, + causal, + qv, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + num_splits, + pack_gqa, + deterministic, + sm_margin, + return_attn_probs, + ) + + +def flash_attn_varlen_func( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + seqused_q=None, + seqused_k=None, + softmax_scale=None, + causal=False, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + return_attn_probs=False, +): + return FlashAttnVarlenFunc.apply( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + qv, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + num_splits, + pack_gqa, + deterministic, + sm_margin, + return_attn_probs, + ) + + +def flash_attn_combine(out_partial, lse_partial, out=None, out_dtype=None): + return flash_attn_3_cuda.fwd_combine(out_partial, lse_partial, out, out_dtype) + + +def flash_attn_with_kvcache( + q, + k_cache, + v_cache, + k=None, + v=None, + qv=None, + rotary_cos=None, + rotary_sin=None, + cache_seqlens: Optional[Union[(int, torch.Tensor)]] = None, + cache_batch_idx: Optional[torch.Tensor] = None, + cache_leftpad: Optional[torch.Tensor] = None, + page_table: Optional[torch.Tensor] = None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + rotary_seqlens: Optional[torch.Tensor] = None, + q_descale: Optional[torch.Tensor] = None, + k_descale: Optional[torch.Tensor] = None, + v_descale: Optional[torch.Tensor] = None, + softmax_scale=None, + causal=False, + window_size=(-1, -1), # -1 means infinite context window + attention_chunk=0, + softcap=0.0, # 0.0 means deactivated + rotary_interleaved=True, + scheduler_metadata=None, + num_splits=0, # Can be tuned for speed + pack_gqa=None, # Can be tuned for speed + sm_margin=0, # Can be tuned if some SMs are used for communication + return_softmax_lse=False, +): + """ + If k and v are not None, k_cache and v_cache will be updated *inplace* with the new values from + k and v. This is useful for incremental decoding: you can pass in the cached keys/values from + the previous step, and update them with the new keys/values from the current step, and do + attention with the updated cache, all in 1 kernel. + + If you pass in k / v, you must make sure that the cache is large enough to hold the new values. + For example, the KV cache could be pre-allocated with the max sequence length, and you can use + cache_seqlens to keep track of the current sequence lengths of each sequence in the batch. + + Also apply rotary embedding if rotary_cos and rotary_sin are passed in. The key @k will be + rotated by rotary_cos and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc. + If causal or local (i.e., window_size != (-1, -1)), the query @q will be rotated by rotary_cos + and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc. + If not causal and not local, the query @q will be rotated by rotary_cos and rotary_sin at + indices cache_seqlens only (i.e. we consider all tokens in @q to be at position cache_seqlens). + + See tests/test_flash_attn.py::test_flash_attn_kvcache for examples of how to use this function. + + Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads + than Q. Note that the number of heads in Q must be divisible by the number of heads in KV. + For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head + 0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V. + + If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix. + For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is: + 1 1 1 1 0 + 1 1 1 1 1 + If seqlen_q = 5 and seqlen_k = 2, the causal mask is: + 0 0 + 0 0 + 0 0 + 1 0 + 1 1 + If the row of the mask is all zero, the output will be zero. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between + [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive. + + Note: Does not support backward pass. + + Arguments: + q: (batch_size, seqlen, nheads, headdim) + k_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim) if there's no page_table, + or (num_blocks, page_block_size, nheads_k, headdim) if there's a page_table (i.e. paged KV cache) + page_block_size can be arbitrary (e.g, 1, 2, 3, 64, etc.). + v_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim_v) if there's no page_table, + or (num_blocks, page_block_size, nheads_k, headdim_v) if there's a page_table (i.e. paged KV cache) + k [optional]: (batch_size, seqlen_new, nheads_k, headdim). If not None, we concatenate + k with k_cache, starting at the indices specified by cache_seqlens. + v [optional]: (batch_size, seqlen_new, nheads_k, headdim_v). Similar to k. + qv [optional]: (batch_size, seqlen, nheads, headdim_v) + rotary_cos [optional]: (seqlen_ro, rotary_dim / 2). If not None, we apply rotary embedding + to k and q. Only applicable if k and v are passed in. rotary_dim must be divisible by 16. + rotary_sin [optional]: (seqlen_ro, rotary_dim / 2). Similar to rotary_cos. + cache_seqlens: int, or (batch_size,), dtype torch.int32. The sequence lengths of the + KV cache. + cache_batch_idx: (batch_size,), dtype torch.int32. The indices used to index into the KV cache. + If None, we assume that the batch indices are [0, 1, 2, ..., batch_size - 1]. + If the indices are not distinct, and k and v are provided, the values updated in the cache + might come from any of the duplicate indices. + cache_leftpad: (batch_size,), dtype torch.int32. The index that the KV cache starts. If None, assume 0. + page_table [optional]: (batch_size, max_num_blocks_per_seq), dtype torch.int32. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + softcap: float. Anything > 0 activates softcapping attention. + rotary_interleaved: bool. Only applicable if rotary_cos and rotary_sin are passed in. + If True, rotary embedding will combine dimensions 0 & 1, 2 & 3, etc. If False, + rotary embedding will combine dimensions 0 & rotary_dim / 2, 1 & rotary_dim / 2 + 1 + (i.e. GPT-NeoX style). + num_splits: int. If > 1, split the key/value into this many chunks along the sequence. + If num_splits == 1, we don't split the key/value. If num_splits == 0, we use a heuristic + to automatically determine the number of splits. + Don't change this unless you know what you are doing. + return_softmax_lse: bool. Whether to return the logsumexp of the attention scores. + + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_softmax_lse=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + """ + assert k_cache.stride(-1) == 1, "k_cache must have contiguous last dimension" + assert v_cache.stride(-1) == 1, "v_cache must have contiguous last dimension" + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + if cache_seqlens is not None and isinstance(cache_seqlens, int): + cache_seqlens = torch.full( + (q.shape[0],), cache_seqlens, dtype=torch.int32, device=k_cache.device + ) + cache_seqlens = maybe_contiguous(cache_seqlens) + out, softmax_lse, *rest = _flash_attn_forward( + q, + k_cache, + v_cache, + k, + v, + qv, + None, # out + cu_seqlens_q, + None, # cu_seqlens_k + cu_seqlens_k_new, + None, # seqused_q + cache_seqlens, + max_seqlen_q, + None, # max_seqlen_k + page_table, + cache_batch_idx, + cache_leftpad, + rotary_cos, + rotary_sin, + rotary_seqlens, + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size_left=window_size[0], + window_size_right=window_size[1], + attention_chunk=attention_chunk, + softcap=softcap, + rotary_interleaved=rotary_interleaved, + scheduler_metadata=scheduler_metadata, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # return (out, softmax_lse) if return_softmax_lse else out + return (out, softmax_lse, *rest) if return_softmax_lse else out + + +def get_scheduler_metadata( + batch_size, max_seqlen_q, max_seqlen_k, num_heads_q, num_heads_kv, headdim, + cache_seqlens: torch.Tensor, + qkv_dtype=torch.bfloat16, + headdim_v=None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + cache_leftpad: Optional[torch.Tensor] = None, + page_size: Optional[int] = None, + max_seqlen_k_new=0, + causal=False, + window_size=(-1, -1), # -1 means infinite context window + attention_chunk=0, + has_softcap=False, + num_splits=0, # Can be tuned for speed + pack_gqa=None, # Can be tuned for speed + sm_margin=0, # Can be tuned if some SMs are used for communication +): + cache_seqlens = maybe_contiguous(cache_seqlens) + if headdim_v is None: + headdim_v = headdim + scheduler_metadata = flash_attn_3_cuda.get_scheduler_metadata( + batch_size, max_seqlen_q, max_seqlen_k, num_heads_q, num_heads_kv, headdim, headdim_v, + qkv_dtype, + cache_seqlens, + cu_seqlens_q, + None, # cu_seqlens_k + cu_seqlens_k_new, + None, # seqused_q + cache_leftpad, + page_size, + max_seqlen_k_new, + causal, + window_size[0], window_size[1], + attention_chunk, + has_softcap, + num_splits, + pack_gqa, + sm_margin, + ) + return scheduler_metadata diff --git a/build/torch210-cxx11-cu128-x86_64-linux/metadata.json b/build/torch210-cxx11-cu128-x86_64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..e0cc0c4f14abd10162c22499be39c75e6d922f94 --- /dev/null +++ b/build/torch210-cxx11-cu128-x86_64-linux/metadata.json @@ -0,0 +1,12 @@ +{ + "version": 1, + "license": "BSD-3-Clause", + "python-depends": [], + "backend": { + "type": "cuda", + "archs": [ + "8.0", + "9.0a" + ] + } +} diff --git a/build/torch210-cxx11-cu130-x86_64-linux/__init__.py b/build/torch210-cxx11-cu130-x86_64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3d12edd447e84d6e65c4fa9ab91ae3c3192884c1 --- /dev/null +++ b/build/torch210-cxx11-cu130-x86_64-linux/__init__.py @@ -0,0 +1,17 @@ +from .flash_attn_interface import ( + flash_attn_combine, + flash_attn_func, + flash_attn_qkvpacked_func, + flash_attn_varlen_func, + flash_attn_with_kvcache, + get_scheduler_metadata, +) + +__all__ = [ + "flash_attn_combine", + "flash_attn_func", + "flash_attn_qkvpacked_func", + "flash_attn_varlen_func", + "flash_attn_with_kvcache", + "get_scheduler_metadata", +] diff --git a/build/torch210-cxx11-cu130-x86_64-linux/_flash_attn3_cuda_e1d5be2.abi3.so b/build/torch210-cxx11-cu130-x86_64-linux/_flash_attn3_cuda_e1d5be2.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..24832ab2e17cbfd9841ae8c17425df61895c268d --- /dev/null +++ b/build/torch210-cxx11-cu130-x86_64-linux/_flash_attn3_cuda_e1d5be2.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:745d488c7c96de2e8b2625c6a0331dd7ecf7693054e8500e6a92c00a5a2fd11c +size 823699368 diff --git a/build/torch210-cxx11-cu130-x86_64-linux/_ops.py b/build/torch210-cxx11-cu130-x86_64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..b845f529a7b7c424bbe0db01a32c2cd6db45bb1a --- /dev/null +++ b/build/torch210-cxx11-cu130-x86_64-linux/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _flash_attn3_cuda_e1d5be2 +ops = torch.ops._flash_attn3_cuda_e1d5be2 + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_flash_attn3_cuda_e1d5be2::{op_name}" diff --git a/build/torch210-cxx11-cu130-x86_64-linux/flash_attn3/__init__.py b/build/torch210-cxx11-cu130-x86_64-linux/flash_attn3/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a9b2672c1cd85b74c1b3ded0fc0b2100e1aeac23 --- /dev/null +++ b/build/torch210-cxx11-cu130-x86_64-linux/flash_attn3/__init__.py @@ -0,0 +1,26 @@ +import ctypes +import importlib.util +import sys +from pathlib import Path +from types import ModuleType + + +def _import_from_path(file_path: Path) -> ModuleType: + # We cannot use the module name as-is, after adding it to `sys.modules`, + # it would also be used for other imports. So, we make a module name that + # depends on the path for it to be unique using the hex-encoded hash of + # the path. + path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value) + module_name = path_hash + spec = importlib.util.spec_from_file_location(module_name, file_path) + if spec is None: + raise ImportError(f"Cannot load spec for {module_name} from {file_path}") + module = importlib.util.module_from_spec(spec) + if module is None: + raise ImportError(f"Cannot load module {module_name} from spec") + sys.modules[module_name] = module + spec.loader.exec_module(module) # type: ignore + return module + + +globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py"))) diff --git a/build/torch210-cxx11-cu130-x86_64-linux/flash_attn_config.py b/build/torch210-cxx11-cu130-x86_64-linux/flash_attn_config.py new file mode 100644 index 0000000000000000000000000000000000000000..c4b7130324be2c4cabee7ada7000e170a2a81865 --- /dev/null +++ b/build/torch210-cxx11-cu130-x86_64-linux/flash_attn_config.py @@ -0,0 +1,7 @@ +# Auto-generated by flash attention 3 setup.py +CONFIG = {'build_flags': {'FLASHATTENTION_DISABLE_BACKWARD': False, 'FLASHATTENTION_DISABLE_SPLIT': False, 'FLASHATTENTION_DISABLE_PAGEDKV': False, 'FLASHATTENTION_DISABLE_APPENDKV': False, 'FLASHATTENTION_DISABLE_LOCAL': False, 'FLASHATTENTION_DISABLE_SOFTCAP': False, 'FLASHATTENTION_DISABLE_PACKGQA': False, 'FLASHATTENTION_DISABLE_FP16': False, 'FLASHATTENTION_DISABLE_FP8': False, 'FLASHATTENTION_DISABLE_VARLEN': False, 'FLASHATTENTION_DISABLE_CLUSTER': False, 'FLASHATTENTION_DISABLE_HDIM64': False, 'FLASHATTENTION_DISABLE_HDIM96': False, 'FLASHATTENTION_DISABLE_HDIM128': False, 'FLASHATTENTION_DISABLE_HDIM192': False, 'FLASHATTENTION_DISABLE_HDIM256': False, 'FLASHATTENTION_DISABLE_SM8x': False, 'FLASHATTENTION_ENABLE_VCOLMAJOR': False, 'FLASH_ATTENTION_DISABLE_HDIMDIFF64': False, 'FLASH_ATTENTION_DISABLE_HDIMDIFF192': False}} + +def show(): + from pprint import pprint + pprint(CONFIG) + diff --git a/build/torch210-cxx11-cu130-x86_64-linux/flash_attn_interface.py b/build/torch210-cxx11-cu130-x86_64-linux/flash_attn_interface.py new file mode 100644 index 0000000000000000000000000000000000000000..2e2a3cdd3e89cbc59bf6c956454e2776c305e0e5 --- /dev/null +++ b/build/torch210-cxx11-cu130-x86_64-linux/flash_attn_interface.py @@ -0,0 +1,1127 @@ +# Copyright (c) 2023, Tri Dao. + +from typing import Optional, Union, List, Tuple + +import torch +import torch.nn as nn + +from ._ops import ops as flash_attn_3_cuda +from ._ops import add_op_namespace_prefix + +def maybe_contiguous(x): + return x.contiguous() if x is not None and x.stride(-1) != 1 else x + + +def round_multiple(x, m): + return (x + m - 1) // m * m + + +def round_up_headdim(head_size: int) -> int: + from .flash_attn_config import CONFIG + + if not CONFIG["build_flags"]["FLASHATTENTION_DISABLE_HDIM64"]: + if head_size <= 64: + return 64 + if not CONFIG["build_flags"]["FLASHATTENTION_DISABLE_HDIM96"]: + if head_size <= 96: + return 96 + if not CONFIG["build_flags"]["FLASHATTENTION_DISABLE_HDIM128"]: + if head_size <= 128: + return 128 + if not CONFIG["build_flags"]["FLASHATTENTION_DISABLE_HDIM192"]: + if head_size <= 192: + return 192 + if not CONFIG["build_flags"]["FLASHATTENTION_DISABLE_HDIM256"]: + if head_size <= 256: + return 256 + return 256 + + +@torch.library.custom_op(add_op_namespace_prefix("_flash_attn_forward"), mutates_args=(), device_types="cuda") +def _flash_attn_forward( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_new: Optional[torch.Tensor] = None, + v_new: Optional[torch.Tensor] = None, + qv: Optional[torch.Tensor] = None, + out_: Optional[torch.Tensor] = None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + seqused_q: Optional[torch.Tensor] = None, + seqused_k: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_k: Optional[int] = None, + page_table: Optional[torch.Tensor] = None, + kv_batch_idx: Optional[torch.Tensor] = None, + leftpad_k: Optional[torch.Tensor] = None, + rotary_cos: Optional[torch.Tensor] = None, + rotary_sin: Optional[torch.Tensor] = None, + seqlens_rotary: Optional[torch.Tensor] = None, + q_descale: Optional[torch.Tensor] = None, + k_descale: Optional[torch.Tensor] = None, + v_descale: Optional[torch.Tensor] = None, + softmax_scale: Optional[float] = None, + causal: bool = False, + window_size_left: int = -1, + window_size_right: int = -1, + attention_chunk: int = 0, + softcap: float = 0.0, + rotary_interleaved: bool = True, + scheduler_metadata: Optional[torch.Tensor] = None, + num_splits: int = 1, + pack_gqa: Optional[bool] = None, + sm_margin: int = 0, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + q, k, k_new, v_new = [maybe_contiguous(x) for x in (q, k, k_new, v_new)] + v = v.contiguous() if v.stride(-1) != 1 and v.stride(-3) != 1 else v + cu_seqlens_q, cu_seqlens_k, cu_seqlens_k_new = [ + maybe_contiguous(x) for x in (cu_seqlens_q, cu_seqlens_k, cu_seqlens_k_new) + ] + seqused_q, seqused_k = [maybe_contiguous(x) for x in (seqused_q, seqused_k)] + page_table, kv_batch_idx, leftpad_k = [ + maybe_contiguous(x) for x in (page_table, kv_batch_idx, leftpad_k) + ] + rotary_cos, rotary_sin = [maybe_contiguous(x) for x in (rotary_cos, rotary_sin)] + seqlens_rotary = maybe_contiguous(seqlens_rotary) + out, softmax_lse, out_accum, softmax_lse_accum = flash_attn_3_cuda.fwd( + q, + k, + v, + k_new, + v_new, + qv, + out_, + cu_seqlens_q, + cu_seqlens_k, + cu_seqlens_k_new, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + page_table, + kv_batch_idx, + leftpad_k, + rotary_cos, + rotary_sin, + seqlens_rotary, + q_descale, + k_descale, + v_descale, + softmax_scale, + causal, + window_size_left, + window_size_right, + attention_chunk, + softcap, + rotary_interleaved, + scheduler_metadata, + num_splits, + pack_gqa, + sm_margin, + ) + + if out_accum is None: + out_accum = torch.tensor([], device=out.device) + + if softmax_lse_accum is None: + softmax_lse_accum = torch.tensor([], device=out.device) + + return out, softmax_lse, out_accum, softmax_lse_accum + + +@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_forward")) +def _flash_attn_forward_fake( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_new: Optional[torch.Tensor] = None, + v_new: Optional[torch.Tensor] = None, + qv: Optional[torch.Tensor] = None, + out_: Optional[torch.Tensor] = None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + seqused_q: Optional[torch.Tensor] = None, + seqused_k: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_k: Optional[int] = None, + page_table: Optional[torch.Tensor] = None, + kv_batch_idx: Optional[torch.Tensor] = None, + leftpad_k: Optional[torch.Tensor] = None, + rotary_cos: Optional[torch.Tensor] = None, + rotary_sin: Optional[torch.Tensor] = None, + seqlens_rotary: Optional[torch.Tensor] = None, + q_descale: Optional[torch.Tensor] = None, + k_descale: Optional[torch.Tensor] = None, + v_descale: Optional[torch.Tensor] = None, + softmax_scale: Optional[float] = None, + causal: bool = False, + window_size_left: int = -1, + window_size_right: int = -1, + attention_chunk: int = 0, + softcap: float = 0.0, + rotary_interleaved: bool = True, + scheduler_metadata: Optional[torch.Tensor] = None, + num_splits: int = 1, + pack_gqa: Optional[bool] = None, + sm_margin: int = 0, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Symbolic fake implementation of flash attention forward. + Returns tensors with the correct shapes and dtypes without actual computation. + """ + + # Determine if we're in varlen mode + is_varlen_q = cu_seqlens_q is not None + + # Get dimensions from query tensor + if is_varlen_q: + # varlen mode: q is (total_q, num_heads, head_size) + total_q, num_heads, head_size = q.shape + batch_size = cu_seqlens_q.shape[0] - 1 + + if max_seqlen_q is None: + raise ValueError("max_seqlen_q must be provided if cu_seqlens_q is provided") + seqlen_q = max_seqlen_q + else: + # batch mode: q is (batch_size, seqlen_q, num_heads, head_size) + batch_size, seqlen_q, num_heads, head_size = q.shape + total_q = batch_size * q.shape[1] + # Get value head dimension + head_size_v = v.shape[-1] + + # Determine output dtype (FP8 inputs produce BF16 outputs) + q_type = q.dtype + if q_type == torch.float8_e4m3fn: + out_dtype = torch.bfloat16 + else: + out_dtype = q_type + + # Create output tensor + if out_ is not None: + # If out_ is provided, _flash_attn_forward becomes non-functional + raise TypeError("Tracing (torch.compile/torch.export) with pre-allocated output tensor is not supported.") + + if is_varlen_q: + out = torch.empty((total_q, num_heads, head_size_v), dtype=out_dtype, device=q.device) + else: + out = torch.empty((batch_size, seqlen_q, num_heads, head_size_v), dtype=out_dtype, device=q.device) + + # Create softmax_lse tensor + if is_varlen_q: + softmax_lse = torch.empty((num_heads, total_q), dtype=torch.float32, device=q.device) + else: + softmax_lse = torch.empty((batch_size, num_heads, seqlen_q), dtype=torch.float32, device=q.device) + + # TODO(guilhermeleobas): Implement "get_num_splits" + # There's an heuristic to compute num_splits when "num_splits <= 0" + # assert that num_splits is > 0 for now + if num_splits <= 0: + raise ValueError(f"tracing (torch.compile/torch.export) with num_splits <= 0 not supported. Got {num_splits=}") + + if num_splits > 1: + if is_varlen_q: + out_accum = torch.empty((num_splits, num_heads, total_q, head_size_v), dtype=torch.float32, device=q.device) + softmax_lse_accum = torch.empty((num_splits, num_heads, total_q), dtype=torch.float32, device=q.device) + else: + out_accum = torch.empty((num_splits, batch_size, num_heads, seqlen_q, head_size_v), dtype=torch.float32, device=q.device) + softmax_lse_accum = torch.empty((num_splits, batch_size, num_heads, seqlen_q), dtype=torch.float32, device=q.device) + else: + # Tensors are not set when num_splits < 1 + out_accum = torch.tensor([], device=out.device) + softmax_lse_accum = torch.tensor([], device=out.device) + + return out, softmax_lse, out_accum, softmax_lse_accum + + +@torch.library.custom_op(add_op_namespace_prefix("_flash_attn_backward"), mutates_args=("dq", "dk", "dv"), device_types="cuda") +def _flash_attn_backward( + dout: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + out: torch.Tensor, + softmax_lse: torch.Tensor, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k: Optional[torch.Tensor] = None, + sequed_q: Optional[torch.Tensor] = None, + sequed_k: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_k: Optional[int] = None, + dq: Optional[torch.Tensor] = None, + dk: Optional[torch.Tensor] = None, + dv: Optional[torch.Tensor] = None, + softmax_scale: Optional[float] = None, + is_causal: bool = False, + window_size_left: int = -1, + window_size_right: int = -1, + softcap: float = 0.0, + deterministic: bool = False, + sm_margin: int = 0, +) -> torch.Tensor: + # dq, dk, dv are allocated by us so they should already be contiguous + dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] + softmax_d, *rest = flash_attn_3_cuda.bwd( + dout, + q, + k, + v, + out, + softmax_lse, + dq, + dk, + dv, + cu_seqlens_q, + cu_seqlens_k, + sequed_q, + sequed_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + is_causal, + window_size_left, + window_size_right, + softcap, + deterministic, + sm_margin, + ) + return softmax_d + + +@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_backward")) +def _flash_attn_backward_fake( + dout: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + out: torch.Tensor, + softmax_lse: torch.Tensor, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k: Optional[torch.Tensor] = None, + sequed_q: Optional[torch.Tensor] = None, + sequed_k: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_k: Optional[int] = None, + dq: Optional[torch.Tensor] = None, + dk: Optional[torch.Tensor] = None, + dv: Optional[torch.Tensor] = None, + softmax_scale: Optional[float] = None, + is_causal: bool = False, + window_size_left: int = -1, + window_size_right: int = -1, + softcap: float = 0.0, + deterministic: bool = False, + sm_margin: int = 0, +) -> torch.Tensor: + + is_varlen_q = cu_seqlens_q is not None + is_varlen_k = cu_seqlens_q is not None + is_varlen = is_varlen_q or is_varlen_k or sequed_q is not None or sequed_k is not None + + if not is_varlen_q: + batch_size = q.size(0) + seqlen_q = q.size(1) + seqlen_k = k.size(1) + total_q = batch_size * q.size(1) + else: + batch_size = cu_seqlens_q.size(0) - 1 + total_q = q.size(0) + seqlen_q = max_seqlen_q + seqlen_k = max_seqlen_k + + if window_size_left >= seqlen_k - 1: + window_size_left = -1 + + if window_size_right >= seqlen_q - 1: + window_size_right = -1 + + if is_causal: + window_size_right = 0 + + is_causal = window_size_left < 0 and window_size_right == 0 + + head_size = q.size(-1) + head_size_v = v.size(-1) + head_size_rounded = round_up_headdim(max(head_size, head_size_v)) + + # Hopper gpus uses cuda compute capabilities 9.0 + cap = torch.cuda.get_device_capability(q.device) + arch = cap[0] * 10 + cap[1] + + is_local = (window_size_left >= 0 or window_size_right >= 0) and not is_causal + + if head_size_rounded <= 64: + kBlockM_sm90 = 96 if (is_causal and softcap > 0.0) else 128 + elif head_size_rounded <= 96: + kBlockM_sm90 = 64 + elif head_size_rounded <= 128: + kBlockM_sm90 = 64 if (is_causal or is_local or softcap > 0.0) else 80 + else: + kBlockM_sm90 = 64 + + kBlockM_sm80 = 128 if head_size_rounded <= 64 else 64 + kBlockM_sm86 = 64 if head_size_rounded <= 192 else 32 + + if arch >= 90: + kBlockM = kBlockM_sm90 + elif arch == 86 or arch == 89: + kBlockM = kBlockM_sm86 + else: + kBlockM = kBlockM_sm80 + + num_heads = q.shape[-2] + seqlen_q_rounded = round_multiple(seqlen_q, kBlockM) + + total_q_padded_rounded = round_multiple(total_q + batch_size * kBlockM, kBlockM) + + dq = torch.empty_like(q) if dq is None else dq + dk = torch.empty_like(k) if dk is None else dk + dv = torch.empty_like(v) if dv is None else dv + + if not is_varlen: + softmax_d = torch.empty((batch_size, num_heads, seqlen_q_rounded), dtype=torch.float32, device=q.device) + else: + softmax_d = torch.empty((num_heads, total_q_padded_rounded), dtype=torch.float32, device=q.device) + + return softmax_d + + +def setup_context(ctx, inputs, output): + q, k, v = inputs[:3] + out, softmax_lse, _, _ = output + ctx.save_for_backward(q, k, v, out, softmax_lse) + ctx.softmax_scale = inputs[-11] + ctx.causal = inputs[-10] + ctx.window_size = [inputs[-9], inputs[-8]] + ctx.attention_chunk = inputs[-7] + ctx.softcap = inputs[-6] + ctx.sm_margin = inputs[-1] + + +def _backward(ctx, dout, *grads): + q, k, v, out, softmax_lse = ctx.saved_tensors + dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + None, None, # cu_seqlens_q, cu_seqlens_k, + None, None, # sequed_q, sequed_k, + None, None, # max_seqlen_q, max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size[0], + ctx.window_size[1], + ctx.softcap, + False, # deterministic + ctx.sm_margin, + ) + return dq, dk, dv, *((None,) * 21) + + +_flash_attn_forward.register_autograd(_backward, setup_context=setup_context) + + + +class FlashAttnQKVPackedFunc(torch.autograd.Function): + @staticmethod + def forward( + ctx, + qkv, + softmax_scale, + causal, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + deterministic=False, + num_heads_q=None, + sm_margin=0, + return_softmax=False, + ): + if softmax_scale is None: + softmax_scale = qkv.shape[-1] ** (-0.5) + if qkv.dim() == 5: + assert qkv.shape[-3] == 3 + q, k, v = qkv.unbind(dim=-3) + else: + assert qkv.dim() == 4 + assert num_heads_q is not None + num_heads_k = (qkv.shape[2] - num_heads_q) // 2 + assert num_heads_k * 2 + num_heads_q == qkv.shape[2] + q, k, v = qkv.split([num_heads_q, num_heads_k, num_heads_k], dim=-2) + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + None, # qv + None, # out + None, None, None, # cu_seqlens_q/k/k_new + None, None, # seqused_q/k + None, None, # max_seqlen_q/k + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size_left=window_size[0], + window_size_right=window_size[1], + attention_chunk=attention_chunk, + softcap=softcap, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse) + ctx.save_for_backward(q, k, v, out, softmax_lse) + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.ndim = qkv.dim() + ctx.sm_margin = sm_margin + return (out, softmax_lse) if return_softmax else out + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + if ctx.ndim == 5: + qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) + dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) + dq, dk, dv = dqkv.unbind(dim=-3) + else: + num_heads_q = q.shape[2] + num_heads_k = k.shape[2] + qkv_shape = q.shape[:-2] + (num_heads_q + num_heads_k * 2, *q.shape[-1:]) + dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) + dq, dk, dv = dqkv.split([num_heads_q, num_heads_k, num_heads_k], dim=-2) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + None, None, # cu_seqlens_q, cu_seqlens_k, + None, None, # sequed_q, sequed_k, + None, None, # max_seqlen_q, max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size[0], + ctx.window_size[1], + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dqkv = dqkv[..., : dout.shape[-1]] # We could have padded the head dimension + return dqkv, None, None, None, None, None, None, None, None, None, None, None, None + + +class FlashAttnFunc(torch.autograd.Function): + + @staticmethod + def forward( + ctx, + q, + k, + v, + softmax_scale, + causal, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + return_softmax=False, + ): + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + # out, q, k, v, out_padded, softmax_lse = _flash_attn_forward( + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + qv, # qv + None, # out + None, None, None, # cu_seqlens_q/k/k_new + None, None, # seqused_q/k + None, None, # max_seqlen_q/k + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size_left=window_size[0], + window_size_right=window_size[1], + attention_chunk=attention_chunk, + softcap=softcap, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse) + ctx.save_for_backward(q, k, v, out, softmax_lse) + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.sm_margin = sm_margin + return (out, softmax_lse) if return_softmax else out + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + None, None, # cu_seqlens_q, cu_seqlens_k, + None, None, # sequed_q, sequed_k, + None, None, # max_seqlen_q, max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size[0], + ctx.window_size[1], + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dq = dq[..., : q.shape[-1]] # We could have padded the head dimension + dk = dk[..., : k.shape[-1]] + dv = dv[..., : v.shape[-1]] + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None + + +class FlashAttnVarlenFunc(torch.autograd.Function): + + @staticmethod + def forward( + ctx, + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + return_softmax=False, + ): + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + # out, q, k, v, out_padded, softmax_lse = _flash_attn_varlen_forward( + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + qv, # qv + None, # out + cu_seqlens_q, + cu_seqlens_k, + None, # cu_seqlens_k_new + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size_left=window_size[0], + window_size_right=window_size[1], + attention_chunk=attention_chunk, + softcap=softcap, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k) + ctx.save_for_backward(q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k) + ctx.max_seqlen_q = max_seqlen_q + ctx.max_seqlen_k = max_seqlen_k + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.sm_margin = sm_margin + return (out, softmax_lse) if return_softmax else out + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + ctx.max_seqlen_q, + ctx.max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size[0], + ctx.window_size[1], + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dq = dq[..., : q.shape[-1]] # We could have padded the head dimension + dk = dk[..., : k.shape[-1]] + dv = dv[..., : v.shape[-1]] + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None + + +def flash_attn_qkvpacked_func( + qkv, + softmax_scale=None, + causal=False, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + deterministic=False, + num_heads_q=None, + sm_margin=0, + return_attn_probs=False, +): + """dropout_p should be set to 0.0 during evaluation + If Q, K, V are already stacked into 1 tensor, this function will be faster than + calling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation + of the gradients of Q, K, V. + For multi-query and grouped-query attention (MQA/GQA), please see + flash_attn_kvpacked_func and flash_attn_func. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between [i - window_size[0], i + window_size[1]] inclusive. + + Arguments: + qkv: (batch_size, seqlen, 3, nheads, headdim) + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + softcap: float. Anything > 0 activates softcapping attention. + alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of (-alibi_slope * |i - j|) is added to + the attention score of query i and key j. + deterministic: bool. Whether to use the deterministic implementation of the backward pass, + which is slightly slower and uses more memory. The forward pass is always deterministic. + return_attn_probs: bool. Whether to return the attention probabilities. This option is for + testing only. The returned probabilities are not guaranteed to be correct + (they might not have the right scaling). + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen). + The output of softmax (possibly with different scaling). It also encodes the dropout + pattern (negative means that location was dropped, nonnegative means it was kept). + """ + return FlashAttnQKVPackedFunc.apply( + qkv, + softmax_scale, + causal, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + deterministic, + num_heads_q, + sm_margin, + return_attn_probs, + ) + + +def flash_attn_func( + q, + k, + v, + softmax_scale=None, + causal=False, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + return_attn_probs=False, +): + """dropout_p should be set to 0.0 during evaluation + Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads + than Q. Note that the number of heads in Q must be divisible by the number of heads in KV. + For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head + 0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V. + + If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix. + For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is: + 1 1 1 1 0 + 1 1 1 1 1 + If seqlen_q = 5 and seqlen_k = 2, the causal mask is: + 0 0 + 0 0 + 0 0 + 1 0 + 1 1 + If the row of the mask is all zero, the output will be zero. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between + [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive. + + Arguments: + q: (batch_size, seqlen, nheads, headdim) + k: (batch_size, seqlen, nheads_k, headdim) + v: (batch_size, seqlen, nheads_k, headdim) + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of + (-alibi_slope * |i + seqlen_k - seqlen_q - j|) + is added to the attention score of query i and key j. + deterministic: bool. Whether to use the deterministic implementation of the backward pass, + which is slightly slower and uses more memory. The forward pass is always deterministic. + return_attn_probs: bool. Whether to return the attention probabilities. This option is for + testing only. The returned probabilities are not guaranteed to be correct + (they might not have the right scaling). + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + """ + return FlashAttnFunc.apply( + q, + k, + v, + softmax_scale, + causal, + qv, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + num_splits, + pack_gqa, + deterministic, + sm_margin, + return_attn_probs, + ) + + +def flash_attn_varlen_func( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + seqused_q=None, + seqused_k=None, + softmax_scale=None, + causal=False, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + return_attn_probs=False, +): + return FlashAttnVarlenFunc.apply( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + qv, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + num_splits, + pack_gqa, + deterministic, + sm_margin, + return_attn_probs, + ) + + +def flash_attn_combine(out_partial, lse_partial, out=None, out_dtype=None): + return flash_attn_3_cuda.fwd_combine(out_partial, lse_partial, out, out_dtype) + + +def flash_attn_with_kvcache( + q, + k_cache, + v_cache, + k=None, + v=None, + qv=None, + rotary_cos=None, + rotary_sin=None, + cache_seqlens: Optional[Union[(int, torch.Tensor)]] = None, + cache_batch_idx: Optional[torch.Tensor] = None, + cache_leftpad: Optional[torch.Tensor] = None, + page_table: Optional[torch.Tensor] = None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + rotary_seqlens: Optional[torch.Tensor] = None, + q_descale: Optional[torch.Tensor] = None, + k_descale: Optional[torch.Tensor] = None, + v_descale: Optional[torch.Tensor] = None, + softmax_scale=None, + causal=False, + window_size=(-1, -1), # -1 means infinite context window + attention_chunk=0, + softcap=0.0, # 0.0 means deactivated + rotary_interleaved=True, + scheduler_metadata=None, + num_splits=0, # Can be tuned for speed + pack_gqa=None, # Can be tuned for speed + sm_margin=0, # Can be tuned if some SMs are used for communication + return_softmax_lse=False, +): + """ + If k and v are not None, k_cache and v_cache will be updated *inplace* with the new values from + k and v. This is useful for incremental decoding: you can pass in the cached keys/values from + the previous step, and update them with the new keys/values from the current step, and do + attention with the updated cache, all in 1 kernel. + + If you pass in k / v, you must make sure that the cache is large enough to hold the new values. + For example, the KV cache could be pre-allocated with the max sequence length, and you can use + cache_seqlens to keep track of the current sequence lengths of each sequence in the batch. + + Also apply rotary embedding if rotary_cos and rotary_sin are passed in. The key @k will be + rotated by rotary_cos and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc. + If causal or local (i.e., window_size != (-1, -1)), the query @q will be rotated by rotary_cos + and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc. + If not causal and not local, the query @q will be rotated by rotary_cos and rotary_sin at + indices cache_seqlens only (i.e. we consider all tokens in @q to be at position cache_seqlens). + + See tests/test_flash_attn.py::test_flash_attn_kvcache for examples of how to use this function. + + Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads + than Q. Note that the number of heads in Q must be divisible by the number of heads in KV. + For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head + 0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V. + + If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix. + For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is: + 1 1 1 1 0 + 1 1 1 1 1 + If seqlen_q = 5 and seqlen_k = 2, the causal mask is: + 0 0 + 0 0 + 0 0 + 1 0 + 1 1 + If the row of the mask is all zero, the output will be zero. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between + [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive. + + Note: Does not support backward pass. + + Arguments: + q: (batch_size, seqlen, nheads, headdim) + k_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim) if there's no page_table, + or (num_blocks, page_block_size, nheads_k, headdim) if there's a page_table (i.e. paged KV cache) + page_block_size can be arbitrary (e.g, 1, 2, 3, 64, etc.). + v_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim_v) if there's no page_table, + or (num_blocks, page_block_size, nheads_k, headdim_v) if there's a page_table (i.e. paged KV cache) + k [optional]: (batch_size, seqlen_new, nheads_k, headdim). If not None, we concatenate + k with k_cache, starting at the indices specified by cache_seqlens. + v [optional]: (batch_size, seqlen_new, nheads_k, headdim_v). Similar to k. + qv [optional]: (batch_size, seqlen, nheads, headdim_v) + rotary_cos [optional]: (seqlen_ro, rotary_dim / 2). If not None, we apply rotary embedding + to k and q. Only applicable if k and v are passed in. rotary_dim must be divisible by 16. + rotary_sin [optional]: (seqlen_ro, rotary_dim / 2). Similar to rotary_cos. + cache_seqlens: int, or (batch_size,), dtype torch.int32. The sequence lengths of the + KV cache. + cache_batch_idx: (batch_size,), dtype torch.int32. The indices used to index into the KV cache. + If None, we assume that the batch indices are [0, 1, 2, ..., batch_size - 1]. + If the indices are not distinct, and k and v are provided, the values updated in the cache + might come from any of the duplicate indices. + cache_leftpad: (batch_size,), dtype torch.int32. The index that the KV cache starts. If None, assume 0. + page_table [optional]: (batch_size, max_num_blocks_per_seq), dtype torch.int32. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + softcap: float. Anything > 0 activates softcapping attention. + rotary_interleaved: bool. Only applicable if rotary_cos and rotary_sin are passed in. + If True, rotary embedding will combine dimensions 0 & 1, 2 & 3, etc. If False, + rotary embedding will combine dimensions 0 & rotary_dim / 2, 1 & rotary_dim / 2 + 1 + (i.e. GPT-NeoX style). + num_splits: int. If > 1, split the key/value into this many chunks along the sequence. + If num_splits == 1, we don't split the key/value. If num_splits == 0, we use a heuristic + to automatically determine the number of splits. + Don't change this unless you know what you are doing. + return_softmax_lse: bool. Whether to return the logsumexp of the attention scores. + + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_softmax_lse=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + """ + assert k_cache.stride(-1) == 1, "k_cache must have contiguous last dimension" + assert v_cache.stride(-1) == 1, "v_cache must have contiguous last dimension" + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + if cache_seqlens is not None and isinstance(cache_seqlens, int): + cache_seqlens = torch.full( + (q.shape[0],), cache_seqlens, dtype=torch.int32, device=k_cache.device + ) + cache_seqlens = maybe_contiguous(cache_seqlens) + out, softmax_lse, *rest = _flash_attn_forward( + q, + k_cache, + v_cache, + k, + v, + qv, + None, # out + cu_seqlens_q, + None, # cu_seqlens_k + cu_seqlens_k_new, + None, # seqused_q + cache_seqlens, + max_seqlen_q, + None, # max_seqlen_k + page_table, + cache_batch_idx, + cache_leftpad, + rotary_cos, + rotary_sin, + rotary_seqlens, + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size_left=window_size[0], + window_size_right=window_size[1], + attention_chunk=attention_chunk, + softcap=softcap, + rotary_interleaved=rotary_interleaved, + scheduler_metadata=scheduler_metadata, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # return (out, softmax_lse) if return_softmax_lse else out + return (out, softmax_lse, *rest) if return_softmax_lse else out + + +def get_scheduler_metadata( + batch_size, max_seqlen_q, max_seqlen_k, num_heads_q, num_heads_kv, headdim, + cache_seqlens: torch.Tensor, + qkv_dtype=torch.bfloat16, + headdim_v=None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + cache_leftpad: Optional[torch.Tensor] = None, + page_size: Optional[int] = None, + max_seqlen_k_new=0, + causal=False, + window_size=(-1, -1), # -1 means infinite context window + attention_chunk=0, + has_softcap=False, + num_splits=0, # Can be tuned for speed + pack_gqa=None, # Can be tuned for speed + sm_margin=0, # Can be tuned if some SMs are used for communication +): + cache_seqlens = maybe_contiguous(cache_seqlens) + if headdim_v is None: + headdim_v = headdim + scheduler_metadata = flash_attn_3_cuda.get_scheduler_metadata( + batch_size, max_seqlen_q, max_seqlen_k, num_heads_q, num_heads_kv, headdim, headdim_v, + qkv_dtype, + cache_seqlens, + cu_seqlens_q, + None, # cu_seqlens_k + cu_seqlens_k_new, + None, # seqused_q + cache_leftpad, + page_size, + max_seqlen_k_new, + causal, + window_size[0], window_size[1], + attention_chunk, + has_softcap, + num_splits, + pack_gqa, + sm_margin, + ) + return scheduler_metadata diff --git a/build/torch210-cxx11-cu130-x86_64-linux/metadata.json b/build/torch210-cxx11-cu130-x86_64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..e0cc0c4f14abd10162c22499be39c75e6d922f94 --- /dev/null +++ b/build/torch210-cxx11-cu130-x86_64-linux/metadata.json @@ -0,0 +1,12 @@ +{ + "version": 1, + "license": "BSD-3-Clause", + "python-depends": [], + "backend": { + "type": "cuda", + "archs": [ + "8.0", + "9.0a" + ] + } +} diff --git a/build/torch211-cxx11-cu128-x86_64-linux/__init__.py b/build/torch211-cxx11-cu128-x86_64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3d12edd447e84d6e65c4fa9ab91ae3c3192884c1 --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/__init__.py @@ -0,0 +1,17 @@ +from .flash_attn_interface import ( + flash_attn_combine, + flash_attn_func, + flash_attn_qkvpacked_func, + flash_attn_varlen_func, + flash_attn_with_kvcache, + get_scheduler_metadata, +) + +__all__ = [ + "flash_attn_combine", + "flash_attn_func", + "flash_attn_qkvpacked_func", + "flash_attn_varlen_func", + "flash_attn_with_kvcache", + "get_scheduler_metadata", +] diff --git a/build/torch211-cxx11-cu128-x86_64-linux/_flash_attn3_cuda_e1d5be2.abi3.so b/build/torch211-cxx11-cu128-x86_64-linux/_flash_attn3_cuda_e1d5be2.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..ffb3178dc17d05d0cc6102a7201010e66aa8da9a --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/_flash_attn3_cuda_e1d5be2.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:1cb641764fe0815bab292dab3bab1d93edb5b21cf08fe1e6ee8b3d4e3e4276ff +size 804184120 diff --git a/build/torch211-cxx11-cu128-x86_64-linux/_ops.py b/build/torch211-cxx11-cu128-x86_64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..b845f529a7b7c424bbe0db01a32c2cd6db45bb1a --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _flash_attn3_cuda_e1d5be2 +ops = torch.ops._flash_attn3_cuda_e1d5be2 + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_flash_attn3_cuda_e1d5be2::{op_name}" diff --git a/build/torch211-cxx11-cu128-x86_64-linux/flash_attn3/__init__.py b/build/torch211-cxx11-cu128-x86_64-linux/flash_attn3/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a9b2672c1cd85b74c1b3ded0fc0b2100e1aeac23 --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/flash_attn3/__init__.py @@ -0,0 +1,26 @@ +import ctypes +import importlib.util +import sys +from pathlib import Path +from types import ModuleType + + +def _import_from_path(file_path: Path) -> ModuleType: + # We cannot use the module name as-is, after adding it to `sys.modules`, + # it would also be used for other imports. So, we make a module name that + # depends on the path for it to be unique using the hex-encoded hash of + # the path. + path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value) + module_name = path_hash + spec = importlib.util.spec_from_file_location(module_name, file_path) + if spec is None: + raise ImportError(f"Cannot load spec for {module_name} from {file_path}") + module = importlib.util.module_from_spec(spec) + if module is None: + raise ImportError(f"Cannot load module {module_name} from spec") + sys.modules[module_name] = module + spec.loader.exec_module(module) # type: ignore + return module + + +globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py"))) diff --git a/build/torch211-cxx11-cu128-x86_64-linux/flash_attn_config.py b/build/torch211-cxx11-cu128-x86_64-linux/flash_attn_config.py new file mode 100644 index 0000000000000000000000000000000000000000..c4b7130324be2c4cabee7ada7000e170a2a81865 --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/flash_attn_config.py @@ -0,0 +1,7 @@ +# Auto-generated by flash attention 3 setup.py +CONFIG = {'build_flags': {'FLASHATTENTION_DISABLE_BACKWARD': False, 'FLASHATTENTION_DISABLE_SPLIT': False, 'FLASHATTENTION_DISABLE_PAGEDKV': False, 'FLASHATTENTION_DISABLE_APPENDKV': False, 'FLASHATTENTION_DISABLE_LOCAL': False, 'FLASHATTENTION_DISABLE_SOFTCAP': False, 'FLASHATTENTION_DISABLE_PACKGQA': False, 'FLASHATTENTION_DISABLE_FP16': False, 'FLASHATTENTION_DISABLE_FP8': False, 'FLASHATTENTION_DISABLE_VARLEN': False, 'FLASHATTENTION_DISABLE_CLUSTER': False, 'FLASHATTENTION_DISABLE_HDIM64': False, 'FLASHATTENTION_DISABLE_HDIM96': False, 'FLASHATTENTION_DISABLE_HDIM128': False, 'FLASHATTENTION_DISABLE_HDIM192': False, 'FLASHATTENTION_DISABLE_HDIM256': False, 'FLASHATTENTION_DISABLE_SM8x': False, 'FLASHATTENTION_ENABLE_VCOLMAJOR': False, 'FLASH_ATTENTION_DISABLE_HDIMDIFF64': False, 'FLASH_ATTENTION_DISABLE_HDIMDIFF192': False}} + +def show(): + from pprint import pprint + pprint(CONFIG) + diff --git a/build/torch211-cxx11-cu128-x86_64-linux/flash_attn_interface.py b/build/torch211-cxx11-cu128-x86_64-linux/flash_attn_interface.py new file mode 100644 index 0000000000000000000000000000000000000000..2e2a3cdd3e89cbc59bf6c956454e2776c305e0e5 --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/flash_attn_interface.py @@ -0,0 +1,1127 @@ +# Copyright (c) 2023, Tri Dao. + +from typing import Optional, Union, List, Tuple + +import torch +import torch.nn as nn + +from ._ops import ops as flash_attn_3_cuda +from ._ops import add_op_namespace_prefix + +def maybe_contiguous(x): + return x.contiguous() if x is not None and x.stride(-1) != 1 else x + + +def round_multiple(x, m): + return (x + m - 1) // m * m + + +def round_up_headdim(head_size: int) -> int: + from .flash_attn_config import CONFIG + + if not CONFIG["build_flags"]["FLASHATTENTION_DISABLE_HDIM64"]: + if head_size <= 64: + return 64 + if not CONFIG["build_flags"]["FLASHATTENTION_DISABLE_HDIM96"]: + if head_size <= 96: + return 96 + if not CONFIG["build_flags"]["FLASHATTENTION_DISABLE_HDIM128"]: + if head_size <= 128: + return 128 + if not CONFIG["build_flags"]["FLASHATTENTION_DISABLE_HDIM192"]: + if head_size <= 192: + return 192 + if not CONFIG["build_flags"]["FLASHATTENTION_DISABLE_HDIM256"]: + if head_size <= 256: + return 256 + return 256 + + +@torch.library.custom_op(add_op_namespace_prefix("_flash_attn_forward"), mutates_args=(), device_types="cuda") +def _flash_attn_forward( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_new: Optional[torch.Tensor] = None, + v_new: Optional[torch.Tensor] = None, + qv: Optional[torch.Tensor] = None, + out_: Optional[torch.Tensor] = None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + seqused_q: Optional[torch.Tensor] = None, + seqused_k: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_k: Optional[int] = None, + page_table: Optional[torch.Tensor] = None, + kv_batch_idx: Optional[torch.Tensor] = None, + leftpad_k: Optional[torch.Tensor] = None, + rotary_cos: Optional[torch.Tensor] = None, + rotary_sin: Optional[torch.Tensor] = None, + seqlens_rotary: Optional[torch.Tensor] = None, + q_descale: Optional[torch.Tensor] = None, + k_descale: Optional[torch.Tensor] = None, + v_descale: Optional[torch.Tensor] = None, + softmax_scale: Optional[float] = None, + causal: bool = False, + window_size_left: int = -1, + window_size_right: int = -1, + attention_chunk: int = 0, + softcap: float = 0.0, + rotary_interleaved: bool = True, + scheduler_metadata: Optional[torch.Tensor] = None, + num_splits: int = 1, + pack_gqa: Optional[bool] = None, + sm_margin: int = 0, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + q, k, k_new, v_new = [maybe_contiguous(x) for x in (q, k, k_new, v_new)] + v = v.contiguous() if v.stride(-1) != 1 and v.stride(-3) != 1 else v + cu_seqlens_q, cu_seqlens_k, cu_seqlens_k_new = [ + maybe_contiguous(x) for x in (cu_seqlens_q, cu_seqlens_k, cu_seqlens_k_new) + ] + seqused_q, seqused_k = [maybe_contiguous(x) for x in (seqused_q, seqused_k)] + page_table, kv_batch_idx, leftpad_k = [ + maybe_contiguous(x) for x in (page_table, kv_batch_idx, leftpad_k) + ] + rotary_cos, rotary_sin = [maybe_contiguous(x) for x in (rotary_cos, rotary_sin)] + seqlens_rotary = maybe_contiguous(seqlens_rotary) + out, softmax_lse, out_accum, softmax_lse_accum = flash_attn_3_cuda.fwd( + q, + k, + v, + k_new, + v_new, + qv, + out_, + cu_seqlens_q, + cu_seqlens_k, + cu_seqlens_k_new, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + page_table, + kv_batch_idx, + leftpad_k, + rotary_cos, + rotary_sin, + seqlens_rotary, + q_descale, + k_descale, + v_descale, + softmax_scale, + causal, + window_size_left, + window_size_right, + attention_chunk, + softcap, + rotary_interleaved, + scheduler_metadata, + num_splits, + pack_gqa, + sm_margin, + ) + + if out_accum is None: + out_accum = torch.tensor([], device=out.device) + + if softmax_lse_accum is None: + softmax_lse_accum = torch.tensor([], device=out.device) + + return out, softmax_lse, out_accum, softmax_lse_accum + + +@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_forward")) +def _flash_attn_forward_fake( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_new: Optional[torch.Tensor] = None, + v_new: Optional[torch.Tensor] = None, + qv: Optional[torch.Tensor] = None, + out_: Optional[torch.Tensor] = None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + seqused_q: Optional[torch.Tensor] = None, + seqused_k: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_k: Optional[int] = None, + page_table: Optional[torch.Tensor] = None, + kv_batch_idx: Optional[torch.Tensor] = None, + leftpad_k: Optional[torch.Tensor] = None, + rotary_cos: Optional[torch.Tensor] = None, + rotary_sin: Optional[torch.Tensor] = None, + seqlens_rotary: Optional[torch.Tensor] = None, + q_descale: Optional[torch.Tensor] = None, + k_descale: Optional[torch.Tensor] = None, + v_descale: Optional[torch.Tensor] = None, + softmax_scale: Optional[float] = None, + causal: bool = False, + window_size_left: int = -1, + window_size_right: int = -1, + attention_chunk: int = 0, + softcap: float = 0.0, + rotary_interleaved: bool = True, + scheduler_metadata: Optional[torch.Tensor] = None, + num_splits: int = 1, + pack_gqa: Optional[bool] = None, + sm_margin: int = 0, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Symbolic fake implementation of flash attention forward. + Returns tensors with the correct shapes and dtypes without actual computation. + """ + + # Determine if we're in varlen mode + is_varlen_q = cu_seqlens_q is not None + + # Get dimensions from query tensor + if is_varlen_q: + # varlen mode: q is (total_q, num_heads, head_size) + total_q, num_heads, head_size = q.shape + batch_size = cu_seqlens_q.shape[0] - 1 + + if max_seqlen_q is None: + raise ValueError("max_seqlen_q must be provided if cu_seqlens_q is provided") + seqlen_q = max_seqlen_q + else: + # batch mode: q is (batch_size, seqlen_q, num_heads, head_size) + batch_size, seqlen_q, num_heads, head_size = q.shape + total_q = batch_size * q.shape[1] + # Get value head dimension + head_size_v = v.shape[-1] + + # Determine output dtype (FP8 inputs produce BF16 outputs) + q_type = q.dtype + if q_type == torch.float8_e4m3fn: + out_dtype = torch.bfloat16 + else: + out_dtype = q_type + + # Create output tensor + if out_ is not None: + # If out_ is provided, _flash_attn_forward becomes non-functional + raise TypeError("Tracing (torch.compile/torch.export) with pre-allocated output tensor is not supported.") + + if is_varlen_q: + out = torch.empty((total_q, num_heads, head_size_v), dtype=out_dtype, device=q.device) + else: + out = torch.empty((batch_size, seqlen_q, num_heads, head_size_v), dtype=out_dtype, device=q.device) + + # Create softmax_lse tensor + if is_varlen_q: + softmax_lse = torch.empty((num_heads, total_q), dtype=torch.float32, device=q.device) + else: + softmax_lse = torch.empty((batch_size, num_heads, seqlen_q), dtype=torch.float32, device=q.device) + + # TODO(guilhermeleobas): Implement "get_num_splits" + # There's an heuristic to compute num_splits when "num_splits <= 0" + # assert that num_splits is > 0 for now + if num_splits <= 0: + raise ValueError(f"tracing (torch.compile/torch.export) with num_splits <= 0 not supported. Got {num_splits=}") + + if num_splits > 1: + if is_varlen_q: + out_accum = torch.empty((num_splits, num_heads, total_q, head_size_v), dtype=torch.float32, device=q.device) + softmax_lse_accum = torch.empty((num_splits, num_heads, total_q), dtype=torch.float32, device=q.device) + else: + out_accum = torch.empty((num_splits, batch_size, num_heads, seqlen_q, head_size_v), dtype=torch.float32, device=q.device) + softmax_lse_accum = torch.empty((num_splits, batch_size, num_heads, seqlen_q), dtype=torch.float32, device=q.device) + else: + # Tensors are not set when num_splits < 1 + out_accum = torch.tensor([], device=out.device) + softmax_lse_accum = torch.tensor([], device=out.device) + + return out, softmax_lse, out_accum, softmax_lse_accum + + +@torch.library.custom_op(add_op_namespace_prefix("_flash_attn_backward"), mutates_args=("dq", "dk", "dv"), device_types="cuda") +def _flash_attn_backward( + dout: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + out: torch.Tensor, + softmax_lse: torch.Tensor, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k: Optional[torch.Tensor] = None, + sequed_q: Optional[torch.Tensor] = None, + sequed_k: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_k: Optional[int] = None, + dq: Optional[torch.Tensor] = None, + dk: Optional[torch.Tensor] = None, + dv: Optional[torch.Tensor] = None, + softmax_scale: Optional[float] = None, + is_causal: bool = False, + window_size_left: int = -1, + window_size_right: int = -1, + softcap: float = 0.0, + deterministic: bool = False, + sm_margin: int = 0, +) -> torch.Tensor: + # dq, dk, dv are allocated by us so they should already be contiguous + dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] + softmax_d, *rest = flash_attn_3_cuda.bwd( + dout, + q, + k, + v, + out, + softmax_lse, + dq, + dk, + dv, + cu_seqlens_q, + cu_seqlens_k, + sequed_q, + sequed_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + is_causal, + window_size_left, + window_size_right, + softcap, + deterministic, + sm_margin, + ) + return softmax_d + + +@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_backward")) +def _flash_attn_backward_fake( + dout: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + out: torch.Tensor, + softmax_lse: torch.Tensor, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k: Optional[torch.Tensor] = None, + sequed_q: Optional[torch.Tensor] = None, + sequed_k: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_k: Optional[int] = None, + dq: Optional[torch.Tensor] = None, + dk: Optional[torch.Tensor] = None, + dv: Optional[torch.Tensor] = None, + softmax_scale: Optional[float] = None, + is_causal: bool = False, + window_size_left: int = -1, + window_size_right: int = -1, + softcap: float = 0.0, + deterministic: bool = False, + sm_margin: int = 0, +) -> torch.Tensor: + + is_varlen_q = cu_seqlens_q is not None + is_varlen_k = cu_seqlens_q is not None + is_varlen = is_varlen_q or is_varlen_k or sequed_q is not None or sequed_k is not None + + if not is_varlen_q: + batch_size = q.size(0) + seqlen_q = q.size(1) + seqlen_k = k.size(1) + total_q = batch_size * q.size(1) + else: + batch_size = cu_seqlens_q.size(0) - 1 + total_q = q.size(0) + seqlen_q = max_seqlen_q + seqlen_k = max_seqlen_k + + if window_size_left >= seqlen_k - 1: + window_size_left = -1 + + if window_size_right >= seqlen_q - 1: + window_size_right = -1 + + if is_causal: + window_size_right = 0 + + is_causal = window_size_left < 0 and window_size_right == 0 + + head_size = q.size(-1) + head_size_v = v.size(-1) + head_size_rounded = round_up_headdim(max(head_size, head_size_v)) + + # Hopper gpus uses cuda compute capabilities 9.0 + cap = torch.cuda.get_device_capability(q.device) + arch = cap[0] * 10 + cap[1] + + is_local = (window_size_left >= 0 or window_size_right >= 0) and not is_causal + + if head_size_rounded <= 64: + kBlockM_sm90 = 96 if (is_causal and softcap > 0.0) else 128 + elif head_size_rounded <= 96: + kBlockM_sm90 = 64 + elif head_size_rounded <= 128: + kBlockM_sm90 = 64 if (is_causal or is_local or softcap > 0.0) else 80 + else: + kBlockM_sm90 = 64 + + kBlockM_sm80 = 128 if head_size_rounded <= 64 else 64 + kBlockM_sm86 = 64 if head_size_rounded <= 192 else 32 + + if arch >= 90: + kBlockM = kBlockM_sm90 + elif arch == 86 or arch == 89: + kBlockM = kBlockM_sm86 + else: + kBlockM = kBlockM_sm80 + + num_heads = q.shape[-2] + seqlen_q_rounded = round_multiple(seqlen_q, kBlockM) + + total_q_padded_rounded = round_multiple(total_q + batch_size * kBlockM, kBlockM) + + dq = torch.empty_like(q) if dq is None else dq + dk = torch.empty_like(k) if dk is None else dk + dv = torch.empty_like(v) if dv is None else dv + + if not is_varlen: + softmax_d = torch.empty((batch_size, num_heads, seqlen_q_rounded), dtype=torch.float32, device=q.device) + else: + softmax_d = torch.empty((num_heads, total_q_padded_rounded), dtype=torch.float32, device=q.device) + + return softmax_d + + +def setup_context(ctx, inputs, output): + q, k, v = inputs[:3] + out, softmax_lse, _, _ = output + ctx.save_for_backward(q, k, v, out, softmax_lse) + ctx.softmax_scale = inputs[-11] + ctx.causal = inputs[-10] + ctx.window_size = [inputs[-9], inputs[-8]] + ctx.attention_chunk = inputs[-7] + ctx.softcap = inputs[-6] + ctx.sm_margin = inputs[-1] + + +def _backward(ctx, dout, *grads): + q, k, v, out, softmax_lse = ctx.saved_tensors + dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + None, None, # cu_seqlens_q, cu_seqlens_k, + None, None, # sequed_q, sequed_k, + None, None, # max_seqlen_q, max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size[0], + ctx.window_size[1], + ctx.softcap, + False, # deterministic + ctx.sm_margin, + ) + return dq, dk, dv, *((None,) * 21) + + +_flash_attn_forward.register_autograd(_backward, setup_context=setup_context) + + + +class FlashAttnQKVPackedFunc(torch.autograd.Function): + @staticmethod + def forward( + ctx, + qkv, + softmax_scale, + causal, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + deterministic=False, + num_heads_q=None, + sm_margin=0, + return_softmax=False, + ): + if softmax_scale is None: + softmax_scale = qkv.shape[-1] ** (-0.5) + if qkv.dim() == 5: + assert qkv.shape[-3] == 3 + q, k, v = qkv.unbind(dim=-3) + else: + assert qkv.dim() == 4 + assert num_heads_q is not None + num_heads_k = (qkv.shape[2] - num_heads_q) // 2 + assert num_heads_k * 2 + num_heads_q == qkv.shape[2] + q, k, v = qkv.split([num_heads_q, num_heads_k, num_heads_k], dim=-2) + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + None, # qv + None, # out + None, None, None, # cu_seqlens_q/k/k_new + None, None, # seqused_q/k + None, None, # max_seqlen_q/k + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size_left=window_size[0], + window_size_right=window_size[1], + attention_chunk=attention_chunk, + softcap=softcap, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse) + ctx.save_for_backward(q, k, v, out, softmax_lse) + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.ndim = qkv.dim() + ctx.sm_margin = sm_margin + return (out, softmax_lse) if return_softmax else out + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + if ctx.ndim == 5: + qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) + dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) + dq, dk, dv = dqkv.unbind(dim=-3) + else: + num_heads_q = q.shape[2] + num_heads_k = k.shape[2] + qkv_shape = q.shape[:-2] + (num_heads_q + num_heads_k * 2, *q.shape[-1:]) + dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) + dq, dk, dv = dqkv.split([num_heads_q, num_heads_k, num_heads_k], dim=-2) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + None, None, # cu_seqlens_q, cu_seqlens_k, + None, None, # sequed_q, sequed_k, + None, None, # max_seqlen_q, max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size[0], + ctx.window_size[1], + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dqkv = dqkv[..., : dout.shape[-1]] # We could have padded the head dimension + return dqkv, None, None, None, None, None, None, None, None, None, None, None, None + + +class FlashAttnFunc(torch.autograd.Function): + + @staticmethod + def forward( + ctx, + q, + k, + v, + softmax_scale, + causal, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + return_softmax=False, + ): + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + # out, q, k, v, out_padded, softmax_lse = _flash_attn_forward( + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + qv, # qv + None, # out + None, None, None, # cu_seqlens_q/k/k_new + None, None, # seqused_q/k + None, None, # max_seqlen_q/k + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size_left=window_size[0], + window_size_right=window_size[1], + attention_chunk=attention_chunk, + softcap=softcap, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse) + ctx.save_for_backward(q, k, v, out, softmax_lse) + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.sm_margin = sm_margin + return (out, softmax_lse) if return_softmax else out + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + None, None, # cu_seqlens_q, cu_seqlens_k, + None, None, # sequed_q, sequed_k, + None, None, # max_seqlen_q, max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size[0], + ctx.window_size[1], + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dq = dq[..., : q.shape[-1]] # We could have padded the head dimension + dk = dk[..., : k.shape[-1]] + dv = dv[..., : v.shape[-1]] + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None + + +class FlashAttnVarlenFunc(torch.autograd.Function): + + @staticmethod + def forward( + ctx, + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + return_softmax=False, + ): + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + # out, q, k, v, out_padded, softmax_lse = _flash_attn_varlen_forward( + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + qv, # qv + None, # out + cu_seqlens_q, + cu_seqlens_k, + None, # cu_seqlens_k_new + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size_left=window_size[0], + window_size_right=window_size[1], + attention_chunk=attention_chunk, + softcap=softcap, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k) + ctx.save_for_backward(q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k) + ctx.max_seqlen_q = max_seqlen_q + ctx.max_seqlen_k = max_seqlen_k + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.sm_margin = sm_margin + return (out, softmax_lse) if return_softmax else out + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + ctx.max_seqlen_q, + ctx.max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size[0], + ctx.window_size[1], + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dq = dq[..., : q.shape[-1]] # We could have padded the head dimension + dk = dk[..., : k.shape[-1]] + dv = dv[..., : v.shape[-1]] + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None + + +def flash_attn_qkvpacked_func( + qkv, + softmax_scale=None, + causal=False, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + deterministic=False, + num_heads_q=None, + sm_margin=0, + return_attn_probs=False, +): + """dropout_p should be set to 0.0 during evaluation + If Q, K, V are already stacked into 1 tensor, this function will be faster than + calling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation + of the gradients of Q, K, V. + For multi-query and grouped-query attention (MQA/GQA), please see + flash_attn_kvpacked_func and flash_attn_func. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between [i - window_size[0], i + window_size[1]] inclusive. + + Arguments: + qkv: (batch_size, seqlen, 3, nheads, headdim) + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + softcap: float. Anything > 0 activates softcapping attention. + alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of (-alibi_slope * |i - j|) is added to + the attention score of query i and key j. + deterministic: bool. Whether to use the deterministic implementation of the backward pass, + which is slightly slower and uses more memory. The forward pass is always deterministic. + return_attn_probs: bool. Whether to return the attention probabilities. This option is for + testing only. The returned probabilities are not guaranteed to be correct + (they might not have the right scaling). + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen). + The output of softmax (possibly with different scaling). It also encodes the dropout + pattern (negative means that location was dropped, nonnegative means it was kept). + """ + return FlashAttnQKVPackedFunc.apply( + qkv, + softmax_scale, + causal, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + deterministic, + num_heads_q, + sm_margin, + return_attn_probs, + ) + + +def flash_attn_func( + q, + k, + v, + softmax_scale=None, + causal=False, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + return_attn_probs=False, +): + """dropout_p should be set to 0.0 during evaluation + Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads + than Q. Note that the number of heads in Q must be divisible by the number of heads in KV. + For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head + 0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V. + + If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix. + For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is: + 1 1 1 1 0 + 1 1 1 1 1 + If seqlen_q = 5 and seqlen_k = 2, the causal mask is: + 0 0 + 0 0 + 0 0 + 1 0 + 1 1 + If the row of the mask is all zero, the output will be zero. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between + [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive. + + Arguments: + q: (batch_size, seqlen, nheads, headdim) + k: (batch_size, seqlen, nheads_k, headdim) + v: (batch_size, seqlen, nheads_k, headdim) + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of + (-alibi_slope * |i + seqlen_k - seqlen_q - j|) + is added to the attention score of query i and key j. + deterministic: bool. Whether to use the deterministic implementation of the backward pass, + which is slightly slower and uses more memory. The forward pass is always deterministic. + return_attn_probs: bool. Whether to return the attention probabilities. This option is for + testing only. The returned probabilities are not guaranteed to be correct + (they might not have the right scaling). + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + """ + return FlashAttnFunc.apply( + q, + k, + v, + softmax_scale, + causal, + qv, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + num_splits, + pack_gqa, + deterministic, + sm_margin, + return_attn_probs, + ) + + +def flash_attn_varlen_func( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + seqused_q=None, + seqused_k=None, + softmax_scale=None, + causal=False, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + return_attn_probs=False, +): + return FlashAttnVarlenFunc.apply( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + qv, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + num_splits, + pack_gqa, + deterministic, + sm_margin, + return_attn_probs, + ) + + +def flash_attn_combine(out_partial, lse_partial, out=None, out_dtype=None): + return flash_attn_3_cuda.fwd_combine(out_partial, lse_partial, out, out_dtype) + + +def flash_attn_with_kvcache( + q, + k_cache, + v_cache, + k=None, + v=None, + qv=None, + rotary_cos=None, + rotary_sin=None, + cache_seqlens: Optional[Union[(int, torch.Tensor)]] = None, + cache_batch_idx: Optional[torch.Tensor] = None, + cache_leftpad: Optional[torch.Tensor] = None, + page_table: Optional[torch.Tensor] = None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + rotary_seqlens: Optional[torch.Tensor] = None, + q_descale: Optional[torch.Tensor] = None, + k_descale: Optional[torch.Tensor] = None, + v_descale: Optional[torch.Tensor] = None, + softmax_scale=None, + causal=False, + window_size=(-1, -1), # -1 means infinite context window + attention_chunk=0, + softcap=0.0, # 0.0 means deactivated + rotary_interleaved=True, + scheduler_metadata=None, + num_splits=0, # Can be tuned for speed + pack_gqa=None, # Can be tuned for speed + sm_margin=0, # Can be tuned if some SMs are used for communication + return_softmax_lse=False, +): + """ + If k and v are not None, k_cache and v_cache will be updated *inplace* with the new values from + k and v. This is useful for incremental decoding: you can pass in the cached keys/values from + the previous step, and update them with the new keys/values from the current step, and do + attention with the updated cache, all in 1 kernel. + + If you pass in k / v, you must make sure that the cache is large enough to hold the new values. + For example, the KV cache could be pre-allocated with the max sequence length, and you can use + cache_seqlens to keep track of the current sequence lengths of each sequence in the batch. + + Also apply rotary embedding if rotary_cos and rotary_sin are passed in. The key @k will be + rotated by rotary_cos and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc. + If causal or local (i.e., window_size != (-1, -1)), the query @q will be rotated by rotary_cos + and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc. + If not causal and not local, the query @q will be rotated by rotary_cos and rotary_sin at + indices cache_seqlens only (i.e. we consider all tokens in @q to be at position cache_seqlens). + + See tests/test_flash_attn.py::test_flash_attn_kvcache for examples of how to use this function. + + Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads + than Q. Note that the number of heads in Q must be divisible by the number of heads in KV. + For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head + 0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V. + + If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix. + For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is: + 1 1 1 1 0 + 1 1 1 1 1 + If seqlen_q = 5 and seqlen_k = 2, the causal mask is: + 0 0 + 0 0 + 0 0 + 1 0 + 1 1 + If the row of the mask is all zero, the output will be zero. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between + [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive. + + Note: Does not support backward pass. + + Arguments: + q: (batch_size, seqlen, nheads, headdim) + k_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim) if there's no page_table, + or (num_blocks, page_block_size, nheads_k, headdim) if there's a page_table (i.e. paged KV cache) + page_block_size can be arbitrary (e.g, 1, 2, 3, 64, etc.). + v_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim_v) if there's no page_table, + or (num_blocks, page_block_size, nheads_k, headdim_v) if there's a page_table (i.e. paged KV cache) + k [optional]: (batch_size, seqlen_new, nheads_k, headdim). If not None, we concatenate + k with k_cache, starting at the indices specified by cache_seqlens. + v [optional]: (batch_size, seqlen_new, nheads_k, headdim_v). Similar to k. + qv [optional]: (batch_size, seqlen, nheads, headdim_v) + rotary_cos [optional]: (seqlen_ro, rotary_dim / 2). If not None, we apply rotary embedding + to k and q. Only applicable if k and v are passed in. rotary_dim must be divisible by 16. + rotary_sin [optional]: (seqlen_ro, rotary_dim / 2). Similar to rotary_cos. + cache_seqlens: int, or (batch_size,), dtype torch.int32. The sequence lengths of the + KV cache. + cache_batch_idx: (batch_size,), dtype torch.int32. The indices used to index into the KV cache. + If None, we assume that the batch indices are [0, 1, 2, ..., batch_size - 1]. + If the indices are not distinct, and k and v are provided, the values updated in the cache + might come from any of the duplicate indices. + cache_leftpad: (batch_size,), dtype torch.int32. The index that the KV cache starts. If None, assume 0. + page_table [optional]: (batch_size, max_num_blocks_per_seq), dtype torch.int32. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + softcap: float. Anything > 0 activates softcapping attention. + rotary_interleaved: bool. Only applicable if rotary_cos and rotary_sin are passed in. + If True, rotary embedding will combine dimensions 0 & 1, 2 & 3, etc. If False, + rotary embedding will combine dimensions 0 & rotary_dim / 2, 1 & rotary_dim / 2 + 1 + (i.e. GPT-NeoX style). + num_splits: int. If > 1, split the key/value into this many chunks along the sequence. + If num_splits == 1, we don't split the key/value. If num_splits == 0, we use a heuristic + to automatically determine the number of splits. + Don't change this unless you know what you are doing. + return_softmax_lse: bool. Whether to return the logsumexp of the attention scores. + + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_softmax_lse=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + """ + assert k_cache.stride(-1) == 1, "k_cache must have contiguous last dimension" + assert v_cache.stride(-1) == 1, "v_cache must have contiguous last dimension" + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + if cache_seqlens is not None and isinstance(cache_seqlens, int): + cache_seqlens = torch.full( + (q.shape[0],), cache_seqlens, dtype=torch.int32, device=k_cache.device + ) + cache_seqlens = maybe_contiguous(cache_seqlens) + out, softmax_lse, *rest = _flash_attn_forward( + q, + k_cache, + v_cache, + k, + v, + qv, + None, # out + cu_seqlens_q, + None, # cu_seqlens_k + cu_seqlens_k_new, + None, # seqused_q + cache_seqlens, + max_seqlen_q, + None, # max_seqlen_k + page_table, + cache_batch_idx, + cache_leftpad, + rotary_cos, + rotary_sin, + rotary_seqlens, + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size_left=window_size[0], + window_size_right=window_size[1], + attention_chunk=attention_chunk, + softcap=softcap, + rotary_interleaved=rotary_interleaved, + scheduler_metadata=scheduler_metadata, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # return (out, softmax_lse) if return_softmax_lse else out + return (out, softmax_lse, *rest) if return_softmax_lse else out + + +def get_scheduler_metadata( + batch_size, max_seqlen_q, max_seqlen_k, num_heads_q, num_heads_kv, headdim, + cache_seqlens: torch.Tensor, + qkv_dtype=torch.bfloat16, + headdim_v=None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + cache_leftpad: Optional[torch.Tensor] = None, + page_size: Optional[int] = None, + max_seqlen_k_new=0, + causal=False, + window_size=(-1, -1), # -1 means infinite context window + attention_chunk=0, + has_softcap=False, + num_splits=0, # Can be tuned for speed + pack_gqa=None, # Can be tuned for speed + sm_margin=0, # Can be tuned if some SMs are used for communication +): + cache_seqlens = maybe_contiguous(cache_seqlens) + if headdim_v is None: + headdim_v = headdim + scheduler_metadata = flash_attn_3_cuda.get_scheduler_metadata( + batch_size, max_seqlen_q, max_seqlen_k, num_heads_q, num_heads_kv, headdim, headdim_v, + qkv_dtype, + cache_seqlens, + cu_seqlens_q, + None, # cu_seqlens_k + cu_seqlens_k_new, + None, # seqused_q + cache_leftpad, + page_size, + max_seqlen_k_new, + causal, + window_size[0], window_size[1], + attention_chunk, + has_softcap, + num_splits, + pack_gqa, + sm_margin, + ) + return scheduler_metadata diff --git a/build/torch211-cxx11-cu128-x86_64-linux/metadata.json b/build/torch211-cxx11-cu128-x86_64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..e0cc0c4f14abd10162c22499be39c75e6d922f94 --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/metadata.json @@ -0,0 +1,12 @@ +{ + "version": 1, + "license": "BSD-3-Clause", + "python-depends": [], + "backend": { + "type": "cuda", + "archs": [ + "8.0", + "9.0a" + ] + } +} diff --git a/build/torch211-cxx11-cu130-x86_64-linux/__init__.py b/build/torch211-cxx11-cu130-x86_64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3d12edd447e84d6e65c4fa9ab91ae3c3192884c1 --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/__init__.py @@ -0,0 +1,17 @@ +from .flash_attn_interface import ( + flash_attn_combine, + flash_attn_func, + flash_attn_qkvpacked_func, + flash_attn_varlen_func, + flash_attn_with_kvcache, + get_scheduler_metadata, +) + +__all__ = [ + "flash_attn_combine", + "flash_attn_func", + "flash_attn_qkvpacked_func", + "flash_attn_varlen_func", + "flash_attn_with_kvcache", + "get_scheduler_metadata", +] diff --git a/build/torch211-cxx11-cu130-x86_64-linux/_flash_attn3_cuda_e1d5be2.abi3.so b/build/torch211-cxx11-cu130-x86_64-linux/_flash_attn3_cuda_e1d5be2.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..55b21af6d75fc5da4d4f5856e6bc8f500884b374 --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/_flash_attn3_cuda_e1d5be2.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:bd55c8f7b848dce30414fd77aee49d9167b6d0a8f73e6832c47de618ac58b24f +size 823692344 diff --git a/build/torch211-cxx11-cu130-x86_64-linux/_ops.py b/build/torch211-cxx11-cu130-x86_64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..b845f529a7b7c424bbe0db01a32c2cd6db45bb1a --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _flash_attn3_cuda_e1d5be2 +ops = torch.ops._flash_attn3_cuda_e1d5be2 + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_flash_attn3_cuda_e1d5be2::{op_name}" diff --git a/build/torch211-cxx11-cu130-x86_64-linux/flash_attn3/__init__.py b/build/torch211-cxx11-cu130-x86_64-linux/flash_attn3/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a9b2672c1cd85b74c1b3ded0fc0b2100e1aeac23 --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/flash_attn3/__init__.py @@ -0,0 +1,26 @@ +import ctypes +import importlib.util +import sys +from pathlib import Path +from types import ModuleType + + +def _import_from_path(file_path: Path) -> ModuleType: + # We cannot use the module name as-is, after adding it to `sys.modules`, + # it would also be used for other imports. So, we make a module name that + # depends on the path for it to be unique using the hex-encoded hash of + # the path. + path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value) + module_name = path_hash + spec = importlib.util.spec_from_file_location(module_name, file_path) + if spec is None: + raise ImportError(f"Cannot load spec for {module_name} from {file_path}") + module = importlib.util.module_from_spec(spec) + if module is None: + raise ImportError(f"Cannot load module {module_name} from spec") + sys.modules[module_name] = module + spec.loader.exec_module(module) # type: ignore + return module + + +globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py"))) diff --git a/build/torch211-cxx11-cu130-x86_64-linux/flash_attn_config.py b/build/torch211-cxx11-cu130-x86_64-linux/flash_attn_config.py new file mode 100644 index 0000000000000000000000000000000000000000..c4b7130324be2c4cabee7ada7000e170a2a81865 --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/flash_attn_config.py @@ -0,0 +1,7 @@ +# Auto-generated by flash attention 3 setup.py +CONFIG = {'build_flags': {'FLASHATTENTION_DISABLE_BACKWARD': False, 'FLASHATTENTION_DISABLE_SPLIT': False, 'FLASHATTENTION_DISABLE_PAGEDKV': False, 'FLASHATTENTION_DISABLE_APPENDKV': False, 'FLASHATTENTION_DISABLE_LOCAL': False, 'FLASHATTENTION_DISABLE_SOFTCAP': False, 'FLASHATTENTION_DISABLE_PACKGQA': False, 'FLASHATTENTION_DISABLE_FP16': False, 'FLASHATTENTION_DISABLE_FP8': False, 'FLASHATTENTION_DISABLE_VARLEN': False, 'FLASHATTENTION_DISABLE_CLUSTER': False, 'FLASHATTENTION_DISABLE_HDIM64': False, 'FLASHATTENTION_DISABLE_HDIM96': False, 'FLASHATTENTION_DISABLE_HDIM128': False, 'FLASHATTENTION_DISABLE_HDIM192': False, 'FLASHATTENTION_DISABLE_HDIM256': False, 'FLASHATTENTION_DISABLE_SM8x': False, 'FLASHATTENTION_ENABLE_VCOLMAJOR': False, 'FLASH_ATTENTION_DISABLE_HDIMDIFF64': False, 'FLASH_ATTENTION_DISABLE_HDIMDIFF192': False}} + +def show(): + from pprint import pprint + pprint(CONFIG) + diff --git a/build/torch211-cxx11-cu130-x86_64-linux/flash_attn_interface.py b/build/torch211-cxx11-cu130-x86_64-linux/flash_attn_interface.py new file mode 100644 index 0000000000000000000000000000000000000000..2e2a3cdd3e89cbc59bf6c956454e2776c305e0e5 --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/flash_attn_interface.py @@ -0,0 +1,1127 @@ +# Copyright (c) 2023, Tri Dao. + +from typing import Optional, Union, List, Tuple + +import torch +import torch.nn as nn + +from ._ops import ops as flash_attn_3_cuda +from ._ops import add_op_namespace_prefix + +def maybe_contiguous(x): + return x.contiguous() if x is not None and x.stride(-1) != 1 else x + + +def round_multiple(x, m): + return (x + m - 1) // m * m + + +def round_up_headdim(head_size: int) -> int: + from .flash_attn_config import CONFIG + + if not CONFIG["build_flags"]["FLASHATTENTION_DISABLE_HDIM64"]: + if head_size <= 64: + return 64 + if not CONFIG["build_flags"]["FLASHATTENTION_DISABLE_HDIM96"]: + if head_size <= 96: + return 96 + if not CONFIG["build_flags"]["FLASHATTENTION_DISABLE_HDIM128"]: + if head_size <= 128: + return 128 + if not CONFIG["build_flags"]["FLASHATTENTION_DISABLE_HDIM192"]: + if head_size <= 192: + return 192 + if not CONFIG["build_flags"]["FLASHATTENTION_DISABLE_HDIM256"]: + if head_size <= 256: + return 256 + return 256 + + +@torch.library.custom_op(add_op_namespace_prefix("_flash_attn_forward"), mutates_args=(), device_types="cuda") +def _flash_attn_forward( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_new: Optional[torch.Tensor] = None, + v_new: Optional[torch.Tensor] = None, + qv: Optional[torch.Tensor] = None, + out_: Optional[torch.Tensor] = None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + seqused_q: Optional[torch.Tensor] = None, + seqused_k: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_k: Optional[int] = None, + page_table: Optional[torch.Tensor] = None, + kv_batch_idx: Optional[torch.Tensor] = None, + leftpad_k: Optional[torch.Tensor] = None, + rotary_cos: Optional[torch.Tensor] = None, + rotary_sin: Optional[torch.Tensor] = None, + seqlens_rotary: Optional[torch.Tensor] = None, + q_descale: Optional[torch.Tensor] = None, + k_descale: Optional[torch.Tensor] = None, + v_descale: Optional[torch.Tensor] = None, + softmax_scale: Optional[float] = None, + causal: bool = False, + window_size_left: int = -1, + window_size_right: int = -1, + attention_chunk: int = 0, + softcap: float = 0.0, + rotary_interleaved: bool = True, + scheduler_metadata: Optional[torch.Tensor] = None, + num_splits: int = 1, + pack_gqa: Optional[bool] = None, + sm_margin: int = 0, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + q, k, k_new, v_new = [maybe_contiguous(x) for x in (q, k, k_new, v_new)] + v = v.contiguous() if v.stride(-1) != 1 and v.stride(-3) != 1 else v + cu_seqlens_q, cu_seqlens_k, cu_seqlens_k_new = [ + maybe_contiguous(x) for x in (cu_seqlens_q, cu_seqlens_k, cu_seqlens_k_new) + ] + seqused_q, seqused_k = [maybe_contiguous(x) for x in (seqused_q, seqused_k)] + page_table, kv_batch_idx, leftpad_k = [ + maybe_contiguous(x) for x in (page_table, kv_batch_idx, leftpad_k) + ] + rotary_cos, rotary_sin = [maybe_contiguous(x) for x in (rotary_cos, rotary_sin)] + seqlens_rotary = maybe_contiguous(seqlens_rotary) + out, softmax_lse, out_accum, softmax_lse_accum = flash_attn_3_cuda.fwd( + q, + k, + v, + k_new, + v_new, + qv, + out_, + cu_seqlens_q, + cu_seqlens_k, + cu_seqlens_k_new, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + page_table, + kv_batch_idx, + leftpad_k, + rotary_cos, + rotary_sin, + seqlens_rotary, + q_descale, + k_descale, + v_descale, + softmax_scale, + causal, + window_size_left, + window_size_right, + attention_chunk, + softcap, + rotary_interleaved, + scheduler_metadata, + num_splits, + pack_gqa, + sm_margin, + ) + + if out_accum is None: + out_accum = torch.tensor([], device=out.device) + + if softmax_lse_accum is None: + softmax_lse_accum = torch.tensor([], device=out.device) + + return out, softmax_lse, out_accum, softmax_lse_accum + + +@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_forward")) +def _flash_attn_forward_fake( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_new: Optional[torch.Tensor] = None, + v_new: Optional[torch.Tensor] = None, + qv: Optional[torch.Tensor] = None, + out_: Optional[torch.Tensor] = None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + seqused_q: Optional[torch.Tensor] = None, + seqused_k: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_k: Optional[int] = None, + page_table: Optional[torch.Tensor] = None, + kv_batch_idx: Optional[torch.Tensor] = None, + leftpad_k: Optional[torch.Tensor] = None, + rotary_cos: Optional[torch.Tensor] = None, + rotary_sin: Optional[torch.Tensor] = None, + seqlens_rotary: Optional[torch.Tensor] = None, + q_descale: Optional[torch.Tensor] = None, + k_descale: Optional[torch.Tensor] = None, + v_descale: Optional[torch.Tensor] = None, + softmax_scale: Optional[float] = None, + causal: bool = False, + window_size_left: int = -1, + window_size_right: int = -1, + attention_chunk: int = 0, + softcap: float = 0.0, + rotary_interleaved: bool = True, + scheduler_metadata: Optional[torch.Tensor] = None, + num_splits: int = 1, + pack_gqa: Optional[bool] = None, + sm_margin: int = 0, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Symbolic fake implementation of flash attention forward. + Returns tensors with the correct shapes and dtypes without actual computation. + """ + + # Determine if we're in varlen mode + is_varlen_q = cu_seqlens_q is not None + + # Get dimensions from query tensor + if is_varlen_q: + # varlen mode: q is (total_q, num_heads, head_size) + total_q, num_heads, head_size = q.shape + batch_size = cu_seqlens_q.shape[0] - 1 + + if max_seqlen_q is None: + raise ValueError("max_seqlen_q must be provided if cu_seqlens_q is provided") + seqlen_q = max_seqlen_q + else: + # batch mode: q is (batch_size, seqlen_q, num_heads, head_size) + batch_size, seqlen_q, num_heads, head_size = q.shape + total_q = batch_size * q.shape[1] + # Get value head dimension + head_size_v = v.shape[-1] + + # Determine output dtype (FP8 inputs produce BF16 outputs) + q_type = q.dtype + if q_type == torch.float8_e4m3fn: + out_dtype = torch.bfloat16 + else: + out_dtype = q_type + + # Create output tensor + if out_ is not None: + # If out_ is provided, _flash_attn_forward becomes non-functional + raise TypeError("Tracing (torch.compile/torch.export) with pre-allocated output tensor is not supported.") + + if is_varlen_q: + out = torch.empty((total_q, num_heads, head_size_v), dtype=out_dtype, device=q.device) + else: + out = torch.empty((batch_size, seqlen_q, num_heads, head_size_v), dtype=out_dtype, device=q.device) + + # Create softmax_lse tensor + if is_varlen_q: + softmax_lse = torch.empty((num_heads, total_q), dtype=torch.float32, device=q.device) + else: + softmax_lse = torch.empty((batch_size, num_heads, seqlen_q), dtype=torch.float32, device=q.device) + + # TODO(guilhermeleobas): Implement "get_num_splits" + # There's an heuristic to compute num_splits when "num_splits <= 0" + # assert that num_splits is > 0 for now + if num_splits <= 0: + raise ValueError(f"tracing (torch.compile/torch.export) with num_splits <= 0 not supported. Got {num_splits=}") + + if num_splits > 1: + if is_varlen_q: + out_accum = torch.empty((num_splits, num_heads, total_q, head_size_v), dtype=torch.float32, device=q.device) + softmax_lse_accum = torch.empty((num_splits, num_heads, total_q), dtype=torch.float32, device=q.device) + else: + out_accum = torch.empty((num_splits, batch_size, num_heads, seqlen_q, head_size_v), dtype=torch.float32, device=q.device) + softmax_lse_accum = torch.empty((num_splits, batch_size, num_heads, seqlen_q), dtype=torch.float32, device=q.device) + else: + # Tensors are not set when num_splits < 1 + out_accum = torch.tensor([], device=out.device) + softmax_lse_accum = torch.tensor([], device=out.device) + + return out, softmax_lse, out_accum, softmax_lse_accum + + +@torch.library.custom_op(add_op_namespace_prefix("_flash_attn_backward"), mutates_args=("dq", "dk", "dv"), device_types="cuda") +def _flash_attn_backward( + dout: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + out: torch.Tensor, + softmax_lse: torch.Tensor, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k: Optional[torch.Tensor] = None, + sequed_q: Optional[torch.Tensor] = None, + sequed_k: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_k: Optional[int] = None, + dq: Optional[torch.Tensor] = None, + dk: Optional[torch.Tensor] = None, + dv: Optional[torch.Tensor] = None, + softmax_scale: Optional[float] = None, + is_causal: bool = False, + window_size_left: int = -1, + window_size_right: int = -1, + softcap: float = 0.0, + deterministic: bool = False, + sm_margin: int = 0, +) -> torch.Tensor: + # dq, dk, dv are allocated by us so they should already be contiguous + dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] + softmax_d, *rest = flash_attn_3_cuda.bwd( + dout, + q, + k, + v, + out, + softmax_lse, + dq, + dk, + dv, + cu_seqlens_q, + cu_seqlens_k, + sequed_q, + sequed_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + is_causal, + window_size_left, + window_size_right, + softcap, + deterministic, + sm_margin, + ) + return softmax_d + + +@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_backward")) +def _flash_attn_backward_fake( + dout: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + out: torch.Tensor, + softmax_lse: torch.Tensor, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k: Optional[torch.Tensor] = None, + sequed_q: Optional[torch.Tensor] = None, + sequed_k: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_k: Optional[int] = None, + dq: Optional[torch.Tensor] = None, + dk: Optional[torch.Tensor] = None, + dv: Optional[torch.Tensor] = None, + softmax_scale: Optional[float] = None, + is_causal: bool = False, + window_size_left: int = -1, + window_size_right: int = -1, + softcap: float = 0.0, + deterministic: bool = False, + sm_margin: int = 0, +) -> torch.Tensor: + + is_varlen_q = cu_seqlens_q is not None + is_varlen_k = cu_seqlens_q is not None + is_varlen = is_varlen_q or is_varlen_k or sequed_q is not None or sequed_k is not None + + if not is_varlen_q: + batch_size = q.size(0) + seqlen_q = q.size(1) + seqlen_k = k.size(1) + total_q = batch_size * q.size(1) + else: + batch_size = cu_seqlens_q.size(0) - 1 + total_q = q.size(0) + seqlen_q = max_seqlen_q + seqlen_k = max_seqlen_k + + if window_size_left >= seqlen_k - 1: + window_size_left = -1 + + if window_size_right >= seqlen_q - 1: + window_size_right = -1 + + if is_causal: + window_size_right = 0 + + is_causal = window_size_left < 0 and window_size_right == 0 + + head_size = q.size(-1) + head_size_v = v.size(-1) + head_size_rounded = round_up_headdim(max(head_size, head_size_v)) + + # Hopper gpus uses cuda compute capabilities 9.0 + cap = torch.cuda.get_device_capability(q.device) + arch = cap[0] * 10 + cap[1] + + is_local = (window_size_left >= 0 or window_size_right >= 0) and not is_causal + + if head_size_rounded <= 64: + kBlockM_sm90 = 96 if (is_causal and softcap > 0.0) else 128 + elif head_size_rounded <= 96: + kBlockM_sm90 = 64 + elif head_size_rounded <= 128: + kBlockM_sm90 = 64 if (is_causal or is_local or softcap > 0.0) else 80 + else: + kBlockM_sm90 = 64 + + kBlockM_sm80 = 128 if head_size_rounded <= 64 else 64 + kBlockM_sm86 = 64 if head_size_rounded <= 192 else 32 + + if arch >= 90: + kBlockM = kBlockM_sm90 + elif arch == 86 or arch == 89: + kBlockM = kBlockM_sm86 + else: + kBlockM = kBlockM_sm80 + + num_heads = q.shape[-2] + seqlen_q_rounded = round_multiple(seqlen_q, kBlockM) + + total_q_padded_rounded = round_multiple(total_q + batch_size * kBlockM, kBlockM) + + dq = torch.empty_like(q) if dq is None else dq + dk = torch.empty_like(k) if dk is None else dk + dv = torch.empty_like(v) if dv is None else dv + + if not is_varlen: + softmax_d = torch.empty((batch_size, num_heads, seqlen_q_rounded), dtype=torch.float32, device=q.device) + else: + softmax_d = torch.empty((num_heads, total_q_padded_rounded), dtype=torch.float32, device=q.device) + + return softmax_d + + +def setup_context(ctx, inputs, output): + q, k, v = inputs[:3] + out, softmax_lse, _, _ = output + ctx.save_for_backward(q, k, v, out, softmax_lse) + ctx.softmax_scale = inputs[-11] + ctx.causal = inputs[-10] + ctx.window_size = [inputs[-9], inputs[-8]] + ctx.attention_chunk = inputs[-7] + ctx.softcap = inputs[-6] + ctx.sm_margin = inputs[-1] + + +def _backward(ctx, dout, *grads): + q, k, v, out, softmax_lse = ctx.saved_tensors + dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + None, None, # cu_seqlens_q, cu_seqlens_k, + None, None, # sequed_q, sequed_k, + None, None, # max_seqlen_q, max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size[0], + ctx.window_size[1], + ctx.softcap, + False, # deterministic + ctx.sm_margin, + ) + return dq, dk, dv, *((None,) * 21) + + +_flash_attn_forward.register_autograd(_backward, setup_context=setup_context) + + + +class FlashAttnQKVPackedFunc(torch.autograd.Function): + @staticmethod + def forward( + ctx, + qkv, + softmax_scale, + causal, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + deterministic=False, + num_heads_q=None, + sm_margin=0, + return_softmax=False, + ): + if softmax_scale is None: + softmax_scale = qkv.shape[-1] ** (-0.5) + if qkv.dim() == 5: + assert qkv.shape[-3] == 3 + q, k, v = qkv.unbind(dim=-3) + else: + assert qkv.dim() == 4 + assert num_heads_q is not None + num_heads_k = (qkv.shape[2] - num_heads_q) // 2 + assert num_heads_k * 2 + num_heads_q == qkv.shape[2] + q, k, v = qkv.split([num_heads_q, num_heads_k, num_heads_k], dim=-2) + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + None, # qv + None, # out + None, None, None, # cu_seqlens_q/k/k_new + None, None, # seqused_q/k + None, None, # max_seqlen_q/k + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size_left=window_size[0], + window_size_right=window_size[1], + attention_chunk=attention_chunk, + softcap=softcap, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse) + ctx.save_for_backward(q, k, v, out, softmax_lse) + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.ndim = qkv.dim() + ctx.sm_margin = sm_margin + return (out, softmax_lse) if return_softmax else out + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + if ctx.ndim == 5: + qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) + dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) + dq, dk, dv = dqkv.unbind(dim=-3) + else: + num_heads_q = q.shape[2] + num_heads_k = k.shape[2] + qkv_shape = q.shape[:-2] + (num_heads_q + num_heads_k * 2, *q.shape[-1:]) + dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) + dq, dk, dv = dqkv.split([num_heads_q, num_heads_k, num_heads_k], dim=-2) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + None, None, # cu_seqlens_q, cu_seqlens_k, + None, None, # sequed_q, sequed_k, + None, None, # max_seqlen_q, max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size[0], + ctx.window_size[1], + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dqkv = dqkv[..., : dout.shape[-1]] # We could have padded the head dimension + return dqkv, None, None, None, None, None, None, None, None, None, None, None, None + + +class FlashAttnFunc(torch.autograd.Function): + + @staticmethod + def forward( + ctx, + q, + k, + v, + softmax_scale, + causal, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + return_softmax=False, + ): + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + # out, q, k, v, out_padded, softmax_lse = _flash_attn_forward( + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + qv, # qv + None, # out + None, None, None, # cu_seqlens_q/k/k_new + None, None, # seqused_q/k + None, None, # max_seqlen_q/k + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size_left=window_size[0], + window_size_right=window_size[1], + attention_chunk=attention_chunk, + softcap=softcap, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse) + ctx.save_for_backward(q, k, v, out, softmax_lse) + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.sm_margin = sm_margin + return (out, softmax_lse) if return_softmax else out + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + None, None, # cu_seqlens_q, cu_seqlens_k, + None, None, # sequed_q, sequed_k, + None, None, # max_seqlen_q, max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size[0], + ctx.window_size[1], + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dq = dq[..., : q.shape[-1]] # We could have padded the head dimension + dk = dk[..., : k.shape[-1]] + dv = dv[..., : v.shape[-1]] + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None + + +class FlashAttnVarlenFunc(torch.autograd.Function): + + @staticmethod + def forward( + ctx, + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + return_softmax=False, + ): + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + # out, q, k, v, out_padded, softmax_lse = _flash_attn_varlen_forward( + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + qv, # qv + None, # out + cu_seqlens_q, + cu_seqlens_k, + None, # cu_seqlens_k_new + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size_left=window_size[0], + window_size_right=window_size[1], + attention_chunk=attention_chunk, + softcap=softcap, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k) + ctx.save_for_backward(q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k) + ctx.max_seqlen_q = max_seqlen_q + ctx.max_seqlen_k = max_seqlen_k + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.sm_margin = sm_margin + return (out, softmax_lse) if return_softmax else out + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + ctx.max_seqlen_q, + ctx.max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size[0], + ctx.window_size[1], + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dq = dq[..., : q.shape[-1]] # We could have padded the head dimension + dk = dk[..., : k.shape[-1]] + dv = dv[..., : v.shape[-1]] + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None + + +def flash_attn_qkvpacked_func( + qkv, + softmax_scale=None, + causal=False, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + deterministic=False, + num_heads_q=None, + sm_margin=0, + return_attn_probs=False, +): + """dropout_p should be set to 0.0 during evaluation + If Q, K, V are already stacked into 1 tensor, this function will be faster than + calling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation + of the gradients of Q, K, V. + For multi-query and grouped-query attention (MQA/GQA), please see + flash_attn_kvpacked_func and flash_attn_func. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between [i - window_size[0], i + window_size[1]] inclusive. + + Arguments: + qkv: (batch_size, seqlen, 3, nheads, headdim) + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + softcap: float. Anything > 0 activates softcapping attention. + alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of (-alibi_slope * |i - j|) is added to + the attention score of query i and key j. + deterministic: bool. Whether to use the deterministic implementation of the backward pass, + which is slightly slower and uses more memory. The forward pass is always deterministic. + return_attn_probs: bool. Whether to return the attention probabilities. This option is for + testing only. The returned probabilities are not guaranteed to be correct + (they might not have the right scaling). + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen). + The output of softmax (possibly with different scaling). It also encodes the dropout + pattern (negative means that location was dropped, nonnegative means it was kept). + """ + return FlashAttnQKVPackedFunc.apply( + qkv, + softmax_scale, + causal, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + deterministic, + num_heads_q, + sm_margin, + return_attn_probs, + ) + + +def flash_attn_func( + q, + k, + v, + softmax_scale=None, + causal=False, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + return_attn_probs=False, +): + """dropout_p should be set to 0.0 during evaluation + Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads + than Q. Note that the number of heads in Q must be divisible by the number of heads in KV. + For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head + 0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V. + + If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix. + For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is: + 1 1 1 1 0 + 1 1 1 1 1 + If seqlen_q = 5 and seqlen_k = 2, the causal mask is: + 0 0 + 0 0 + 0 0 + 1 0 + 1 1 + If the row of the mask is all zero, the output will be zero. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between + [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive. + + Arguments: + q: (batch_size, seqlen, nheads, headdim) + k: (batch_size, seqlen, nheads_k, headdim) + v: (batch_size, seqlen, nheads_k, headdim) + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of + (-alibi_slope * |i + seqlen_k - seqlen_q - j|) + is added to the attention score of query i and key j. + deterministic: bool. Whether to use the deterministic implementation of the backward pass, + which is slightly slower and uses more memory. The forward pass is always deterministic. + return_attn_probs: bool. Whether to return the attention probabilities. This option is for + testing only. The returned probabilities are not guaranteed to be correct + (they might not have the right scaling). + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + """ + return FlashAttnFunc.apply( + q, + k, + v, + softmax_scale, + causal, + qv, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + num_splits, + pack_gqa, + deterministic, + sm_margin, + return_attn_probs, + ) + + +def flash_attn_varlen_func( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + seqused_q=None, + seqused_k=None, + softmax_scale=None, + causal=False, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + return_attn_probs=False, +): + return FlashAttnVarlenFunc.apply( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + qv, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + num_splits, + pack_gqa, + deterministic, + sm_margin, + return_attn_probs, + ) + + +def flash_attn_combine(out_partial, lse_partial, out=None, out_dtype=None): + return flash_attn_3_cuda.fwd_combine(out_partial, lse_partial, out, out_dtype) + + +def flash_attn_with_kvcache( + q, + k_cache, + v_cache, + k=None, + v=None, + qv=None, + rotary_cos=None, + rotary_sin=None, + cache_seqlens: Optional[Union[(int, torch.Tensor)]] = None, + cache_batch_idx: Optional[torch.Tensor] = None, + cache_leftpad: Optional[torch.Tensor] = None, + page_table: Optional[torch.Tensor] = None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + rotary_seqlens: Optional[torch.Tensor] = None, + q_descale: Optional[torch.Tensor] = None, + k_descale: Optional[torch.Tensor] = None, + v_descale: Optional[torch.Tensor] = None, + softmax_scale=None, + causal=False, + window_size=(-1, -1), # -1 means infinite context window + attention_chunk=0, + softcap=0.0, # 0.0 means deactivated + rotary_interleaved=True, + scheduler_metadata=None, + num_splits=0, # Can be tuned for speed + pack_gqa=None, # Can be tuned for speed + sm_margin=0, # Can be tuned if some SMs are used for communication + return_softmax_lse=False, +): + """ + If k and v are not None, k_cache and v_cache will be updated *inplace* with the new values from + k and v. This is useful for incremental decoding: you can pass in the cached keys/values from + the previous step, and update them with the new keys/values from the current step, and do + attention with the updated cache, all in 1 kernel. + + If you pass in k / v, you must make sure that the cache is large enough to hold the new values. + For example, the KV cache could be pre-allocated with the max sequence length, and you can use + cache_seqlens to keep track of the current sequence lengths of each sequence in the batch. + + Also apply rotary embedding if rotary_cos and rotary_sin are passed in. The key @k will be + rotated by rotary_cos and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc. + If causal or local (i.e., window_size != (-1, -1)), the query @q will be rotated by rotary_cos + and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc. + If not causal and not local, the query @q will be rotated by rotary_cos and rotary_sin at + indices cache_seqlens only (i.e. we consider all tokens in @q to be at position cache_seqlens). + + See tests/test_flash_attn.py::test_flash_attn_kvcache for examples of how to use this function. + + Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads + than Q. Note that the number of heads in Q must be divisible by the number of heads in KV. + For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head + 0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V. + + If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix. + For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is: + 1 1 1 1 0 + 1 1 1 1 1 + If seqlen_q = 5 and seqlen_k = 2, the causal mask is: + 0 0 + 0 0 + 0 0 + 1 0 + 1 1 + If the row of the mask is all zero, the output will be zero. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between + [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive. + + Note: Does not support backward pass. + + Arguments: + q: (batch_size, seqlen, nheads, headdim) + k_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim) if there's no page_table, + or (num_blocks, page_block_size, nheads_k, headdim) if there's a page_table (i.e. paged KV cache) + page_block_size can be arbitrary (e.g, 1, 2, 3, 64, etc.). + v_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim_v) if there's no page_table, + or (num_blocks, page_block_size, nheads_k, headdim_v) if there's a page_table (i.e. paged KV cache) + k [optional]: (batch_size, seqlen_new, nheads_k, headdim). If not None, we concatenate + k with k_cache, starting at the indices specified by cache_seqlens. + v [optional]: (batch_size, seqlen_new, nheads_k, headdim_v). Similar to k. + qv [optional]: (batch_size, seqlen, nheads, headdim_v) + rotary_cos [optional]: (seqlen_ro, rotary_dim / 2). If not None, we apply rotary embedding + to k and q. Only applicable if k and v are passed in. rotary_dim must be divisible by 16. + rotary_sin [optional]: (seqlen_ro, rotary_dim / 2). Similar to rotary_cos. + cache_seqlens: int, or (batch_size,), dtype torch.int32. The sequence lengths of the + KV cache. + cache_batch_idx: (batch_size,), dtype torch.int32. The indices used to index into the KV cache. + If None, we assume that the batch indices are [0, 1, 2, ..., batch_size - 1]. + If the indices are not distinct, and k and v are provided, the values updated in the cache + might come from any of the duplicate indices. + cache_leftpad: (batch_size,), dtype torch.int32. The index that the KV cache starts. If None, assume 0. + page_table [optional]: (batch_size, max_num_blocks_per_seq), dtype torch.int32. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + softcap: float. Anything > 0 activates softcapping attention. + rotary_interleaved: bool. Only applicable if rotary_cos and rotary_sin are passed in. + If True, rotary embedding will combine dimensions 0 & 1, 2 & 3, etc. If False, + rotary embedding will combine dimensions 0 & rotary_dim / 2, 1 & rotary_dim / 2 + 1 + (i.e. GPT-NeoX style). + num_splits: int. If > 1, split the key/value into this many chunks along the sequence. + If num_splits == 1, we don't split the key/value. If num_splits == 0, we use a heuristic + to automatically determine the number of splits. + Don't change this unless you know what you are doing. + return_softmax_lse: bool. Whether to return the logsumexp of the attention scores. + + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_softmax_lse=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + """ + assert k_cache.stride(-1) == 1, "k_cache must have contiguous last dimension" + assert v_cache.stride(-1) == 1, "v_cache must have contiguous last dimension" + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + if cache_seqlens is not None and isinstance(cache_seqlens, int): + cache_seqlens = torch.full( + (q.shape[0],), cache_seqlens, dtype=torch.int32, device=k_cache.device + ) + cache_seqlens = maybe_contiguous(cache_seqlens) + out, softmax_lse, *rest = _flash_attn_forward( + q, + k_cache, + v_cache, + k, + v, + qv, + None, # out + cu_seqlens_q, + None, # cu_seqlens_k + cu_seqlens_k_new, + None, # seqused_q + cache_seqlens, + max_seqlen_q, + None, # max_seqlen_k + page_table, + cache_batch_idx, + cache_leftpad, + rotary_cos, + rotary_sin, + rotary_seqlens, + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size_left=window_size[0], + window_size_right=window_size[1], + attention_chunk=attention_chunk, + softcap=softcap, + rotary_interleaved=rotary_interleaved, + scheduler_metadata=scheduler_metadata, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # return (out, softmax_lse) if return_softmax_lse else out + return (out, softmax_lse, *rest) if return_softmax_lse else out + + +def get_scheduler_metadata( + batch_size, max_seqlen_q, max_seqlen_k, num_heads_q, num_heads_kv, headdim, + cache_seqlens: torch.Tensor, + qkv_dtype=torch.bfloat16, + headdim_v=None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + cache_leftpad: Optional[torch.Tensor] = None, + page_size: Optional[int] = None, + max_seqlen_k_new=0, + causal=False, + window_size=(-1, -1), # -1 means infinite context window + attention_chunk=0, + has_softcap=False, + num_splits=0, # Can be tuned for speed + pack_gqa=None, # Can be tuned for speed + sm_margin=0, # Can be tuned if some SMs are used for communication +): + cache_seqlens = maybe_contiguous(cache_seqlens) + if headdim_v is None: + headdim_v = headdim + scheduler_metadata = flash_attn_3_cuda.get_scheduler_metadata( + batch_size, max_seqlen_q, max_seqlen_k, num_heads_q, num_heads_kv, headdim, headdim_v, + qkv_dtype, + cache_seqlens, + cu_seqlens_q, + None, # cu_seqlens_k + cu_seqlens_k_new, + None, # seqused_q + cache_leftpad, + page_size, + max_seqlen_k_new, + causal, + window_size[0], window_size[1], + attention_chunk, + has_softcap, + num_splits, + pack_gqa, + sm_margin, + ) + return scheduler_metadata diff --git a/build/torch211-cxx11-cu130-x86_64-linux/metadata.json b/build/torch211-cxx11-cu130-x86_64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..e0cc0c4f14abd10162c22499be39c75e6d922f94 --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/metadata.json @@ -0,0 +1,12 @@ +{ + "version": 1, + "license": "BSD-3-Clause", + "python-depends": [], + "backend": { + "type": "cuda", + "archs": [ + "8.0", + "9.0a" + ] + } +} diff --git a/build/torch26-cxx11-cu124-x86_64-linux/flash_attn3/__init__.py b/build/torch26-cxx11-cu124-x86_64-linux/flash_attn3/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3d12edd447e84d6e65c4fa9ab91ae3c3192884c1 --- /dev/null +++ b/build/torch26-cxx11-cu124-x86_64-linux/flash_attn3/__init__.py @@ -0,0 +1,17 @@ +from .flash_attn_interface import ( + flash_attn_combine, + flash_attn_func, + flash_attn_qkvpacked_func, + flash_attn_varlen_func, + flash_attn_with_kvcache, + get_scheduler_metadata, +) + +__all__ = [ + "flash_attn_combine", + "flash_attn_func", + "flash_attn_qkvpacked_func", + "flash_attn_varlen_func", + "flash_attn_with_kvcache", + "get_scheduler_metadata", +] diff --git a/build/torch26-cxx11-cu124-x86_64-linux/flash_attn3/_flash_attn3_2e75662.abi3.so b/build/torch26-cxx11-cu124-x86_64-linux/flash_attn3/_flash_attn3_2e75662.abi3.so new file mode 100755 index 0000000000000000000000000000000000000000..9fab1a524613c478a8720bc4d58fda0574cfc225 --- /dev/null +++ b/build/torch26-cxx11-cu124-x86_64-linux/flash_attn3/_flash_attn3_2e75662.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:21b44e8e5e447a8b8ee051d347f0e32a3446a750f79d0bd1755e553f2119aa3b +size 838459656 diff --git a/build/torch26-cxx11-cu124-x86_64-linux/flash_attn3/_flash_attn3_557701f.abi3.so b/build/torch26-cxx11-cu124-x86_64-linux/flash_attn3/_flash_attn3_557701f.abi3.so new file mode 100755 index 0000000000000000000000000000000000000000..67c719849cf59ed70335a7ab4d13ea28c41c17a7 --- /dev/null +++ b/build/torch26-cxx11-cu124-x86_64-linux/flash_attn3/_flash_attn3_557701f.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:12d4ff964085fd02252777a2008f5ca47c90ea6a93da590e2fc5065dd5330207 +size 838459656 diff --git a/build/torch26-cxx11-cu124-x86_64-linux/flash_attn3/_ops.py b/build/torch26-cxx11-cu124-x86_64-linux/flash_attn3/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..0c0daa483e5a131df44714bd9cb4c4c8143bdcff --- /dev/null +++ b/build/torch26-cxx11-cu124-x86_64-linux/flash_attn3/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _flash_attn3_557701f +ops = torch.ops._flash_attn3_557701f + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_flash_attn3_557701f::{op_name}" \ No newline at end of file diff --git a/build/torch26-cxx11-cu124-x86_64-linux/flash_attn3/flash_attn_interface.py b/build/torch26-cxx11-cu124-x86_64-linux/flash_attn3/flash_attn_interface.py new file mode 100644 index 0000000000000000000000000000000000000000..7572297784a73fdf891af6aafee88428120b72f2 --- /dev/null +++ b/build/torch26-cxx11-cu124-x86_64-linux/flash_attn3/flash_attn_interface.py @@ -0,0 +1,828 @@ +# Copyright (c) 2023, Tri Dao. + +from typing import Optional, Union + +import torch +import torch.nn as nn + +from ._ops import ops as flash_attn_3_cuda + +def maybe_contiguous(x): + return x.contiguous() if x is not None and x.stride(-1) != 1 else x + + +def _flash_attn_forward( + q, + k, + v, + k_new, + v_new, + qv, + out, + cu_seqlens_q, + cu_seqlens_k, + cu_seqlens_k_new, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + page_table, + kv_batch_idx, + leftpad_k, + rotary_cos, + rotary_sin, + seqlens_rotary, + q_descale, + k_descale, + v_descale, + softmax_scale, + causal, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + rotary_interleaved=True, + scheduler_metadata=None, + num_splits=1, + pack_gqa=None, + sm_margin=0): + q, k, k_new, v_new = [maybe_contiguous(x) for x in (q, k, k_new, v_new)] + v = v.contiguous() if v.stride(-1) != 1 and v.stride(-3) != 1 else v + cu_seqlens_q, cu_seqlens_k, cu_seqlens_k_new = [ + maybe_contiguous(x) for x in (cu_seqlens_q, cu_seqlens_k, cu_seqlens_k_new) + ] + seqused_q, seqused_k = [maybe_contiguous(x) for x in (seqused_q, seqused_k)] + page_table, kv_batch_idx, leftpad_k = [ + maybe_contiguous(x) for x in (page_table, kv_batch_idx, leftpad_k) + ] + rotary_cos, rotary_sin = [maybe_contiguous(x) for x in (rotary_cos, rotary_sin)] + seqlens_rotary = maybe_contiguous(seqlens_rotary) + out, softmax_lse, *rest = flash_attn_3_cuda.fwd( + q, + k, + v, + k_new, + v_new, + qv, + out, + cu_seqlens_q, + cu_seqlens_k, + cu_seqlens_k_new, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + page_table, + kv_batch_idx, + leftpad_k, + rotary_cos, + rotary_sin, + seqlens_rotary, + q_descale, + k_descale, + v_descale, + softmax_scale, + causal, + window_size[0], + window_size[1], + attention_chunk, + softcap, + rotary_interleaved, + scheduler_metadata, + num_splits, + pack_gqa, + sm_margin, + ) + return out, softmax_lse, *rest + + +def _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + cu_seqlens_q, + cu_seqlens_k, + sequed_q, + sequed_k, + max_seqlen_q, + max_seqlen_k, + dq, + dk, + dv, + softmax_scale, + causal, + window_size=(-1, -1), + softcap=0.0, + deterministic=False, + sm_margin=0, +): + # dq, dk, dv are allocated by us so they should already be contiguous + dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] + dq, dk, dv, softmax_d, *rest = flash_attn_3_cuda.bwd( + dout, + q, + k, + v, + out, + softmax_lse, + dq, + dk, + dv, + cu_seqlens_q, + cu_seqlens_k, + sequed_q, + sequed_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + window_size[0], + window_size[1], + softcap, + deterministic, + sm_margin, + ) + return dq, dk, dv, softmax_d + + +class FlashAttnQKVPackedFunc(torch.autograd.Function): + @staticmethod + def forward( + ctx, + qkv, + softmax_scale, + causal, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + deterministic=False, + num_heads_q=None, + sm_margin=0, + ): + if softmax_scale is None: + softmax_scale = qkv.shape[-1] ** (-0.5) + if qkv.dim() == 5: + assert qkv.shape[-3] == 3 + q, k, v = qkv.unbind(dim=-3) + else: + assert qkv.dim() == 4 + assert num_heads_q is not None + num_heads_k = (qkv.shape[2] - num_heads_q) // 2 + assert num_heads_k * 2 + num_heads_q == qkv.shape[2] + q, k, v = qkv.split([num_heads_q, num_heads_k, num_heads_k], dim=-2) + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + None, # qv + None, # out + None, None, None, # cu_seqlens_q/k/k_new + None, None, # seqused_q/k + None, None, # max_seqlen_q/k + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size=window_size, + attention_chunk=attention_chunk, + softcap=softcap, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse) + ctx.save_for_backward(q, k, v, out, softmax_lse) + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.ndim = qkv.dim() + ctx.sm_margin = sm_margin + # return out, softmax_lse + return out + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + if ctx.ndim == 5: + qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) + dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) + dq, dk, dv = dqkv.unbind(dim=-3) + else: + num_heads_q = q.shape[2] + num_heads_k = k.shape[2] + qkv_shape = q.shape[:-2] + (num_heads_q + num_heads_k * 2, *q.shape[-1:]) + dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) + dq, dk, dv = dqkv.split([num_heads_q, num_heads_k, num_heads_k], dim=-2) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + None, None, # cu_seqlens_q, cu_seqlens_k, + None, None, # sequed_q, sequed_k, + None, None, # max_seqlen_q, max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size, + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dqkv = dqkv[..., : dout.shape[-1]] # We could have padded the head dimension + return dqkv, None, None, None, None, None, None, None, None, None, None, None + + +class FlashAttnFunc(torch.autograd.Function): + + @staticmethod + def forward( + ctx, + q, + k, + v, + softmax_scale, + causal, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + ): + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + # out, q, k, v, out_padded, softmax_lse = _flash_attn_forward( + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + qv, # qv + None, # out + None, None, None, # cu_seqlens_q/k/k_new + None, None, # seqused_q/k + None, None, # max_seqlen_q/k + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size=window_size, + attention_chunk=attention_chunk, + softcap=softcap, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse) + ctx.save_for_backward(q, k, v, out, softmax_lse) + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.sm_margin = sm_margin + return out, softmax_lse + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + None, None, # cu_seqlens_q, cu_seqlens_k, + None, None, # sequed_q, sequed_k, + None, None, # max_seqlen_q, max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size, + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dq = dq[..., : q.shape[-1]] # We could have padded the head dimension + dk = dk[..., : k.shape[-1]] + dv = dv[..., : v.shape[-1]] + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None + + +class FlashAttnVarlenFunc(torch.autograd.Function): + + @staticmethod + def forward( + ctx, + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + ): + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + # out, q, k, v, out_padded, softmax_lse = _flash_attn_varlen_forward( + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + qv, # qv + None, # out + cu_seqlens_q, + cu_seqlens_k, + None, # cu_seqlens_k_new + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size=window_size, + attention_chunk=attention_chunk, + softcap=softcap, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k) + ctx.save_for_backward(q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k) + ctx.max_seqlen_q = max_seqlen_q + ctx.max_seqlen_k = max_seqlen_k + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.sm_margin = sm_margin + return out, softmax_lse + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + ctx.max_seqlen_q, + ctx.max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size, + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dq = dq[..., : q.shape[-1]] # We could have padded the head dimension + dk = dk[..., : k.shape[-1]] + dv = dv[..., : v.shape[-1]] + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None + + +def flash_attn_qkvpacked_func( + qkv, + softmax_scale=None, + causal=False, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + deterministic=False, + num_heads_q=None, + sm_margin=0, +): + """dropout_p should be set to 0.0 during evaluation + If Q, K, V are already stacked into 1 tensor, this function will be faster than + calling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation + of the gradients of Q, K, V. + For multi-query and grouped-query attention (MQA/GQA), please see + flash_attn_kvpacked_func and flash_attn_func. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between [i - window_size[0], i + window_size[1]] inclusive. + + Arguments: + qkv: (batch_size, seqlen, 3, nheads, headdim) + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + softcap: float. Anything > 0 activates softcapping attention. + alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of (-alibi_slope * |i - j|) is added to + the attention score of query i and key j. + deterministic: bool. Whether to use the deterministic implementation of the backward pass, + which is slightly slower and uses more memory. The forward pass is always deterministic. + return_attn_probs: bool. Whether to return the attention probabilities. This option is for + testing only. The returned probabilities are not guaranteed to be correct + (they might not have the right scaling). + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen). + The output of softmax (possibly with different scaling). It also encodes the dropout + pattern (negative means that location was dropped, nonnegative means it was kept). + """ + return FlashAttnQKVPackedFunc.apply( + qkv, + softmax_scale, + causal, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + deterministic, + num_heads_q, + sm_margin, + ) + + +def flash_attn_func( + q, + k, + v, + softmax_scale=None, + causal=False, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, +): + """dropout_p should be set to 0.0 during evaluation + Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads + than Q. Note that the number of heads in Q must be divisible by the number of heads in KV. + For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head + 0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V. + + If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix. + For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is: + 1 1 1 1 0 + 1 1 1 1 1 + If seqlen_q = 5 and seqlen_k = 2, the causal mask is: + 0 0 + 0 0 + 0 0 + 1 0 + 1 1 + If the row of the mask is all zero, the output will be zero. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between + [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive. + + Arguments: + q: (batch_size, seqlen, nheads, headdim) + k: (batch_size, seqlen, nheads_k, headdim) + v: (batch_size, seqlen, nheads_k, headdim) + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of + (-alibi_slope * |i + seqlen_k - seqlen_q - j|) + is added to the attention score of query i and key j. + deterministic: bool. Whether to use the deterministic implementation of the backward pass, + which is slightly slower and uses more memory. The forward pass is always deterministic. + return_attn_probs: bool. Whether to return the attention probabilities. This option is for + testing only. The returned probabilities are not guaranteed to be correct + (they might not have the right scaling). + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + """ + return FlashAttnFunc.apply( + q, + k, + v, + softmax_scale, + causal, + qv, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + num_splits, + pack_gqa, + deterministic, + sm_margin, + ) + + +def flash_attn_varlen_func( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + seqused_q=None, + seqused_k=None, + softmax_scale=None, + causal=False, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, +): + return FlashAttnVarlenFunc.apply( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + qv, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + num_splits, + pack_gqa, + deterministic, + sm_margin, + ) + + +def flash_attn_combine(out_partial, lse_partial, out=None, out_dtype=None): + return flash_attn_3_cuda.fwd_combine(out_partial, lse_partial, out, out_dtype) + + +def flash_attn_with_kvcache( + q, + k_cache, + v_cache, + k=None, + v=None, + qv=None, + rotary_cos=None, + rotary_sin=None, + cache_seqlens: Optional[Union[(int, torch.Tensor)]] = None, + cache_batch_idx: Optional[torch.Tensor] = None, + cache_leftpad: Optional[torch.Tensor] = None, + page_table: Optional[torch.Tensor] = None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + rotary_seqlens: Optional[torch.Tensor] = None, + q_descale: Optional[torch.Tensor] = None, + k_descale: Optional[torch.Tensor] = None, + v_descale: Optional[torch.Tensor] = None, + softmax_scale=None, + causal=False, + window_size=(-1, -1), # -1 means infinite context window + attention_chunk=0, + softcap=0.0, # 0.0 means deactivated + rotary_interleaved=True, + scheduler_metadata=None, + num_splits=0, # Can be tuned for speed + pack_gqa=None, # Can be tuned for speed + sm_margin=0, # Can be tuned if some SMs are used for communication + return_softmax_lse=False, +): + """ + If k and v are not None, k_cache and v_cache will be updated *inplace* with the new values from + k and v. This is useful for incremental decoding: you can pass in the cached keys/values from + the previous step, and update them with the new keys/values from the current step, and do + attention with the updated cache, all in 1 kernel. + + If you pass in k / v, you must make sure that the cache is large enough to hold the new values. + For example, the KV cache could be pre-allocated with the max sequence length, and you can use + cache_seqlens to keep track of the current sequence lengths of each sequence in the batch. + + Also apply rotary embedding if rotary_cos and rotary_sin are passed in. The key @k will be + rotated by rotary_cos and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc. + If causal or local (i.e., window_size != (-1, -1)), the query @q will be rotated by rotary_cos + and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc. + If not causal and not local, the query @q will be rotated by rotary_cos and rotary_sin at + indices cache_seqlens only (i.e. we consider all tokens in @q to be at position cache_seqlens). + + See tests/test_flash_attn.py::test_flash_attn_kvcache for examples of how to use this function. + + Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads + than Q. Note that the number of heads in Q must be divisible by the number of heads in KV. + For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head + 0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V. + + If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix. + For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is: + 1 1 1 1 0 + 1 1 1 1 1 + If seqlen_q = 5 and seqlen_k = 2, the causal mask is: + 0 0 + 0 0 + 0 0 + 1 0 + 1 1 + If the row of the mask is all zero, the output will be zero. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between + [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive. + + Note: Does not support backward pass. + + Arguments: + q: (batch_size, seqlen, nheads, headdim) + k_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim) if there's no page_table, + or (num_blocks, page_block_size, nheads_k, headdim) if there's a page_table (i.e. paged KV cache) + page_block_size must be a multiple of 256. + v_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim_v) if there's no page_table, + or (num_blocks, page_block_size, nheads_k, headdim_v) if there's a page_table (i.e. paged KV cache) + k [optional]: (batch_size, seqlen_new, nheads_k, headdim). If not None, we concatenate + k with k_cache, starting at the indices specified by cache_seqlens. + v [optional]: (batch_size, seqlen_new, nheads_k, headdim_v). Similar to k. + qv [optional]: (batch_size, seqlen, nheads, headdim_v) + rotary_cos [optional]: (seqlen_ro, rotary_dim / 2). If not None, we apply rotary embedding + to k and q. Only applicable if k and v are passed in. rotary_dim must be divisible by 16. + rotary_sin [optional]: (seqlen_ro, rotary_dim / 2). Similar to rotary_cos. + cache_seqlens: int, or (batch_size,), dtype torch.int32. The sequence lengths of the + KV cache. + cache_batch_idx: (batch_size,), dtype torch.int32. The indices used to index into the KV cache. + If None, we assume that the batch indices are [0, 1, 2, ..., batch_size - 1]. + If the indices are not distinct, and k and v are provided, the values updated in the cache + might come from any of the duplicate indices. + cache_leftpad: (batch_size,), dtype torch.int32. The index that the KV cache starts. If None, assume 0. + page_table [optional]: (batch_size, max_num_blocks_per_seq), dtype torch.int32. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + softcap: float. Anything > 0 activates softcapping attention. + rotary_interleaved: bool. Only applicable if rotary_cos and rotary_sin are passed in. + If True, rotary embedding will combine dimensions 0 & 1, 2 & 3, etc. If False, + rotary embedding will combine dimensions 0 & rotary_dim / 2, 1 & rotary_dim / 2 + 1 + (i.e. GPT-NeoX style). + num_splits: int. If > 1, split the key/value into this many chunks along the sequence. + If num_splits == 1, we don't split the key/value. If num_splits == 0, we use a heuristic + to automatically determine the number of splits. + Don't change this unless you know what you are doing. + return_softmax_lse: bool. Whether to return the logsumexp of the attention scores. + + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_softmax_lse=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + """ + assert k_cache.stride(-1) == 1, "k_cache must have contiguous last dimension" + assert v_cache.stride(-1) == 1, "v_cache must have contiguous last dimension" + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + if cache_seqlens is not None and isinstance(cache_seqlens, int): + cache_seqlens = torch.full( + (k_cache.shape[0],), cache_seqlens, dtype=torch.int32, device=k_cache.device + ) + cache_seqlens = maybe_contiguous(cache_seqlens) + out, softmax_lse, *rest = _flash_attn_forward( + q, + k_cache, + v_cache, + k, + v, + qv, + None, # out + cu_seqlens_q, + None, # cu_seqlens_k + cu_seqlens_k_new, + None, # seqused_q + cache_seqlens, + max_seqlen_q, + None, # max_seqlen_k + page_table, + cache_batch_idx, + cache_leftpad, + rotary_cos, + rotary_sin, + rotary_seqlens, + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size=window_size, + attention_chunk=attention_chunk, + softcap=softcap, + rotary_interleaved=rotary_interleaved, + scheduler_metadata=scheduler_metadata, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # return (out, softmax_lse) if return_softmax_lse else out + return (out, softmax_lse, *rest) if return_softmax_lse else out + + +def get_scheduler_metadata( + batch_size, max_seqlen_q, max_seqlen_k, num_heads_q, num_heads_kv, headdim, + cache_seqlens: torch.Tensor, + qkv_dtype=torch.bfloat16, + headdim_v=None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + cache_leftpad: Optional[torch.Tensor] = None, + page_size: Optional[int] = None, + max_seqlen_k_new=0, + causal=False, + window_size=(-1, -1), # -1 means infinite context window + attention_chunk=0, + has_softcap=False, + num_splits=0, # Can be tuned for speed + pack_gqa=None, # Can be tuned for speed + sm_margin=0, # Can be tuned if some SMs are used for communication +): + cache_seqlens = maybe_contiguous(cache_seqlens) + if headdim_v is None: + headdim_v = headdim + scheduler_metadata = flash_attn_3_cuda.get_scheduler_metadata( + batch_size, max_seqlen_q, max_seqlen_k, num_heads_q, num_heads_kv, headdim, headdim_v, + qkv_dtype, + cache_seqlens, + cu_seqlens_q, + None, # cu_seqlens_k + cu_seqlens_k_new, + None, # seqused_q + cache_leftpad, + page_size, + max_seqlen_k_new, + causal, + window_size[0], window_size[1], + attention_chunk, + has_softcap, + num_splits, + pack_gqa, + sm_margin, + ) + return scheduler_metadata diff --git a/build/torch26-cxx11-cu126-x86_64-linux/flash_attn3/__init__.py b/build/torch26-cxx11-cu126-x86_64-linux/flash_attn3/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3d12edd447e84d6e65c4fa9ab91ae3c3192884c1 --- /dev/null +++ b/build/torch26-cxx11-cu126-x86_64-linux/flash_attn3/__init__.py @@ -0,0 +1,17 @@ +from .flash_attn_interface import ( + flash_attn_combine, + flash_attn_func, + flash_attn_qkvpacked_func, + flash_attn_varlen_func, + flash_attn_with_kvcache, + get_scheduler_metadata, +) + +__all__ = [ + "flash_attn_combine", + "flash_attn_func", + "flash_attn_qkvpacked_func", + "flash_attn_varlen_func", + "flash_attn_with_kvcache", + "get_scheduler_metadata", +] diff --git a/build/torch26-cxx11-cu126-x86_64-linux/flash_attn3/_flash_attn3_2e75662.abi3.so b/build/torch26-cxx11-cu126-x86_64-linux/flash_attn3/_flash_attn3_2e75662.abi3.so new file mode 100755 index 0000000000000000000000000000000000000000..9fab1a524613c478a8720bc4d58fda0574cfc225 --- /dev/null +++ b/build/torch26-cxx11-cu126-x86_64-linux/flash_attn3/_flash_attn3_2e75662.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:21b44e8e5e447a8b8ee051d347f0e32a3446a750f79d0bd1755e553f2119aa3b +size 838459656 diff --git a/build/torch26-cxx11-cu126-x86_64-linux/flash_attn3/_flash_attn3_557701f.abi3.so b/build/torch26-cxx11-cu126-x86_64-linux/flash_attn3/_flash_attn3_557701f.abi3.so new file mode 100755 index 0000000000000000000000000000000000000000..67c719849cf59ed70335a7ab4d13ea28c41c17a7 --- /dev/null +++ b/build/torch26-cxx11-cu126-x86_64-linux/flash_attn3/_flash_attn3_557701f.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:12d4ff964085fd02252777a2008f5ca47c90ea6a93da590e2fc5065dd5330207 +size 838459656 diff --git a/build/torch26-cxx11-cu126-x86_64-linux/flash_attn3/_ops.py b/build/torch26-cxx11-cu126-x86_64-linux/flash_attn3/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..0c0daa483e5a131df44714bd9cb4c4c8143bdcff --- /dev/null +++ b/build/torch26-cxx11-cu126-x86_64-linux/flash_attn3/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _flash_attn3_557701f +ops = torch.ops._flash_attn3_557701f + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_flash_attn3_557701f::{op_name}" \ No newline at end of file diff --git a/build/torch26-cxx11-cu126-x86_64-linux/flash_attn3/flash_attn_interface.py b/build/torch26-cxx11-cu126-x86_64-linux/flash_attn3/flash_attn_interface.py new file mode 100644 index 0000000000000000000000000000000000000000..7572297784a73fdf891af6aafee88428120b72f2 --- /dev/null +++ b/build/torch26-cxx11-cu126-x86_64-linux/flash_attn3/flash_attn_interface.py @@ -0,0 +1,828 @@ +# Copyright (c) 2023, Tri Dao. + +from typing import Optional, Union + +import torch +import torch.nn as nn + +from ._ops import ops as flash_attn_3_cuda + +def maybe_contiguous(x): + return x.contiguous() if x is not None and x.stride(-1) != 1 else x + + +def _flash_attn_forward( + q, + k, + v, + k_new, + v_new, + qv, + out, + cu_seqlens_q, + cu_seqlens_k, + cu_seqlens_k_new, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + page_table, + kv_batch_idx, + leftpad_k, + rotary_cos, + rotary_sin, + seqlens_rotary, + q_descale, + k_descale, + v_descale, + softmax_scale, + causal, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + rotary_interleaved=True, + scheduler_metadata=None, + num_splits=1, + pack_gqa=None, + sm_margin=0): + q, k, k_new, v_new = [maybe_contiguous(x) for x in (q, k, k_new, v_new)] + v = v.contiguous() if v.stride(-1) != 1 and v.stride(-3) != 1 else v + cu_seqlens_q, cu_seqlens_k, cu_seqlens_k_new = [ + maybe_contiguous(x) for x in (cu_seqlens_q, cu_seqlens_k, cu_seqlens_k_new) + ] + seqused_q, seqused_k = [maybe_contiguous(x) for x in (seqused_q, seqused_k)] + page_table, kv_batch_idx, leftpad_k = [ + maybe_contiguous(x) for x in (page_table, kv_batch_idx, leftpad_k) + ] + rotary_cos, rotary_sin = [maybe_contiguous(x) for x in (rotary_cos, rotary_sin)] + seqlens_rotary = maybe_contiguous(seqlens_rotary) + out, softmax_lse, *rest = flash_attn_3_cuda.fwd( + q, + k, + v, + k_new, + v_new, + qv, + out, + cu_seqlens_q, + cu_seqlens_k, + cu_seqlens_k_new, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + page_table, + kv_batch_idx, + leftpad_k, + rotary_cos, + rotary_sin, + seqlens_rotary, + q_descale, + k_descale, + v_descale, + softmax_scale, + causal, + window_size[0], + window_size[1], + attention_chunk, + softcap, + rotary_interleaved, + scheduler_metadata, + num_splits, + pack_gqa, + sm_margin, + ) + return out, softmax_lse, *rest + + +def _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + cu_seqlens_q, + cu_seqlens_k, + sequed_q, + sequed_k, + max_seqlen_q, + max_seqlen_k, + dq, + dk, + dv, + softmax_scale, + causal, + window_size=(-1, -1), + softcap=0.0, + deterministic=False, + sm_margin=0, +): + # dq, dk, dv are allocated by us so they should already be contiguous + dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] + dq, dk, dv, softmax_d, *rest = flash_attn_3_cuda.bwd( + dout, + q, + k, + v, + out, + softmax_lse, + dq, + dk, + dv, + cu_seqlens_q, + cu_seqlens_k, + sequed_q, + sequed_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + window_size[0], + window_size[1], + softcap, + deterministic, + sm_margin, + ) + return dq, dk, dv, softmax_d + + +class FlashAttnQKVPackedFunc(torch.autograd.Function): + @staticmethod + def forward( + ctx, + qkv, + softmax_scale, + causal, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + deterministic=False, + num_heads_q=None, + sm_margin=0, + ): + if softmax_scale is None: + softmax_scale = qkv.shape[-1] ** (-0.5) + if qkv.dim() == 5: + assert qkv.shape[-3] == 3 + q, k, v = qkv.unbind(dim=-3) + else: + assert qkv.dim() == 4 + assert num_heads_q is not None + num_heads_k = (qkv.shape[2] - num_heads_q) // 2 + assert num_heads_k * 2 + num_heads_q == qkv.shape[2] + q, k, v = qkv.split([num_heads_q, num_heads_k, num_heads_k], dim=-2) + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + None, # qv + None, # out + None, None, None, # cu_seqlens_q/k/k_new + None, None, # seqused_q/k + None, None, # max_seqlen_q/k + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size=window_size, + attention_chunk=attention_chunk, + softcap=softcap, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse) + ctx.save_for_backward(q, k, v, out, softmax_lse) + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.ndim = qkv.dim() + ctx.sm_margin = sm_margin + # return out, softmax_lse + return out + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + if ctx.ndim == 5: + qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) + dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) + dq, dk, dv = dqkv.unbind(dim=-3) + else: + num_heads_q = q.shape[2] + num_heads_k = k.shape[2] + qkv_shape = q.shape[:-2] + (num_heads_q + num_heads_k * 2, *q.shape[-1:]) + dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) + dq, dk, dv = dqkv.split([num_heads_q, num_heads_k, num_heads_k], dim=-2) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + None, None, # cu_seqlens_q, cu_seqlens_k, + None, None, # sequed_q, sequed_k, + None, None, # max_seqlen_q, max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size, + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dqkv = dqkv[..., : dout.shape[-1]] # We could have padded the head dimension + return dqkv, None, None, None, None, None, None, None, None, None, None, None + + +class FlashAttnFunc(torch.autograd.Function): + + @staticmethod + def forward( + ctx, + q, + k, + v, + softmax_scale, + causal, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + ): + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + # out, q, k, v, out_padded, softmax_lse = _flash_attn_forward( + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + qv, # qv + None, # out + None, None, None, # cu_seqlens_q/k/k_new + None, None, # seqused_q/k + None, None, # max_seqlen_q/k + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size=window_size, + attention_chunk=attention_chunk, + softcap=softcap, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse) + ctx.save_for_backward(q, k, v, out, softmax_lse) + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.sm_margin = sm_margin + return out, softmax_lse + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + None, None, # cu_seqlens_q, cu_seqlens_k, + None, None, # sequed_q, sequed_k, + None, None, # max_seqlen_q, max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size, + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dq = dq[..., : q.shape[-1]] # We could have padded the head dimension + dk = dk[..., : k.shape[-1]] + dv = dv[..., : v.shape[-1]] + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None + + +class FlashAttnVarlenFunc(torch.autograd.Function): + + @staticmethod + def forward( + ctx, + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + ): + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + # out, q, k, v, out_padded, softmax_lse = _flash_attn_varlen_forward( + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + qv, # qv + None, # out + cu_seqlens_q, + cu_seqlens_k, + None, # cu_seqlens_k_new + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size=window_size, + attention_chunk=attention_chunk, + softcap=softcap, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k) + ctx.save_for_backward(q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k) + ctx.max_seqlen_q = max_seqlen_q + ctx.max_seqlen_k = max_seqlen_k + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.sm_margin = sm_margin + return out, softmax_lse + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + ctx.max_seqlen_q, + ctx.max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size, + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dq = dq[..., : q.shape[-1]] # We could have padded the head dimension + dk = dk[..., : k.shape[-1]] + dv = dv[..., : v.shape[-1]] + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None + + +def flash_attn_qkvpacked_func( + qkv, + softmax_scale=None, + causal=False, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + deterministic=False, + num_heads_q=None, + sm_margin=0, +): + """dropout_p should be set to 0.0 during evaluation + If Q, K, V are already stacked into 1 tensor, this function will be faster than + calling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation + of the gradients of Q, K, V. + For multi-query and grouped-query attention (MQA/GQA), please see + flash_attn_kvpacked_func and flash_attn_func. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between [i - window_size[0], i + window_size[1]] inclusive. + + Arguments: + qkv: (batch_size, seqlen, 3, nheads, headdim) + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + softcap: float. Anything > 0 activates softcapping attention. + alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of (-alibi_slope * |i - j|) is added to + the attention score of query i and key j. + deterministic: bool. Whether to use the deterministic implementation of the backward pass, + which is slightly slower and uses more memory. The forward pass is always deterministic. + return_attn_probs: bool. Whether to return the attention probabilities. This option is for + testing only. The returned probabilities are not guaranteed to be correct + (they might not have the right scaling). + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen). + The output of softmax (possibly with different scaling). It also encodes the dropout + pattern (negative means that location was dropped, nonnegative means it was kept). + """ + return FlashAttnQKVPackedFunc.apply( + qkv, + softmax_scale, + causal, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + deterministic, + num_heads_q, + sm_margin, + ) + + +def flash_attn_func( + q, + k, + v, + softmax_scale=None, + causal=False, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, +): + """dropout_p should be set to 0.0 during evaluation + Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads + than Q. Note that the number of heads in Q must be divisible by the number of heads in KV. + For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head + 0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V. + + If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix. + For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is: + 1 1 1 1 0 + 1 1 1 1 1 + If seqlen_q = 5 and seqlen_k = 2, the causal mask is: + 0 0 + 0 0 + 0 0 + 1 0 + 1 1 + If the row of the mask is all zero, the output will be zero. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between + [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive. + + Arguments: + q: (batch_size, seqlen, nheads, headdim) + k: (batch_size, seqlen, nheads_k, headdim) + v: (batch_size, seqlen, nheads_k, headdim) + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of + (-alibi_slope * |i + seqlen_k - seqlen_q - j|) + is added to the attention score of query i and key j. + deterministic: bool. Whether to use the deterministic implementation of the backward pass, + which is slightly slower and uses more memory. The forward pass is always deterministic. + return_attn_probs: bool. Whether to return the attention probabilities. This option is for + testing only. The returned probabilities are not guaranteed to be correct + (they might not have the right scaling). + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + """ + return FlashAttnFunc.apply( + q, + k, + v, + softmax_scale, + causal, + qv, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + num_splits, + pack_gqa, + deterministic, + sm_margin, + ) + + +def flash_attn_varlen_func( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + seqused_q=None, + seqused_k=None, + softmax_scale=None, + causal=False, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, +): + return FlashAttnVarlenFunc.apply( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + qv, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + num_splits, + pack_gqa, + deterministic, + sm_margin, + ) + + +def flash_attn_combine(out_partial, lse_partial, out=None, out_dtype=None): + return flash_attn_3_cuda.fwd_combine(out_partial, lse_partial, out, out_dtype) + + +def flash_attn_with_kvcache( + q, + k_cache, + v_cache, + k=None, + v=None, + qv=None, + rotary_cos=None, + rotary_sin=None, + cache_seqlens: Optional[Union[(int, torch.Tensor)]] = None, + cache_batch_idx: Optional[torch.Tensor] = None, + cache_leftpad: Optional[torch.Tensor] = None, + page_table: Optional[torch.Tensor] = None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + rotary_seqlens: Optional[torch.Tensor] = None, + q_descale: Optional[torch.Tensor] = None, + k_descale: Optional[torch.Tensor] = None, + v_descale: Optional[torch.Tensor] = None, + softmax_scale=None, + causal=False, + window_size=(-1, -1), # -1 means infinite context window + attention_chunk=0, + softcap=0.0, # 0.0 means deactivated + rotary_interleaved=True, + scheduler_metadata=None, + num_splits=0, # Can be tuned for speed + pack_gqa=None, # Can be tuned for speed + sm_margin=0, # Can be tuned if some SMs are used for communication + return_softmax_lse=False, +): + """ + If k and v are not None, k_cache and v_cache will be updated *inplace* with the new values from + k and v. This is useful for incremental decoding: you can pass in the cached keys/values from + the previous step, and update them with the new keys/values from the current step, and do + attention with the updated cache, all in 1 kernel. + + If you pass in k / v, you must make sure that the cache is large enough to hold the new values. + For example, the KV cache could be pre-allocated with the max sequence length, and you can use + cache_seqlens to keep track of the current sequence lengths of each sequence in the batch. + + Also apply rotary embedding if rotary_cos and rotary_sin are passed in. The key @k will be + rotated by rotary_cos and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc. + If causal or local (i.e., window_size != (-1, -1)), the query @q will be rotated by rotary_cos + and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc. + If not causal and not local, the query @q will be rotated by rotary_cos and rotary_sin at + indices cache_seqlens only (i.e. we consider all tokens in @q to be at position cache_seqlens). + + See tests/test_flash_attn.py::test_flash_attn_kvcache for examples of how to use this function. + + Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads + than Q. Note that the number of heads in Q must be divisible by the number of heads in KV. + For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head + 0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V. + + If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix. + For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is: + 1 1 1 1 0 + 1 1 1 1 1 + If seqlen_q = 5 and seqlen_k = 2, the causal mask is: + 0 0 + 0 0 + 0 0 + 1 0 + 1 1 + If the row of the mask is all zero, the output will be zero. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between + [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive. + + Note: Does not support backward pass. + + Arguments: + q: (batch_size, seqlen, nheads, headdim) + k_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim) if there's no page_table, + or (num_blocks, page_block_size, nheads_k, headdim) if there's a page_table (i.e. paged KV cache) + page_block_size must be a multiple of 256. + v_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim_v) if there's no page_table, + or (num_blocks, page_block_size, nheads_k, headdim_v) if there's a page_table (i.e. paged KV cache) + k [optional]: (batch_size, seqlen_new, nheads_k, headdim). If not None, we concatenate + k with k_cache, starting at the indices specified by cache_seqlens. + v [optional]: (batch_size, seqlen_new, nheads_k, headdim_v). Similar to k. + qv [optional]: (batch_size, seqlen, nheads, headdim_v) + rotary_cos [optional]: (seqlen_ro, rotary_dim / 2). If not None, we apply rotary embedding + to k and q. Only applicable if k and v are passed in. rotary_dim must be divisible by 16. + rotary_sin [optional]: (seqlen_ro, rotary_dim / 2). Similar to rotary_cos. + cache_seqlens: int, or (batch_size,), dtype torch.int32. The sequence lengths of the + KV cache. + cache_batch_idx: (batch_size,), dtype torch.int32. The indices used to index into the KV cache. + If None, we assume that the batch indices are [0, 1, 2, ..., batch_size - 1]. + If the indices are not distinct, and k and v are provided, the values updated in the cache + might come from any of the duplicate indices. + cache_leftpad: (batch_size,), dtype torch.int32. The index that the KV cache starts. If None, assume 0. + page_table [optional]: (batch_size, max_num_blocks_per_seq), dtype torch.int32. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + softcap: float. Anything > 0 activates softcapping attention. + rotary_interleaved: bool. Only applicable if rotary_cos and rotary_sin are passed in. + If True, rotary embedding will combine dimensions 0 & 1, 2 & 3, etc. If False, + rotary embedding will combine dimensions 0 & rotary_dim / 2, 1 & rotary_dim / 2 + 1 + (i.e. GPT-NeoX style). + num_splits: int. If > 1, split the key/value into this many chunks along the sequence. + If num_splits == 1, we don't split the key/value. If num_splits == 0, we use a heuristic + to automatically determine the number of splits. + Don't change this unless you know what you are doing. + return_softmax_lse: bool. Whether to return the logsumexp of the attention scores. + + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_softmax_lse=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + """ + assert k_cache.stride(-1) == 1, "k_cache must have contiguous last dimension" + assert v_cache.stride(-1) == 1, "v_cache must have contiguous last dimension" + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + if cache_seqlens is not None and isinstance(cache_seqlens, int): + cache_seqlens = torch.full( + (k_cache.shape[0],), cache_seqlens, dtype=torch.int32, device=k_cache.device + ) + cache_seqlens = maybe_contiguous(cache_seqlens) + out, softmax_lse, *rest = _flash_attn_forward( + q, + k_cache, + v_cache, + k, + v, + qv, + None, # out + cu_seqlens_q, + None, # cu_seqlens_k + cu_seqlens_k_new, + None, # seqused_q + cache_seqlens, + max_seqlen_q, + None, # max_seqlen_k + page_table, + cache_batch_idx, + cache_leftpad, + rotary_cos, + rotary_sin, + rotary_seqlens, + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size=window_size, + attention_chunk=attention_chunk, + softcap=softcap, + rotary_interleaved=rotary_interleaved, + scheduler_metadata=scheduler_metadata, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # return (out, softmax_lse) if return_softmax_lse else out + return (out, softmax_lse, *rest) if return_softmax_lse else out + + +def get_scheduler_metadata( + batch_size, max_seqlen_q, max_seqlen_k, num_heads_q, num_heads_kv, headdim, + cache_seqlens: torch.Tensor, + qkv_dtype=torch.bfloat16, + headdim_v=None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + cache_leftpad: Optional[torch.Tensor] = None, + page_size: Optional[int] = None, + max_seqlen_k_new=0, + causal=False, + window_size=(-1, -1), # -1 means infinite context window + attention_chunk=0, + has_softcap=False, + num_splits=0, # Can be tuned for speed + pack_gqa=None, # Can be tuned for speed + sm_margin=0, # Can be tuned if some SMs are used for communication +): + cache_seqlens = maybe_contiguous(cache_seqlens) + if headdim_v is None: + headdim_v = headdim + scheduler_metadata = flash_attn_3_cuda.get_scheduler_metadata( + batch_size, max_seqlen_q, max_seqlen_k, num_heads_q, num_heads_kv, headdim, headdim_v, + qkv_dtype, + cache_seqlens, + cu_seqlens_q, + None, # cu_seqlens_k + cu_seqlens_k_new, + None, # seqused_q + cache_leftpad, + page_size, + max_seqlen_k_new, + causal, + window_size[0], window_size[1], + attention_chunk, + has_softcap, + num_splits, + pack_gqa, + sm_margin, + ) + return scheduler_metadata diff --git a/build/torch26-cxx98-cu124-x86_64-linux/flash_attn3/__init__.py b/build/torch26-cxx98-cu124-x86_64-linux/flash_attn3/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3d12edd447e84d6e65c4fa9ab91ae3c3192884c1 --- /dev/null +++ b/build/torch26-cxx98-cu124-x86_64-linux/flash_attn3/__init__.py @@ -0,0 +1,17 @@ +from .flash_attn_interface import ( + flash_attn_combine, + flash_attn_func, + flash_attn_qkvpacked_func, + flash_attn_varlen_func, + flash_attn_with_kvcache, + get_scheduler_metadata, +) + +__all__ = [ + "flash_attn_combine", + "flash_attn_func", + "flash_attn_qkvpacked_func", + "flash_attn_varlen_func", + "flash_attn_with_kvcache", + "get_scheduler_metadata", +] diff --git a/build/torch26-cxx98-cu124-x86_64-linux/flash_attn3/_flash_attn3_2e75662.abi3.so b/build/torch26-cxx98-cu124-x86_64-linux/flash_attn3/_flash_attn3_2e75662.abi3.so new file mode 100755 index 0000000000000000000000000000000000000000..ff91e1b41158d1d2f3622599e5486d7b3644d7a8 --- /dev/null +++ b/build/torch26-cxx98-cu124-x86_64-linux/flash_attn3/_flash_attn3_2e75662.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9627e08ec8778d2a409a2a0477572edb3e03eaca2b45e7b4810ee0a9126d6547 +size 838456048 diff --git a/build/torch26-cxx98-cu124-x86_64-linux/flash_attn3/_flash_attn3_557701f.abi3.so b/build/torch26-cxx98-cu124-x86_64-linux/flash_attn3/_flash_attn3_557701f.abi3.so new file mode 100755 index 0000000000000000000000000000000000000000..683aae960afda0e72ff1b0e1cf1d57c4759606d8 --- /dev/null +++ b/build/torch26-cxx98-cu124-x86_64-linux/flash_attn3/_flash_attn3_557701f.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:07fe025ba95671f6ff957991f74c66063bfb10ab6737641c88f88116c9f83718 +size 838456048 diff --git a/build/torch26-cxx98-cu124-x86_64-linux/flash_attn3/_ops.py b/build/torch26-cxx98-cu124-x86_64-linux/flash_attn3/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..0c0daa483e5a131df44714bd9cb4c4c8143bdcff --- /dev/null +++ b/build/torch26-cxx98-cu124-x86_64-linux/flash_attn3/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _flash_attn3_557701f +ops = torch.ops._flash_attn3_557701f + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_flash_attn3_557701f::{op_name}" \ No newline at end of file diff --git a/build/torch26-cxx98-cu124-x86_64-linux/flash_attn3/flash_attn_interface.py b/build/torch26-cxx98-cu124-x86_64-linux/flash_attn3/flash_attn_interface.py new file mode 100644 index 0000000000000000000000000000000000000000..7572297784a73fdf891af6aafee88428120b72f2 --- /dev/null +++ b/build/torch26-cxx98-cu124-x86_64-linux/flash_attn3/flash_attn_interface.py @@ -0,0 +1,828 @@ +# Copyright (c) 2023, Tri Dao. + +from typing import Optional, Union + +import torch +import torch.nn as nn + +from ._ops import ops as flash_attn_3_cuda + +def maybe_contiguous(x): + return x.contiguous() if x is not None and x.stride(-1) != 1 else x + + +def _flash_attn_forward( + q, + k, + v, + k_new, + v_new, + qv, + out, + cu_seqlens_q, + cu_seqlens_k, + cu_seqlens_k_new, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + page_table, + kv_batch_idx, + leftpad_k, + rotary_cos, + rotary_sin, + seqlens_rotary, + q_descale, + k_descale, + v_descale, + softmax_scale, + causal, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + rotary_interleaved=True, + scheduler_metadata=None, + num_splits=1, + pack_gqa=None, + sm_margin=0): + q, k, k_new, v_new = [maybe_contiguous(x) for x in (q, k, k_new, v_new)] + v = v.contiguous() if v.stride(-1) != 1 and v.stride(-3) != 1 else v + cu_seqlens_q, cu_seqlens_k, cu_seqlens_k_new = [ + maybe_contiguous(x) for x in (cu_seqlens_q, cu_seqlens_k, cu_seqlens_k_new) + ] + seqused_q, seqused_k = [maybe_contiguous(x) for x in (seqused_q, seqused_k)] + page_table, kv_batch_idx, leftpad_k = [ + maybe_contiguous(x) for x in (page_table, kv_batch_idx, leftpad_k) + ] + rotary_cos, rotary_sin = [maybe_contiguous(x) for x in (rotary_cos, rotary_sin)] + seqlens_rotary = maybe_contiguous(seqlens_rotary) + out, softmax_lse, *rest = flash_attn_3_cuda.fwd( + q, + k, + v, + k_new, + v_new, + qv, + out, + cu_seqlens_q, + cu_seqlens_k, + cu_seqlens_k_new, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + page_table, + kv_batch_idx, + leftpad_k, + rotary_cos, + rotary_sin, + seqlens_rotary, + q_descale, + k_descale, + v_descale, + softmax_scale, + causal, + window_size[0], + window_size[1], + attention_chunk, + softcap, + rotary_interleaved, + scheduler_metadata, + num_splits, + pack_gqa, + sm_margin, + ) + return out, softmax_lse, *rest + + +def _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + cu_seqlens_q, + cu_seqlens_k, + sequed_q, + sequed_k, + max_seqlen_q, + max_seqlen_k, + dq, + dk, + dv, + softmax_scale, + causal, + window_size=(-1, -1), + softcap=0.0, + deterministic=False, + sm_margin=0, +): + # dq, dk, dv are allocated by us so they should already be contiguous + dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] + dq, dk, dv, softmax_d, *rest = flash_attn_3_cuda.bwd( + dout, + q, + k, + v, + out, + softmax_lse, + dq, + dk, + dv, + cu_seqlens_q, + cu_seqlens_k, + sequed_q, + sequed_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + window_size[0], + window_size[1], + softcap, + deterministic, + sm_margin, + ) + return dq, dk, dv, softmax_d + + +class FlashAttnQKVPackedFunc(torch.autograd.Function): + @staticmethod + def forward( + ctx, + qkv, + softmax_scale, + causal, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + deterministic=False, + num_heads_q=None, + sm_margin=0, + ): + if softmax_scale is None: + softmax_scale = qkv.shape[-1] ** (-0.5) + if qkv.dim() == 5: + assert qkv.shape[-3] == 3 + q, k, v = qkv.unbind(dim=-3) + else: + assert qkv.dim() == 4 + assert num_heads_q is not None + num_heads_k = (qkv.shape[2] - num_heads_q) // 2 + assert num_heads_k * 2 + num_heads_q == qkv.shape[2] + q, k, v = qkv.split([num_heads_q, num_heads_k, num_heads_k], dim=-2) + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + None, # qv + None, # out + None, None, None, # cu_seqlens_q/k/k_new + None, None, # seqused_q/k + None, None, # max_seqlen_q/k + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size=window_size, + attention_chunk=attention_chunk, + softcap=softcap, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse) + ctx.save_for_backward(q, k, v, out, softmax_lse) + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.ndim = qkv.dim() + ctx.sm_margin = sm_margin + # return out, softmax_lse + return out + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + if ctx.ndim == 5: + qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) + dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) + dq, dk, dv = dqkv.unbind(dim=-3) + else: + num_heads_q = q.shape[2] + num_heads_k = k.shape[2] + qkv_shape = q.shape[:-2] + (num_heads_q + num_heads_k * 2, *q.shape[-1:]) + dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) + dq, dk, dv = dqkv.split([num_heads_q, num_heads_k, num_heads_k], dim=-2) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + None, None, # cu_seqlens_q, cu_seqlens_k, + None, None, # sequed_q, sequed_k, + None, None, # max_seqlen_q, max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size, + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dqkv = dqkv[..., : dout.shape[-1]] # We could have padded the head dimension + return dqkv, None, None, None, None, None, None, None, None, None, None, None + + +class FlashAttnFunc(torch.autograd.Function): + + @staticmethod + def forward( + ctx, + q, + k, + v, + softmax_scale, + causal, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + ): + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + # out, q, k, v, out_padded, softmax_lse = _flash_attn_forward( + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + qv, # qv + None, # out + None, None, None, # cu_seqlens_q/k/k_new + None, None, # seqused_q/k + None, None, # max_seqlen_q/k + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size=window_size, + attention_chunk=attention_chunk, + softcap=softcap, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse) + ctx.save_for_backward(q, k, v, out, softmax_lse) + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.sm_margin = sm_margin + return out, softmax_lse + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + None, None, # cu_seqlens_q, cu_seqlens_k, + None, None, # sequed_q, sequed_k, + None, None, # max_seqlen_q, max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size, + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dq = dq[..., : q.shape[-1]] # We could have padded the head dimension + dk = dk[..., : k.shape[-1]] + dv = dv[..., : v.shape[-1]] + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None + + +class FlashAttnVarlenFunc(torch.autograd.Function): + + @staticmethod + def forward( + ctx, + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + ): + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + # out, q, k, v, out_padded, softmax_lse = _flash_attn_varlen_forward( + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + qv, # qv + None, # out + cu_seqlens_q, + cu_seqlens_k, + None, # cu_seqlens_k_new + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size=window_size, + attention_chunk=attention_chunk, + softcap=softcap, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k) + ctx.save_for_backward(q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k) + ctx.max_seqlen_q = max_seqlen_q + ctx.max_seqlen_k = max_seqlen_k + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.sm_margin = sm_margin + return out, softmax_lse + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + ctx.max_seqlen_q, + ctx.max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size, + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dq = dq[..., : q.shape[-1]] # We could have padded the head dimension + dk = dk[..., : k.shape[-1]] + dv = dv[..., : v.shape[-1]] + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None + + +def flash_attn_qkvpacked_func( + qkv, + softmax_scale=None, + causal=False, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + deterministic=False, + num_heads_q=None, + sm_margin=0, +): + """dropout_p should be set to 0.0 during evaluation + If Q, K, V are already stacked into 1 tensor, this function will be faster than + calling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation + of the gradients of Q, K, V. + For multi-query and grouped-query attention (MQA/GQA), please see + flash_attn_kvpacked_func and flash_attn_func. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between [i - window_size[0], i + window_size[1]] inclusive. + + Arguments: + qkv: (batch_size, seqlen, 3, nheads, headdim) + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + softcap: float. Anything > 0 activates softcapping attention. + alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of (-alibi_slope * |i - j|) is added to + the attention score of query i and key j. + deterministic: bool. Whether to use the deterministic implementation of the backward pass, + which is slightly slower and uses more memory. The forward pass is always deterministic. + return_attn_probs: bool. Whether to return the attention probabilities. This option is for + testing only. The returned probabilities are not guaranteed to be correct + (they might not have the right scaling). + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen). + The output of softmax (possibly with different scaling). It also encodes the dropout + pattern (negative means that location was dropped, nonnegative means it was kept). + """ + return FlashAttnQKVPackedFunc.apply( + qkv, + softmax_scale, + causal, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + deterministic, + num_heads_q, + sm_margin, + ) + + +def flash_attn_func( + q, + k, + v, + softmax_scale=None, + causal=False, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, +): + """dropout_p should be set to 0.0 during evaluation + Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads + than Q. Note that the number of heads in Q must be divisible by the number of heads in KV. + For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head + 0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V. + + If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix. + For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is: + 1 1 1 1 0 + 1 1 1 1 1 + If seqlen_q = 5 and seqlen_k = 2, the causal mask is: + 0 0 + 0 0 + 0 0 + 1 0 + 1 1 + If the row of the mask is all zero, the output will be zero. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between + [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive. + + Arguments: + q: (batch_size, seqlen, nheads, headdim) + k: (batch_size, seqlen, nheads_k, headdim) + v: (batch_size, seqlen, nheads_k, headdim) + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of + (-alibi_slope * |i + seqlen_k - seqlen_q - j|) + is added to the attention score of query i and key j. + deterministic: bool. Whether to use the deterministic implementation of the backward pass, + which is slightly slower and uses more memory. The forward pass is always deterministic. + return_attn_probs: bool. Whether to return the attention probabilities. This option is for + testing only. The returned probabilities are not guaranteed to be correct + (they might not have the right scaling). + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + """ + return FlashAttnFunc.apply( + q, + k, + v, + softmax_scale, + causal, + qv, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + num_splits, + pack_gqa, + deterministic, + sm_margin, + ) + + +def flash_attn_varlen_func( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + seqused_q=None, + seqused_k=None, + softmax_scale=None, + causal=False, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, +): + return FlashAttnVarlenFunc.apply( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + qv, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + num_splits, + pack_gqa, + deterministic, + sm_margin, + ) + + +def flash_attn_combine(out_partial, lse_partial, out=None, out_dtype=None): + return flash_attn_3_cuda.fwd_combine(out_partial, lse_partial, out, out_dtype) + + +def flash_attn_with_kvcache( + q, + k_cache, + v_cache, + k=None, + v=None, + qv=None, + rotary_cos=None, + rotary_sin=None, + cache_seqlens: Optional[Union[(int, torch.Tensor)]] = None, + cache_batch_idx: Optional[torch.Tensor] = None, + cache_leftpad: Optional[torch.Tensor] = None, + page_table: Optional[torch.Tensor] = None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + rotary_seqlens: Optional[torch.Tensor] = None, + q_descale: Optional[torch.Tensor] = None, + k_descale: Optional[torch.Tensor] = None, + v_descale: Optional[torch.Tensor] = None, + softmax_scale=None, + causal=False, + window_size=(-1, -1), # -1 means infinite context window + attention_chunk=0, + softcap=0.0, # 0.0 means deactivated + rotary_interleaved=True, + scheduler_metadata=None, + num_splits=0, # Can be tuned for speed + pack_gqa=None, # Can be tuned for speed + sm_margin=0, # Can be tuned if some SMs are used for communication + return_softmax_lse=False, +): + """ + If k and v are not None, k_cache and v_cache will be updated *inplace* with the new values from + k and v. This is useful for incremental decoding: you can pass in the cached keys/values from + the previous step, and update them with the new keys/values from the current step, and do + attention with the updated cache, all in 1 kernel. + + If you pass in k / v, you must make sure that the cache is large enough to hold the new values. + For example, the KV cache could be pre-allocated with the max sequence length, and you can use + cache_seqlens to keep track of the current sequence lengths of each sequence in the batch. + + Also apply rotary embedding if rotary_cos and rotary_sin are passed in. The key @k will be + rotated by rotary_cos and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc. + If causal or local (i.e., window_size != (-1, -1)), the query @q will be rotated by rotary_cos + and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc. + If not causal and not local, the query @q will be rotated by rotary_cos and rotary_sin at + indices cache_seqlens only (i.e. we consider all tokens in @q to be at position cache_seqlens). + + See tests/test_flash_attn.py::test_flash_attn_kvcache for examples of how to use this function. + + Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads + than Q. Note that the number of heads in Q must be divisible by the number of heads in KV. + For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head + 0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V. + + If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix. + For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is: + 1 1 1 1 0 + 1 1 1 1 1 + If seqlen_q = 5 and seqlen_k = 2, the causal mask is: + 0 0 + 0 0 + 0 0 + 1 0 + 1 1 + If the row of the mask is all zero, the output will be zero. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between + [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive. + + Note: Does not support backward pass. + + Arguments: + q: (batch_size, seqlen, nheads, headdim) + k_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim) if there's no page_table, + or (num_blocks, page_block_size, nheads_k, headdim) if there's a page_table (i.e. paged KV cache) + page_block_size must be a multiple of 256. + v_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim_v) if there's no page_table, + or (num_blocks, page_block_size, nheads_k, headdim_v) if there's a page_table (i.e. paged KV cache) + k [optional]: (batch_size, seqlen_new, nheads_k, headdim). If not None, we concatenate + k with k_cache, starting at the indices specified by cache_seqlens. + v [optional]: (batch_size, seqlen_new, nheads_k, headdim_v). Similar to k. + qv [optional]: (batch_size, seqlen, nheads, headdim_v) + rotary_cos [optional]: (seqlen_ro, rotary_dim / 2). If not None, we apply rotary embedding + to k and q. Only applicable if k and v are passed in. rotary_dim must be divisible by 16. + rotary_sin [optional]: (seqlen_ro, rotary_dim / 2). Similar to rotary_cos. + cache_seqlens: int, or (batch_size,), dtype torch.int32. The sequence lengths of the + KV cache. + cache_batch_idx: (batch_size,), dtype torch.int32. The indices used to index into the KV cache. + If None, we assume that the batch indices are [0, 1, 2, ..., batch_size - 1]. + If the indices are not distinct, and k and v are provided, the values updated in the cache + might come from any of the duplicate indices. + cache_leftpad: (batch_size,), dtype torch.int32. The index that the KV cache starts. If None, assume 0. + page_table [optional]: (batch_size, max_num_blocks_per_seq), dtype torch.int32. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + softcap: float. Anything > 0 activates softcapping attention. + rotary_interleaved: bool. Only applicable if rotary_cos and rotary_sin are passed in. + If True, rotary embedding will combine dimensions 0 & 1, 2 & 3, etc. If False, + rotary embedding will combine dimensions 0 & rotary_dim / 2, 1 & rotary_dim / 2 + 1 + (i.e. GPT-NeoX style). + num_splits: int. If > 1, split the key/value into this many chunks along the sequence. + If num_splits == 1, we don't split the key/value. If num_splits == 0, we use a heuristic + to automatically determine the number of splits. + Don't change this unless you know what you are doing. + return_softmax_lse: bool. Whether to return the logsumexp of the attention scores. + + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_softmax_lse=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + """ + assert k_cache.stride(-1) == 1, "k_cache must have contiguous last dimension" + assert v_cache.stride(-1) == 1, "v_cache must have contiguous last dimension" + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + if cache_seqlens is not None and isinstance(cache_seqlens, int): + cache_seqlens = torch.full( + (k_cache.shape[0],), cache_seqlens, dtype=torch.int32, device=k_cache.device + ) + cache_seqlens = maybe_contiguous(cache_seqlens) + out, softmax_lse, *rest = _flash_attn_forward( + q, + k_cache, + v_cache, + k, + v, + qv, + None, # out + cu_seqlens_q, + None, # cu_seqlens_k + cu_seqlens_k_new, + None, # seqused_q + cache_seqlens, + max_seqlen_q, + None, # max_seqlen_k + page_table, + cache_batch_idx, + cache_leftpad, + rotary_cos, + rotary_sin, + rotary_seqlens, + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size=window_size, + attention_chunk=attention_chunk, + softcap=softcap, + rotary_interleaved=rotary_interleaved, + scheduler_metadata=scheduler_metadata, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # return (out, softmax_lse) if return_softmax_lse else out + return (out, softmax_lse, *rest) if return_softmax_lse else out + + +def get_scheduler_metadata( + batch_size, max_seqlen_q, max_seqlen_k, num_heads_q, num_heads_kv, headdim, + cache_seqlens: torch.Tensor, + qkv_dtype=torch.bfloat16, + headdim_v=None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + cache_leftpad: Optional[torch.Tensor] = None, + page_size: Optional[int] = None, + max_seqlen_k_new=0, + causal=False, + window_size=(-1, -1), # -1 means infinite context window + attention_chunk=0, + has_softcap=False, + num_splits=0, # Can be tuned for speed + pack_gqa=None, # Can be tuned for speed + sm_margin=0, # Can be tuned if some SMs are used for communication +): + cache_seqlens = maybe_contiguous(cache_seqlens) + if headdim_v is None: + headdim_v = headdim + scheduler_metadata = flash_attn_3_cuda.get_scheduler_metadata( + batch_size, max_seqlen_q, max_seqlen_k, num_heads_q, num_heads_kv, headdim, headdim_v, + qkv_dtype, + cache_seqlens, + cu_seqlens_q, + None, # cu_seqlens_k + cu_seqlens_k_new, + None, # seqused_q + cache_leftpad, + page_size, + max_seqlen_k_new, + causal, + window_size[0], window_size[1], + attention_chunk, + has_softcap, + num_splits, + pack_gqa, + sm_margin, + ) + return scheduler_metadata diff --git a/build/torch26-cxx98-cu126-x86_64-linux/flash_attn3/__init__.py b/build/torch26-cxx98-cu126-x86_64-linux/flash_attn3/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3d12edd447e84d6e65c4fa9ab91ae3c3192884c1 --- /dev/null +++ b/build/torch26-cxx98-cu126-x86_64-linux/flash_attn3/__init__.py @@ -0,0 +1,17 @@ +from .flash_attn_interface import ( + flash_attn_combine, + flash_attn_func, + flash_attn_qkvpacked_func, + flash_attn_varlen_func, + flash_attn_with_kvcache, + get_scheduler_metadata, +) + +__all__ = [ + "flash_attn_combine", + "flash_attn_func", + "flash_attn_qkvpacked_func", + "flash_attn_varlen_func", + "flash_attn_with_kvcache", + "get_scheduler_metadata", +] diff --git a/build/torch26-cxx98-cu126-x86_64-linux/flash_attn3/_flash_attn3_2e75662.abi3.so b/build/torch26-cxx98-cu126-x86_64-linux/flash_attn3/_flash_attn3_2e75662.abi3.so new file mode 100755 index 0000000000000000000000000000000000000000..ff91e1b41158d1d2f3622599e5486d7b3644d7a8 --- /dev/null +++ b/build/torch26-cxx98-cu126-x86_64-linux/flash_attn3/_flash_attn3_2e75662.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9627e08ec8778d2a409a2a0477572edb3e03eaca2b45e7b4810ee0a9126d6547 +size 838456048 diff --git a/build/torch26-cxx98-cu126-x86_64-linux/flash_attn3/_flash_attn3_557701f.abi3.so b/build/torch26-cxx98-cu126-x86_64-linux/flash_attn3/_flash_attn3_557701f.abi3.so new file mode 100755 index 0000000000000000000000000000000000000000..683aae960afda0e72ff1b0e1cf1d57c4759606d8 --- /dev/null +++ b/build/torch26-cxx98-cu126-x86_64-linux/flash_attn3/_flash_attn3_557701f.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:07fe025ba95671f6ff957991f74c66063bfb10ab6737641c88f88116c9f83718 +size 838456048 diff --git a/build/torch26-cxx98-cu126-x86_64-linux/flash_attn3/_ops.py b/build/torch26-cxx98-cu126-x86_64-linux/flash_attn3/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..0c0daa483e5a131df44714bd9cb4c4c8143bdcff --- /dev/null +++ b/build/torch26-cxx98-cu126-x86_64-linux/flash_attn3/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _flash_attn3_557701f +ops = torch.ops._flash_attn3_557701f + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_flash_attn3_557701f::{op_name}" \ No newline at end of file diff --git a/build/torch26-cxx98-cu126-x86_64-linux/flash_attn3/flash_attn_interface.py b/build/torch26-cxx98-cu126-x86_64-linux/flash_attn3/flash_attn_interface.py new file mode 100644 index 0000000000000000000000000000000000000000..7572297784a73fdf891af6aafee88428120b72f2 --- /dev/null +++ b/build/torch26-cxx98-cu126-x86_64-linux/flash_attn3/flash_attn_interface.py @@ -0,0 +1,828 @@ +# Copyright (c) 2023, Tri Dao. + +from typing import Optional, Union + +import torch +import torch.nn as nn + +from ._ops import ops as flash_attn_3_cuda + +def maybe_contiguous(x): + return x.contiguous() if x is not None and x.stride(-1) != 1 else x + + +def _flash_attn_forward( + q, + k, + v, + k_new, + v_new, + qv, + out, + cu_seqlens_q, + cu_seqlens_k, + cu_seqlens_k_new, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + page_table, + kv_batch_idx, + leftpad_k, + rotary_cos, + rotary_sin, + seqlens_rotary, + q_descale, + k_descale, + v_descale, + softmax_scale, + causal, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + rotary_interleaved=True, + scheduler_metadata=None, + num_splits=1, + pack_gqa=None, + sm_margin=0): + q, k, k_new, v_new = [maybe_contiguous(x) for x in (q, k, k_new, v_new)] + v = v.contiguous() if v.stride(-1) != 1 and v.stride(-3) != 1 else v + cu_seqlens_q, cu_seqlens_k, cu_seqlens_k_new = [ + maybe_contiguous(x) for x in (cu_seqlens_q, cu_seqlens_k, cu_seqlens_k_new) + ] + seqused_q, seqused_k = [maybe_contiguous(x) for x in (seqused_q, seqused_k)] + page_table, kv_batch_idx, leftpad_k = [ + maybe_contiguous(x) for x in (page_table, kv_batch_idx, leftpad_k) + ] + rotary_cos, rotary_sin = [maybe_contiguous(x) for x in (rotary_cos, rotary_sin)] + seqlens_rotary = maybe_contiguous(seqlens_rotary) + out, softmax_lse, *rest = flash_attn_3_cuda.fwd( + q, + k, + v, + k_new, + v_new, + qv, + out, + cu_seqlens_q, + cu_seqlens_k, + cu_seqlens_k_new, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + page_table, + kv_batch_idx, + leftpad_k, + rotary_cos, + rotary_sin, + seqlens_rotary, + q_descale, + k_descale, + v_descale, + softmax_scale, + causal, + window_size[0], + window_size[1], + attention_chunk, + softcap, + rotary_interleaved, + scheduler_metadata, + num_splits, + pack_gqa, + sm_margin, + ) + return out, softmax_lse, *rest + + +def _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + cu_seqlens_q, + cu_seqlens_k, + sequed_q, + sequed_k, + max_seqlen_q, + max_seqlen_k, + dq, + dk, + dv, + softmax_scale, + causal, + window_size=(-1, -1), + softcap=0.0, + deterministic=False, + sm_margin=0, +): + # dq, dk, dv are allocated by us so they should already be contiguous + dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] + dq, dk, dv, softmax_d, *rest = flash_attn_3_cuda.bwd( + dout, + q, + k, + v, + out, + softmax_lse, + dq, + dk, + dv, + cu_seqlens_q, + cu_seqlens_k, + sequed_q, + sequed_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + window_size[0], + window_size[1], + softcap, + deterministic, + sm_margin, + ) + return dq, dk, dv, softmax_d + + +class FlashAttnQKVPackedFunc(torch.autograd.Function): + @staticmethod + def forward( + ctx, + qkv, + softmax_scale, + causal, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + deterministic=False, + num_heads_q=None, + sm_margin=0, + ): + if softmax_scale is None: + softmax_scale = qkv.shape[-1] ** (-0.5) + if qkv.dim() == 5: + assert qkv.shape[-3] == 3 + q, k, v = qkv.unbind(dim=-3) + else: + assert qkv.dim() == 4 + assert num_heads_q is not None + num_heads_k = (qkv.shape[2] - num_heads_q) // 2 + assert num_heads_k * 2 + num_heads_q == qkv.shape[2] + q, k, v = qkv.split([num_heads_q, num_heads_k, num_heads_k], dim=-2) + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + None, # qv + None, # out + None, None, None, # cu_seqlens_q/k/k_new + None, None, # seqused_q/k + None, None, # max_seqlen_q/k + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size=window_size, + attention_chunk=attention_chunk, + softcap=softcap, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse) + ctx.save_for_backward(q, k, v, out, softmax_lse) + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.ndim = qkv.dim() + ctx.sm_margin = sm_margin + # return out, softmax_lse + return out + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + if ctx.ndim == 5: + qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) + dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) + dq, dk, dv = dqkv.unbind(dim=-3) + else: + num_heads_q = q.shape[2] + num_heads_k = k.shape[2] + qkv_shape = q.shape[:-2] + (num_heads_q + num_heads_k * 2, *q.shape[-1:]) + dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) + dq, dk, dv = dqkv.split([num_heads_q, num_heads_k, num_heads_k], dim=-2) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + None, None, # cu_seqlens_q, cu_seqlens_k, + None, None, # sequed_q, sequed_k, + None, None, # max_seqlen_q, max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size, + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dqkv = dqkv[..., : dout.shape[-1]] # We could have padded the head dimension + return dqkv, None, None, None, None, None, None, None, None, None, None, None + + +class FlashAttnFunc(torch.autograd.Function): + + @staticmethod + def forward( + ctx, + q, + k, + v, + softmax_scale, + causal, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + ): + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + # out, q, k, v, out_padded, softmax_lse = _flash_attn_forward( + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + qv, # qv + None, # out + None, None, None, # cu_seqlens_q/k/k_new + None, None, # seqused_q/k + None, None, # max_seqlen_q/k + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size=window_size, + attention_chunk=attention_chunk, + softcap=softcap, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse) + ctx.save_for_backward(q, k, v, out, softmax_lse) + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.sm_margin = sm_margin + return out, softmax_lse + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + None, None, # cu_seqlens_q, cu_seqlens_k, + None, None, # sequed_q, sequed_k, + None, None, # max_seqlen_q, max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size, + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dq = dq[..., : q.shape[-1]] # We could have padded the head dimension + dk = dk[..., : k.shape[-1]] + dv = dv[..., : v.shape[-1]] + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None + + +class FlashAttnVarlenFunc(torch.autograd.Function): + + @staticmethod + def forward( + ctx, + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + ): + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + # out, q, k, v, out_padded, softmax_lse = _flash_attn_varlen_forward( + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + qv, # qv + None, # out + cu_seqlens_q, + cu_seqlens_k, + None, # cu_seqlens_k_new + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size=window_size, + attention_chunk=attention_chunk, + softcap=softcap, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k) + ctx.save_for_backward(q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k) + ctx.max_seqlen_q = max_seqlen_q + ctx.max_seqlen_k = max_seqlen_k + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.sm_margin = sm_margin + return out, softmax_lse + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + ctx.max_seqlen_q, + ctx.max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size, + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dq = dq[..., : q.shape[-1]] # We could have padded the head dimension + dk = dk[..., : k.shape[-1]] + dv = dv[..., : v.shape[-1]] + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None + + +def flash_attn_qkvpacked_func( + qkv, + softmax_scale=None, + causal=False, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + deterministic=False, + num_heads_q=None, + sm_margin=0, +): + """dropout_p should be set to 0.0 during evaluation + If Q, K, V are already stacked into 1 tensor, this function will be faster than + calling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation + of the gradients of Q, K, V. + For multi-query and grouped-query attention (MQA/GQA), please see + flash_attn_kvpacked_func and flash_attn_func. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between [i - window_size[0], i + window_size[1]] inclusive. + + Arguments: + qkv: (batch_size, seqlen, 3, nheads, headdim) + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + softcap: float. Anything > 0 activates softcapping attention. + alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of (-alibi_slope * |i - j|) is added to + the attention score of query i and key j. + deterministic: bool. Whether to use the deterministic implementation of the backward pass, + which is slightly slower and uses more memory. The forward pass is always deterministic. + return_attn_probs: bool. Whether to return the attention probabilities. This option is for + testing only. The returned probabilities are not guaranteed to be correct + (they might not have the right scaling). + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen). + The output of softmax (possibly with different scaling). It also encodes the dropout + pattern (negative means that location was dropped, nonnegative means it was kept). + """ + return FlashAttnQKVPackedFunc.apply( + qkv, + softmax_scale, + causal, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + deterministic, + num_heads_q, + sm_margin, + ) + + +def flash_attn_func( + q, + k, + v, + softmax_scale=None, + causal=False, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, +): + """dropout_p should be set to 0.0 during evaluation + Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads + than Q. Note that the number of heads in Q must be divisible by the number of heads in KV. + For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head + 0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V. + + If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix. + For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is: + 1 1 1 1 0 + 1 1 1 1 1 + If seqlen_q = 5 and seqlen_k = 2, the causal mask is: + 0 0 + 0 0 + 0 0 + 1 0 + 1 1 + If the row of the mask is all zero, the output will be zero. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between + [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive. + + Arguments: + q: (batch_size, seqlen, nheads, headdim) + k: (batch_size, seqlen, nheads_k, headdim) + v: (batch_size, seqlen, nheads_k, headdim) + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of + (-alibi_slope * |i + seqlen_k - seqlen_q - j|) + is added to the attention score of query i and key j. + deterministic: bool. Whether to use the deterministic implementation of the backward pass, + which is slightly slower and uses more memory. The forward pass is always deterministic. + return_attn_probs: bool. Whether to return the attention probabilities. This option is for + testing only. The returned probabilities are not guaranteed to be correct + (they might not have the right scaling). + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + """ + return FlashAttnFunc.apply( + q, + k, + v, + softmax_scale, + causal, + qv, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + num_splits, + pack_gqa, + deterministic, + sm_margin, + ) + + +def flash_attn_varlen_func( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + seqused_q=None, + seqused_k=None, + softmax_scale=None, + causal=False, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, +): + return FlashAttnVarlenFunc.apply( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + qv, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + num_splits, + pack_gqa, + deterministic, + sm_margin, + ) + + +def flash_attn_combine(out_partial, lse_partial, out=None, out_dtype=None): + return flash_attn_3_cuda.fwd_combine(out_partial, lse_partial, out, out_dtype) + + +def flash_attn_with_kvcache( + q, + k_cache, + v_cache, + k=None, + v=None, + qv=None, + rotary_cos=None, + rotary_sin=None, + cache_seqlens: Optional[Union[(int, torch.Tensor)]] = None, + cache_batch_idx: Optional[torch.Tensor] = None, + cache_leftpad: Optional[torch.Tensor] = None, + page_table: Optional[torch.Tensor] = None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + rotary_seqlens: Optional[torch.Tensor] = None, + q_descale: Optional[torch.Tensor] = None, + k_descale: Optional[torch.Tensor] = None, + v_descale: Optional[torch.Tensor] = None, + softmax_scale=None, + causal=False, + window_size=(-1, -1), # -1 means infinite context window + attention_chunk=0, + softcap=0.0, # 0.0 means deactivated + rotary_interleaved=True, + scheduler_metadata=None, + num_splits=0, # Can be tuned for speed + pack_gqa=None, # Can be tuned for speed + sm_margin=0, # Can be tuned if some SMs are used for communication + return_softmax_lse=False, +): + """ + If k and v are not None, k_cache and v_cache will be updated *inplace* with the new values from + k and v. This is useful for incremental decoding: you can pass in the cached keys/values from + the previous step, and update them with the new keys/values from the current step, and do + attention with the updated cache, all in 1 kernel. + + If you pass in k / v, you must make sure that the cache is large enough to hold the new values. + For example, the KV cache could be pre-allocated with the max sequence length, and you can use + cache_seqlens to keep track of the current sequence lengths of each sequence in the batch. + + Also apply rotary embedding if rotary_cos and rotary_sin are passed in. The key @k will be + rotated by rotary_cos and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc. + If causal or local (i.e., window_size != (-1, -1)), the query @q will be rotated by rotary_cos + and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc. + If not causal and not local, the query @q will be rotated by rotary_cos and rotary_sin at + indices cache_seqlens only (i.e. we consider all tokens in @q to be at position cache_seqlens). + + See tests/test_flash_attn.py::test_flash_attn_kvcache for examples of how to use this function. + + Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads + than Q. Note that the number of heads in Q must be divisible by the number of heads in KV. + For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head + 0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V. + + If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix. + For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is: + 1 1 1 1 0 + 1 1 1 1 1 + If seqlen_q = 5 and seqlen_k = 2, the causal mask is: + 0 0 + 0 0 + 0 0 + 1 0 + 1 1 + If the row of the mask is all zero, the output will be zero. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between + [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive. + + Note: Does not support backward pass. + + Arguments: + q: (batch_size, seqlen, nheads, headdim) + k_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim) if there's no page_table, + or (num_blocks, page_block_size, nheads_k, headdim) if there's a page_table (i.e. paged KV cache) + page_block_size must be a multiple of 256. + v_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim_v) if there's no page_table, + or (num_blocks, page_block_size, nheads_k, headdim_v) if there's a page_table (i.e. paged KV cache) + k [optional]: (batch_size, seqlen_new, nheads_k, headdim). If not None, we concatenate + k with k_cache, starting at the indices specified by cache_seqlens. + v [optional]: (batch_size, seqlen_new, nheads_k, headdim_v). Similar to k. + qv [optional]: (batch_size, seqlen, nheads, headdim_v) + rotary_cos [optional]: (seqlen_ro, rotary_dim / 2). If not None, we apply rotary embedding + to k and q. Only applicable if k and v are passed in. rotary_dim must be divisible by 16. + rotary_sin [optional]: (seqlen_ro, rotary_dim / 2). Similar to rotary_cos. + cache_seqlens: int, or (batch_size,), dtype torch.int32. The sequence lengths of the + KV cache. + cache_batch_idx: (batch_size,), dtype torch.int32. The indices used to index into the KV cache. + If None, we assume that the batch indices are [0, 1, 2, ..., batch_size - 1]. + If the indices are not distinct, and k and v are provided, the values updated in the cache + might come from any of the duplicate indices. + cache_leftpad: (batch_size,), dtype torch.int32. The index that the KV cache starts. If None, assume 0. + page_table [optional]: (batch_size, max_num_blocks_per_seq), dtype torch.int32. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + softcap: float. Anything > 0 activates softcapping attention. + rotary_interleaved: bool. Only applicable if rotary_cos and rotary_sin are passed in. + If True, rotary embedding will combine dimensions 0 & 1, 2 & 3, etc. If False, + rotary embedding will combine dimensions 0 & rotary_dim / 2, 1 & rotary_dim / 2 + 1 + (i.e. GPT-NeoX style). + num_splits: int. If > 1, split the key/value into this many chunks along the sequence. + If num_splits == 1, we don't split the key/value. If num_splits == 0, we use a heuristic + to automatically determine the number of splits. + Don't change this unless you know what you are doing. + return_softmax_lse: bool. Whether to return the logsumexp of the attention scores. + + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_softmax_lse=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + """ + assert k_cache.stride(-1) == 1, "k_cache must have contiguous last dimension" + assert v_cache.stride(-1) == 1, "v_cache must have contiguous last dimension" + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + if cache_seqlens is not None and isinstance(cache_seqlens, int): + cache_seqlens = torch.full( + (k_cache.shape[0],), cache_seqlens, dtype=torch.int32, device=k_cache.device + ) + cache_seqlens = maybe_contiguous(cache_seqlens) + out, softmax_lse, *rest = _flash_attn_forward( + q, + k_cache, + v_cache, + k, + v, + qv, + None, # out + cu_seqlens_q, + None, # cu_seqlens_k + cu_seqlens_k_new, + None, # seqused_q + cache_seqlens, + max_seqlen_q, + None, # max_seqlen_k + page_table, + cache_batch_idx, + cache_leftpad, + rotary_cos, + rotary_sin, + rotary_seqlens, + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size=window_size, + attention_chunk=attention_chunk, + softcap=softcap, + rotary_interleaved=rotary_interleaved, + scheduler_metadata=scheduler_metadata, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # return (out, softmax_lse) if return_softmax_lse else out + return (out, softmax_lse, *rest) if return_softmax_lse else out + + +def get_scheduler_metadata( + batch_size, max_seqlen_q, max_seqlen_k, num_heads_q, num_heads_kv, headdim, + cache_seqlens: torch.Tensor, + qkv_dtype=torch.bfloat16, + headdim_v=None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + cache_leftpad: Optional[torch.Tensor] = None, + page_size: Optional[int] = None, + max_seqlen_k_new=0, + causal=False, + window_size=(-1, -1), # -1 means infinite context window + attention_chunk=0, + has_softcap=False, + num_splits=0, # Can be tuned for speed + pack_gqa=None, # Can be tuned for speed + sm_margin=0, # Can be tuned if some SMs are used for communication +): + cache_seqlens = maybe_contiguous(cache_seqlens) + if headdim_v is None: + headdim_v = headdim + scheduler_metadata = flash_attn_3_cuda.get_scheduler_metadata( + batch_size, max_seqlen_q, max_seqlen_k, num_heads_q, num_heads_kv, headdim, headdim_v, + qkv_dtype, + cache_seqlens, + cu_seqlens_q, + None, # cu_seqlens_k + cu_seqlens_k_new, + None, # seqused_q + cache_leftpad, + page_size, + max_seqlen_k_new, + causal, + window_size[0], window_size[1], + attention_chunk, + has_softcap, + num_splits, + pack_gqa, + sm_margin, + ) + return scheduler_metadata diff --git a/build/torch27-cxx11-cu126-x86_64-linux/flash_attn3/__init__.py b/build/torch27-cxx11-cu126-x86_64-linux/flash_attn3/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3d12edd447e84d6e65c4fa9ab91ae3c3192884c1 --- /dev/null +++ b/build/torch27-cxx11-cu126-x86_64-linux/flash_attn3/__init__.py @@ -0,0 +1,17 @@ +from .flash_attn_interface import ( + flash_attn_combine, + flash_attn_func, + flash_attn_qkvpacked_func, + flash_attn_varlen_func, + flash_attn_with_kvcache, + get_scheduler_metadata, +) + +__all__ = [ + "flash_attn_combine", + "flash_attn_func", + "flash_attn_qkvpacked_func", + "flash_attn_varlen_func", + "flash_attn_with_kvcache", + "get_scheduler_metadata", +] diff --git a/build/torch27-cxx11-cu126-x86_64-linux/flash_attn3/_flash_attn3_2e75662.abi3.so b/build/torch27-cxx11-cu126-x86_64-linux/flash_attn3/_flash_attn3_2e75662.abi3.so new file mode 100755 index 0000000000000000000000000000000000000000..5dbb976b32fb12d25582d764b4a0f0c1050c2d51 --- /dev/null +++ b/build/torch27-cxx11-cu126-x86_64-linux/flash_attn3/_flash_attn3_2e75662.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c0302224ac29ba4773d926d4cb16c01c45a374c6dd61286aae1f423f2bf495ea +size 838459544 diff --git a/build/torch27-cxx11-cu126-x86_64-linux/flash_attn3/_ops.py b/build/torch27-cxx11-cu126-x86_64-linux/flash_attn3/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..ca7833b9555e0f0a24a98f83825fc1b58f3d1089 --- /dev/null +++ b/build/torch27-cxx11-cu126-x86_64-linux/flash_attn3/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _flash_attn3_2e75662 +ops = torch.ops._flash_attn3_2e75662 + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_flash_attn3_2e75662::{op_name}" \ No newline at end of file diff --git a/build/torch27-cxx11-cu126-x86_64-linux/flash_attn3/flash_attn_interface.py b/build/torch27-cxx11-cu126-x86_64-linux/flash_attn3/flash_attn_interface.py new file mode 100644 index 0000000000000000000000000000000000000000..7572297784a73fdf891af6aafee88428120b72f2 --- /dev/null +++ b/build/torch27-cxx11-cu126-x86_64-linux/flash_attn3/flash_attn_interface.py @@ -0,0 +1,828 @@ +# Copyright (c) 2023, Tri Dao. + +from typing import Optional, Union + +import torch +import torch.nn as nn + +from ._ops import ops as flash_attn_3_cuda + +def maybe_contiguous(x): + return x.contiguous() if x is not None and x.stride(-1) != 1 else x + + +def _flash_attn_forward( + q, + k, + v, + k_new, + v_new, + qv, + out, + cu_seqlens_q, + cu_seqlens_k, + cu_seqlens_k_new, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + page_table, + kv_batch_idx, + leftpad_k, + rotary_cos, + rotary_sin, + seqlens_rotary, + q_descale, + k_descale, + v_descale, + softmax_scale, + causal, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + rotary_interleaved=True, + scheduler_metadata=None, + num_splits=1, + pack_gqa=None, + sm_margin=0): + q, k, k_new, v_new = [maybe_contiguous(x) for x in (q, k, k_new, v_new)] + v = v.contiguous() if v.stride(-1) != 1 and v.stride(-3) != 1 else v + cu_seqlens_q, cu_seqlens_k, cu_seqlens_k_new = [ + maybe_contiguous(x) for x in (cu_seqlens_q, cu_seqlens_k, cu_seqlens_k_new) + ] + seqused_q, seqused_k = [maybe_contiguous(x) for x in (seqused_q, seqused_k)] + page_table, kv_batch_idx, leftpad_k = [ + maybe_contiguous(x) for x in (page_table, kv_batch_idx, leftpad_k) + ] + rotary_cos, rotary_sin = [maybe_contiguous(x) for x in (rotary_cos, rotary_sin)] + seqlens_rotary = maybe_contiguous(seqlens_rotary) + out, softmax_lse, *rest = flash_attn_3_cuda.fwd( + q, + k, + v, + k_new, + v_new, + qv, + out, + cu_seqlens_q, + cu_seqlens_k, + cu_seqlens_k_new, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + page_table, + kv_batch_idx, + leftpad_k, + rotary_cos, + rotary_sin, + seqlens_rotary, + q_descale, + k_descale, + v_descale, + softmax_scale, + causal, + window_size[0], + window_size[1], + attention_chunk, + softcap, + rotary_interleaved, + scheduler_metadata, + num_splits, + pack_gqa, + sm_margin, + ) + return out, softmax_lse, *rest + + +def _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + cu_seqlens_q, + cu_seqlens_k, + sequed_q, + sequed_k, + max_seqlen_q, + max_seqlen_k, + dq, + dk, + dv, + softmax_scale, + causal, + window_size=(-1, -1), + softcap=0.0, + deterministic=False, + sm_margin=0, +): + # dq, dk, dv are allocated by us so they should already be contiguous + dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] + dq, dk, dv, softmax_d, *rest = flash_attn_3_cuda.bwd( + dout, + q, + k, + v, + out, + softmax_lse, + dq, + dk, + dv, + cu_seqlens_q, + cu_seqlens_k, + sequed_q, + sequed_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + window_size[0], + window_size[1], + softcap, + deterministic, + sm_margin, + ) + return dq, dk, dv, softmax_d + + +class FlashAttnQKVPackedFunc(torch.autograd.Function): + @staticmethod + def forward( + ctx, + qkv, + softmax_scale, + causal, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + deterministic=False, + num_heads_q=None, + sm_margin=0, + ): + if softmax_scale is None: + softmax_scale = qkv.shape[-1] ** (-0.5) + if qkv.dim() == 5: + assert qkv.shape[-3] == 3 + q, k, v = qkv.unbind(dim=-3) + else: + assert qkv.dim() == 4 + assert num_heads_q is not None + num_heads_k = (qkv.shape[2] - num_heads_q) // 2 + assert num_heads_k * 2 + num_heads_q == qkv.shape[2] + q, k, v = qkv.split([num_heads_q, num_heads_k, num_heads_k], dim=-2) + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + None, # qv + None, # out + None, None, None, # cu_seqlens_q/k/k_new + None, None, # seqused_q/k + None, None, # max_seqlen_q/k + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size=window_size, + attention_chunk=attention_chunk, + softcap=softcap, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse) + ctx.save_for_backward(q, k, v, out, softmax_lse) + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.ndim = qkv.dim() + ctx.sm_margin = sm_margin + # return out, softmax_lse + return out + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + if ctx.ndim == 5: + qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) + dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) + dq, dk, dv = dqkv.unbind(dim=-3) + else: + num_heads_q = q.shape[2] + num_heads_k = k.shape[2] + qkv_shape = q.shape[:-2] + (num_heads_q + num_heads_k * 2, *q.shape[-1:]) + dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) + dq, dk, dv = dqkv.split([num_heads_q, num_heads_k, num_heads_k], dim=-2) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + None, None, # cu_seqlens_q, cu_seqlens_k, + None, None, # sequed_q, sequed_k, + None, None, # max_seqlen_q, max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size, + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dqkv = dqkv[..., : dout.shape[-1]] # We could have padded the head dimension + return dqkv, None, None, None, None, None, None, None, None, None, None, None + + +class FlashAttnFunc(torch.autograd.Function): + + @staticmethod + def forward( + ctx, + q, + k, + v, + softmax_scale, + causal, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + ): + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + # out, q, k, v, out_padded, softmax_lse = _flash_attn_forward( + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + qv, # qv + None, # out + None, None, None, # cu_seqlens_q/k/k_new + None, None, # seqused_q/k + None, None, # max_seqlen_q/k + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size=window_size, + attention_chunk=attention_chunk, + softcap=softcap, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse) + ctx.save_for_backward(q, k, v, out, softmax_lse) + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.sm_margin = sm_margin + return out, softmax_lse + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + None, None, # cu_seqlens_q, cu_seqlens_k, + None, None, # sequed_q, sequed_k, + None, None, # max_seqlen_q, max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size, + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dq = dq[..., : q.shape[-1]] # We could have padded the head dimension + dk = dk[..., : k.shape[-1]] + dv = dv[..., : v.shape[-1]] + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None + + +class FlashAttnVarlenFunc(torch.autograd.Function): + + @staticmethod + def forward( + ctx, + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + ): + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + # out, q, k, v, out_padded, softmax_lse = _flash_attn_varlen_forward( + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + qv, # qv + None, # out + cu_seqlens_q, + cu_seqlens_k, + None, # cu_seqlens_k_new + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size=window_size, + attention_chunk=attention_chunk, + softcap=softcap, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k) + ctx.save_for_backward(q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k) + ctx.max_seqlen_q = max_seqlen_q + ctx.max_seqlen_k = max_seqlen_k + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.sm_margin = sm_margin + return out, softmax_lse + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + ctx.max_seqlen_q, + ctx.max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size, + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dq = dq[..., : q.shape[-1]] # We could have padded the head dimension + dk = dk[..., : k.shape[-1]] + dv = dv[..., : v.shape[-1]] + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None + + +def flash_attn_qkvpacked_func( + qkv, + softmax_scale=None, + causal=False, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + deterministic=False, + num_heads_q=None, + sm_margin=0, +): + """dropout_p should be set to 0.0 during evaluation + If Q, K, V are already stacked into 1 tensor, this function will be faster than + calling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation + of the gradients of Q, K, V. + For multi-query and grouped-query attention (MQA/GQA), please see + flash_attn_kvpacked_func and flash_attn_func. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between [i - window_size[0], i + window_size[1]] inclusive. + + Arguments: + qkv: (batch_size, seqlen, 3, nheads, headdim) + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + softcap: float. Anything > 0 activates softcapping attention. + alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of (-alibi_slope * |i - j|) is added to + the attention score of query i and key j. + deterministic: bool. Whether to use the deterministic implementation of the backward pass, + which is slightly slower and uses more memory. The forward pass is always deterministic. + return_attn_probs: bool. Whether to return the attention probabilities. This option is for + testing only. The returned probabilities are not guaranteed to be correct + (they might not have the right scaling). + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen). + The output of softmax (possibly with different scaling). It also encodes the dropout + pattern (negative means that location was dropped, nonnegative means it was kept). + """ + return FlashAttnQKVPackedFunc.apply( + qkv, + softmax_scale, + causal, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + deterministic, + num_heads_q, + sm_margin, + ) + + +def flash_attn_func( + q, + k, + v, + softmax_scale=None, + causal=False, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, +): + """dropout_p should be set to 0.0 during evaluation + Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads + than Q. Note that the number of heads in Q must be divisible by the number of heads in KV. + For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head + 0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V. + + If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix. + For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is: + 1 1 1 1 0 + 1 1 1 1 1 + If seqlen_q = 5 and seqlen_k = 2, the causal mask is: + 0 0 + 0 0 + 0 0 + 1 0 + 1 1 + If the row of the mask is all zero, the output will be zero. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between + [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive. + + Arguments: + q: (batch_size, seqlen, nheads, headdim) + k: (batch_size, seqlen, nheads_k, headdim) + v: (batch_size, seqlen, nheads_k, headdim) + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of + (-alibi_slope * |i + seqlen_k - seqlen_q - j|) + is added to the attention score of query i and key j. + deterministic: bool. Whether to use the deterministic implementation of the backward pass, + which is slightly slower and uses more memory. The forward pass is always deterministic. + return_attn_probs: bool. Whether to return the attention probabilities. This option is for + testing only. The returned probabilities are not guaranteed to be correct + (they might not have the right scaling). + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + """ + return FlashAttnFunc.apply( + q, + k, + v, + softmax_scale, + causal, + qv, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + num_splits, + pack_gqa, + deterministic, + sm_margin, + ) + + +def flash_attn_varlen_func( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + seqused_q=None, + seqused_k=None, + softmax_scale=None, + causal=False, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, +): + return FlashAttnVarlenFunc.apply( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + qv, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + num_splits, + pack_gqa, + deterministic, + sm_margin, + ) + + +def flash_attn_combine(out_partial, lse_partial, out=None, out_dtype=None): + return flash_attn_3_cuda.fwd_combine(out_partial, lse_partial, out, out_dtype) + + +def flash_attn_with_kvcache( + q, + k_cache, + v_cache, + k=None, + v=None, + qv=None, + rotary_cos=None, + rotary_sin=None, + cache_seqlens: Optional[Union[(int, torch.Tensor)]] = None, + cache_batch_idx: Optional[torch.Tensor] = None, + cache_leftpad: Optional[torch.Tensor] = None, + page_table: Optional[torch.Tensor] = None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + rotary_seqlens: Optional[torch.Tensor] = None, + q_descale: Optional[torch.Tensor] = None, + k_descale: Optional[torch.Tensor] = None, + v_descale: Optional[torch.Tensor] = None, + softmax_scale=None, + causal=False, + window_size=(-1, -1), # -1 means infinite context window + attention_chunk=0, + softcap=0.0, # 0.0 means deactivated + rotary_interleaved=True, + scheduler_metadata=None, + num_splits=0, # Can be tuned for speed + pack_gqa=None, # Can be tuned for speed + sm_margin=0, # Can be tuned if some SMs are used for communication + return_softmax_lse=False, +): + """ + If k and v are not None, k_cache and v_cache will be updated *inplace* with the new values from + k and v. This is useful for incremental decoding: you can pass in the cached keys/values from + the previous step, and update them with the new keys/values from the current step, and do + attention with the updated cache, all in 1 kernel. + + If you pass in k / v, you must make sure that the cache is large enough to hold the new values. + For example, the KV cache could be pre-allocated with the max sequence length, and you can use + cache_seqlens to keep track of the current sequence lengths of each sequence in the batch. + + Also apply rotary embedding if rotary_cos and rotary_sin are passed in. The key @k will be + rotated by rotary_cos and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc. + If causal or local (i.e., window_size != (-1, -1)), the query @q will be rotated by rotary_cos + and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc. + If not causal and not local, the query @q will be rotated by rotary_cos and rotary_sin at + indices cache_seqlens only (i.e. we consider all tokens in @q to be at position cache_seqlens). + + See tests/test_flash_attn.py::test_flash_attn_kvcache for examples of how to use this function. + + Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads + than Q. Note that the number of heads in Q must be divisible by the number of heads in KV. + For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head + 0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V. + + If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix. + For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is: + 1 1 1 1 0 + 1 1 1 1 1 + If seqlen_q = 5 and seqlen_k = 2, the causal mask is: + 0 0 + 0 0 + 0 0 + 1 0 + 1 1 + If the row of the mask is all zero, the output will be zero. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between + [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive. + + Note: Does not support backward pass. + + Arguments: + q: (batch_size, seqlen, nheads, headdim) + k_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim) if there's no page_table, + or (num_blocks, page_block_size, nheads_k, headdim) if there's a page_table (i.e. paged KV cache) + page_block_size must be a multiple of 256. + v_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim_v) if there's no page_table, + or (num_blocks, page_block_size, nheads_k, headdim_v) if there's a page_table (i.e. paged KV cache) + k [optional]: (batch_size, seqlen_new, nheads_k, headdim). If not None, we concatenate + k with k_cache, starting at the indices specified by cache_seqlens. + v [optional]: (batch_size, seqlen_new, nheads_k, headdim_v). Similar to k. + qv [optional]: (batch_size, seqlen, nheads, headdim_v) + rotary_cos [optional]: (seqlen_ro, rotary_dim / 2). If not None, we apply rotary embedding + to k and q. Only applicable if k and v are passed in. rotary_dim must be divisible by 16. + rotary_sin [optional]: (seqlen_ro, rotary_dim / 2). Similar to rotary_cos. + cache_seqlens: int, or (batch_size,), dtype torch.int32. The sequence lengths of the + KV cache. + cache_batch_idx: (batch_size,), dtype torch.int32. The indices used to index into the KV cache. + If None, we assume that the batch indices are [0, 1, 2, ..., batch_size - 1]. + If the indices are not distinct, and k and v are provided, the values updated in the cache + might come from any of the duplicate indices. + cache_leftpad: (batch_size,), dtype torch.int32. The index that the KV cache starts. If None, assume 0. + page_table [optional]: (batch_size, max_num_blocks_per_seq), dtype torch.int32. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + softcap: float. Anything > 0 activates softcapping attention. + rotary_interleaved: bool. Only applicable if rotary_cos and rotary_sin are passed in. + If True, rotary embedding will combine dimensions 0 & 1, 2 & 3, etc. If False, + rotary embedding will combine dimensions 0 & rotary_dim / 2, 1 & rotary_dim / 2 + 1 + (i.e. GPT-NeoX style). + num_splits: int. If > 1, split the key/value into this many chunks along the sequence. + If num_splits == 1, we don't split the key/value. If num_splits == 0, we use a heuristic + to automatically determine the number of splits. + Don't change this unless you know what you are doing. + return_softmax_lse: bool. Whether to return the logsumexp of the attention scores. + + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_softmax_lse=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + """ + assert k_cache.stride(-1) == 1, "k_cache must have contiguous last dimension" + assert v_cache.stride(-1) == 1, "v_cache must have contiguous last dimension" + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + if cache_seqlens is not None and isinstance(cache_seqlens, int): + cache_seqlens = torch.full( + (k_cache.shape[0],), cache_seqlens, dtype=torch.int32, device=k_cache.device + ) + cache_seqlens = maybe_contiguous(cache_seqlens) + out, softmax_lse, *rest = _flash_attn_forward( + q, + k_cache, + v_cache, + k, + v, + qv, + None, # out + cu_seqlens_q, + None, # cu_seqlens_k + cu_seqlens_k_new, + None, # seqused_q + cache_seqlens, + max_seqlen_q, + None, # max_seqlen_k + page_table, + cache_batch_idx, + cache_leftpad, + rotary_cos, + rotary_sin, + rotary_seqlens, + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size=window_size, + attention_chunk=attention_chunk, + softcap=softcap, + rotary_interleaved=rotary_interleaved, + scheduler_metadata=scheduler_metadata, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # return (out, softmax_lse) if return_softmax_lse else out + return (out, softmax_lse, *rest) if return_softmax_lse else out + + +def get_scheduler_metadata( + batch_size, max_seqlen_q, max_seqlen_k, num_heads_q, num_heads_kv, headdim, + cache_seqlens: torch.Tensor, + qkv_dtype=torch.bfloat16, + headdim_v=None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + cache_leftpad: Optional[torch.Tensor] = None, + page_size: Optional[int] = None, + max_seqlen_k_new=0, + causal=False, + window_size=(-1, -1), # -1 means infinite context window + attention_chunk=0, + has_softcap=False, + num_splits=0, # Can be tuned for speed + pack_gqa=None, # Can be tuned for speed + sm_margin=0, # Can be tuned if some SMs are used for communication +): + cache_seqlens = maybe_contiguous(cache_seqlens) + if headdim_v is None: + headdim_v = headdim + scheduler_metadata = flash_attn_3_cuda.get_scheduler_metadata( + batch_size, max_seqlen_q, max_seqlen_k, num_heads_q, num_heads_kv, headdim, headdim_v, + qkv_dtype, + cache_seqlens, + cu_seqlens_q, + None, # cu_seqlens_k + cu_seqlens_k_new, + None, # seqused_q + cache_leftpad, + page_size, + max_seqlen_k_new, + causal, + window_size[0], window_size[1], + attention_chunk, + has_softcap, + num_splits, + pack_gqa, + sm_margin, + ) + return scheduler_metadata diff --git a/build/torch27-cxx11-cu128-x86_64-linux/flash_attn3/__init__.py b/build/torch27-cxx11-cu128-x86_64-linux/flash_attn3/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3d12edd447e84d6e65c4fa9ab91ae3c3192884c1 --- /dev/null +++ b/build/torch27-cxx11-cu128-x86_64-linux/flash_attn3/__init__.py @@ -0,0 +1,17 @@ +from .flash_attn_interface import ( + flash_attn_combine, + flash_attn_func, + flash_attn_qkvpacked_func, + flash_attn_varlen_func, + flash_attn_with_kvcache, + get_scheduler_metadata, +) + +__all__ = [ + "flash_attn_combine", + "flash_attn_func", + "flash_attn_qkvpacked_func", + "flash_attn_varlen_func", + "flash_attn_with_kvcache", + "get_scheduler_metadata", +] diff --git a/build/torch27-cxx11-cu128-x86_64-linux/flash_attn3/_flash_attn3_2e75662.abi3.so b/build/torch27-cxx11-cu128-x86_64-linux/flash_attn3/_flash_attn3_2e75662.abi3.so new file mode 100755 index 0000000000000000000000000000000000000000..5dbb976b32fb12d25582d764b4a0f0c1050c2d51 --- /dev/null +++ b/build/torch27-cxx11-cu128-x86_64-linux/flash_attn3/_flash_attn3_2e75662.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c0302224ac29ba4773d926d4cb16c01c45a374c6dd61286aae1f423f2bf495ea +size 838459544 diff --git a/build/torch27-cxx11-cu128-x86_64-linux/flash_attn3/_ops.py b/build/torch27-cxx11-cu128-x86_64-linux/flash_attn3/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..ca7833b9555e0f0a24a98f83825fc1b58f3d1089 --- /dev/null +++ b/build/torch27-cxx11-cu128-x86_64-linux/flash_attn3/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _flash_attn3_2e75662 +ops = torch.ops._flash_attn3_2e75662 + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_flash_attn3_2e75662::{op_name}" \ No newline at end of file diff --git a/build/torch27-cxx11-cu128-x86_64-linux/flash_attn3/flash_attn_interface.py b/build/torch27-cxx11-cu128-x86_64-linux/flash_attn3/flash_attn_interface.py new file mode 100644 index 0000000000000000000000000000000000000000..7572297784a73fdf891af6aafee88428120b72f2 --- /dev/null +++ b/build/torch27-cxx11-cu128-x86_64-linux/flash_attn3/flash_attn_interface.py @@ -0,0 +1,828 @@ +# Copyright (c) 2023, Tri Dao. + +from typing import Optional, Union + +import torch +import torch.nn as nn + +from ._ops import ops as flash_attn_3_cuda + +def maybe_contiguous(x): + return x.contiguous() if x is not None and x.stride(-1) != 1 else x + + +def _flash_attn_forward( + q, + k, + v, + k_new, + v_new, + qv, + out, + cu_seqlens_q, + cu_seqlens_k, + cu_seqlens_k_new, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + page_table, + kv_batch_idx, + leftpad_k, + rotary_cos, + rotary_sin, + seqlens_rotary, + q_descale, + k_descale, + v_descale, + softmax_scale, + causal, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + rotary_interleaved=True, + scheduler_metadata=None, + num_splits=1, + pack_gqa=None, + sm_margin=0): + q, k, k_new, v_new = [maybe_contiguous(x) for x in (q, k, k_new, v_new)] + v = v.contiguous() if v.stride(-1) != 1 and v.stride(-3) != 1 else v + cu_seqlens_q, cu_seqlens_k, cu_seqlens_k_new = [ + maybe_contiguous(x) for x in (cu_seqlens_q, cu_seqlens_k, cu_seqlens_k_new) + ] + seqused_q, seqused_k = [maybe_contiguous(x) for x in (seqused_q, seqused_k)] + page_table, kv_batch_idx, leftpad_k = [ + maybe_contiguous(x) for x in (page_table, kv_batch_idx, leftpad_k) + ] + rotary_cos, rotary_sin = [maybe_contiguous(x) for x in (rotary_cos, rotary_sin)] + seqlens_rotary = maybe_contiguous(seqlens_rotary) + out, softmax_lse, *rest = flash_attn_3_cuda.fwd( + q, + k, + v, + k_new, + v_new, + qv, + out, + cu_seqlens_q, + cu_seqlens_k, + cu_seqlens_k_new, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + page_table, + kv_batch_idx, + leftpad_k, + rotary_cos, + rotary_sin, + seqlens_rotary, + q_descale, + k_descale, + v_descale, + softmax_scale, + causal, + window_size[0], + window_size[1], + attention_chunk, + softcap, + rotary_interleaved, + scheduler_metadata, + num_splits, + pack_gqa, + sm_margin, + ) + return out, softmax_lse, *rest + + +def _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + cu_seqlens_q, + cu_seqlens_k, + sequed_q, + sequed_k, + max_seqlen_q, + max_seqlen_k, + dq, + dk, + dv, + softmax_scale, + causal, + window_size=(-1, -1), + softcap=0.0, + deterministic=False, + sm_margin=0, +): + # dq, dk, dv are allocated by us so they should already be contiguous + dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] + dq, dk, dv, softmax_d, *rest = flash_attn_3_cuda.bwd( + dout, + q, + k, + v, + out, + softmax_lse, + dq, + dk, + dv, + cu_seqlens_q, + cu_seqlens_k, + sequed_q, + sequed_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + window_size[0], + window_size[1], + softcap, + deterministic, + sm_margin, + ) + return dq, dk, dv, softmax_d + + +class FlashAttnQKVPackedFunc(torch.autograd.Function): + @staticmethod + def forward( + ctx, + qkv, + softmax_scale, + causal, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + deterministic=False, + num_heads_q=None, + sm_margin=0, + ): + if softmax_scale is None: + softmax_scale = qkv.shape[-1] ** (-0.5) + if qkv.dim() == 5: + assert qkv.shape[-3] == 3 + q, k, v = qkv.unbind(dim=-3) + else: + assert qkv.dim() == 4 + assert num_heads_q is not None + num_heads_k = (qkv.shape[2] - num_heads_q) // 2 + assert num_heads_k * 2 + num_heads_q == qkv.shape[2] + q, k, v = qkv.split([num_heads_q, num_heads_k, num_heads_k], dim=-2) + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + None, # qv + None, # out + None, None, None, # cu_seqlens_q/k/k_new + None, None, # seqused_q/k + None, None, # max_seqlen_q/k + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size=window_size, + attention_chunk=attention_chunk, + softcap=softcap, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse) + ctx.save_for_backward(q, k, v, out, softmax_lse) + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.ndim = qkv.dim() + ctx.sm_margin = sm_margin + # return out, softmax_lse + return out + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + if ctx.ndim == 5: + qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) + dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) + dq, dk, dv = dqkv.unbind(dim=-3) + else: + num_heads_q = q.shape[2] + num_heads_k = k.shape[2] + qkv_shape = q.shape[:-2] + (num_heads_q + num_heads_k * 2, *q.shape[-1:]) + dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) + dq, dk, dv = dqkv.split([num_heads_q, num_heads_k, num_heads_k], dim=-2) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + None, None, # cu_seqlens_q, cu_seqlens_k, + None, None, # sequed_q, sequed_k, + None, None, # max_seqlen_q, max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size, + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dqkv = dqkv[..., : dout.shape[-1]] # We could have padded the head dimension + return dqkv, None, None, None, None, None, None, None, None, None, None, None + + +class FlashAttnFunc(torch.autograd.Function): + + @staticmethod + def forward( + ctx, + q, + k, + v, + softmax_scale, + causal, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + ): + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + # out, q, k, v, out_padded, softmax_lse = _flash_attn_forward( + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + qv, # qv + None, # out + None, None, None, # cu_seqlens_q/k/k_new + None, None, # seqused_q/k + None, None, # max_seqlen_q/k + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size=window_size, + attention_chunk=attention_chunk, + softcap=softcap, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse) + ctx.save_for_backward(q, k, v, out, softmax_lse) + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.sm_margin = sm_margin + return out, softmax_lse + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + None, None, # cu_seqlens_q, cu_seqlens_k, + None, None, # sequed_q, sequed_k, + None, None, # max_seqlen_q, max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size, + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dq = dq[..., : q.shape[-1]] # We could have padded the head dimension + dk = dk[..., : k.shape[-1]] + dv = dv[..., : v.shape[-1]] + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None + + +class FlashAttnVarlenFunc(torch.autograd.Function): + + @staticmethod + def forward( + ctx, + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + ): + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + # out, q, k, v, out_padded, softmax_lse = _flash_attn_varlen_forward( + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + qv, # qv + None, # out + cu_seqlens_q, + cu_seqlens_k, + None, # cu_seqlens_k_new + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size=window_size, + attention_chunk=attention_chunk, + softcap=softcap, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k) + ctx.save_for_backward(q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k) + ctx.max_seqlen_q = max_seqlen_q + ctx.max_seqlen_k = max_seqlen_k + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.sm_margin = sm_margin + return out, softmax_lse + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + ctx.max_seqlen_q, + ctx.max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size, + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dq = dq[..., : q.shape[-1]] # We could have padded the head dimension + dk = dk[..., : k.shape[-1]] + dv = dv[..., : v.shape[-1]] + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None + + +def flash_attn_qkvpacked_func( + qkv, + softmax_scale=None, + causal=False, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + deterministic=False, + num_heads_q=None, + sm_margin=0, +): + """dropout_p should be set to 0.0 during evaluation + If Q, K, V are already stacked into 1 tensor, this function will be faster than + calling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation + of the gradients of Q, K, V. + For multi-query and grouped-query attention (MQA/GQA), please see + flash_attn_kvpacked_func and flash_attn_func. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between [i - window_size[0], i + window_size[1]] inclusive. + + Arguments: + qkv: (batch_size, seqlen, 3, nheads, headdim) + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + softcap: float. Anything > 0 activates softcapping attention. + alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of (-alibi_slope * |i - j|) is added to + the attention score of query i and key j. + deterministic: bool. Whether to use the deterministic implementation of the backward pass, + which is slightly slower and uses more memory. The forward pass is always deterministic. + return_attn_probs: bool. Whether to return the attention probabilities. This option is for + testing only. The returned probabilities are not guaranteed to be correct + (they might not have the right scaling). + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen). + The output of softmax (possibly with different scaling). It also encodes the dropout + pattern (negative means that location was dropped, nonnegative means it was kept). + """ + return FlashAttnQKVPackedFunc.apply( + qkv, + softmax_scale, + causal, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + deterministic, + num_heads_q, + sm_margin, + ) + + +def flash_attn_func( + q, + k, + v, + softmax_scale=None, + causal=False, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, +): + """dropout_p should be set to 0.0 during evaluation + Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads + than Q. Note that the number of heads in Q must be divisible by the number of heads in KV. + For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head + 0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V. + + If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix. + For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is: + 1 1 1 1 0 + 1 1 1 1 1 + If seqlen_q = 5 and seqlen_k = 2, the causal mask is: + 0 0 + 0 0 + 0 0 + 1 0 + 1 1 + If the row of the mask is all zero, the output will be zero. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between + [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive. + + Arguments: + q: (batch_size, seqlen, nheads, headdim) + k: (batch_size, seqlen, nheads_k, headdim) + v: (batch_size, seqlen, nheads_k, headdim) + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of + (-alibi_slope * |i + seqlen_k - seqlen_q - j|) + is added to the attention score of query i and key j. + deterministic: bool. Whether to use the deterministic implementation of the backward pass, + which is slightly slower and uses more memory. The forward pass is always deterministic. + return_attn_probs: bool. Whether to return the attention probabilities. This option is for + testing only. The returned probabilities are not guaranteed to be correct + (they might not have the right scaling). + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + """ + return FlashAttnFunc.apply( + q, + k, + v, + softmax_scale, + causal, + qv, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + num_splits, + pack_gqa, + deterministic, + sm_margin, + ) + + +def flash_attn_varlen_func( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + seqused_q=None, + seqused_k=None, + softmax_scale=None, + causal=False, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, +): + return FlashAttnVarlenFunc.apply( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + qv, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + num_splits, + pack_gqa, + deterministic, + sm_margin, + ) + + +def flash_attn_combine(out_partial, lse_partial, out=None, out_dtype=None): + return flash_attn_3_cuda.fwd_combine(out_partial, lse_partial, out, out_dtype) + + +def flash_attn_with_kvcache( + q, + k_cache, + v_cache, + k=None, + v=None, + qv=None, + rotary_cos=None, + rotary_sin=None, + cache_seqlens: Optional[Union[(int, torch.Tensor)]] = None, + cache_batch_idx: Optional[torch.Tensor] = None, + cache_leftpad: Optional[torch.Tensor] = None, + page_table: Optional[torch.Tensor] = None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + rotary_seqlens: Optional[torch.Tensor] = None, + q_descale: Optional[torch.Tensor] = None, + k_descale: Optional[torch.Tensor] = None, + v_descale: Optional[torch.Tensor] = None, + softmax_scale=None, + causal=False, + window_size=(-1, -1), # -1 means infinite context window + attention_chunk=0, + softcap=0.0, # 0.0 means deactivated + rotary_interleaved=True, + scheduler_metadata=None, + num_splits=0, # Can be tuned for speed + pack_gqa=None, # Can be tuned for speed + sm_margin=0, # Can be tuned if some SMs are used for communication + return_softmax_lse=False, +): + """ + If k and v are not None, k_cache and v_cache will be updated *inplace* with the new values from + k and v. This is useful for incremental decoding: you can pass in the cached keys/values from + the previous step, and update them with the new keys/values from the current step, and do + attention with the updated cache, all in 1 kernel. + + If you pass in k / v, you must make sure that the cache is large enough to hold the new values. + For example, the KV cache could be pre-allocated with the max sequence length, and you can use + cache_seqlens to keep track of the current sequence lengths of each sequence in the batch. + + Also apply rotary embedding if rotary_cos and rotary_sin are passed in. The key @k will be + rotated by rotary_cos and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc. + If causal or local (i.e., window_size != (-1, -1)), the query @q will be rotated by rotary_cos + and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc. + If not causal and not local, the query @q will be rotated by rotary_cos and rotary_sin at + indices cache_seqlens only (i.e. we consider all tokens in @q to be at position cache_seqlens). + + See tests/test_flash_attn.py::test_flash_attn_kvcache for examples of how to use this function. + + Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads + than Q. Note that the number of heads in Q must be divisible by the number of heads in KV. + For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head + 0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V. + + If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix. + For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is: + 1 1 1 1 0 + 1 1 1 1 1 + If seqlen_q = 5 and seqlen_k = 2, the causal mask is: + 0 0 + 0 0 + 0 0 + 1 0 + 1 1 + If the row of the mask is all zero, the output will be zero. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between + [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive. + + Note: Does not support backward pass. + + Arguments: + q: (batch_size, seqlen, nheads, headdim) + k_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim) if there's no page_table, + or (num_blocks, page_block_size, nheads_k, headdim) if there's a page_table (i.e. paged KV cache) + page_block_size must be a multiple of 256. + v_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim_v) if there's no page_table, + or (num_blocks, page_block_size, nheads_k, headdim_v) if there's a page_table (i.e. paged KV cache) + k [optional]: (batch_size, seqlen_new, nheads_k, headdim). If not None, we concatenate + k with k_cache, starting at the indices specified by cache_seqlens. + v [optional]: (batch_size, seqlen_new, nheads_k, headdim_v). Similar to k. + qv [optional]: (batch_size, seqlen, nheads, headdim_v) + rotary_cos [optional]: (seqlen_ro, rotary_dim / 2). If not None, we apply rotary embedding + to k and q. Only applicable if k and v are passed in. rotary_dim must be divisible by 16. + rotary_sin [optional]: (seqlen_ro, rotary_dim / 2). Similar to rotary_cos. + cache_seqlens: int, or (batch_size,), dtype torch.int32. The sequence lengths of the + KV cache. + cache_batch_idx: (batch_size,), dtype torch.int32. The indices used to index into the KV cache. + If None, we assume that the batch indices are [0, 1, 2, ..., batch_size - 1]. + If the indices are not distinct, and k and v are provided, the values updated in the cache + might come from any of the duplicate indices. + cache_leftpad: (batch_size,), dtype torch.int32. The index that the KV cache starts. If None, assume 0. + page_table [optional]: (batch_size, max_num_blocks_per_seq), dtype torch.int32. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + softcap: float. Anything > 0 activates softcapping attention. + rotary_interleaved: bool. Only applicable if rotary_cos and rotary_sin are passed in. + If True, rotary embedding will combine dimensions 0 & 1, 2 & 3, etc. If False, + rotary embedding will combine dimensions 0 & rotary_dim / 2, 1 & rotary_dim / 2 + 1 + (i.e. GPT-NeoX style). + num_splits: int. If > 1, split the key/value into this many chunks along the sequence. + If num_splits == 1, we don't split the key/value. If num_splits == 0, we use a heuristic + to automatically determine the number of splits. + Don't change this unless you know what you are doing. + return_softmax_lse: bool. Whether to return the logsumexp of the attention scores. + + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_softmax_lse=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + """ + assert k_cache.stride(-1) == 1, "k_cache must have contiguous last dimension" + assert v_cache.stride(-1) == 1, "v_cache must have contiguous last dimension" + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + if cache_seqlens is not None and isinstance(cache_seqlens, int): + cache_seqlens = torch.full( + (k_cache.shape[0],), cache_seqlens, dtype=torch.int32, device=k_cache.device + ) + cache_seqlens = maybe_contiguous(cache_seqlens) + out, softmax_lse, *rest = _flash_attn_forward( + q, + k_cache, + v_cache, + k, + v, + qv, + None, # out + cu_seqlens_q, + None, # cu_seqlens_k + cu_seqlens_k_new, + None, # seqused_q + cache_seqlens, + max_seqlen_q, + None, # max_seqlen_k + page_table, + cache_batch_idx, + cache_leftpad, + rotary_cos, + rotary_sin, + rotary_seqlens, + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size=window_size, + attention_chunk=attention_chunk, + softcap=softcap, + rotary_interleaved=rotary_interleaved, + scheduler_metadata=scheduler_metadata, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # return (out, softmax_lse) if return_softmax_lse else out + return (out, softmax_lse, *rest) if return_softmax_lse else out + + +def get_scheduler_metadata( + batch_size, max_seqlen_q, max_seqlen_k, num_heads_q, num_heads_kv, headdim, + cache_seqlens: torch.Tensor, + qkv_dtype=torch.bfloat16, + headdim_v=None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + cache_leftpad: Optional[torch.Tensor] = None, + page_size: Optional[int] = None, + max_seqlen_k_new=0, + causal=False, + window_size=(-1, -1), # -1 means infinite context window + attention_chunk=0, + has_softcap=False, + num_splits=0, # Can be tuned for speed + pack_gqa=None, # Can be tuned for speed + sm_margin=0, # Can be tuned if some SMs are used for communication +): + cache_seqlens = maybe_contiguous(cache_seqlens) + if headdim_v is None: + headdim_v = headdim + scheduler_metadata = flash_attn_3_cuda.get_scheduler_metadata( + batch_size, max_seqlen_q, max_seqlen_k, num_heads_q, num_heads_kv, headdim, headdim_v, + qkv_dtype, + cache_seqlens, + cu_seqlens_q, + None, # cu_seqlens_k + cu_seqlens_k_new, + None, # seqused_q + cache_leftpad, + page_size, + max_seqlen_k_new, + causal, + window_size[0], window_size[1], + attention_chunk, + has_softcap, + num_splits, + pack_gqa, + sm_margin, + ) + return scheduler_metadata diff --git a/build/torch28-cxx11-cu126-aarch64-linux/flash_attn3/__init__.py b/build/torch28-cxx11-cu126-aarch64-linux/flash_attn3/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3d12edd447e84d6e65c4fa9ab91ae3c3192884c1 --- /dev/null +++ b/build/torch28-cxx11-cu126-aarch64-linux/flash_attn3/__init__.py @@ -0,0 +1,17 @@ +from .flash_attn_interface import ( + flash_attn_combine, + flash_attn_func, + flash_attn_qkvpacked_func, + flash_attn_varlen_func, + flash_attn_with_kvcache, + get_scheduler_metadata, +) + +__all__ = [ + "flash_attn_combine", + "flash_attn_func", + "flash_attn_qkvpacked_func", + "flash_attn_varlen_func", + "flash_attn_with_kvcache", + "get_scheduler_metadata", +] diff --git a/build/torch28-cxx11-cu126-aarch64-linux/flash_attn3/__pycache__/__init__.cpython-313.pyc b/build/torch28-cxx11-cu126-aarch64-linux/flash_attn3/__pycache__/__init__.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..67000c0f5e7fdc01e96d411dda7f2cd337af428a Binary files /dev/null and b/build/torch28-cxx11-cu126-aarch64-linux/flash_attn3/__pycache__/__init__.cpython-313.pyc differ diff --git a/build/torch28-cxx11-cu126-aarch64-linux/flash_attn3/__pycache__/_ops.cpython-313.pyc b/build/torch28-cxx11-cu126-aarch64-linux/flash_attn3/__pycache__/_ops.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9b6208136ffb53d517974b01db22044d4865af7f Binary files /dev/null and b/build/torch28-cxx11-cu126-aarch64-linux/flash_attn3/__pycache__/_ops.cpython-313.pyc differ diff --git a/build/torch28-cxx11-cu126-aarch64-linux/flash_attn3/__pycache__/flash_attn_interface.cpython-313.pyc b/build/torch28-cxx11-cu126-aarch64-linux/flash_attn3/__pycache__/flash_attn_interface.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0abb8046346f3ec15bc4c5c3ac9e0de4c0ee1a93 Binary files /dev/null and b/build/torch28-cxx11-cu126-aarch64-linux/flash_attn3/__pycache__/flash_attn_interface.cpython-313.pyc differ diff --git a/build/torch28-cxx11-cu126-aarch64-linux/flash_attn3/_flash_attn3_8d4f83f.abi3.so b/build/torch28-cxx11-cu126-aarch64-linux/flash_attn3/_flash_attn3_8d4f83f.abi3.so new file mode 100755 index 0000000000000000000000000000000000000000..9b2ccbd3ec7cebe80f62a99cc6ac6d228710393d --- /dev/null +++ b/build/torch28-cxx11-cu126-aarch64-linux/flash_attn3/_flash_attn3_8d4f83f.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e9aef52109e5974778e3ccc2f697c4e6050b365624c843a675ce894b938341cc +size 822395576 diff --git a/build/torch28-cxx11-cu126-aarch64-linux/flash_attn3/_ops.py b/build/torch28-cxx11-cu126-aarch64-linux/flash_attn3/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..31f16e5de836523820e69c0fefef2a57bae2073a --- /dev/null +++ b/build/torch28-cxx11-cu126-aarch64-linux/flash_attn3/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _flash_attn3_8d4f83f +ops = torch.ops._flash_attn3_8d4f83f + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_flash_attn3_8d4f83f::{op_name}" \ No newline at end of file diff --git a/build/torch28-cxx11-cu126-aarch64-linux/flash_attn3/flash_attn_interface.py b/build/torch28-cxx11-cu126-aarch64-linux/flash_attn3/flash_attn_interface.py new file mode 100644 index 0000000000000000000000000000000000000000..7572297784a73fdf891af6aafee88428120b72f2 --- /dev/null +++ b/build/torch28-cxx11-cu126-aarch64-linux/flash_attn3/flash_attn_interface.py @@ -0,0 +1,828 @@ +# Copyright (c) 2023, Tri Dao. + +from typing import Optional, Union + +import torch +import torch.nn as nn + +from ._ops import ops as flash_attn_3_cuda + +def maybe_contiguous(x): + return x.contiguous() if x is not None and x.stride(-1) != 1 else x + + +def _flash_attn_forward( + q, + k, + v, + k_new, + v_new, + qv, + out, + cu_seqlens_q, + cu_seqlens_k, + cu_seqlens_k_new, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + page_table, + kv_batch_idx, + leftpad_k, + rotary_cos, + rotary_sin, + seqlens_rotary, + q_descale, + k_descale, + v_descale, + softmax_scale, + causal, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + rotary_interleaved=True, + scheduler_metadata=None, + num_splits=1, + pack_gqa=None, + sm_margin=0): + q, k, k_new, v_new = [maybe_contiguous(x) for x in (q, k, k_new, v_new)] + v = v.contiguous() if v.stride(-1) != 1 and v.stride(-3) != 1 else v + cu_seqlens_q, cu_seqlens_k, cu_seqlens_k_new = [ + maybe_contiguous(x) for x in (cu_seqlens_q, cu_seqlens_k, cu_seqlens_k_new) + ] + seqused_q, seqused_k = [maybe_contiguous(x) for x in (seqused_q, seqused_k)] + page_table, kv_batch_idx, leftpad_k = [ + maybe_contiguous(x) for x in (page_table, kv_batch_idx, leftpad_k) + ] + rotary_cos, rotary_sin = [maybe_contiguous(x) for x in (rotary_cos, rotary_sin)] + seqlens_rotary = maybe_contiguous(seqlens_rotary) + out, softmax_lse, *rest = flash_attn_3_cuda.fwd( + q, + k, + v, + k_new, + v_new, + qv, + out, + cu_seqlens_q, + cu_seqlens_k, + cu_seqlens_k_new, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + page_table, + kv_batch_idx, + leftpad_k, + rotary_cos, + rotary_sin, + seqlens_rotary, + q_descale, + k_descale, + v_descale, + softmax_scale, + causal, + window_size[0], + window_size[1], + attention_chunk, + softcap, + rotary_interleaved, + scheduler_metadata, + num_splits, + pack_gqa, + sm_margin, + ) + return out, softmax_lse, *rest + + +def _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + cu_seqlens_q, + cu_seqlens_k, + sequed_q, + sequed_k, + max_seqlen_q, + max_seqlen_k, + dq, + dk, + dv, + softmax_scale, + causal, + window_size=(-1, -1), + softcap=0.0, + deterministic=False, + sm_margin=0, +): + # dq, dk, dv are allocated by us so they should already be contiguous + dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] + dq, dk, dv, softmax_d, *rest = flash_attn_3_cuda.bwd( + dout, + q, + k, + v, + out, + softmax_lse, + dq, + dk, + dv, + cu_seqlens_q, + cu_seqlens_k, + sequed_q, + sequed_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + window_size[0], + window_size[1], + softcap, + deterministic, + sm_margin, + ) + return dq, dk, dv, softmax_d + + +class FlashAttnQKVPackedFunc(torch.autograd.Function): + @staticmethod + def forward( + ctx, + qkv, + softmax_scale, + causal, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + deterministic=False, + num_heads_q=None, + sm_margin=0, + ): + if softmax_scale is None: + softmax_scale = qkv.shape[-1] ** (-0.5) + if qkv.dim() == 5: + assert qkv.shape[-3] == 3 + q, k, v = qkv.unbind(dim=-3) + else: + assert qkv.dim() == 4 + assert num_heads_q is not None + num_heads_k = (qkv.shape[2] - num_heads_q) // 2 + assert num_heads_k * 2 + num_heads_q == qkv.shape[2] + q, k, v = qkv.split([num_heads_q, num_heads_k, num_heads_k], dim=-2) + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + None, # qv + None, # out + None, None, None, # cu_seqlens_q/k/k_new + None, None, # seqused_q/k + None, None, # max_seqlen_q/k + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size=window_size, + attention_chunk=attention_chunk, + softcap=softcap, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse) + ctx.save_for_backward(q, k, v, out, softmax_lse) + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.ndim = qkv.dim() + ctx.sm_margin = sm_margin + # return out, softmax_lse + return out + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + if ctx.ndim == 5: + qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) + dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) + dq, dk, dv = dqkv.unbind(dim=-3) + else: + num_heads_q = q.shape[2] + num_heads_k = k.shape[2] + qkv_shape = q.shape[:-2] + (num_heads_q + num_heads_k * 2, *q.shape[-1:]) + dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) + dq, dk, dv = dqkv.split([num_heads_q, num_heads_k, num_heads_k], dim=-2) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + None, None, # cu_seqlens_q, cu_seqlens_k, + None, None, # sequed_q, sequed_k, + None, None, # max_seqlen_q, max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size, + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dqkv = dqkv[..., : dout.shape[-1]] # We could have padded the head dimension + return dqkv, None, None, None, None, None, None, None, None, None, None, None + + +class FlashAttnFunc(torch.autograd.Function): + + @staticmethod + def forward( + ctx, + q, + k, + v, + softmax_scale, + causal, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + ): + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + # out, q, k, v, out_padded, softmax_lse = _flash_attn_forward( + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + qv, # qv + None, # out + None, None, None, # cu_seqlens_q/k/k_new + None, None, # seqused_q/k + None, None, # max_seqlen_q/k + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size=window_size, + attention_chunk=attention_chunk, + softcap=softcap, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse) + ctx.save_for_backward(q, k, v, out, softmax_lse) + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.sm_margin = sm_margin + return out, softmax_lse + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + None, None, # cu_seqlens_q, cu_seqlens_k, + None, None, # sequed_q, sequed_k, + None, None, # max_seqlen_q, max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size, + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dq = dq[..., : q.shape[-1]] # We could have padded the head dimension + dk = dk[..., : k.shape[-1]] + dv = dv[..., : v.shape[-1]] + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None + + +class FlashAttnVarlenFunc(torch.autograd.Function): + + @staticmethod + def forward( + ctx, + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + ): + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + # out, q, k, v, out_padded, softmax_lse = _flash_attn_varlen_forward( + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + qv, # qv + None, # out + cu_seqlens_q, + cu_seqlens_k, + None, # cu_seqlens_k_new + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size=window_size, + attention_chunk=attention_chunk, + softcap=softcap, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k) + ctx.save_for_backward(q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k) + ctx.max_seqlen_q = max_seqlen_q + ctx.max_seqlen_k = max_seqlen_k + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.sm_margin = sm_margin + return out, softmax_lse + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + ctx.max_seqlen_q, + ctx.max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size, + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dq = dq[..., : q.shape[-1]] # We could have padded the head dimension + dk = dk[..., : k.shape[-1]] + dv = dv[..., : v.shape[-1]] + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None + + +def flash_attn_qkvpacked_func( + qkv, + softmax_scale=None, + causal=False, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + deterministic=False, + num_heads_q=None, + sm_margin=0, +): + """dropout_p should be set to 0.0 during evaluation + If Q, K, V are already stacked into 1 tensor, this function will be faster than + calling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation + of the gradients of Q, K, V. + For multi-query and grouped-query attention (MQA/GQA), please see + flash_attn_kvpacked_func and flash_attn_func. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between [i - window_size[0], i + window_size[1]] inclusive. + + Arguments: + qkv: (batch_size, seqlen, 3, nheads, headdim) + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + softcap: float. Anything > 0 activates softcapping attention. + alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of (-alibi_slope * |i - j|) is added to + the attention score of query i and key j. + deterministic: bool. Whether to use the deterministic implementation of the backward pass, + which is slightly slower and uses more memory. The forward pass is always deterministic. + return_attn_probs: bool. Whether to return the attention probabilities. This option is for + testing only. The returned probabilities are not guaranteed to be correct + (they might not have the right scaling). + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen). + The output of softmax (possibly with different scaling). It also encodes the dropout + pattern (negative means that location was dropped, nonnegative means it was kept). + """ + return FlashAttnQKVPackedFunc.apply( + qkv, + softmax_scale, + causal, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + deterministic, + num_heads_q, + sm_margin, + ) + + +def flash_attn_func( + q, + k, + v, + softmax_scale=None, + causal=False, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, +): + """dropout_p should be set to 0.0 during evaluation + Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads + than Q. Note that the number of heads in Q must be divisible by the number of heads in KV. + For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head + 0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V. + + If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix. + For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is: + 1 1 1 1 0 + 1 1 1 1 1 + If seqlen_q = 5 and seqlen_k = 2, the causal mask is: + 0 0 + 0 0 + 0 0 + 1 0 + 1 1 + If the row of the mask is all zero, the output will be zero. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between + [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive. + + Arguments: + q: (batch_size, seqlen, nheads, headdim) + k: (batch_size, seqlen, nheads_k, headdim) + v: (batch_size, seqlen, nheads_k, headdim) + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of + (-alibi_slope * |i + seqlen_k - seqlen_q - j|) + is added to the attention score of query i and key j. + deterministic: bool. Whether to use the deterministic implementation of the backward pass, + which is slightly slower and uses more memory. The forward pass is always deterministic. + return_attn_probs: bool. Whether to return the attention probabilities. This option is for + testing only. The returned probabilities are not guaranteed to be correct + (they might not have the right scaling). + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + """ + return FlashAttnFunc.apply( + q, + k, + v, + softmax_scale, + causal, + qv, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + num_splits, + pack_gqa, + deterministic, + sm_margin, + ) + + +def flash_attn_varlen_func( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + seqused_q=None, + seqused_k=None, + softmax_scale=None, + causal=False, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, +): + return FlashAttnVarlenFunc.apply( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + qv, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + num_splits, + pack_gqa, + deterministic, + sm_margin, + ) + + +def flash_attn_combine(out_partial, lse_partial, out=None, out_dtype=None): + return flash_attn_3_cuda.fwd_combine(out_partial, lse_partial, out, out_dtype) + + +def flash_attn_with_kvcache( + q, + k_cache, + v_cache, + k=None, + v=None, + qv=None, + rotary_cos=None, + rotary_sin=None, + cache_seqlens: Optional[Union[(int, torch.Tensor)]] = None, + cache_batch_idx: Optional[torch.Tensor] = None, + cache_leftpad: Optional[torch.Tensor] = None, + page_table: Optional[torch.Tensor] = None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + rotary_seqlens: Optional[torch.Tensor] = None, + q_descale: Optional[torch.Tensor] = None, + k_descale: Optional[torch.Tensor] = None, + v_descale: Optional[torch.Tensor] = None, + softmax_scale=None, + causal=False, + window_size=(-1, -1), # -1 means infinite context window + attention_chunk=0, + softcap=0.0, # 0.0 means deactivated + rotary_interleaved=True, + scheduler_metadata=None, + num_splits=0, # Can be tuned for speed + pack_gqa=None, # Can be tuned for speed + sm_margin=0, # Can be tuned if some SMs are used for communication + return_softmax_lse=False, +): + """ + If k and v are not None, k_cache and v_cache will be updated *inplace* with the new values from + k and v. This is useful for incremental decoding: you can pass in the cached keys/values from + the previous step, and update them with the new keys/values from the current step, and do + attention with the updated cache, all in 1 kernel. + + If you pass in k / v, you must make sure that the cache is large enough to hold the new values. + For example, the KV cache could be pre-allocated with the max sequence length, and you can use + cache_seqlens to keep track of the current sequence lengths of each sequence in the batch. + + Also apply rotary embedding if rotary_cos and rotary_sin are passed in. The key @k will be + rotated by rotary_cos and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc. + If causal or local (i.e., window_size != (-1, -1)), the query @q will be rotated by rotary_cos + and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc. + If not causal and not local, the query @q will be rotated by rotary_cos and rotary_sin at + indices cache_seqlens only (i.e. we consider all tokens in @q to be at position cache_seqlens). + + See tests/test_flash_attn.py::test_flash_attn_kvcache for examples of how to use this function. + + Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads + than Q. Note that the number of heads in Q must be divisible by the number of heads in KV. + For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head + 0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V. + + If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix. + For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is: + 1 1 1 1 0 + 1 1 1 1 1 + If seqlen_q = 5 and seqlen_k = 2, the causal mask is: + 0 0 + 0 0 + 0 0 + 1 0 + 1 1 + If the row of the mask is all zero, the output will be zero. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between + [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive. + + Note: Does not support backward pass. + + Arguments: + q: (batch_size, seqlen, nheads, headdim) + k_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim) if there's no page_table, + or (num_blocks, page_block_size, nheads_k, headdim) if there's a page_table (i.e. paged KV cache) + page_block_size must be a multiple of 256. + v_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim_v) if there's no page_table, + or (num_blocks, page_block_size, nheads_k, headdim_v) if there's a page_table (i.e. paged KV cache) + k [optional]: (batch_size, seqlen_new, nheads_k, headdim). If not None, we concatenate + k with k_cache, starting at the indices specified by cache_seqlens. + v [optional]: (batch_size, seqlen_new, nheads_k, headdim_v). Similar to k. + qv [optional]: (batch_size, seqlen, nheads, headdim_v) + rotary_cos [optional]: (seqlen_ro, rotary_dim / 2). If not None, we apply rotary embedding + to k and q. Only applicable if k and v are passed in. rotary_dim must be divisible by 16. + rotary_sin [optional]: (seqlen_ro, rotary_dim / 2). Similar to rotary_cos. + cache_seqlens: int, or (batch_size,), dtype torch.int32. The sequence lengths of the + KV cache. + cache_batch_idx: (batch_size,), dtype torch.int32. The indices used to index into the KV cache. + If None, we assume that the batch indices are [0, 1, 2, ..., batch_size - 1]. + If the indices are not distinct, and k and v are provided, the values updated in the cache + might come from any of the duplicate indices. + cache_leftpad: (batch_size,), dtype torch.int32. The index that the KV cache starts. If None, assume 0. + page_table [optional]: (batch_size, max_num_blocks_per_seq), dtype torch.int32. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + softcap: float. Anything > 0 activates softcapping attention. + rotary_interleaved: bool. Only applicable if rotary_cos and rotary_sin are passed in. + If True, rotary embedding will combine dimensions 0 & 1, 2 & 3, etc. If False, + rotary embedding will combine dimensions 0 & rotary_dim / 2, 1 & rotary_dim / 2 + 1 + (i.e. GPT-NeoX style). + num_splits: int. If > 1, split the key/value into this many chunks along the sequence. + If num_splits == 1, we don't split the key/value. If num_splits == 0, we use a heuristic + to automatically determine the number of splits. + Don't change this unless you know what you are doing. + return_softmax_lse: bool. Whether to return the logsumexp of the attention scores. + + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_softmax_lse=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + """ + assert k_cache.stride(-1) == 1, "k_cache must have contiguous last dimension" + assert v_cache.stride(-1) == 1, "v_cache must have contiguous last dimension" + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + if cache_seqlens is not None and isinstance(cache_seqlens, int): + cache_seqlens = torch.full( + (k_cache.shape[0],), cache_seqlens, dtype=torch.int32, device=k_cache.device + ) + cache_seqlens = maybe_contiguous(cache_seqlens) + out, softmax_lse, *rest = _flash_attn_forward( + q, + k_cache, + v_cache, + k, + v, + qv, + None, # out + cu_seqlens_q, + None, # cu_seqlens_k + cu_seqlens_k_new, + None, # seqused_q + cache_seqlens, + max_seqlen_q, + None, # max_seqlen_k + page_table, + cache_batch_idx, + cache_leftpad, + rotary_cos, + rotary_sin, + rotary_seqlens, + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size=window_size, + attention_chunk=attention_chunk, + softcap=softcap, + rotary_interleaved=rotary_interleaved, + scheduler_metadata=scheduler_metadata, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # return (out, softmax_lse) if return_softmax_lse else out + return (out, softmax_lse, *rest) if return_softmax_lse else out + + +def get_scheduler_metadata( + batch_size, max_seqlen_q, max_seqlen_k, num_heads_q, num_heads_kv, headdim, + cache_seqlens: torch.Tensor, + qkv_dtype=torch.bfloat16, + headdim_v=None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + cache_leftpad: Optional[torch.Tensor] = None, + page_size: Optional[int] = None, + max_seqlen_k_new=0, + causal=False, + window_size=(-1, -1), # -1 means infinite context window + attention_chunk=0, + has_softcap=False, + num_splits=0, # Can be tuned for speed + pack_gqa=None, # Can be tuned for speed + sm_margin=0, # Can be tuned if some SMs are used for communication +): + cache_seqlens = maybe_contiguous(cache_seqlens) + if headdim_v is None: + headdim_v = headdim + scheduler_metadata = flash_attn_3_cuda.get_scheduler_metadata( + batch_size, max_seqlen_q, max_seqlen_k, num_heads_q, num_heads_kv, headdim, headdim_v, + qkv_dtype, + cache_seqlens, + cu_seqlens_q, + None, # cu_seqlens_k + cu_seqlens_k_new, + None, # seqused_q + cache_leftpad, + page_size, + max_seqlen_k_new, + causal, + window_size[0], window_size[1], + attention_chunk, + has_softcap, + num_splits, + pack_gqa, + sm_margin, + ) + return scheduler_metadata diff --git a/build/torch28-cxx11-cu126-x86_64-linux/flash_attn3/__init__.py b/build/torch28-cxx11-cu126-x86_64-linux/flash_attn3/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3d12edd447e84d6e65c4fa9ab91ae3c3192884c1 --- /dev/null +++ b/build/torch28-cxx11-cu126-x86_64-linux/flash_attn3/__init__.py @@ -0,0 +1,17 @@ +from .flash_attn_interface import ( + flash_attn_combine, + flash_attn_func, + flash_attn_qkvpacked_func, + flash_attn_varlen_func, + flash_attn_with_kvcache, + get_scheduler_metadata, +) + +__all__ = [ + "flash_attn_combine", + "flash_attn_func", + "flash_attn_qkvpacked_func", + "flash_attn_varlen_func", + "flash_attn_with_kvcache", + "get_scheduler_metadata", +] diff --git a/build/torch28-cxx11-cu126-x86_64-linux/flash_attn3/_flash_attn3_48fe103_dirty.abi3.so b/build/torch28-cxx11-cu126-x86_64-linux/flash_attn3/_flash_attn3_48fe103_dirty.abi3.so new file mode 100755 index 0000000000000000000000000000000000000000..1ba2a3c163fabd5ab28c2dbb1ecab7237423d64b --- /dev/null +++ b/build/torch28-cxx11-cu126-x86_64-linux/flash_attn3/_flash_attn3_48fe103_dirty.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:bc32b815563bc9051986a333a362ff61e37cbd967893212243292fef03b461a5 +size 838544688 diff --git a/build/torch28-cxx11-cu126-x86_64-linux/flash_attn3/_ops.py b/build/torch28-cxx11-cu126-x86_64-linux/flash_attn3/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..a0f25282c78ae764fdfa0f8e251d6bb8f1c0c4eb --- /dev/null +++ b/build/torch28-cxx11-cu126-x86_64-linux/flash_attn3/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _flash_attn3_48fe103_dirty +ops = torch.ops._flash_attn3_48fe103_dirty + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_flash_attn3_48fe103_dirty::{op_name}" \ No newline at end of file diff --git a/build/torch28-cxx11-cu126-x86_64-linux/flash_attn3/flash_attn_interface.py b/build/torch28-cxx11-cu126-x86_64-linux/flash_attn3/flash_attn_interface.py new file mode 100644 index 0000000000000000000000000000000000000000..7572297784a73fdf891af6aafee88428120b72f2 --- /dev/null +++ b/build/torch28-cxx11-cu126-x86_64-linux/flash_attn3/flash_attn_interface.py @@ -0,0 +1,828 @@ +# Copyright (c) 2023, Tri Dao. + +from typing import Optional, Union + +import torch +import torch.nn as nn + +from ._ops import ops as flash_attn_3_cuda + +def maybe_contiguous(x): + return x.contiguous() if x is not None and x.stride(-1) != 1 else x + + +def _flash_attn_forward( + q, + k, + v, + k_new, + v_new, + qv, + out, + cu_seqlens_q, + cu_seqlens_k, + cu_seqlens_k_new, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + page_table, + kv_batch_idx, + leftpad_k, + rotary_cos, + rotary_sin, + seqlens_rotary, + q_descale, + k_descale, + v_descale, + softmax_scale, + causal, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + rotary_interleaved=True, + scheduler_metadata=None, + num_splits=1, + pack_gqa=None, + sm_margin=0): + q, k, k_new, v_new = [maybe_contiguous(x) for x in (q, k, k_new, v_new)] + v = v.contiguous() if v.stride(-1) != 1 and v.stride(-3) != 1 else v + cu_seqlens_q, cu_seqlens_k, cu_seqlens_k_new = [ + maybe_contiguous(x) for x in (cu_seqlens_q, cu_seqlens_k, cu_seqlens_k_new) + ] + seqused_q, seqused_k = [maybe_contiguous(x) for x in (seqused_q, seqused_k)] + page_table, kv_batch_idx, leftpad_k = [ + maybe_contiguous(x) for x in (page_table, kv_batch_idx, leftpad_k) + ] + rotary_cos, rotary_sin = [maybe_contiguous(x) for x in (rotary_cos, rotary_sin)] + seqlens_rotary = maybe_contiguous(seqlens_rotary) + out, softmax_lse, *rest = flash_attn_3_cuda.fwd( + q, + k, + v, + k_new, + v_new, + qv, + out, + cu_seqlens_q, + cu_seqlens_k, + cu_seqlens_k_new, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + page_table, + kv_batch_idx, + leftpad_k, + rotary_cos, + rotary_sin, + seqlens_rotary, + q_descale, + k_descale, + v_descale, + softmax_scale, + causal, + window_size[0], + window_size[1], + attention_chunk, + softcap, + rotary_interleaved, + scheduler_metadata, + num_splits, + pack_gqa, + sm_margin, + ) + return out, softmax_lse, *rest + + +def _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + cu_seqlens_q, + cu_seqlens_k, + sequed_q, + sequed_k, + max_seqlen_q, + max_seqlen_k, + dq, + dk, + dv, + softmax_scale, + causal, + window_size=(-1, -1), + softcap=0.0, + deterministic=False, + sm_margin=0, +): + # dq, dk, dv are allocated by us so they should already be contiguous + dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] + dq, dk, dv, softmax_d, *rest = flash_attn_3_cuda.bwd( + dout, + q, + k, + v, + out, + softmax_lse, + dq, + dk, + dv, + cu_seqlens_q, + cu_seqlens_k, + sequed_q, + sequed_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + window_size[0], + window_size[1], + softcap, + deterministic, + sm_margin, + ) + return dq, dk, dv, softmax_d + + +class FlashAttnQKVPackedFunc(torch.autograd.Function): + @staticmethod + def forward( + ctx, + qkv, + softmax_scale, + causal, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + deterministic=False, + num_heads_q=None, + sm_margin=0, + ): + if softmax_scale is None: + softmax_scale = qkv.shape[-1] ** (-0.5) + if qkv.dim() == 5: + assert qkv.shape[-3] == 3 + q, k, v = qkv.unbind(dim=-3) + else: + assert qkv.dim() == 4 + assert num_heads_q is not None + num_heads_k = (qkv.shape[2] - num_heads_q) // 2 + assert num_heads_k * 2 + num_heads_q == qkv.shape[2] + q, k, v = qkv.split([num_heads_q, num_heads_k, num_heads_k], dim=-2) + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + None, # qv + None, # out + None, None, None, # cu_seqlens_q/k/k_new + None, None, # seqused_q/k + None, None, # max_seqlen_q/k + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size=window_size, + attention_chunk=attention_chunk, + softcap=softcap, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse) + ctx.save_for_backward(q, k, v, out, softmax_lse) + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.ndim = qkv.dim() + ctx.sm_margin = sm_margin + # return out, softmax_lse + return out + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + if ctx.ndim == 5: + qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) + dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) + dq, dk, dv = dqkv.unbind(dim=-3) + else: + num_heads_q = q.shape[2] + num_heads_k = k.shape[2] + qkv_shape = q.shape[:-2] + (num_heads_q + num_heads_k * 2, *q.shape[-1:]) + dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) + dq, dk, dv = dqkv.split([num_heads_q, num_heads_k, num_heads_k], dim=-2) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + None, None, # cu_seqlens_q, cu_seqlens_k, + None, None, # sequed_q, sequed_k, + None, None, # max_seqlen_q, max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size, + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dqkv = dqkv[..., : dout.shape[-1]] # We could have padded the head dimension + return dqkv, None, None, None, None, None, None, None, None, None, None, None + + +class FlashAttnFunc(torch.autograd.Function): + + @staticmethod + def forward( + ctx, + q, + k, + v, + softmax_scale, + causal, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + ): + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + # out, q, k, v, out_padded, softmax_lse = _flash_attn_forward( + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + qv, # qv + None, # out + None, None, None, # cu_seqlens_q/k/k_new + None, None, # seqused_q/k + None, None, # max_seqlen_q/k + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size=window_size, + attention_chunk=attention_chunk, + softcap=softcap, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse) + ctx.save_for_backward(q, k, v, out, softmax_lse) + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.sm_margin = sm_margin + return out, softmax_lse + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + None, None, # cu_seqlens_q, cu_seqlens_k, + None, None, # sequed_q, sequed_k, + None, None, # max_seqlen_q, max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size, + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dq = dq[..., : q.shape[-1]] # We could have padded the head dimension + dk = dk[..., : k.shape[-1]] + dv = dv[..., : v.shape[-1]] + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None + + +class FlashAttnVarlenFunc(torch.autograd.Function): + + @staticmethod + def forward( + ctx, + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + ): + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + # out, q, k, v, out_padded, softmax_lse = _flash_attn_varlen_forward( + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + qv, # qv + None, # out + cu_seqlens_q, + cu_seqlens_k, + None, # cu_seqlens_k_new + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size=window_size, + attention_chunk=attention_chunk, + softcap=softcap, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k) + ctx.save_for_backward(q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k) + ctx.max_seqlen_q = max_seqlen_q + ctx.max_seqlen_k = max_seqlen_k + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.sm_margin = sm_margin + return out, softmax_lse + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + ctx.max_seqlen_q, + ctx.max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size, + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dq = dq[..., : q.shape[-1]] # We could have padded the head dimension + dk = dk[..., : k.shape[-1]] + dv = dv[..., : v.shape[-1]] + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None + + +def flash_attn_qkvpacked_func( + qkv, + softmax_scale=None, + causal=False, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + deterministic=False, + num_heads_q=None, + sm_margin=0, +): + """dropout_p should be set to 0.0 during evaluation + If Q, K, V are already stacked into 1 tensor, this function will be faster than + calling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation + of the gradients of Q, K, V. + For multi-query and grouped-query attention (MQA/GQA), please see + flash_attn_kvpacked_func and flash_attn_func. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between [i - window_size[0], i + window_size[1]] inclusive. + + Arguments: + qkv: (batch_size, seqlen, 3, nheads, headdim) + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + softcap: float. Anything > 0 activates softcapping attention. + alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of (-alibi_slope * |i - j|) is added to + the attention score of query i and key j. + deterministic: bool. Whether to use the deterministic implementation of the backward pass, + which is slightly slower and uses more memory. The forward pass is always deterministic. + return_attn_probs: bool. Whether to return the attention probabilities. This option is for + testing only. The returned probabilities are not guaranteed to be correct + (they might not have the right scaling). + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen). + The output of softmax (possibly with different scaling). It also encodes the dropout + pattern (negative means that location was dropped, nonnegative means it was kept). + """ + return FlashAttnQKVPackedFunc.apply( + qkv, + softmax_scale, + causal, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + deterministic, + num_heads_q, + sm_margin, + ) + + +def flash_attn_func( + q, + k, + v, + softmax_scale=None, + causal=False, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, +): + """dropout_p should be set to 0.0 during evaluation + Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads + than Q. Note that the number of heads in Q must be divisible by the number of heads in KV. + For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head + 0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V. + + If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix. + For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is: + 1 1 1 1 0 + 1 1 1 1 1 + If seqlen_q = 5 and seqlen_k = 2, the causal mask is: + 0 0 + 0 0 + 0 0 + 1 0 + 1 1 + If the row of the mask is all zero, the output will be zero. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between + [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive. + + Arguments: + q: (batch_size, seqlen, nheads, headdim) + k: (batch_size, seqlen, nheads_k, headdim) + v: (batch_size, seqlen, nheads_k, headdim) + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of + (-alibi_slope * |i + seqlen_k - seqlen_q - j|) + is added to the attention score of query i and key j. + deterministic: bool. Whether to use the deterministic implementation of the backward pass, + which is slightly slower and uses more memory. The forward pass is always deterministic. + return_attn_probs: bool. Whether to return the attention probabilities. This option is for + testing only. The returned probabilities are not guaranteed to be correct + (they might not have the right scaling). + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + """ + return FlashAttnFunc.apply( + q, + k, + v, + softmax_scale, + causal, + qv, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + num_splits, + pack_gqa, + deterministic, + sm_margin, + ) + + +def flash_attn_varlen_func( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + seqused_q=None, + seqused_k=None, + softmax_scale=None, + causal=False, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, +): + return FlashAttnVarlenFunc.apply( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + qv, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + num_splits, + pack_gqa, + deterministic, + sm_margin, + ) + + +def flash_attn_combine(out_partial, lse_partial, out=None, out_dtype=None): + return flash_attn_3_cuda.fwd_combine(out_partial, lse_partial, out, out_dtype) + + +def flash_attn_with_kvcache( + q, + k_cache, + v_cache, + k=None, + v=None, + qv=None, + rotary_cos=None, + rotary_sin=None, + cache_seqlens: Optional[Union[(int, torch.Tensor)]] = None, + cache_batch_idx: Optional[torch.Tensor] = None, + cache_leftpad: Optional[torch.Tensor] = None, + page_table: Optional[torch.Tensor] = None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + rotary_seqlens: Optional[torch.Tensor] = None, + q_descale: Optional[torch.Tensor] = None, + k_descale: Optional[torch.Tensor] = None, + v_descale: Optional[torch.Tensor] = None, + softmax_scale=None, + causal=False, + window_size=(-1, -1), # -1 means infinite context window + attention_chunk=0, + softcap=0.0, # 0.0 means deactivated + rotary_interleaved=True, + scheduler_metadata=None, + num_splits=0, # Can be tuned for speed + pack_gqa=None, # Can be tuned for speed + sm_margin=0, # Can be tuned if some SMs are used for communication + return_softmax_lse=False, +): + """ + If k and v are not None, k_cache and v_cache will be updated *inplace* with the new values from + k and v. This is useful for incremental decoding: you can pass in the cached keys/values from + the previous step, and update them with the new keys/values from the current step, and do + attention with the updated cache, all in 1 kernel. + + If you pass in k / v, you must make sure that the cache is large enough to hold the new values. + For example, the KV cache could be pre-allocated with the max sequence length, and you can use + cache_seqlens to keep track of the current sequence lengths of each sequence in the batch. + + Also apply rotary embedding if rotary_cos and rotary_sin are passed in. The key @k will be + rotated by rotary_cos and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc. + If causal or local (i.e., window_size != (-1, -1)), the query @q will be rotated by rotary_cos + and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc. + If not causal and not local, the query @q will be rotated by rotary_cos and rotary_sin at + indices cache_seqlens only (i.e. we consider all tokens in @q to be at position cache_seqlens). + + See tests/test_flash_attn.py::test_flash_attn_kvcache for examples of how to use this function. + + Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads + than Q. Note that the number of heads in Q must be divisible by the number of heads in KV. + For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head + 0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V. + + If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix. + For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is: + 1 1 1 1 0 + 1 1 1 1 1 + If seqlen_q = 5 and seqlen_k = 2, the causal mask is: + 0 0 + 0 0 + 0 0 + 1 0 + 1 1 + If the row of the mask is all zero, the output will be zero. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between + [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive. + + Note: Does not support backward pass. + + Arguments: + q: (batch_size, seqlen, nheads, headdim) + k_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim) if there's no page_table, + or (num_blocks, page_block_size, nheads_k, headdim) if there's a page_table (i.e. paged KV cache) + page_block_size must be a multiple of 256. + v_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim_v) if there's no page_table, + or (num_blocks, page_block_size, nheads_k, headdim_v) if there's a page_table (i.e. paged KV cache) + k [optional]: (batch_size, seqlen_new, nheads_k, headdim). If not None, we concatenate + k with k_cache, starting at the indices specified by cache_seqlens. + v [optional]: (batch_size, seqlen_new, nheads_k, headdim_v). Similar to k. + qv [optional]: (batch_size, seqlen, nheads, headdim_v) + rotary_cos [optional]: (seqlen_ro, rotary_dim / 2). If not None, we apply rotary embedding + to k and q. Only applicable if k and v are passed in. rotary_dim must be divisible by 16. + rotary_sin [optional]: (seqlen_ro, rotary_dim / 2). Similar to rotary_cos. + cache_seqlens: int, or (batch_size,), dtype torch.int32. The sequence lengths of the + KV cache. + cache_batch_idx: (batch_size,), dtype torch.int32. The indices used to index into the KV cache. + If None, we assume that the batch indices are [0, 1, 2, ..., batch_size - 1]. + If the indices are not distinct, and k and v are provided, the values updated in the cache + might come from any of the duplicate indices. + cache_leftpad: (batch_size,), dtype torch.int32. The index that the KV cache starts. If None, assume 0. + page_table [optional]: (batch_size, max_num_blocks_per_seq), dtype torch.int32. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + softcap: float. Anything > 0 activates softcapping attention. + rotary_interleaved: bool. Only applicable if rotary_cos and rotary_sin are passed in. + If True, rotary embedding will combine dimensions 0 & 1, 2 & 3, etc. If False, + rotary embedding will combine dimensions 0 & rotary_dim / 2, 1 & rotary_dim / 2 + 1 + (i.e. GPT-NeoX style). + num_splits: int. If > 1, split the key/value into this many chunks along the sequence. + If num_splits == 1, we don't split the key/value. If num_splits == 0, we use a heuristic + to automatically determine the number of splits. + Don't change this unless you know what you are doing. + return_softmax_lse: bool. Whether to return the logsumexp of the attention scores. + + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_softmax_lse=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + """ + assert k_cache.stride(-1) == 1, "k_cache must have contiguous last dimension" + assert v_cache.stride(-1) == 1, "v_cache must have contiguous last dimension" + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + if cache_seqlens is not None and isinstance(cache_seqlens, int): + cache_seqlens = torch.full( + (k_cache.shape[0],), cache_seqlens, dtype=torch.int32, device=k_cache.device + ) + cache_seqlens = maybe_contiguous(cache_seqlens) + out, softmax_lse, *rest = _flash_attn_forward( + q, + k_cache, + v_cache, + k, + v, + qv, + None, # out + cu_seqlens_q, + None, # cu_seqlens_k + cu_seqlens_k_new, + None, # seqused_q + cache_seqlens, + max_seqlen_q, + None, # max_seqlen_k + page_table, + cache_batch_idx, + cache_leftpad, + rotary_cos, + rotary_sin, + rotary_seqlens, + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size=window_size, + attention_chunk=attention_chunk, + softcap=softcap, + rotary_interleaved=rotary_interleaved, + scheduler_metadata=scheduler_metadata, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # return (out, softmax_lse) if return_softmax_lse else out + return (out, softmax_lse, *rest) if return_softmax_lse else out + + +def get_scheduler_metadata( + batch_size, max_seqlen_q, max_seqlen_k, num_heads_q, num_heads_kv, headdim, + cache_seqlens: torch.Tensor, + qkv_dtype=torch.bfloat16, + headdim_v=None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + cache_leftpad: Optional[torch.Tensor] = None, + page_size: Optional[int] = None, + max_seqlen_k_new=0, + causal=False, + window_size=(-1, -1), # -1 means infinite context window + attention_chunk=0, + has_softcap=False, + num_splits=0, # Can be tuned for speed + pack_gqa=None, # Can be tuned for speed + sm_margin=0, # Can be tuned if some SMs are used for communication +): + cache_seqlens = maybe_contiguous(cache_seqlens) + if headdim_v is None: + headdim_v = headdim + scheduler_metadata = flash_attn_3_cuda.get_scheduler_metadata( + batch_size, max_seqlen_q, max_seqlen_k, num_heads_q, num_heads_kv, headdim, headdim_v, + qkv_dtype, + cache_seqlens, + cu_seqlens_q, + None, # cu_seqlens_k + cu_seqlens_k_new, + None, # seqused_q + cache_leftpad, + page_size, + max_seqlen_k_new, + causal, + window_size[0], window_size[1], + attention_chunk, + has_softcap, + num_splits, + pack_gqa, + sm_margin, + ) + return scheduler_metadata diff --git a/build/torch28-cxx11-cu126-x86_64-linux/flash_attn_config.py b/build/torch28-cxx11-cu126-x86_64-linux/flash_attn_config.py new file mode 100644 index 0000000000000000000000000000000000000000..c4b7130324be2c4cabee7ada7000e170a2a81865 --- /dev/null +++ b/build/torch28-cxx11-cu126-x86_64-linux/flash_attn_config.py @@ -0,0 +1,7 @@ +# Auto-generated by flash attention 3 setup.py +CONFIG = {'build_flags': {'FLASHATTENTION_DISABLE_BACKWARD': False, 'FLASHATTENTION_DISABLE_SPLIT': False, 'FLASHATTENTION_DISABLE_PAGEDKV': False, 'FLASHATTENTION_DISABLE_APPENDKV': False, 'FLASHATTENTION_DISABLE_LOCAL': False, 'FLASHATTENTION_DISABLE_SOFTCAP': False, 'FLASHATTENTION_DISABLE_PACKGQA': False, 'FLASHATTENTION_DISABLE_FP16': False, 'FLASHATTENTION_DISABLE_FP8': False, 'FLASHATTENTION_DISABLE_VARLEN': False, 'FLASHATTENTION_DISABLE_CLUSTER': False, 'FLASHATTENTION_DISABLE_HDIM64': False, 'FLASHATTENTION_DISABLE_HDIM96': False, 'FLASHATTENTION_DISABLE_HDIM128': False, 'FLASHATTENTION_DISABLE_HDIM192': False, 'FLASHATTENTION_DISABLE_HDIM256': False, 'FLASHATTENTION_DISABLE_SM8x': False, 'FLASHATTENTION_ENABLE_VCOLMAJOR': False, 'FLASH_ATTENTION_DISABLE_HDIMDIFF64': False, 'FLASH_ATTENTION_DISABLE_HDIMDIFF192': False}} + +def show(): + from pprint import pprint + pprint(CONFIG) + diff --git a/build/torch28-cxx11-cu128-aarch64-linux/flash_attn3/__init__.py b/build/torch28-cxx11-cu128-aarch64-linux/flash_attn3/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3d12edd447e84d6e65c4fa9ab91ae3c3192884c1 --- /dev/null +++ b/build/torch28-cxx11-cu128-aarch64-linux/flash_attn3/__init__.py @@ -0,0 +1,17 @@ +from .flash_attn_interface import ( + flash_attn_combine, + flash_attn_func, + flash_attn_qkvpacked_func, + flash_attn_varlen_func, + flash_attn_with_kvcache, + get_scheduler_metadata, +) + +__all__ = [ + "flash_attn_combine", + "flash_attn_func", + "flash_attn_qkvpacked_func", + "flash_attn_varlen_func", + "flash_attn_with_kvcache", + "get_scheduler_metadata", +] diff --git a/build/torch28-cxx11-cu128-aarch64-linux/flash_attn3/__pycache__/__init__.cpython-313.pyc b/build/torch28-cxx11-cu128-aarch64-linux/flash_attn3/__pycache__/__init__.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..67000c0f5e7fdc01e96d411dda7f2cd337af428a Binary files /dev/null and b/build/torch28-cxx11-cu128-aarch64-linux/flash_attn3/__pycache__/__init__.cpython-313.pyc differ diff --git a/build/torch28-cxx11-cu128-aarch64-linux/flash_attn3/__pycache__/_ops.cpython-313.pyc b/build/torch28-cxx11-cu128-aarch64-linux/flash_attn3/__pycache__/_ops.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9b6208136ffb53d517974b01db22044d4865af7f Binary files /dev/null and b/build/torch28-cxx11-cu128-aarch64-linux/flash_attn3/__pycache__/_ops.cpython-313.pyc differ diff --git a/build/torch28-cxx11-cu128-aarch64-linux/flash_attn3/__pycache__/flash_attn_interface.cpython-313.pyc b/build/torch28-cxx11-cu128-aarch64-linux/flash_attn3/__pycache__/flash_attn_interface.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0abb8046346f3ec15bc4c5c3ac9e0de4c0ee1a93 Binary files /dev/null and b/build/torch28-cxx11-cu128-aarch64-linux/flash_attn3/__pycache__/flash_attn_interface.cpython-313.pyc differ diff --git a/build/torch28-cxx11-cu128-aarch64-linux/flash_attn3/_flash_attn3_8d4f83f.abi3.so b/build/torch28-cxx11-cu128-aarch64-linux/flash_attn3/_flash_attn3_8d4f83f.abi3.so new file mode 100755 index 0000000000000000000000000000000000000000..9b2ccbd3ec7cebe80f62a99cc6ac6d228710393d --- /dev/null +++ b/build/torch28-cxx11-cu128-aarch64-linux/flash_attn3/_flash_attn3_8d4f83f.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e9aef52109e5974778e3ccc2f697c4e6050b365624c843a675ce894b938341cc +size 822395576 diff --git a/build/torch28-cxx11-cu128-aarch64-linux/flash_attn3/_ops.py b/build/torch28-cxx11-cu128-aarch64-linux/flash_attn3/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..31f16e5de836523820e69c0fefef2a57bae2073a --- /dev/null +++ b/build/torch28-cxx11-cu128-aarch64-linux/flash_attn3/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _flash_attn3_8d4f83f +ops = torch.ops._flash_attn3_8d4f83f + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_flash_attn3_8d4f83f::{op_name}" \ No newline at end of file diff --git a/build/torch28-cxx11-cu128-aarch64-linux/flash_attn3/flash_attn_interface.py b/build/torch28-cxx11-cu128-aarch64-linux/flash_attn3/flash_attn_interface.py new file mode 100644 index 0000000000000000000000000000000000000000..7572297784a73fdf891af6aafee88428120b72f2 --- /dev/null +++ b/build/torch28-cxx11-cu128-aarch64-linux/flash_attn3/flash_attn_interface.py @@ -0,0 +1,828 @@ +# Copyright (c) 2023, Tri Dao. + +from typing import Optional, Union + +import torch +import torch.nn as nn + +from ._ops import ops as flash_attn_3_cuda + +def maybe_contiguous(x): + return x.contiguous() if x is not None and x.stride(-1) != 1 else x + + +def _flash_attn_forward( + q, + k, + v, + k_new, + v_new, + qv, + out, + cu_seqlens_q, + cu_seqlens_k, + cu_seqlens_k_new, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + page_table, + kv_batch_idx, + leftpad_k, + rotary_cos, + rotary_sin, + seqlens_rotary, + q_descale, + k_descale, + v_descale, + softmax_scale, + causal, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + rotary_interleaved=True, + scheduler_metadata=None, + num_splits=1, + pack_gqa=None, + sm_margin=0): + q, k, k_new, v_new = [maybe_contiguous(x) for x in (q, k, k_new, v_new)] + v = v.contiguous() if v.stride(-1) != 1 and v.stride(-3) != 1 else v + cu_seqlens_q, cu_seqlens_k, cu_seqlens_k_new = [ + maybe_contiguous(x) for x in (cu_seqlens_q, cu_seqlens_k, cu_seqlens_k_new) + ] + seqused_q, seqused_k = [maybe_contiguous(x) for x in (seqused_q, seqused_k)] + page_table, kv_batch_idx, leftpad_k = [ + maybe_contiguous(x) for x in (page_table, kv_batch_idx, leftpad_k) + ] + rotary_cos, rotary_sin = [maybe_contiguous(x) for x in (rotary_cos, rotary_sin)] + seqlens_rotary = maybe_contiguous(seqlens_rotary) + out, softmax_lse, *rest = flash_attn_3_cuda.fwd( + q, + k, + v, + k_new, + v_new, + qv, + out, + cu_seqlens_q, + cu_seqlens_k, + cu_seqlens_k_new, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + page_table, + kv_batch_idx, + leftpad_k, + rotary_cos, + rotary_sin, + seqlens_rotary, + q_descale, + k_descale, + v_descale, + softmax_scale, + causal, + window_size[0], + window_size[1], + attention_chunk, + softcap, + rotary_interleaved, + scheduler_metadata, + num_splits, + pack_gqa, + sm_margin, + ) + return out, softmax_lse, *rest + + +def _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + cu_seqlens_q, + cu_seqlens_k, + sequed_q, + sequed_k, + max_seqlen_q, + max_seqlen_k, + dq, + dk, + dv, + softmax_scale, + causal, + window_size=(-1, -1), + softcap=0.0, + deterministic=False, + sm_margin=0, +): + # dq, dk, dv are allocated by us so they should already be contiguous + dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] + dq, dk, dv, softmax_d, *rest = flash_attn_3_cuda.bwd( + dout, + q, + k, + v, + out, + softmax_lse, + dq, + dk, + dv, + cu_seqlens_q, + cu_seqlens_k, + sequed_q, + sequed_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + window_size[0], + window_size[1], + softcap, + deterministic, + sm_margin, + ) + return dq, dk, dv, softmax_d + + +class FlashAttnQKVPackedFunc(torch.autograd.Function): + @staticmethod + def forward( + ctx, + qkv, + softmax_scale, + causal, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + deterministic=False, + num_heads_q=None, + sm_margin=0, + ): + if softmax_scale is None: + softmax_scale = qkv.shape[-1] ** (-0.5) + if qkv.dim() == 5: + assert qkv.shape[-3] == 3 + q, k, v = qkv.unbind(dim=-3) + else: + assert qkv.dim() == 4 + assert num_heads_q is not None + num_heads_k = (qkv.shape[2] - num_heads_q) // 2 + assert num_heads_k * 2 + num_heads_q == qkv.shape[2] + q, k, v = qkv.split([num_heads_q, num_heads_k, num_heads_k], dim=-2) + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + None, # qv + None, # out + None, None, None, # cu_seqlens_q/k/k_new + None, None, # seqused_q/k + None, None, # max_seqlen_q/k + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size=window_size, + attention_chunk=attention_chunk, + softcap=softcap, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse) + ctx.save_for_backward(q, k, v, out, softmax_lse) + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.ndim = qkv.dim() + ctx.sm_margin = sm_margin + # return out, softmax_lse + return out + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + if ctx.ndim == 5: + qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) + dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) + dq, dk, dv = dqkv.unbind(dim=-3) + else: + num_heads_q = q.shape[2] + num_heads_k = k.shape[2] + qkv_shape = q.shape[:-2] + (num_heads_q + num_heads_k * 2, *q.shape[-1:]) + dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) + dq, dk, dv = dqkv.split([num_heads_q, num_heads_k, num_heads_k], dim=-2) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + None, None, # cu_seqlens_q, cu_seqlens_k, + None, None, # sequed_q, sequed_k, + None, None, # max_seqlen_q, max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size, + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dqkv = dqkv[..., : dout.shape[-1]] # We could have padded the head dimension + return dqkv, None, None, None, None, None, None, None, None, None, None, None + + +class FlashAttnFunc(torch.autograd.Function): + + @staticmethod + def forward( + ctx, + q, + k, + v, + softmax_scale, + causal, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + ): + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + # out, q, k, v, out_padded, softmax_lse = _flash_attn_forward( + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + qv, # qv + None, # out + None, None, None, # cu_seqlens_q/k/k_new + None, None, # seqused_q/k + None, None, # max_seqlen_q/k + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size=window_size, + attention_chunk=attention_chunk, + softcap=softcap, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse) + ctx.save_for_backward(q, k, v, out, softmax_lse) + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.sm_margin = sm_margin + return out, softmax_lse + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + None, None, # cu_seqlens_q, cu_seqlens_k, + None, None, # sequed_q, sequed_k, + None, None, # max_seqlen_q, max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size, + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dq = dq[..., : q.shape[-1]] # We could have padded the head dimension + dk = dk[..., : k.shape[-1]] + dv = dv[..., : v.shape[-1]] + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None + + +class FlashAttnVarlenFunc(torch.autograd.Function): + + @staticmethod + def forward( + ctx, + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + ): + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + # out, q, k, v, out_padded, softmax_lse = _flash_attn_varlen_forward( + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + qv, # qv + None, # out + cu_seqlens_q, + cu_seqlens_k, + None, # cu_seqlens_k_new + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size=window_size, + attention_chunk=attention_chunk, + softcap=softcap, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k) + ctx.save_for_backward(q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k) + ctx.max_seqlen_q = max_seqlen_q + ctx.max_seqlen_k = max_seqlen_k + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.sm_margin = sm_margin + return out, softmax_lse + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + ctx.max_seqlen_q, + ctx.max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size, + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dq = dq[..., : q.shape[-1]] # We could have padded the head dimension + dk = dk[..., : k.shape[-1]] + dv = dv[..., : v.shape[-1]] + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None + + +def flash_attn_qkvpacked_func( + qkv, + softmax_scale=None, + causal=False, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + deterministic=False, + num_heads_q=None, + sm_margin=0, +): + """dropout_p should be set to 0.0 during evaluation + If Q, K, V are already stacked into 1 tensor, this function will be faster than + calling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation + of the gradients of Q, K, V. + For multi-query and grouped-query attention (MQA/GQA), please see + flash_attn_kvpacked_func and flash_attn_func. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between [i - window_size[0], i + window_size[1]] inclusive. + + Arguments: + qkv: (batch_size, seqlen, 3, nheads, headdim) + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + softcap: float. Anything > 0 activates softcapping attention. + alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of (-alibi_slope * |i - j|) is added to + the attention score of query i and key j. + deterministic: bool. Whether to use the deterministic implementation of the backward pass, + which is slightly slower and uses more memory. The forward pass is always deterministic. + return_attn_probs: bool. Whether to return the attention probabilities. This option is for + testing only. The returned probabilities are not guaranteed to be correct + (they might not have the right scaling). + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen). + The output of softmax (possibly with different scaling). It also encodes the dropout + pattern (negative means that location was dropped, nonnegative means it was kept). + """ + return FlashAttnQKVPackedFunc.apply( + qkv, + softmax_scale, + causal, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + deterministic, + num_heads_q, + sm_margin, + ) + + +def flash_attn_func( + q, + k, + v, + softmax_scale=None, + causal=False, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, +): + """dropout_p should be set to 0.0 during evaluation + Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads + than Q. Note that the number of heads in Q must be divisible by the number of heads in KV. + For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head + 0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V. + + If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix. + For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is: + 1 1 1 1 0 + 1 1 1 1 1 + If seqlen_q = 5 and seqlen_k = 2, the causal mask is: + 0 0 + 0 0 + 0 0 + 1 0 + 1 1 + If the row of the mask is all zero, the output will be zero. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between + [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive. + + Arguments: + q: (batch_size, seqlen, nheads, headdim) + k: (batch_size, seqlen, nheads_k, headdim) + v: (batch_size, seqlen, nheads_k, headdim) + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of + (-alibi_slope * |i + seqlen_k - seqlen_q - j|) + is added to the attention score of query i and key j. + deterministic: bool. Whether to use the deterministic implementation of the backward pass, + which is slightly slower and uses more memory. The forward pass is always deterministic. + return_attn_probs: bool. Whether to return the attention probabilities. This option is for + testing only. The returned probabilities are not guaranteed to be correct + (they might not have the right scaling). + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + """ + return FlashAttnFunc.apply( + q, + k, + v, + softmax_scale, + causal, + qv, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + num_splits, + pack_gqa, + deterministic, + sm_margin, + ) + + +def flash_attn_varlen_func( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + seqused_q=None, + seqused_k=None, + softmax_scale=None, + causal=False, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, +): + return FlashAttnVarlenFunc.apply( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + qv, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + num_splits, + pack_gqa, + deterministic, + sm_margin, + ) + + +def flash_attn_combine(out_partial, lse_partial, out=None, out_dtype=None): + return flash_attn_3_cuda.fwd_combine(out_partial, lse_partial, out, out_dtype) + + +def flash_attn_with_kvcache( + q, + k_cache, + v_cache, + k=None, + v=None, + qv=None, + rotary_cos=None, + rotary_sin=None, + cache_seqlens: Optional[Union[(int, torch.Tensor)]] = None, + cache_batch_idx: Optional[torch.Tensor] = None, + cache_leftpad: Optional[torch.Tensor] = None, + page_table: Optional[torch.Tensor] = None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + rotary_seqlens: Optional[torch.Tensor] = None, + q_descale: Optional[torch.Tensor] = None, + k_descale: Optional[torch.Tensor] = None, + v_descale: Optional[torch.Tensor] = None, + softmax_scale=None, + causal=False, + window_size=(-1, -1), # -1 means infinite context window + attention_chunk=0, + softcap=0.0, # 0.0 means deactivated + rotary_interleaved=True, + scheduler_metadata=None, + num_splits=0, # Can be tuned for speed + pack_gqa=None, # Can be tuned for speed + sm_margin=0, # Can be tuned if some SMs are used for communication + return_softmax_lse=False, +): + """ + If k and v are not None, k_cache and v_cache will be updated *inplace* with the new values from + k and v. This is useful for incremental decoding: you can pass in the cached keys/values from + the previous step, and update them with the new keys/values from the current step, and do + attention with the updated cache, all in 1 kernel. + + If you pass in k / v, you must make sure that the cache is large enough to hold the new values. + For example, the KV cache could be pre-allocated with the max sequence length, and you can use + cache_seqlens to keep track of the current sequence lengths of each sequence in the batch. + + Also apply rotary embedding if rotary_cos and rotary_sin are passed in. The key @k will be + rotated by rotary_cos and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc. + If causal or local (i.e., window_size != (-1, -1)), the query @q will be rotated by rotary_cos + and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc. + If not causal and not local, the query @q will be rotated by rotary_cos and rotary_sin at + indices cache_seqlens only (i.e. we consider all tokens in @q to be at position cache_seqlens). + + See tests/test_flash_attn.py::test_flash_attn_kvcache for examples of how to use this function. + + Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads + than Q. Note that the number of heads in Q must be divisible by the number of heads in KV. + For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head + 0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V. + + If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix. + For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is: + 1 1 1 1 0 + 1 1 1 1 1 + If seqlen_q = 5 and seqlen_k = 2, the causal mask is: + 0 0 + 0 0 + 0 0 + 1 0 + 1 1 + If the row of the mask is all zero, the output will be zero. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between + [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive. + + Note: Does not support backward pass. + + Arguments: + q: (batch_size, seqlen, nheads, headdim) + k_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim) if there's no page_table, + or (num_blocks, page_block_size, nheads_k, headdim) if there's a page_table (i.e. paged KV cache) + page_block_size must be a multiple of 256. + v_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim_v) if there's no page_table, + or (num_blocks, page_block_size, nheads_k, headdim_v) if there's a page_table (i.e. paged KV cache) + k [optional]: (batch_size, seqlen_new, nheads_k, headdim). If not None, we concatenate + k with k_cache, starting at the indices specified by cache_seqlens. + v [optional]: (batch_size, seqlen_new, nheads_k, headdim_v). Similar to k. + qv [optional]: (batch_size, seqlen, nheads, headdim_v) + rotary_cos [optional]: (seqlen_ro, rotary_dim / 2). If not None, we apply rotary embedding + to k and q. Only applicable if k and v are passed in. rotary_dim must be divisible by 16. + rotary_sin [optional]: (seqlen_ro, rotary_dim / 2). Similar to rotary_cos. + cache_seqlens: int, or (batch_size,), dtype torch.int32. The sequence lengths of the + KV cache. + cache_batch_idx: (batch_size,), dtype torch.int32. The indices used to index into the KV cache. + If None, we assume that the batch indices are [0, 1, 2, ..., batch_size - 1]. + If the indices are not distinct, and k and v are provided, the values updated in the cache + might come from any of the duplicate indices. + cache_leftpad: (batch_size,), dtype torch.int32. The index that the KV cache starts. If None, assume 0. + page_table [optional]: (batch_size, max_num_blocks_per_seq), dtype torch.int32. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + softcap: float. Anything > 0 activates softcapping attention. + rotary_interleaved: bool. Only applicable if rotary_cos and rotary_sin are passed in. + If True, rotary embedding will combine dimensions 0 & 1, 2 & 3, etc. If False, + rotary embedding will combine dimensions 0 & rotary_dim / 2, 1 & rotary_dim / 2 + 1 + (i.e. GPT-NeoX style). + num_splits: int. If > 1, split the key/value into this many chunks along the sequence. + If num_splits == 1, we don't split the key/value. If num_splits == 0, we use a heuristic + to automatically determine the number of splits. + Don't change this unless you know what you are doing. + return_softmax_lse: bool. Whether to return the logsumexp of the attention scores. + + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_softmax_lse=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + """ + assert k_cache.stride(-1) == 1, "k_cache must have contiguous last dimension" + assert v_cache.stride(-1) == 1, "v_cache must have contiguous last dimension" + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + if cache_seqlens is not None and isinstance(cache_seqlens, int): + cache_seqlens = torch.full( + (k_cache.shape[0],), cache_seqlens, dtype=torch.int32, device=k_cache.device + ) + cache_seqlens = maybe_contiguous(cache_seqlens) + out, softmax_lse, *rest = _flash_attn_forward( + q, + k_cache, + v_cache, + k, + v, + qv, + None, # out + cu_seqlens_q, + None, # cu_seqlens_k + cu_seqlens_k_new, + None, # seqused_q + cache_seqlens, + max_seqlen_q, + None, # max_seqlen_k + page_table, + cache_batch_idx, + cache_leftpad, + rotary_cos, + rotary_sin, + rotary_seqlens, + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size=window_size, + attention_chunk=attention_chunk, + softcap=softcap, + rotary_interleaved=rotary_interleaved, + scheduler_metadata=scheduler_metadata, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # return (out, softmax_lse) if return_softmax_lse else out + return (out, softmax_lse, *rest) if return_softmax_lse else out + + +def get_scheduler_metadata( + batch_size, max_seqlen_q, max_seqlen_k, num_heads_q, num_heads_kv, headdim, + cache_seqlens: torch.Tensor, + qkv_dtype=torch.bfloat16, + headdim_v=None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + cache_leftpad: Optional[torch.Tensor] = None, + page_size: Optional[int] = None, + max_seqlen_k_new=0, + causal=False, + window_size=(-1, -1), # -1 means infinite context window + attention_chunk=0, + has_softcap=False, + num_splits=0, # Can be tuned for speed + pack_gqa=None, # Can be tuned for speed + sm_margin=0, # Can be tuned if some SMs are used for communication +): + cache_seqlens = maybe_contiguous(cache_seqlens) + if headdim_v is None: + headdim_v = headdim + scheduler_metadata = flash_attn_3_cuda.get_scheduler_metadata( + batch_size, max_seqlen_q, max_seqlen_k, num_heads_q, num_heads_kv, headdim, headdim_v, + qkv_dtype, + cache_seqlens, + cu_seqlens_q, + None, # cu_seqlens_k + cu_seqlens_k_new, + None, # seqused_q + cache_leftpad, + page_size, + max_seqlen_k_new, + causal, + window_size[0], window_size[1], + attention_chunk, + has_softcap, + num_splits, + pack_gqa, + sm_margin, + ) + return scheduler_metadata diff --git a/build/torch28-cxx11-cu128-x86_64-linux/__init__.py b/build/torch28-cxx11-cu128-x86_64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3d12edd447e84d6e65c4fa9ab91ae3c3192884c1 --- /dev/null +++ b/build/torch28-cxx11-cu128-x86_64-linux/__init__.py @@ -0,0 +1,17 @@ +from .flash_attn_interface import ( + flash_attn_combine, + flash_attn_func, + flash_attn_qkvpacked_func, + flash_attn_varlen_func, + flash_attn_with_kvcache, + get_scheduler_metadata, +) + +__all__ = [ + "flash_attn_combine", + "flash_attn_func", + "flash_attn_qkvpacked_func", + "flash_attn_varlen_func", + "flash_attn_with_kvcache", + "get_scheduler_metadata", +] diff --git a/build/torch28-cxx11-cu128-x86_64-linux/_flash_attn3_5630d32.abi3.so b/build/torch28-cxx11-cu128-x86_64-linux/_flash_attn3_5630d32.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..b37dc5b1a79eeffa97b40e3e47f3e667bfce3fcc --- /dev/null +++ b/build/torch28-cxx11-cu128-x86_64-linux/_flash_attn3_5630d32.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2391775caf0ae40d0fd5d06a011b60d47380ce42527ea84ef0cff3e88330a95c +size 804184512 diff --git a/build/torch28-cxx11-cu128-x86_64-linux/_ops.py b/build/torch28-cxx11-cu128-x86_64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..02dbf81201d16b6791a086b87a8b46dee7f45122 --- /dev/null +++ b/build/torch28-cxx11-cu128-x86_64-linux/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _flash_attn3_5630d32 +ops = torch.ops._flash_attn3_5630d32 + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_flash_attn3_5630d32::{op_name}" \ No newline at end of file diff --git a/build/torch28-cxx11-cu128-x86_64-linux/flash_attn3/__init__.py b/build/torch28-cxx11-cu128-x86_64-linux/flash_attn3/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..03dbc1afe1cf156661a2b1b22003cd5f599a0309 --- /dev/null +++ b/build/torch28-cxx11-cu128-x86_64-linux/flash_attn3/__init__.py @@ -0,0 +1,26 @@ +import ctypes +import sys + +import importlib +from pathlib import Path +from types import ModuleType + +def _import_from_path(file_path: Path) -> ModuleType: + # We cannot use the module name as-is, after adding it to `sys.modules`, + # it would also be used for other imports. So, we make a module name that + # depends on the path for it to be unique using the hex-encoded hash of + # the path. + path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value) + module_name = path_hash + spec = importlib.util.spec_from_file_location(module_name, file_path) + if spec is None: + raise ImportError(f"Cannot load spec for {module_name} from {file_path}") + module = importlib.util.module_from_spec(spec) + if module is None: + raise ImportError(f"Cannot load module {module_name} from spec") + sys.modules[module_name] = module + spec.loader.exec_module(module) # type: ignore + return module + + +globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py"))) diff --git a/build/torch28-cxx11-cu128-x86_64-linux/flash_attn_config.py b/build/torch28-cxx11-cu128-x86_64-linux/flash_attn_config.py new file mode 100644 index 0000000000000000000000000000000000000000..c4b7130324be2c4cabee7ada7000e170a2a81865 --- /dev/null +++ b/build/torch28-cxx11-cu128-x86_64-linux/flash_attn_config.py @@ -0,0 +1,7 @@ +# Auto-generated by flash attention 3 setup.py +CONFIG = {'build_flags': {'FLASHATTENTION_DISABLE_BACKWARD': False, 'FLASHATTENTION_DISABLE_SPLIT': False, 'FLASHATTENTION_DISABLE_PAGEDKV': False, 'FLASHATTENTION_DISABLE_APPENDKV': False, 'FLASHATTENTION_DISABLE_LOCAL': False, 'FLASHATTENTION_DISABLE_SOFTCAP': False, 'FLASHATTENTION_DISABLE_PACKGQA': False, 'FLASHATTENTION_DISABLE_FP16': False, 'FLASHATTENTION_DISABLE_FP8': False, 'FLASHATTENTION_DISABLE_VARLEN': False, 'FLASHATTENTION_DISABLE_CLUSTER': False, 'FLASHATTENTION_DISABLE_HDIM64': False, 'FLASHATTENTION_DISABLE_HDIM96': False, 'FLASHATTENTION_DISABLE_HDIM128': False, 'FLASHATTENTION_DISABLE_HDIM192': False, 'FLASHATTENTION_DISABLE_HDIM256': False, 'FLASHATTENTION_DISABLE_SM8x': False, 'FLASHATTENTION_ENABLE_VCOLMAJOR': False, 'FLASH_ATTENTION_DISABLE_HDIMDIFF64': False, 'FLASH_ATTENTION_DISABLE_HDIMDIFF192': False}} + +def show(): + from pprint import pprint + pprint(CONFIG) + diff --git a/build/torch28-cxx11-cu128-x86_64-linux/flash_attn_interface.py b/build/torch28-cxx11-cu128-x86_64-linux/flash_attn_interface.py new file mode 100644 index 0000000000000000000000000000000000000000..2e2a3cdd3e89cbc59bf6c956454e2776c305e0e5 --- /dev/null +++ b/build/torch28-cxx11-cu128-x86_64-linux/flash_attn_interface.py @@ -0,0 +1,1127 @@ +# Copyright (c) 2023, Tri Dao. + +from typing import Optional, Union, List, Tuple + +import torch +import torch.nn as nn + +from ._ops import ops as flash_attn_3_cuda +from ._ops import add_op_namespace_prefix + +def maybe_contiguous(x): + return x.contiguous() if x is not None and x.stride(-1) != 1 else x + + +def round_multiple(x, m): + return (x + m - 1) // m * m + + +def round_up_headdim(head_size: int) -> int: + from .flash_attn_config import CONFIG + + if not CONFIG["build_flags"]["FLASHATTENTION_DISABLE_HDIM64"]: + if head_size <= 64: + return 64 + if not CONFIG["build_flags"]["FLASHATTENTION_DISABLE_HDIM96"]: + if head_size <= 96: + return 96 + if not CONFIG["build_flags"]["FLASHATTENTION_DISABLE_HDIM128"]: + if head_size <= 128: + return 128 + if not CONFIG["build_flags"]["FLASHATTENTION_DISABLE_HDIM192"]: + if head_size <= 192: + return 192 + if not CONFIG["build_flags"]["FLASHATTENTION_DISABLE_HDIM256"]: + if head_size <= 256: + return 256 + return 256 + + +@torch.library.custom_op(add_op_namespace_prefix("_flash_attn_forward"), mutates_args=(), device_types="cuda") +def _flash_attn_forward( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_new: Optional[torch.Tensor] = None, + v_new: Optional[torch.Tensor] = None, + qv: Optional[torch.Tensor] = None, + out_: Optional[torch.Tensor] = None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + seqused_q: Optional[torch.Tensor] = None, + seqused_k: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_k: Optional[int] = None, + page_table: Optional[torch.Tensor] = None, + kv_batch_idx: Optional[torch.Tensor] = None, + leftpad_k: Optional[torch.Tensor] = None, + rotary_cos: Optional[torch.Tensor] = None, + rotary_sin: Optional[torch.Tensor] = None, + seqlens_rotary: Optional[torch.Tensor] = None, + q_descale: Optional[torch.Tensor] = None, + k_descale: Optional[torch.Tensor] = None, + v_descale: Optional[torch.Tensor] = None, + softmax_scale: Optional[float] = None, + causal: bool = False, + window_size_left: int = -1, + window_size_right: int = -1, + attention_chunk: int = 0, + softcap: float = 0.0, + rotary_interleaved: bool = True, + scheduler_metadata: Optional[torch.Tensor] = None, + num_splits: int = 1, + pack_gqa: Optional[bool] = None, + sm_margin: int = 0, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + q, k, k_new, v_new = [maybe_contiguous(x) for x in (q, k, k_new, v_new)] + v = v.contiguous() if v.stride(-1) != 1 and v.stride(-3) != 1 else v + cu_seqlens_q, cu_seqlens_k, cu_seqlens_k_new = [ + maybe_contiguous(x) for x in (cu_seqlens_q, cu_seqlens_k, cu_seqlens_k_new) + ] + seqused_q, seqused_k = [maybe_contiguous(x) for x in (seqused_q, seqused_k)] + page_table, kv_batch_idx, leftpad_k = [ + maybe_contiguous(x) for x in (page_table, kv_batch_idx, leftpad_k) + ] + rotary_cos, rotary_sin = [maybe_contiguous(x) for x in (rotary_cos, rotary_sin)] + seqlens_rotary = maybe_contiguous(seqlens_rotary) + out, softmax_lse, out_accum, softmax_lse_accum = flash_attn_3_cuda.fwd( + q, + k, + v, + k_new, + v_new, + qv, + out_, + cu_seqlens_q, + cu_seqlens_k, + cu_seqlens_k_new, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + page_table, + kv_batch_idx, + leftpad_k, + rotary_cos, + rotary_sin, + seqlens_rotary, + q_descale, + k_descale, + v_descale, + softmax_scale, + causal, + window_size_left, + window_size_right, + attention_chunk, + softcap, + rotary_interleaved, + scheduler_metadata, + num_splits, + pack_gqa, + sm_margin, + ) + + if out_accum is None: + out_accum = torch.tensor([], device=out.device) + + if softmax_lse_accum is None: + softmax_lse_accum = torch.tensor([], device=out.device) + + return out, softmax_lse, out_accum, softmax_lse_accum + + +@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_forward")) +def _flash_attn_forward_fake( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_new: Optional[torch.Tensor] = None, + v_new: Optional[torch.Tensor] = None, + qv: Optional[torch.Tensor] = None, + out_: Optional[torch.Tensor] = None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + seqused_q: Optional[torch.Tensor] = None, + seqused_k: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_k: Optional[int] = None, + page_table: Optional[torch.Tensor] = None, + kv_batch_idx: Optional[torch.Tensor] = None, + leftpad_k: Optional[torch.Tensor] = None, + rotary_cos: Optional[torch.Tensor] = None, + rotary_sin: Optional[torch.Tensor] = None, + seqlens_rotary: Optional[torch.Tensor] = None, + q_descale: Optional[torch.Tensor] = None, + k_descale: Optional[torch.Tensor] = None, + v_descale: Optional[torch.Tensor] = None, + softmax_scale: Optional[float] = None, + causal: bool = False, + window_size_left: int = -1, + window_size_right: int = -1, + attention_chunk: int = 0, + softcap: float = 0.0, + rotary_interleaved: bool = True, + scheduler_metadata: Optional[torch.Tensor] = None, + num_splits: int = 1, + pack_gqa: Optional[bool] = None, + sm_margin: int = 0, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Symbolic fake implementation of flash attention forward. + Returns tensors with the correct shapes and dtypes without actual computation. + """ + + # Determine if we're in varlen mode + is_varlen_q = cu_seqlens_q is not None + + # Get dimensions from query tensor + if is_varlen_q: + # varlen mode: q is (total_q, num_heads, head_size) + total_q, num_heads, head_size = q.shape + batch_size = cu_seqlens_q.shape[0] - 1 + + if max_seqlen_q is None: + raise ValueError("max_seqlen_q must be provided if cu_seqlens_q is provided") + seqlen_q = max_seqlen_q + else: + # batch mode: q is (batch_size, seqlen_q, num_heads, head_size) + batch_size, seqlen_q, num_heads, head_size = q.shape + total_q = batch_size * q.shape[1] + # Get value head dimension + head_size_v = v.shape[-1] + + # Determine output dtype (FP8 inputs produce BF16 outputs) + q_type = q.dtype + if q_type == torch.float8_e4m3fn: + out_dtype = torch.bfloat16 + else: + out_dtype = q_type + + # Create output tensor + if out_ is not None: + # If out_ is provided, _flash_attn_forward becomes non-functional + raise TypeError("Tracing (torch.compile/torch.export) with pre-allocated output tensor is not supported.") + + if is_varlen_q: + out = torch.empty((total_q, num_heads, head_size_v), dtype=out_dtype, device=q.device) + else: + out = torch.empty((batch_size, seqlen_q, num_heads, head_size_v), dtype=out_dtype, device=q.device) + + # Create softmax_lse tensor + if is_varlen_q: + softmax_lse = torch.empty((num_heads, total_q), dtype=torch.float32, device=q.device) + else: + softmax_lse = torch.empty((batch_size, num_heads, seqlen_q), dtype=torch.float32, device=q.device) + + # TODO(guilhermeleobas): Implement "get_num_splits" + # There's an heuristic to compute num_splits when "num_splits <= 0" + # assert that num_splits is > 0 for now + if num_splits <= 0: + raise ValueError(f"tracing (torch.compile/torch.export) with num_splits <= 0 not supported. Got {num_splits=}") + + if num_splits > 1: + if is_varlen_q: + out_accum = torch.empty((num_splits, num_heads, total_q, head_size_v), dtype=torch.float32, device=q.device) + softmax_lse_accum = torch.empty((num_splits, num_heads, total_q), dtype=torch.float32, device=q.device) + else: + out_accum = torch.empty((num_splits, batch_size, num_heads, seqlen_q, head_size_v), dtype=torch.float32, device=q.device) + softmax_lse_accum = torch.empty((num_splits, batch_size, num_heads, seqlen_q), dtype=torch.float32, device=q.device) + else: + # Tensors are not set when num_splits < 1 + out_accum = torch.tensor([], device=out.device) + softmax_lse_accum = torch.tensor([], device=out.device) + + return out, softmax_lse, out_accum, softmax_lse_accum + + +@torch.library.custom_op(add_op_namespace_prefix("_flash_attn_backward"), mutates_args=("dq", "dk", "dv"), device_types="cuda") +def _flash_attn_backward( + dout: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + out: torch.Tensor, + softmax_lse: torch.Tensor, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k: Optional[torch.Tensor] = None, + sequed_q: Optional[torch.Tensor] = None, + sequed_k: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_k: Optional[int] = None, + dq: Optional[torch.Tensor] = None, + dk: Optional[torch.Tensor] = None, + dv: Optional[torch.Tensor] = None, + softmax_scale: Optional[float] = None, + is_causal: bool = False, + window_size_left: int = -1, + window_size_right: int = -1, + softcap: float = 0.0, + deterministic: bool = False, + sm_margin: int = 0, +) -> torch.Tensor: + # dq, dk, dv are allocated by us so they should already be contiguous + dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] + softmax_d, *rest = flash_attn_3_cuda.bwd( + dout, + q, + k, + v, + out, + softmax_lse, + dq, + dk, + dv, + cu_seqlens_q, + cu_seqlens_k, + sequed_q, + sequed_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + is_causal, + window_size_left, + window_size_right, + softcap, + deterministic, + sm_margin, + ) + return softmax_d + + +@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_backward")) +def _flash_attn_backward_fake( + dout: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + out: torch.Tensor, + softmax_lse: torch.Tensor, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k: Optional[torch.Tensor] = None, + sequed_q: Optional[torch.Tensor] = None, + sequed_k: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_k: Optional[int] = None, + dq: Optional[torch.Tensor] = None, + dk: Optional[torch.Tensor] = None, + dv: Optional[torch.Tensor] = None, + softmax_scale: Optional[float] = None, + is_causal: bool = False, + window_size_left: int = -1, + window_size_right: int = -1, + softcap: float = 0.0, + deterministic: bool = False, + sm_margin: int = 0, +) -> torch.Tensor: + + is_varlen_q = cu_seqlens_q is not None + is_varlen_k = cu_seqlens_q is not None + is_varlen = is_varlen_q or is_varlen_k or sequed_q is not None or sequed_k is not None + + if not is_varlen_q: + batch_size = q.size(0) + seqlen_q = q.size(1) + seqlen_k = k.size(1) + total_q = batch_size * q.size(1) + else: + batch_size = cu_seqlens_q.size(0) - 1 + total_q = q.size(0) + seqlen_q = max_seqlen_q + seqlen_k = max_seqlen_k + + if window_size_left >= seqlen_k - 1: + window_size_left = -1 + + if window_size_right >= seqlen_q - 1: + window_size_right = -1 + + if is_causal: + window_size_right = 0 + + is_causal = window_size_left < 0 and window_size_right == 0 + + head_size = q.size(-1) + head_size_v = v.size(-1) + head_size_rounded = round_up_headdim(max(head_size, head_size_v)) + + # Hopper gpus uses cuda compute capabilities 9.0 + cap = torch.cuda.get_device_capability(q.device) + arch = cap[0] * 10 + cap[1] + + is_local = (window_size_left >= 0 or window_size_right >= 0) and not is_causal + + if head_size_rounded <= 64: + kBlockM_sm90 = 96 if (is_causal and softcap > 0.0) else 128 + elif head_size_rounded <= 96: + kBlockM_sm90 = 64 + elif head_size_rounded <= 128: + kBlockM_sm90 = 64 if (is_causal or is_local or softcap > 0.0) else 80 + else: + kBlockM_sm90 = 64 + + kBlockM_sm80 = 128 if head_size_rounded <= 64 else 64 + kBlockM_sm86 = 64 if head_size_rounded <= 192 else 32 + + if arch >= 90: + kBlockM = kBlockM_sm90 + elif arch == 86 or arch == 89: + kBlockM = kBlockM_sm86 + else: + kBlockM = kBlockM_sm80 + + num_heads = q.shape[-2] + seqlen_q_rounded = round_multiple(seqlen_q, kBlockM) + + total_q_padded_rounded = round_multiple(total_q + batch_size * kBlockM, kBlockM) + + dq = torch.empty_like(q) if dq is None else dq + dk = torch.empty_like(k) if dk is None else dk + dv = torch.empty_like(v) if dv is None else dv + + if not is_varlen: + softmax_d = torch.empty((batch_size, num_heads, seqlen_q_rounded), dtype=torch.float32, device=q.device) + else: + softmax_d = torch.empty((num_heads, total_q_padded_rounded), dtype=torch.float32, device=q.device) + + return softmax_d + + +def setup_context(ctx, inputs, output): + q, k, v = inputs[:3] + out, softmax_lse, _, _ = output + ctx.save_for_backward(q, k, v, out, softmax_lse) + ctx.softmax_scale = inputs[-11] + ctx.causal = inputs[-10] + ctx.window_size = [inputs[-9], inputs[-8]] + ctx.attention_chunk = inputs[-7] + ctx.softcap = inputs[-6] + ctx.sm_margin = inputs[-1] + + +def _backward(ctx, dout, *grads): + q, k, v, out, softmax_lse = ctx.saved_tensors + dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + None, None, # cu_seqlens_q, cu_seqlens_k, + None, None, # sequed_q, sequed_k, + None, None, # max_seqlen_q, max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size[0], + ctx.window_size[1], + ctx.softcap, + False, # deterministic + ctx.sm_margin, + ) + return dq, dk, dv, *((None,) * 21) + + +_flash_attn_forward.register_autograd(_backward, setup_context=setup_context) + + + +class FlashAttnQKVPackedFunc(torch.autograd.Function): + @staticmethod + def forward( + ctx, + qkv, + softmax_scale, + causal, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + deterministic=False, + num_heads_q=None, + sm_margin=0, + return_softmax=False, + ): + if softmax_scale is None: + softmax_scale = qkv.shape[-1] ** (-0.5) + if qkv.dim() == 5: + assert qkv.shape[-3] == 3 + q, k, v = qkv.unbind(dim=-3) + else: + assert qkv.dim() == 4 + assert num_heads_q is not None + num_heads_k = (qkv.shape[2] - num_heads_q) // 2 + assert num_heads_k * 2 + num_heads_q == qkv.shape[2] + q, k, v = qkv.split([num_heads_q, num_heads_k, num_heads_k], dim=-2) + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + None, # qv + None, # out + None, None, None, # cu_seqlens_q/k/k_new + None, None, # seqused_q/k + None, None, # max_seqlen_q/k + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size_left=window_size[0], + window_size_right=window_size[1], + attention_chunk=attention_chunk, + softcap=softcap, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse) + ctx.save_for_backward(q, k, v, out, softmax_lse) + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.ndim = qkv.dim() + ctx.sm_margin = sm_margin + return (out, softmax_lse) if return_softmax else out + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + if ctx.ndim == 5: + qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) + dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) + dq, dk, dv = dqkv.unbind(dim=-3) + else: + num_heads_q = q.shape[2] + num_heads_k = k.shape[2] + qkv_shape = q.shape[:-2] + (num_heads_q + num_heads_k * 2, *q.shape[-1:]) + dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) + dq, dk, dv = dqkv.split([num_heads_q, num_heads_k, num_heads_k], dim=-2) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + None, None, # cu_seqlens_q, cu_seqlens_k, + None, None, # sequed_q, sequed_k, + None, None, # max_seqlen_q, max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size[0], + ctx.window_size[1], + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dqkv = dqkv[..., : dout.shape[-1]] # We could have padded the head dimension + return dqkv, None, None, None, None, None, None, None, None, None, None, None, None + + +class FlashAttnFunc(torch.autograd.Function): + + @staticmethod + def forward( + ctx, + q, + k, + v, + softmax_scale, + causal, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + return_softmax=False, + ): + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + # out, q, k, v, out_padded, softmax_lse = _flash_attn_forward( + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + qv, # qv + None, # out + None, None, None, # cu_seqlens_q/k/k_new + None, None, # seqused_q/k + None, None, # max_seqlen_q/k + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size_left=window_size[0], + window_size_right=window_size[1], + attention_chunk=attention_chunk, + softcap=softcap, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse) + ctx.save_for_backward(q, k, v, out, softmax_lse) + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.sm_margin = sm_margin + return (out, softmax_lse) if return_softmax else out + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + None, None, # cu_seqlens_q, cu_seqlens_k, + None, None, # sequed_q, sequed_k, + None, None, # max_seqlen_q, max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size[0], + ctx.window_size[1], + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dq = dq[..., : q.shape[-1]] # We could have padded the head dimension + dk = dk[..., : k.shape[-1]] + dv = dv[..., : v.shape[-1]] + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None + + +class FlashAttnVarlenFunc(torch.autograd.Function): + + @staticmethod + def forward( + ctx, + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + return_softmax=False, + ): + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + # out, q, k, v, out_padded, softmax_lse = _flash_attn_varlen_forward( + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + qv, # qv + None, # out + cu_seqlens_q, + cu_seqlens_k, + None, # cu_seqlens_k_new + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size_left=window_size[0], + window_size_right=window_size[1], + attention_chunk=attention_chunk, + softcap=softcap, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k) + ctx.save_for_backward(q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k) + ctx.max_seqlen_q = max_seqlen_q + ctx.max_seqlen_k = max_seqlen_k + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.sm_margin = sm_margin + return (out, softmax_lse) if return_softmax else out + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + ctx.max_seqlen_q, + ctx.max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size[0], + ctx.window_size[1], + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dq = dq[..., : q.shape[-1]] # We could have padded the head dimension + dk = dk[..., : k.shape[-1]] + dv = dv[..., : v.shape[-1]] + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None + + +def flash_attn_qkvpacked_func( + qkv, + softmax_scale=None, + causal=False, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + deterministic=False, + num_heads_q=None, + sm_margin=0, + return_attn_probs=False, +): + """dropout_p should be set to 0.0 during evaluation + If Q, K, V are already stacked into 1 tensor, this function will be faster than + calling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation + of the gradients of Q, K, V. + For multi-query and grouped-query attention (MQA/GQA), please see + flash_attn_kvpacked_func and flash_attn_func. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between [i - window_size[0], i + window_size[1]] inclusive. + + Arguments: + qkv: (batch_size, seqlen, 3, nheads, headdim) + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + softcap: float. Anything > 0 activates softcapping attention. + alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of (-alibi_slope * |i - j|) is added to + the attention score of query i and key j. + deterministic: bool. Whether to use the deterministic implementation of the backward pass, + which is slightly slower and uses more memory. The forward pass is always deterministic. + return_attn_probs: bool. Whether to return the attention probabilities. This option is for + testing only. The returned probabilities are not guaranteed to be correct + (they might not have the right scaling). + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen). + The output of softmax (possibly with different scaling). It also encodes the dropout + pattern (negative means that location was dropped, nonnegative means it was kept). + """ + return FlashAttnQKVPackedFunc.apply( + qkv, + softmax_scale, + causal, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + deterministic, + num_heads_q, + sm_margin, + return_attn_probs, + ) + + +def flash_attn_func( + q, + k, + v, + softmax_scale=None, + causal=False, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + return_attn_probs=False, +): + """dropout_p should be set to 0.0 during evaluation + Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads + than Q. Note that the number of heads in Q must be divisible by the number of heads in KV. + For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head + 0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V. + + If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix. + For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is: + 1 1 1 1 0 + 1 1 1 1 1 + If seqlen_q = 5 and seqlen_k = 2, the causal mask is: + 0 0 + 0 0 + 0 0 + 1 0 + 1 1 + If the row of the mask is all zero, the output will be zero. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between + [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive. + + Arguments: + q: (batch_size, seqlen, nheads, headdim) + k: (batch_size, seqlen, nheads_k, headdim) + v: (batch_size, seqlen, nheads_k, headdim) + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of + (-alibi_slope * |i + seqlen_k - seqlen_q - j|) + is added to the attention score of query i and key j. + deterministic: bool. Whether to use the deterministic implementation of the backward pass, + which is slightly slower and uses more memory. The forward pass is always deterministic. + return_attn_probs: bool. Whether to return the attention probabilities. This option is for + testing only. The returned probabilities are not guaranteed to be correct + (they might not have the right scaling). + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + """ + return FlashAttnFunc.apply( + q, + k, + v, + softmax_scale, + causal, + qv, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + num_splits, + pack_gqa, + deterministic, + sm_margin, + return_attn_probs, + ) + + +def flash_attn_varlen_func( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + seqused_q=None, + seqused_k=None, + softmax_scale=None, + causal=False, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + return_attn_probs=False, +): + return FlashAttnVarlenFunc.apply( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + qv, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + num_splits, + pack_gqa, + deterministic, + sm_margin, + return_attn_probs, + ) + + +def flash_attn_combine(out_partial, lse_partial, out=None, out_dtype=None): + return flash_attn_3_cuda.fwd_combine(out_partial, lse_partial, out, out_dtype) + + +def flash_attn_with_kvcache( + q, + k_cache, + v_cache, + k=None, + v=None, + qv=None, + rotary_cos=None, + rotary_sin=None, + cache_seqlens: Optional[Union[(int, torch.Tensor)]] = None, + cache_batch_idx: Optional[torch.Tensor] = None, + cache_leftpad: Optional[torch.Tensor] = None, + page_table: Optional[torch.Tensor] = None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + rotary_seqlens: Optional[torch.Tensor] = None, + q_descale: Optional[torch.Tensor] = None, + k_descale: Optional[torch.Tensor] = None, + v_descale: Optional[torch.Tensor] = None, + softmax_scale=None, + causal=False, + window_size=(-1, -1), # -1 means infinite context window + attention_chunk=0, + softcap=0.0, # 0.0 means deactivated + rotary_interleaved=True, + scheduler_metadata=None, + num_splits=0, # Can be tuned for speed + pack_gqa=None, # Can be tuned for speed + sm_margin=0, # Can be tuned if some SMs are used for communication + return_softmax_lse=False, +): + """ + If k and v are not None, k_cache and v_cache will be updated *inplace* with the new values from + k and v. This is useful for incremental decoding: you can pass in the cached keys/values from + the previous step, and update them with the new keys/values from the current step, and do + attention with the updated cache, all in 1 kernel. + + If you pass in k / v, you must make sure that the cache is large enough to hold the new values. + For example, the KV cache could be pre-allocated with the max sequence length, and you can use + cache_seqlens to keep track of the current sequence lengths of each sequence in the batch. + + Also apply rotary embedding if rotary_cos and rotary_sin are passed in. The key @k will be + rotated by rotary_cos and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc. + If causal or local (i.e., window_size != (-1, -1)), the query @q will be rotated by rotary_cos + and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc. + If not causal and not local, the query @q will be rotated by rotary_cos and rotary_sin at + indices cache_seqlens only (i.e. we consider all tokens in @q to be at position cache_seqlens). + + See tests/test_flash_attn.py::test_flash_attn_kvcache for examples of how to use this function. + + Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads + than Q. Note that the number of heads in Q must be divisible by the number of heads in KV. + For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head + 0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V. + + If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix. + For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is: + 1 1 1 1 0 + 1 1 1 1 1 + If seqlen_q = 5 and seqlen_k = 2, the causal mask is: + 0 0 + 0 0 + 0 0 + 1 0 + 1 1 + If the row of the mask is all zero, the output will be zero. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between + [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive. + + Note: Does not support backward pass. + + Arguments: + q: (batch_size, seqlen, nheads, headdim) + k_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim) if there's no page_table, + or (num_blocks, page_block_size, nheads_k, headdim) if there's a page_table (i.e. paged KV cache) + page_block_size can be arbitrary (e.g, 1, 2, 3, 64, etc.). + v_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim_v) if there's no page_table, + or (num_blocks, page_block_size, nheads_k, headdim_v) if there's a page_table (i.e. paged KV cache) + k [optional]: (batch_size, seqlen_new, nheads_k, headdim). If not None, we concatenate + k with k_cache, starting at the indices specified by cache_seqlens. + v [optional]: (batch_size, seqlen_new, nheads_k, headdim_v). Similar to k. + qv [optional]: (batch_size, seqlen, nheads, headdim_v) + rotary_cos [optional]: (seqlen_ro, rotary_dim / 2). If not None, we apply rotary embedding + to k and q. Only applicable if k and v are passed in. rotary_dim must be divisible by 16. + rotary_sin [optional]: (seqlen_ro, rotary_dim / 2). Similar to rotary_cos. + cache_seqlens: int, or (batch_size,), dtype torch.int32. The sequence lengths of the + KV cache. + cache_batch_idx: (batch_size,), dtype torch.int32. The indices used to index into the KV cache. + If None, we assume that the batch indices are [0, 1, 2, ..., batch_size - 1]. + If the indices are not distinct, and k and v are provided, the values updated in the cache + might come from any of the duplicate indices. + cache_leftpad: (batch_size,), dtype torch.int32. The index that the KV cache starts. If None, assume 0. + page_table [optional]: (batch_size, max_num_blocks_per_seq), dtype torch.int32. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + softcap: float. Anything > 0 activates softcapping attention. + rotary_interleaved: bool. Only applicable if rotary_cos and rotary_sin are passed in. + If True, rotary embedding will combine dimensions 0 & 1, 2 & 3, etc. If False, + rotary embedding will combine dimensions 0 & rotary_dim / 2, 1 & rotary_dim / 2 + 1 + (i.e. GPT-NeoX style). + num_splits: int. If > 1, split the key/value into this many chunks along the sequence. + If num_splits == 1, we don't split the key/value. If num_splits == 0, we use a heuristic + to automatically determine the number of splits. + Don't change this unless you know what you are doing. + return_softmax_lse: bool. Whether to return the logsumexp of the attention scores. + + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_softmax_lse=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + """ + assert k_cache.stride(-1) == 1, "k_cache must have contiguous last dimension" + assert v_cache.stride(-1) == 1, "v_cache must have contiguous last dimension" + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + if cache_seqlens is not None and isinstance(cache_seqlens, int): + cache_seqlens = torch.full( + (q.shape[0],), cache_seqlens, dtype=torch.int32, device=k_cache.device + ) + cache_seqlens = maybe_contiguous(cache_seqlens) + out, softmax_lse, *rest = _flash_attn_forward( + q, + k_cache, + v_cache, + k, + v, + qv, + None, # out + cu_seqlens_q, + None, # cu_seqlens_k + cu_seqlens_k_new, + None, # seqused_q + cache_seqlens, + max_seqlen_q, + None, # max_seqlen_k + page_table, + cache_batch_idx, + cache_leftpad, + rotary_cos, + rotary_sin, + rotary_seqlens, + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size_left=window_size[0], + window_size_right=window_size[1], + attention_chunk=attention_chunk, + softcap=softcap, + rotary_interleaved=rotary_interleaved, + scheduler_metadata=scheduler_metadata, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # return (out, softmax_lse) if return_softmax_lse else out + return (out, softmax_lse, *rest) if return_softmax_lse else out + + +def get_scheduler_metadata( + batch_size, max_seqlen_q, max_seqlen_k, num_heads_q, num_heads_kv, headdim, + cache_seqlens: torch.Tensor, + qkv_dtype=torch.bfloat16, + headdim_v=None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + cache_leftpad: Optional[torch.Tensor] = None, + page_size: Optional[int] = None, + max_seqlen_k_new=0, + causal=False, + window_size=(-1, -1), # -1 means infinite context window + attention_chunk=0, + has_softcap=False, + num_splits=0, # Can be tuned for speed + pack_gqa=None, # Can be tuned for speed + sm_margin=0, # Can be tuned if some SMs are used for communication +): + cache_seqlens = maybe_contiguous(cache_seqlens) + if headdim_v is None: + headdim_v = headdim + scheduler_metadata = flash_attn_3_cuda.get_scheduler_metadata( + batch_size, max_seqlen_q, max_seqlen_k, num_heads_q, num_heads_kv, headdim, headdim_v, + qkv_dtype, + cache_seqlens, + cu_seqlens_q, + None, # cu_seqlens_k + cu_seqlens_k_new, + None, # seqused_q + cache_leftpad, + page_size, + max_seqlen_k_new, + causal, + window_size[0], window_size[1], + attention_chunk, + has_softcap, + num_splits, + pack_gqa, + sm_margin, + ) + return scheduler_metadata diff --git a/build/torch28-cxx11-cu128-x86_64-linux/metadata.json b/build/torch28-cxx11-cu128-x86_64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..9cf5deed9898dce769f4cc73913d3530b92a0bd8 --- /dev/null +++ b/build/torch28-cxx11-cu128-x86_64-linux/metadata.json @@ -0,0 +1,4 @@ +{ + "version": 1, + "python-depends": [] +} \ No newline at end of file diff --git a/build/torch28-cxx11-cu129-aarch64-linux/flash_attn3/__init__.py b/build/torch28-cxx11-cu129-aarch64-linux/flash_attn3/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3d12edd447e84d6e65c4fa9ab91ae3c3192884c1 --- /dev/null +++ b/build/torch28-cxx11-cu129-aarch64-linux/flash_attn3/__init__.py @@ -0,0 +1,17 @@ +from .flash_attn_interface import ( + flash_attn_combine, + flash_attn_func, + flash_attn_qkvpacked_func, + flash_attn_varlen_func, + flash_attn_with_kvcache, + get_scheduler_metadata, +) + +__all__ = [ + "flash_attn_combine", + "flash_attn_func", + "flash_attn_qkvpacked_func", + "flash_attn_varlen_func", + "flash_attn_with_kvcache", + "get_scheduler_metadata", +] diff --git a/build/torch28-cxx11-cu129-aarch64-linux/flash_attn3/__pycache__/__init__.cpython-313.pyc b/build/torch28-cxx11-cu129-aarch64-linux/flash_attn3/__pycache__/__init__.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..67000c0f5e7fdc01e96d411dda7f2cd337af428a Binary files /dev/null and b/build/torch28-cxx11-cu129-aarch64-linux/flash_attn3/__pycache__/__init__.cpython-313.pyc differ diff --git a/build/torch28-cxx11-cu129-aarch64-linux/flash_attn3/__pycache__/_ops.cpython-313.pyc b/build/torch28-cxx11-cu129-aarch64-linux/flash_attn3/__pycache__/_ops.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9b6208136ffb53d517974b01db22044d4865af7f Binary files /dev/null and b/build/torch28-cxx11-cu129-aarch64-linux/flash_attn3/__pycache__/_ops.cpython-313.pyc differ diff --git a/build/torch28-cxx11-cu129-aarch64-linux/flash_attn3/__pycache__/flash_attn_interface.cpython-313.pyc b/build/torch28-cxx11-cu129-aarch64-linux/flash_attn3/__pycache__/flash_attn_interface.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0abb8046346f3ec15bc4c5c3ac9e0de4c0ee1a93 Binary files /dev/null and b/build/torch28-cxx11-cu129-aarch64-linux/flash_attn3/__pycache__/flash_attn_interface.cpython-313.pyc differ diff --git a/build/torch28-cxx11-cu129-aarch64-linux/flash_attn3/_flash_attn3_8d4f83f.abi3.so b/build/torch28-cxx11-cu129-aarch64-linux/flash_attn3/_flash_attn3_8d4f83f.abi3.so new file mode 100755 index 0000000000000000000000000000000000000000..9b2ccbd3ec7cebe80f62a99cc6ac6d228710393d --- /dev/null +++ b/build/torch28-cxx11-cu129-aarch64-linux/flash_attn3/_flash_attn3_8d4f83f.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e9aef52109e5974778e3ccc2f697c4e6050b365624c843a675ce894b938341cc +size 822395576 diff --git a/build/torch28-cxx11-cu129-aarch64-linux/flash_attn3/_ops.py b/build/torch28-cxx11-cu129-aarch64-linux/flash_attn3/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..31f16e5de836523820e69c0fefef2a57bae2073a --- /dev/null +++ b/build/torch28-cxx11-cu129-aarch64-linux/flash_attn3/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _flash_attn3_8d4f83f +ops = torch.ops._flash_attn3_8d4f83f + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_flash_attn3_8d4f83f::{op_name}" \ No newline at end of file diff --git a/build/torch28-cxx11-cu129-aarch64-linux/flash_attn3/flash_attn_interface.py b/build/torch28-cxx11-cu129-aarch64-linux/flash_attn3/flash_attn_interface.py new file mode 100644 index 0000000000000000000000000000000000000000..7572297784a73fdf891af6aafee88428120b72f2 --- /dev/null +++ b/build/torch28-cxx11-cu129-aarch64-linux/flash_attn3/flash_attn_interface.py @@ -0,0 +1,828 @@ +# Copyright (c) 2023, Tri Dao. + +from typing import Optional, Union + +import torch +import torch.nn as nn + +from ._ops import ops as flash_attn_3_cuda + +def maybe_contiguous(x): + return x.contiguous() if x is not None and x.stride(-1) != 1 else x + + +def _flash_attn_forward( + q, + k, + v, + k_new, + v_new, + qv, + out, + cu_seqlens_q, + cu_seqlens_k, + cu_seqlens_k_new, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + page_table, + kv_batch_idx, + leftpad_k, + rotary_cos, + rotary_sin, + seqlens_rotary, + q_descale, + k_descale, + v_descale, + softmax_scale, + causal, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + rotary_interleaved=True, + scheduler_metadata=None, + num_splits=1, + pack_gqa=None, + sm_margin=0): + q, k, k_new, v_new = [maybe_contiguous(x) for x in (q, k, k_new, v_new)] + v = v.contiguous() if v.stride(-1) != 1 and v.stride(-3) != 1 else v + cu_seqlens_q, cu_seqlens_k, cu_seqlens_k_new = [ + maybe_contiguous(x) for x in (cu_seqlens_q, cu_seqlens_k, cu_seqlens_k_new) + ] + seqused_q, seqused_k = [maybe_contiguous(x) for x in (seqused_q, seqused_k)] + page_table, kv_batch_idx, leftpad_k = [ + maybe_contiguous(x) for x in (page_table, kv_batch_idx, leftpad_k) + ] + rotary_cos, rotary_sin = [maybe_contiguous(x) for x in (rotary_cos, rotary_sin)] + seqlens_rotary = maybe_contiguous(seqlens_rotary) + out, softmax_lse, *rest = flash_attn_3_cuda.fwd( + q, + k, + v, + k_new, + v_new, + qv, + out, + cu_seqlens_q, + cu_seqlens_k, + cu_seqlens_k_new, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + page_table, + kv_batch_idx, + leftpad_k, + rotary_cos, + rotary_sin, + seqlens_rotary, + q_descale, + k_descale, + v_descale, + softmax_scale, + causal, + window_size[0], + window_size[1], + attention_chunk, + softcap, + rotary_interleaved, + scheduler_metadata, + num_splits, + pack_gqa, + sm_margin, + ) + return out, softmax_lse, *rest + + +def _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + cu_seqlens_q, + cu_seqlens_k, + sequed_q, + sequed_k, + max_seqlen_q, + max_seqlen_k, + dq, + dk, + dv, + softmax_scale, + causal, + window_size=(-1, -1), + softcap=0.0, + deterministic=False, + sm_margin=0, +): + # dq, dk, dv are allocated by us so they should already be contiguous + dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] + dq, dk, dv, softmax_d, *rest = flash_attn_3_cuda.bwd( + dout, + q, + k, + v, + out, + softmax_lse, + dq, + dk, + dv, + cu_seqlens_q, + cu_seqlens_k, + sequed_q, + sequed_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + window_size[0], + window_size[1], + softcap, + deterministic, + sm_margin, + ) + return dq, dk, dv, softmax_d + + +class FlashAttnQKVPackedFunc(torch.autograd.Function): + @staticmethod + def forward( + ctx, + qkv, + softmax_scale, + causal, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + deterministic=False, + num_heads_q=None, + sm_margin=0, + ): + if softmax_scale is None: + softmax_scale = qkv.shape[-1] ** (-0.5) + if qkv.dim() == 5: + assert qkv.shape[-3] == 3 + q, k, v = qkv.unbind(dim=-3) + else: + assert qkv.dim() == 4 + assert num_heads_q is not None + num_heads_k = (qkv.shape[2] - num_heads_q) // 2 + assert num_heads_k * 2 + num_heads_q == qkv.shape[2] + q, k, v = qkv.split([num_heads_q, num_heads_k, num_heads_k], dim=-2) + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + None, # qv + None, # out + None, None, None, # cu_seqlens_q/k/k_new + None, None, # seqused_q/k + None, None, # max_seqlen_q/k + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size=window_size, + attention_chunk=attention_chunk, + softcap=softcap, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse) + ctx.save_for_backward(q, k, v, out, softmax_lse) + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.ndim = qkv.dim() + ctx.sm_margin = sm_margin + # return out, softmax_lse + return out + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + if ctx.ndim == 5: + qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) + dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) + dq, dk, dv = dqkv.unbind(dim=-3) + else: + num_heads_q = q.shape[2] + num_heads_k = k.shape[2] + qkv_shape = q.shape[:-2] + (num_heads_q + num_heads_k * 2, *q.shape[-1:]) + dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) + dq, dk, dv = dqkv.split([num_heads_q, num_heads_k, num_heads_k], dim=-2) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + None, None, # cu_seqlens_q, cu_seqlens_k, + None, None, # sequed_q, sequed_k, + None, None, # max_seqlen_q, max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size, + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dqkv = dqkv[..., : dout.shape[-1]] # We could have padded the head dimension + return dqkv, None, None, None, None, None, None, None, None, None, None, None + + +class FlashAttnFunc(torch.autograd.Function): + + @staticmethod + def forward( + ctx, + q, + k, + v, + softmax_scale, + causal, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + ): + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + # out, q, k, v, out_padded, softmax_lse = _flash_attn_forward( + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + qv, # qv + None, # out + None, None, None, # cu_seqlens_q/k/k_new + None, None, # seqused_q/k + None, None, # max_seqlen_q/k + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size=window_size, + attention_chunk=attention_chunk, + softcap=softcap, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse) + ctx.save_for_backward(q, k, v, out, softmax_lse) + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.sm_margin = sm_margin + return out, softmax_lse + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + None, None, # cu_seqlens_q, cu_seqlens_k, + None, None, # sequed_q, sequed_k, + None, None, # max_seqlen_q, max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size, + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dq = dq[..., : q.shape[-1]] # We could have padded the head dimension + dk = dk[..., : k.shape[-1]] + dv = dv[..., : v.shape[-1]] + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None + + +class FlashAttnVarlenFunc(torch.autograd.Function): + + @staticmethod + def forward( + ctx, + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + ): + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + # out, q, k, v, out_padded, softmax_lse = _flash_attn_varlen_forward( + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + qv, # qv + None, # out + cu_seqlens_q, + cu_seqlens_k, + None, # cu_seqlens_k_new + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size=window_size, + attention_chunk=attention_chunk, + softcap=softcap, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k) + ctx.save_for_backward(q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k) + ctx.max_seqlen_q = max_seqlen_q + ctx.max_seqlen_k = max_seqlen_k + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.sm_margin = sm_margin + return out, softmax_lse + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + ctx.max_seqlen_q, + ctx.max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size, + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dq = dq[..., : q.shape[-1]] # We could have padded the head dimension + dk = dk[..., : k.shape[-1]] + dv = dv[..., : v.shape[-1]] + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None + + +def flash_attn_qkvpacked_func( + qkv, + softmax_scale=None, + causal=False, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + deterministic=False, + num_heads_q=None, + sm_margin=0, +): + """dropout_p should be set to 0.0 during evaluation + If Q, K, V are already stacked into 1 tensor, this function will be faster than + calling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation + of the gradients of Q, K, V. + For multi-query and grouped-query attention (MQA/GQA), please see + flash_attn_kvpacked_func and flash_attn_func. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between [i - window_size[0], i + window_size[1]] inclusive. + + Arguments: + qkv: (batch_size, seqlen, 3, nheads, headdim) + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + softcap: float. Anything > 0 activates softcapping attention. + alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of (-alibi_slope * |i - j|) is added to + the attention score of query i and key j. + deterministic: bool. Whether to use the deterministic implementation of the backward pass, + which is slightly slower and uses more memory. The forward pass is always deterministic. + return_attn_probs: bool. Whether to return the attention probabilities. This option is for + testing only. The returned probabilities are not guaranteed to be correct + (they might not have the right scaling). + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen). + The output of softmax (possibly with different scaling). It also encodes the dropout + pattern (negative means that location was dropped, nonnegative means it was kept). + """ + return FlashAttnQKVPackedFunc.apply( + qkv, + softmax_scale, + causal, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + deterministic, + num_heads_q, + sm_margin, + ) + + +def flash_attn_func( + q, + k, + v, + softmax_scale=None, + causal=False, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, +): + """dropout_p should be set to 0.0 during evaluation + Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads + than Q. Note that the number of heads in Q must be divisible by the number of heads in KV. + For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head + 0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V. + + If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix. + For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is: + 1 1 1 1 0 + 1 1 1 1 1 + If seqlen_q = 5 and seqlen_k = 2, the causal mask is: + 0 0 + 0 0 + 0 0 + 1 0 + 1 1 + If the row of the mask is all zero, the output will be zero. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between + [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive. + + Arguments: + q: (batch_size, seqlen, nheads, headdim) + k: (batch_size, seqlen, nheads_k, headdim) + v: (batch_size, seqlen, nheads_k, headdim) + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of + (-alibi_slope * |i + seqlen_k - seqlen_q - j|) + is added to the attention score of query i and key j. + deterministic: bool. Whether to use the deterministic implementation of the backward pass, + which is slightly slower and uses more memory. The forward pass is always deterministic. + return_attn_probs: bool. Whether to return the attention probabilities. This option is for + testing only. The returned probabilities are not guaranteed to be correct + (they might not have the right scaling). + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + """ + return FlashAttnFunc.apply( + q, + k, + v, + softmax_scale, + causal, + qv, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + num_splits, + pack_gqa, + deterministic, + sm_margin, + ) + + +def flash_attn_varlen_func( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + seqused_q=None, + seqused_k=None, + softmax_scale=None, + causal=False, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, +): + return FlashAttnVarlenFunc.apply( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + qv, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + num_splits, + pack_gqa, + deterministic, + sm_margin, + ) + + +def flash_attn_combine(out_partial, lse_partial, out=None, out_dtype=None): + return flash_attn_3_cuda.fwd_combine(out_partial, lse_partial, out, out_dtype) + + +def flash_attn_with_kvcache( + q, + k_cache, + v_cache, + k=None, + v=None, + qv=None, + rotary_cos=None, + rotary_sin=None, + cache_seqlens: Optional[Union[(int, torch.Tensor)]] = None, + cache_batch_idx: Optional[torch.Tensor] = None, + cache_leftpad: Optional[torch.Tensor] = None, + page_table: Optional[torch.Tensor] = None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + rotary_seqlens: Optional[torch.Tensor] = None, + q_descale: Optional[torch.Tensor] = None, + k_descale: Optional[torch.Tensor] = None, + v_descale: Optional[torch.Tensor] = None, + softmax_scale=None, + causal=False, + window_size=(-1, -1), # -1 means infinite context window + attention_chunk=0, + softcap=0.0, # 0.0 means deactivated + rotary_interleaved=True, + scheduler_metadata=None, + num_splits=0, # Can be tuned for speed + pack_gqa=None, # Can be tuned for speed + sm_margin=0, # Can be tuned if some SMs are used for communication + return_softmax_lse=False, +): + """ + If k and v are not None, k_cache and v_cache will be updated *inplace* with the new values from + k and v. This is useful for incremental decoding: you can pass in the cached keys/values from + the previous step, and update them with the new keys/values from the current step, and do + attention with the updated cache, all in 1 kernel. + + If you pass in k / v, you must make sure that the cache is large enough to hold the new values. + For example, the KV cache could be pre-allocated with the max sequence length, and you can use + cache_seqlens to keep track of the current sequence lengths of each sequence in the batch. + + Also apply rotary embedding if rotary_cos and rotary_sin are passed in. The key @k will be + rotated by rotary_cos and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc. + If causal or local (i.e., window_size != (-1, -1)), the query @q will be rotated by rotary_cos + and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc. + If not causal and not local, the query @q will be rotated by rotary_cos and rotary_sin at + indices cache_seqlens only (i.e. we consider all tokens in @q to be at position cache_seqlens). + + See tests/test_flash_attn.py::test_flash_attn_kvcache for examples of how to use this function. + + Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads + than Q. Note that the number of heads in Q must be divisible by the number of heads in KV. + For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head + 0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V. + + If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix. + For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is: + 1 1 1 1 0 + 1 1 1 1 1 + If seqlen_q = 5 and seqlen_k = 2, the causal mask is: + 0 0 + 0 0 + 0 0 + 1 0 + 1 1 + If the row of the mask is all zero, the output will be zero. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between + [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive. + + Note: Does not support backward pass. + + Arguments: + q: (batch_size, seqlen, nheads, headdim) + k_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim) if there's no page_table, + or (num_blocks, page_block_size, nheads_k, headdim) if there's a page_table (i.e. paged KV cache) + page_block_size must be a multiple of 256. + v_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim_v) if there's no page_table, + or (num_blocks, page_block_size, nheads_k, headdim_v) if there's a page_table (i.e. paged KV cache) + k [optional]: (batch_size, seqlen_new, nheads_k, headdim). If not None, we concatenate + k with k_cache, starting at the indices specified by cache_seqlens. + v [optional]: (batch_size, seqlen_new, nheads_k, headdim_v). Similar to k. + qv [optional]: (batch_size, seqlen, nheads, headdim_v) + rotary_cos [optional]: (seqlen_ro, rotary_dim / 2). If not None, we apply rotary embedding + to k and q. Only applicable if k and v are passed in. rotary_dim must be divisible by 16. + rotary_sin [optional]: (seqlen_ro, rotary_dim / 2). Similar to rotary_cos. + cache_seqlens: int, or (batch_size,), dtype torch.int32. The sequence lengths of the + KV cache. + cache_batch_idx: (batch_size,), dtype torch.int32. The indices used to index into the KV cache. + If None, we assume that the batch indices are [0, 1, 2, ..., batch_size - 1]. + If the indices are not distinct, and k and v are provided, the values updated in the cache + might come from any of the duplicate indices. + cache_leftpad: (batch_size,), dtype torch.int32. The index that the KV cache starts. If None, assume 0. + page_table [optional]: (batch_size, max_num_blocks_per_seq), dtype torch.int32. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + softcap: float. Anything > 0 activates softcapping attention. + rotary_interleaved: bool. Only applicable if rotary_cos and rotary_sin are passed in. + If True, rotary embedding will combine dimensions 0 & 1, 2 & 3, etc. If False, + rotary embedding will combine dimensions 0 & rotary_dim / 2, 1 & rotary_dim / 2 + 1 + (i.e. GPT-NeoX style). + num_splits: int. If > 1, split the key/value into this many chunks along the sequence. + If num_splits == 1, we don't split the key/value. If num_splits == 0, we use a heuristic + to automatically determine the number of splits. + Don't change this unless you know what you are doing. + return_softmax_lse: bool. Whether to return the logsumexp of the attention scores. + + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_softmax_lse=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + """ + assert k_cache.stride(-1) == 1, "k_cache must have contiguous last dimension" + assert v_cache.stride(-1) == 1, "v_cache must have contiguous last dimension" + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + if cache_seqlens is not None and isinstance(cache_seqlens, int): + cache_seqlens = torch.full( + (k_cache.shape[0],), cache_seqlens, dtype=torch.int32, device=k_cache.device + ) + cache_seqlens = maybe_contiguous(cache_seqlens) + out, softmax_lse, *rest = _flash_attn_forward( + q, + k_cache, + v_cache, + k, + v, + qv, + None, # out + cu_seqlens_q, + None, # cu_seqlens_k + cu_seqlens_k_new, + None, # seqused_q + cache_seqlens, + max_seqlen_q, + None, # max_seqlen_k + page_table, + cache_batch_idx, + cache_leftpad, + rotary_cos, + rotary_sin, + rotary_seqlens, + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size=window_size, + attention_chunk=attention_chunk, + softcap=softcap, + rotary_interleaved=rotary_interleaved, + scheduler_metadata=scheduler_metadata, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # return (out, softmax_lse) if return_softmax_lse else out + return (out, softmax_lse, *rest) if return_softmax_lse else out + + +def get_scheduler_metadata( + batch_size, max_seqlen_q, max_seqlen_k, num_heads_q, num_heads_kv, headdim, + cache_seqlens: torch.Tensor, + qkv_dtype=torch.bfloat16, + headdim_v=None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + cache_leftpad: Optional[torch.Tensor] = None, + page_size: Optional[int] = None, + max_seqlen_k_new=0, + causal=False, + window_size=(-1, -1), # -1 means infinite context window + attention_chunk=0, + has_softcap=False, + num_splits=0, # Can be tuned for speed + pack_gqa=None, # Can be tuned for speed + sm_margin=0, # Can be tuned if some SMs are used for communication +): + cache_seqlens = maybe_contiguous(cache_seqlens) + if headdim_v is None: + headdim_v = headdim + scheduler_metadata = flash_attn_3_cuda.get_scheduler_metadata( + batch_size, max_seqlen_q, max_seqlen_k, num_heads_q, num_heads_kv, headdim, headdim_v, + qkv_dtype, + cache_seqlens, + cu_seqlens_q, + None, # cu_seqlens_k + cu_seqlens_k_new, + None, # seqused_q + cache_leftpad, + page_size, + max_seqlen_k_new, + causal, + window_size[0], window_size[1], + attention_chunk, + has_softcap, + num_splits, + pack_gqa, + sm_margin, + ) + return scheduler_metadata diff --git a/build/torch28-cxx11-cu129-x86_64-linux/__init__.py b/build/torch28-cxx11-cu129-x86_64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3d12edd447e84d6e65c4fa9ab91ae3c3192884c1 --- /dev/null +++ b/build/torch28-cxx11-cu129-x86_64-linux/__init__.py @@ -0,0 +1,17 @@ +from .flash_attn_interface import ( + flash_attn_combine, + flash_attn_func, + flash_attn_qkvpacked_func, + flash_attn_varlen_func, + flash_attn_with_kvcache, + get_scheduler_metadata, +) + +__all__ = [ + "flash_attn_combine", + "flash_attn_func", + "flash_attn_qkvpacked_func", + "flash_attn_varlen_func", + "flash_attn_with_kvcache", + "get_scheduler_metadata", +] diff --git a/build/torch28-cxx11-cu129-x86_64-linux/_flash_attn3_5630d32.abi3.so b/build/torch28-cxx11-cu129-x86_64-linux/_flash_attn3_5630d32.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..ae939f58ded43044c0403c41767b16dcb94e6ef5 --- /dev/null +++ b/build/torch28-cxx11-cu129-x86_64-linux/_flash_attn3_5630d32.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:14d427bde7b2d40ae4dca976af208c24343fb9a05a3aba5a284d098740274815 +size 804424896 diff --git a/build/torch28-cxx11-cu129-x86_64-linux/_ops.py b/build/torch28-cxx11-cu129-x86_64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..02dbf81201d16b6791a086b87a8b46dee7f45122 --- /dev/null +++ b/build/torch28-cxx11-cu129-x86_64-linux/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _flash_attn3_5630d32 +ops = torch.ops._flash_attn3_5630d32 + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_flash_attn3_5630d32::{op_name}" \ No newline at end of file diff --git a/build/torch28-cxx11-cu129-x86_64-linux/flash_attn3/__init__.py b/build/torch28-cxx11-cu129-x86_64-linux/flash_attn3/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..03dbc1afe1cf156661a2b1b22003cd5f599a0309 --- /dev/null +++ b/build/torch28-cxx11-cu129-x86_64-linux/flash_attn3/__init__.py @@ -0,0 +1,26 @@ +import ctypes +import sys + +import importlib +from pathlib import Path +from types import ModuleType + +def _import_from_path(file_path: Path) -> ModuleType: + # We cannot use the module name as-is, after adding it to `sys.modules`, + # it would also be used for other imports. So, we make a module name that + # depends on the path for it to be unique using the hex-encoded hash of + # the path. + path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value) + module_name = path_hash + spec = importlib.util.spec_from_file_location(module_name, file_path) + if spec is None: + raise ImportError(f"Cannot load spec for {module_name} from {file_path}") + module = importlib.util.module_from_spec(spec) + if module is None: + raise ImportError(f"Cannot load module {module_name} from spec") + sys.modules[module_name] = module + spec.loader.exec_module(module) # type: ignore + return module + + +globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py"))) diff --git a/build/torch28-cxx11-cu129-x86_64-linux/flash_attn_config.py b/build/torch28-cxx11-cu129-x86_64-linux/flash_attn_config.py new file mode 100644 index 0000000000000000000000000000000000000000..c4b7130324be2c4cabee7ada7000e170a2a81865 --- /dev/null +++ b/build/torch28-cxx11-cu129-x86_64-linux/flash_attn_config.py @@ -0,0 +1,7 @@ +# Auto-generated by flash attention 3 setup.py +CONFIG = {'build_flags': {'FLASHATTENTION_DISABLE_BACKWARD': False, 'FLASHATTENTION_DISABLE_SPLIT': False, 'FLASHATTENTION_DISABLE_PAGEDKV': False, 'FLASHATTENTION_DISABLE_APPENDKV': False, 'FLASHATTENTION_DISABLE_LOCAL': False, 'FLASHATTENTION_DISABLE_SOFTCAP': False, 'FLASHATTENTION_DISABLE_PACKGQA': False, 'FLASHATTENTION_DISABLE_FP16': False, 'FLASHATTENTION_DISABLE_FP8': False, 'FLASHATTENTION_DISABLE_VARLEN': False, 'FLASHATTENTION_DISABLE_CLUSTER': False, 'FLASHATTENTION_DISABLE_HDIM64': False, 'FLASHATTENTION_DISABLE_HDIM96': False, 'FLASHATTENTION_DISABLE_HDIM128': False, 'FLASHATTENTION_DISABLE_HDIM192': False, 'FLASHATTENTION_DISABLE_HDIM256': False, 'FLASHATTENTION_DISABLE_SM8x': False, 'FLASHATTENTION_ENABLE_VCOLMAJOR': False, 'FLASH_ATTENTION_DISABLE_HDIMDIFF64': False, 'FLASH_ATTENTION_DISABLE_HDIMDIFF192': False}} + +def show(): + from pprint import pprint + pprint(CONFIG) + diff --git a/build/torch28-cxx11-cu129-x86_64-linux/flash_attn_interface.py b/build/torch28-cxx11-cu129-x86_64-linux/flash_attn_interface.py new file mode 100644 index 0000000000000000000000000000000000000000..2e2a3cdd3e89cbc59bf6c956454e2776c305e0e5 --- /dev/null +++ b/build/torch28-cxx11-cu129-x86_64-linux/flash_attn_interface.py @@ -0,0 +1,1127 @@ +# Copyright (c) 2023, Tri Dao. + +from typing import Optional, Union, List, Tuple + +import torch +import torch.nn as nn + +from ._ops import ops as flash_attn_3_cuda +from ._ops import add_op_namespace_prefix + +def maybe_contiguous(x): + return x.contiguous() if x is not None and x.stride(-1) != 1 else x + + +def round_multiple(x, m): + return (x + m - 1) // m * m + + +def round_up_headdim(head_size: int) -> int: + from .flash_attn_config import CONFIG + + if not CONFIG["build_flags"]["FLASHATTENTION_DISABLE_HDIM64"]: + if head_size <= 64: + return 64 + if not CONFIG["build_flags"]["FLASHATTENTION_DISABLE_HDIM96"]: + if head_size <= 96: + return 96 + if not CONFIG["build_flags"]["FLASHATTENTION_DISABLE_HDIM128"]: + if head_size <= 128: + return 128 + if not CONFIG["build_flags"]["FLASHATTENTION_DISABLE_HDIM192"]: + if head_size <= 192: + return 192 + if not CONFIG["build_flags"]["FLASHATTENTION_DISABLE_HDIM256"]: + if head_size <= 256: + return 256 + return 256 + + +@torch.library.custom_op(add_op_namespace_prefix("_flash_attn_forward"), mutates_args=(), device_types="cuda") +def _flash_attn_forward( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_new: Optional[torch.Tensor] = None, + v_new: Optional[torch.Tensor] = None, + qv: Optional[torch.Tensor] = None, + out_: Optional[torch.Tensor] = None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + seqused_q: Optional[torch.Tensor] = None, + seqused_k: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_k: Optional[int] = None, + page_table: Optional[torch.Tensor] = None, + kv_batch_idx: Optional[torch.Tensor] = None, + leftpad_k: Optional[torch.Tensor] = None, + rotary_cos: Optional[torch.Tensor] = None, + rotary_sin: Optional[torch.Tensor] = None, + seqlens_rotary: Optional[torch.Tensor] = None, + q_descale: Optional[torch.Tensor] = None, + k_descale: Optional[torch.Tensor] = None, + v_descale: Optional[torch.Tensor] = None, + softmax_scale: Optional[float] = None, + causal: bool = False, + window_size_left: int = -1, + window_size_right: int = -1, + attention_chunk: int = 0, + softcap: float = 0.0, + rotary_interleaved: bool = True, + scheduler_metadata: Optional[torch.Tensor] = None, + num_splits: int = 1, + pack_gqa: Optional[bool] = None, + sm_margin: int = 0, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + q, k, k_new, v_new = [maybe_contiguous(x) for x in (q, k, k_new, v_new)] + v = v.contiguous() if v.stride(-1) != 1 and v.stride(-3) != 1 else v + cu_seqlens_q, cu_seqlens_k, cu_seqlens_k_new = [ + maybe_contiguous(x) for x in (cu_seqlens_q, cu_seqlens_k, cu_seqlens_k_new) + ] + seqused_q, seqused_k = [maybe_contiguous(x) for x in (seqused_q, seqused_k)] + page_table, kv_batch_idx, leftpad_k = [ + maybe_contiguous(x) for x in (page_table, kv_batch_idx, leftpad_k) + ] + rotary_cos, rotary_sin = [maybe_contiguous(x) for x in (rotary_cos, rotary_sin)] + seqlens_rotary = maybe_contiguous(seqlens_rotary) + out, softmax_lse, out_accum, softmax_lse_accum = flash_attn_3_cuda.fwd( + q, + k, + v, + k_new, + v_new, + qv, + out_, + cu_seqlens_q, + cu_seqlens_k, + cu_seqlens_k_new, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + page_table, + kv_batch_idx, + leftpad_k, + rotary_cos, + rotary_sin, + seqlens_rotary, + q_descale, + k_descale, + v_descale, + softmax_scale, + causal, + window_size_left, + window_size_right, + attention_chunk, + softcap, + rotary_interleaved, + scheduler_metadata, + num_splits, + pack_gqa, + sm_margin, + ) + + if out_accum is None: + out_accum = torch.tensor([], device=out.device) + + if softmax_lse_accum is None: + softmax_lse_accum = torch.tensor([], device=out.device) + + return out, softmax_lse, out_accum, softmax_lse_accum + + +@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_forward")) +def _flash_attn_forward_fake( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_new: Optional[torch.Tensor] = None, + v_new: Optional[torch.Tensor] = None, + qv: Optional[torch.Tensor] = None, + out_: Optional[torch.Tensor] = None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + seqused_q: Optional[torch.Tensor] = None, + seqused_k: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_k: Optional[int] = None, + page_table: Optional[torch.Tensor] = None, + kv_batch_idx: Optional[torch.Tensor] = None, + leftpad_k: Optional[torch.Tensor] = None, + rotary_cos: Optional[torch.Tensor] = None, + rotary_sin: Optional[torch.Tensor] = None, + seqlens_rotary: Optional[torch.Tensor] = None, + q_descale: Optional[torch.Tensor] = None, + k_descale: Optional[torch.Tensor] = None, + v_descale: Optional[torch.Tensor] = None, + softmax_scale: Optional[float] = None, + causal: bool = False, + window_size_left: int = -1, + window_size_right: int = -1, + attention_chunk: int = 0, + softcap: float = 0.0, + rotary_interleaved: bool = True, + scheduler_metadata: Optional[torch.Tensor] = None, + num_splits: int = 1, + pack_gqa: Optional[bool] = None, + sm_margin: int = 0, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Symbolic fake implementation of flash attention forward. + Returns tensors with the correct shapes and dtypes without actual computation. + """ + + # Determine if we're in varlen mode + is_varlen_q = cu_seqlens_q is not None + + # Get dimensions from query tensor + if is_varlen_q: + # varlen mode: q is (total_q, num_heads, head_size) + total_q, num_heads, head_size = q.shape + batch_size = cu_seqlens_q.shape[0] - 1 + + if max_seqlen_q is None: + raise ValueError("max_seqlen_q must be provided if cu_seqlens_q is provided") + seqlen_q = max_seqlen_q + else: + # batch mode: q is (batch_size, seqlen_q, num_heads, head_size) + batch_size, seqlen_q, num_heads, head_size = q.shape + total_q = batch_size * q.shape[1] + # Get value head dimension + head_size_v = v.shape[-1] + + # Determine output dtype (FP8 inputs produce BF16 outputs) + q_type = q.dtype + if q_type == torch.float8_e4m3fn: + out_dtype = torch.bfloat16 + else: + out_dtype = q_type + + # Create output tensor + if out_ is not None: + # If out_ is provided, _flash_attn_forward becomes non-functional + raise TypeError("Tracing (torch.compile/torch.export) with pre-allocated output tensor is not supported.") + + if is_varlen_q: + out = torch.empty((total_q, num_heads, head_size_v), dtype=out_dtype, device=q.device) + else: + out = torch.empty((batch_size, seqlen_q, num_heads, head_size_v), dtype=out_dtype, device=q.device) + + # Create softmax_lse tensor + if is_varlen_q: + softmax_lse = torch.empty((num_heads, total_q), dtype=torch.float32, device=q.device) + else: + softmax_lse = torch.empty((batch_size, num_heads, seqlen_q), dtype=torch.float32, device=q.device) + + # TODO(guilhermeleobas): Implement "get_num_splits" + # There's an heuristic to compute num_splits when "num_splits <= 0" + # assert that num_splits is > 0 for now + if num_splits <= 0: + raise ValueError(f"tracing (torch.compile/torch.export) with num_splits <= 0 not supported. Got {num_splits=}") + + if num_splits > 1: + if is_varlen_q: + out_accum = torch.empty((num_splits, num_heads, total_q, head_size_v), dtype=torch.float32, device=q.device) + softmax_lse_accum = torch.empty((num_splits, num_heads, total_q), dtype=torch.float32, device=q.device) + else: + out_accum = torch.empty((num_splits, batch_size, num_heads, seqlen_q, head_size_v), dtype=torch.float32, device=q.device) + softmax_lse_accum = torch.empty((num_splits, batch_size, num_heads, seqlen_q), dtype=torch.float32, device=q.device) + else: + # Tensors are not set when num_splits < 1 + out_accum = torch.tensor([], device=out.device) + softmax_lse_accum = torch.tensor([], device=out.device) + + return out, softmax_lse, out_accum, softmax_lse_accum + + +@torch.library.custom_op(add_op_namespace_prefix("_flash_attn_backward"), mutates_args=("dq", "dk", "dv"), device_types="cuda") +def _flash_attn_backward( + dout: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + out: torch.Tensor, + softmax_lse: torch.Tensor, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k: Optional[torch.Tensor] = None, + sequed_q: Optional[torch.Tensor] = None, + sequed_k: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_k: Optional[int] = None, + dq: Optional[torch.Tensor] = None, + dk: Optional[torch.Tensor] = None, + dv: Optional[torch.Tensor] = None, + softmax_scale: Optional[float] = None, + is_causal: bool = False, + window_size_left: int = -1, + window_size_right: int = -1, + softcap: float = 0.0, + deterministic: bool = False, + sm_margin: int = 0, +) -> torch.Tensor: + # dq, dk, dv are allocated by us so they should already be contiguous + dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] + softmax_d, *rest = flash_attn_3_cuda.bwd( + dout, + q, + k, + v, + out, + softmax_lse, + dq, + dk, + dv, + cu_seqlens_q, + cu_seqlens_k, + sequed_q, + sequed_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + is_causal, + window_size_left, + window_size_right, + softcap, + deterministic, + sm_margin, + ) + return softmax_d + + +@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_backward")) +def _flash_attn_backward_fake( + dout: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + out: torch.Tensor, + softmax_lse: torch.Tensor, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k: Optional[torch.Tensor] = None, + sequed_q: Optional[torch.Tensor] = None, + sequed_k: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_k: Optional[int] = None, + dq: Optional[torch.Tensor] = None, + dk: Optional[torch.Tensor] = None, + dv: Optional[torch.Tensor] = None, + softmax_scale: Optional[float] = None, + is_causal: bool = False, + window_size_left: int = -1, + window_size_right: int = -1, + softcap: float = 0.0, + deterministic: bool = False, + sm_margin: int = 0, +) -> torch.Tensor: + + is_varlen_q = cu_seqlens_q is not None + is_varlen_k = cu_seqlens_q is not None + is_varlen = is_varlen_q or is_varlen_k or sequed_q is not None or sequed_k is not None + + if not is_varlen_q: + batch_size = q.size(0) + seqlen_q = q.size(1) + seqlen_k = k.size(1) + total_q = batch_size * q.size(1) + else: + batch_size = cu_seqlens_q.size(0) - 1 + total_q = q.size(0) + seqlen_q = max_seqlen_q + seqlen_k = max_seqlen_k + + if window_size_left >= seqlen_k - 1: + window_size_left = -1 + + if window_size_right >= seqlen_q - 1: + window_size_right = -1 + + if is_causal: + window_size_right = 0 + + is_causal = window_size_left < 0 and window_size_right == 0 + + head_size = q.size(-1) + head_size_v = v.size(-1) + head_size_rounded = round_up_headdim(max(head_size, head_size_v)) + + # Hopper gpus uses cuda compute capabilities 9.0 + cap = torch.cuda.get_device_capability(q.device) + arch = cap[0] * 10 + cap[1] + + is_local = (window_size_left >= 0 or window_size_right >= 0) and not is_causal + + if head_size_rounded <= 64: + kBlockM_sm90 = 96 if (is_causal and softcap > 0.0) else 128 + elif head_size_rounded <= 96: + kBlockM_sm90 = 64 + elif head_size_rounded <= 128: + kBlockM_sm90 = 64 if (is_causal or is_local or softcap > 0.0) else 80 + else: + kBlockM_sm90 = 64 + + kBlockM_sm80 = 128 if head_size_rounded <= 64 else 64 + kBlockM_sm86 = 64 if head_size_rounded <= 192 else 32 + + if arch >= 90: + kBlockM = kBlockM_sm90 + elif arch == 86 or arch == 89: + kBlockM = kBlockM_sm86 + else: + kBlockM = kBlockM_sm80 + + num_heads = q.shape[-2] + seqlen_q_rounded = round_multiple(seqlen_q, kBlockM) + + total_q_padded_rounded = round_multiple(total_q + batch_size * kBlockM, kBlockM) + + dq = torch.empty_like(q) if dq is None else dq + dk = torch.empty_like(k) if dk is None else dk + dv = torch.empty_like(v) if dv is None else dv + + if not is_varlen: + softmax_d = torch.empty((batch_size, num_heads, seqlen_q_rounded), dtype=torch.float32, device=q.device) + else: + softmax_d = torch.empty((num_heads, total_q_padded_rounded), dtype=torch.float32, device=q.device) + + return softmax_d + + +def setup_context(ctx, inputs, output): + q, k, v = inputs[:3] + out, softmax_lse, _, _ = output + ctx.save_for_backward(q, k, v, out, softmax_lse) + ctx.softmax_scale = inputs[-11] + ctx.causal = inputs[-10] + ctx.window_size = [inputs[-9], inputs[-8]] + ctx.attention_chunk = inputs[-7] + ctx.softcap = inputs[-6] + ctx.sm_margin = inputs[-1] + + +def _backward(ctx, dout, *grads): + q, k, v, out, softmax_lse = ctx.saved_tensors + dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + None, None, # cu_seqlens_q, cu_seqlens_k, + None, None, # sequed_q, sequed_k, + None, None, # max_seqlen_q, max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size[0], + ctx.window_size[1], + ctx.softcap, + False, # deterministic + ctx.sm_margin, + ) + return dq, dk, dv, *((None,) * 21) + + +_flash_attn_forward.register_autograd(_backward, setup_context=setup_context) + + + +class FlashAttnQKVPackedFunc(torch.autograd.Function): + @staticmethod + def forward( + ctx, + qkv, + softmax_scale, + causal, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + deterministic=False, + num_heads_q=None, + sm_margin=0, + return_softmax=False, + ): + if softmax_scale is None: + softmax_scale = qkv.shape[-1] ** (-0.5) + if qkv.dim() == 5: + assert qkv.shape[-3] == 3 + q, k, v = qkv.unbind(dim=-3) + else: + assert qkv.dim() == 4 + assert num_heads_q is not None + num_heads_k = (qkv.shape[2] - num_heads_q) // 2 + assert num_heads_k * 2 + num_heads_q == qkv.shape[2] + q, k, v = qkv.split([num_heads_q, num_heads_k, num_heads_k], dim=-2) + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + None, # qv + None, # out + None, None, None, # cu_seqlens_q/k/k_new + None, None, # seqused_q/k + None, None, # max_seqlen_q/k + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size_left=window_size[0], + window_size_right=window_size[1], + attention_chunk=attention_chunk, + softcap=softcap, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse) + ctx.save_for_backward(q, k, v, out, softmax_lse) + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.ndim = qkv.dim() + ctx.sm_margin = sm_margin + return (out, softmax_lse) if return_softmax else out + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + if ctx.ndim == 5: + qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) + dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) + dq, dk, dv = dqkv.unbind(dim=-3) + else: + num_heads_q = q.shape[2] + num_heads_k = k.shape[2] + qkv_shape = q.shape[:-2] + (num_heads_q + num_heads_k * 2, *q.shape[-1:]) + dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) + dq, dk, dv = dqkv.split([num_heads_q, num_heads_k, num_heads_k], dim=-2) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + None, None, # cu_seqlens_q, cu_seqlens_k, + None, None, # sequed_q, sequed_k, + None, None, # max_seqlen_q, max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size[0], + ctx.window_size[1], + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dqkv = dqkv[..., : dout.shape[-1]] # We could have padded the head dimension + return dqkv, None, None, None, None, None, None, None, None, None, None, None, None + + +class FlashAttnFunc(torch.autograd.Function): + + @staticmethod + def forward( + ctx, + q, + k, + v, + softmax_scale, + causal, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + return_softmax=False, + ): + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + # out, q, k, v, out_padded, softmax_lse = _flash_attn_forward( + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + qv, # qv + None, # out + None, None, None, # cu_seqlens_q/k/k_new + None, None, # seqused_q/k + None, None, # max_seqlen_q/k + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size_left=window_size[0], + window_size_right=window_size[1], + attention_chunk=attention_chunk, + softcap=softcap, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse) + ctx.save_for_backward(q, k, v, out, softmax_lse) + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.sm_margin = sm_margin + return (out, softmax_lse) if return_softmax else out + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + None, None, # cu_seqlens_q, cu_seqlens_k, + None, None, # sequed_q, sequed_k, + None, None, # max_seqlen_q, max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size[0], + ctx.window_size[1], + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dq = dq[..., : q.shape[-1]] # We could have padded the head dimension + dk = dk[..., : k.shape[-1]] + dv = dv[..., : v.shape[-1]] + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None + + +class FlashAttnVarlenFunc(torch.autograd.Function): + + @staticmethod + def forward( + ctx, + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + return_softmax=False, + ): + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + # out, q, k, v, out_padded, softmax_lse = _flash_attn_varlen_forward( + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + qv, # qv + None, # out + cu_seqlens_q, + cu_seqlens_k, + None, # cu_seqlens_k_new + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size_left=window_size[0], + window_size_right=window_size[1], + attention_chunk=attention_chunk, + softcap=softcap, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k) + ctx.save_for_backward(q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k) + ctx.max_seqlen_q = max_seqlen_q + ctx.max_seqlen_k = max_seqlen_k + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.sm_margin = sm_margin + return (out, softmax_lse) if return_softmax else out + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + ctx.max_seqlen_q, + ctx.max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size[0], + ctx.window_size[1], + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dq = dq[..., : q.shape[-1]] # We could have padded the head dimension + dk = dk[..., : k.shape[-1]] + dv = dv[..., : v.shape[-1]] + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None + + +def flash_attn_qkvpacked_func( + qkv, + softmax_scale=None, + causal=False, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + deterministic=False, + num_heads_q=None, + sm_margin=0, + return_attn_probs=False, +): + """dropout_p should be set to 0.0 during evaluation + If Q, K, V are already stacked into 1 tensor, this function will be faster than + calling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation + of the gradients of Q, K, V. + For multi-query and grouped-query attention (MQA/GQA), please see + flash_attn_kvpacked_func and flash_attn_func. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between [i - window_size[0], i + window_size[1]] inclusive. + + Arguments: + qkv: (batch_size, seqlen, 3, nheads, headdim) + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + softcap: float. Anything > 0 activates softcapping attention. + alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of (-alibi_slope * |i - j|) is added to + the attention score of query i and key j. + deterministic: bool. Whether to use the deterministic implementation of the backward pass, + which is slightly slower and uses more memory. The forward pass is always deterministic. + return_attn_probs: bool. Whether to return the attention probabilities. This option is for + testing only. The returned probabilities are not guaranteed to be correct + (they might not have the right scaling). + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen). + The output of softmax (possibly with different scaling). It also encodes the dropout + pattern (negative means that location was dropped, nonnegative means it was kept). + """ + return FlashAttnQKVPackedFunc.apply( + qkv, + softmax_scale, + causal, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + deterministic, + num_heads_q, + sm_margin, + return_attn_probs, + ) + + +def flash_attn_func( + q, + k, + v, + softmax_scale=None, + causal=False, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + return_attn_probs=False, +): + """dropout_p should be set to 0.0 during evaluation + Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads + than Q. Note that the number of heads in Q must be divisible by the number of heads in KV. + For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head + 0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V. + + If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix. + For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is: + 1 1 1 1 0 + 1 1 1 1 1 + If seqlen_q = 5 and seqlen_k = 2, the causal mask is: + 0 0 + 0 0 + 0 0 + 1 0 + 1 1 + If the row of the mask is all zero, the output will be zero. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between + [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive. + + Arguments: + q: (batch_size, seqlen, nheads, headdim) + k: (batch_size, seqlen, nheads_k, headdim) + v: (batch_size, seqlen, nheads_k, headdim) + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of + (-alibi_slope * |i + seqlen_k - seqlen_q - j|) + is added to the attention score of query i and key j. + deterministic: bool. Whether to use the deterministic implementation of the backward pass, + which is slightly slower and uses more memory. The forward pass is always deterministic. + return_attn_probs: bool. Whether to return the attention probabilities. This option is for + testing only. The returned probabilities are not guaranteed to be correct + (they might not have the right scaling). + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + """ + return FlashAttnFunc.apply( + q, + k, + v, + softmax_scale, + causal, + qv, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + num_splits, + pack_gqa, + deterministic, + sm_margin, + return_attn_probs, + ) + + +def flash_attn_varlen_func( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + seqused_q=None, + seqused_k=None, + softmax_scale=None, + causal=False, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + return_attn_probs=False, +): + return FlashAttnVarlenFunc.apply( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + qv, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + num_splits, + pack_gqa, + deterministic, + sm_margin, + return_attn_probs, + ) + + +def flash_attn_combine(out_partial, lse_partial, out=None, out_dtype=None): + return flash_attn_3_cuda.fwd_combine(out_partial, lse_partial, out, out_dtype) + + +def flash_attn_with_kvcache( + q, + k_cache, + v_cache, + k=None, + v=None, + qv=None, + rotary_cos=None, + rotary_sin=None, + cache_seqlens: Optional[Union[(int, torch.Tensor)]] = None, + cache_batch_idx: Optional[torch.Tensor] = None, + cache_leftpad: Optional[torch.Tensor] = None, + page_table: Optional[torch.Tensor] = None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + rotary_seqlens: Optional[torch.Tensor] = None, + q_descale: Optional[torch.Tensor] = None, + k_descale: Optional[torch.Tensor] = None, + v_descale: Optional[torch.Tensor] = None, + softmax_scale=None, + causal=False, + window_size=(-1, -1), # -1 means infinite context window + attention_chunk=0, + softcap=0.0, # 0.0 means deactivated + rotary_interleaved=True, + scheduler_metadata=None, + num_splits=0, # Can be tuned for speed + pack_gqa=None, # Can be tuned for speed + sm_margin=0, # Can be tuned if some SMs are used for communication + return_softmax_lse=False, +): + """ + If k and v are not None, k_cache and v_cache will be updated *inplace* with the new values from + k and v. This is useful for incremental decoding: you can pass in the cached keys/values from + the previous step, and update them with the new keys/values from the current step, and do + attention with the updated cache, all in 1 kernel. + + If you pass in k / v, you must make sure that the cache is large enough to hold the new values. + For example, the KV cache could be pre-allocated with the max sequence length, and you can use + cache_seqlens to keep track of the current sequence lengths of each sequence in the batch. + + Also apply rotary embedding if rotary_cos and rotary_sin are passed in. The key @k will be + rotated by rotary_cos and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc. + If causal or local (i.e., window_size != (-1, -1)), the query @q will be rotated by rotary_cos + and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc. + If not causal and not local, the query @q will be rotated by rotary_cos and rotary_sin at + indices cache_seqlens only (i.e. we consider all tokens in @q to be at position cache_seqlens). + + See tests/test_flash_attn.py::test_flash_attn_kvcache for examples of how to use this function. + + Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads + than Q. Note that the number of heads in Q must be divisible by the number of heads in KV. + For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head + 0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V. + + If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix. + For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is: + 1 1 1 1 0 + 1 1 1 1 1 + If seqlen_q = 5 and seqlen_k = 2, the causal mask is: + 0 0 + 0 0 + 0 0 + 1 0 + 1 1 + If the row of the mask is all zero, the output will be zero. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between + [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive. + + Note: Does not support backward pass. + + Arguments: + q: (batch_size, seqlen, nheads, headdim) + k_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim) if there's no page_table, + or (num_blocks, page_block_size, nheads_k, headdim) if there's a page_table (i.e. paged KV cache) + page_block_size can be arbitrary (e.g, 1, 2, 3, 64, etc.). + v_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim_v) if there's no page_table, + or (num_blocks, page_block_size, nheads_k, headdim_v) if there's a page_table (i.e. paged KV cache) + k [optional]: (batch_size, seqlen_new, nheads_k, headdim). If not None, we concatenate + k with k_cache, starting at the indices specified by cache_seqlens. + v [optional]: (batch_size, seqlen_new, nheads_k, headdim_v). Similar to k. + qv [optional]: (batch_size, seqlen, nheads, headdim_v) + rotary_cos [optional]: (seqlen_ro, rotary_dim / 2). If not None, we apply rotary embedding + to k and q. Only applicable if k and v are passed in. rotary_dim must be divisible by 16. + rotary_sin [optional]: (seqlen_ro, rotary_dim / 2). Similar to rotary_cos. + cache_seqlens: int, or (batch_size,), dtype torch.int32. The sequence lengths of the + KV cache. + cache_batch_idx: (batch_size,), dtype torch.int32. The indices used to index into the KV cache. + If None, we assume that the batch indices are [0, 1, 2, ..., batch_size - 1]. + If the indices are not distinct, and k and v are provided, the values updated in the cache + might come from any of the duplicate indices. + cache_leftpad: (batch_size,), dtype torch.int32. The index that the KV cache starts. If None, assume 0. + page_table [optional]: (batch_size, max_num_blocks_per_seq), dtype torch.int32. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + softcap: float. Anything > 0 activates softcapping attention. + rotary_interleaved: bool. Only applicable if rotary_cos and rotary_sin are passed in. + If True, rotary embedding will combine dimensions 0 & 1, 2 & 3, etc. If False, + rotary embedding will combine dimensions 0 & rotary_dim / 2, 1 & rotary_dim / 2 + 1 + (i.e. GPT-NeoX style). + num_splits: int. If > 1, split the key/value into this many chunks along the sequence. + If num_splits == 1, we don't split the key/value. If num_splits == 0, we use a heuristic + to automatically determine the number of splits. + Don't change this unless you know what you are doing. + return_softmax_lse: bool. Whether to return the logsumexp of the attention scores. + + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_softmax_lse=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + """ + assert k_cache.stride(-1) == 1, "k_cache must have contiguous last dimension" + assert v_cache.stride(-1) == 1, "v_cache must have contiguous last dimension" + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + if cache_seqlens is not None and isinstance(cache_seqlens, int): + cache_seqlens = torch.full( + (q.shape[0],), cache_seqlens, dtype=torch.int32, device=k_cache.device + ) + cache_seqlens = maybe_contiguous(cache_seqlens) + out, softmax_lse, *rest = _flash_attn_forward( + q, + k_cache, + v_cache, + k, + v, + qv, + None, # out + cu_seqlens_q, + None, # cu_seqlens_k + cu_seqlens_k_new, + None, # seqused_q + cache_seqlens, + max_seqlen_q, + None, # max_seqlen_k + page_table, + cache_batch_idx, + cache_leftpad, + rotary_cos, + rotary_sin, + rotary_seqlens, + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size_left=window_size[0], + window_size_right=window_size[1], + attention_chunk=attention_chunk, + softcap=softcap, + rotary_interleaved=rotary_interleaved, + scheduler_metadata=scheduler_metadata, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # return (out, softmax_lse) if return_softmax_lse else out + return (out, softmax_lse, *rest) if return_softmax_lse else out + + +def get_scheduler_metadata( + batch_size, max_seqlen_q, max_seqlen_k, num_heads_q, num_heads_kv, headdim, + cache_seqlens: torch.Tensor, + qkv_dtype=torch.bfloat16, + headdim_v=None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + cache_leftpad: Optional[torch.Tensor] = None, + page_size: Optional[int] = None, + max_seqlen_k_new=0, + causal=False, + window_size=(-1, -1), # -1 means infinite context window + attention_chunk=0, + has_softcap=False, + num_splits=0, # Can be tuned for speed + pack_gqa=None, # Can be tuned for speed + sm_margin=0, # Can be tuned if some SMs are used for communication +): + cache_seqlens = maybe_contiguous(cache_seqlens) + if headdim_v is None: + headdim_v = headdim + scheduler_metadata = flash_attn_3_cuda.get_scheduler_metadata( + batch_size, max_seqlen_q, max_seqlen_k, num_heads_q, num_heads_kv, headdim, headdim_v, + qkv_dtype, + cache_seqlens, + cu_seqlens_q, + None, # cu_seqlens_k + cu_seqlens_k_new, + None, # seqused_q + cache_leftpad, + page_size, + max_seqlen_k_new, + causal, + window_size[0], window_size[1], + attention_chunk, + has_softcap, + num_splits, + pack_gqa, + sm_margin, + ) + return scheduler_metadata diff --git a/build/torch28-cxx11-cu129-x86_64-linux/metadata.json b/build/torch28-cxx11-cu129-x86_64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..9cf5deed9898dce769f4cc73913d3530b92a0bd8 --- /dev/null +++ b/build/torch28-cxx11-cu129-x86_64-linux/metadata.json @@ -0,0 +1,4 @@ +{ + "version": 1, + "python-depends": [] +} \ No newline at end of file diff --git a/build/torch29-cxx11-cu126-aarch64-linux/flash_attn3/__init__.py b/build/torch29-cxx11-cu126-aarch64-linux/flash_attn3/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3d12edd447e84d6e65c4fa9ab91ae3c3192884c1 --- /dev/null +++ b/build/torch29-cxx11-cu126-aarch64-linux/flash_attn3/__init__.py @@ -0,0 +1,17 @@ +from .flash_attn_interface import ( + flash_attn_combine, + flash_attn_func, + flash_attn_qkvpacked_func, + flash_attn_varlen_func, + flash_attn_with_kvcache, + get_scheduler_metadata, +) + +__all__ = [ + "flash_attn_combine", + "flash_attn_func", + "flash_attn_qkvpacked_func", + "flash_attn_varlen_func", + "flash_attn_with_kvcache", + "get_scheduler_metadata", +] diff --git a/build/torch29-cxx11-cu126-aarch64-linux/flash_attn3/__pycache__/__init__.cpython-313.pyc b/build/torch29-cxx11-cu126-aarch64-linux/flash_attn3/__pycache__/__init__.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..857fa9ac00bb1e42aa50ce5e49b975e27147ecc5 Binary files /dev/null and b/build/torch29-cxx11-cu126-aarch64-linux/flash_attn3/__pycache__/__init__.cpython-313.pyc differ diff --git a/build/torch29-cxx11-cu126-aarch64-linux/flash_attn3/__pycache__/_ops.cpython-313.pyc b/build/torch29-cxx11-cu126-aarch64-linux/flash_attn3/__pycache__/_ops.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a1b03ee6f27dc2017aa2aade427aeab18dbf8486 Binary files /dev/null and b/build/torch29-cxx11-cu126-aarch64-linux/flash_attn3/__pycache__/_ops.cpython-313.pyc differ diff --git a/build/torch29-cxx11-cu126-aarch64-linux/flash_attn3/__pycache__/flash_attn_interface.cpython-313.pyc b/build/torch29-cxx11-cu126-aarch64-linux/flash_attn3/__pycache__/flash_attn_interface.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..485b2c2591cb28b809904f1769de106ca9e97481 Binary files /dev/null and b/build/torch29-cxx11-cu126-aarch64-linux/flash_attn3/__pycache__/flash_attn_interface.cpython-313.pyc differ diff --git a/build/torch29-cxx11-cu126-aarch64-linux/flash_attn3/_flash_attn3_847092b_dirty.abi3.so b/build/torch29-cxx11-cu126-aarch64-linux/flash_attn3/_flash_attn3_847092b_dirty.abi3.so new file mode 100755 index 0000000000000000000000000000000000000000..dec7879b4b5bca58ed4c772d780c368462632625 --- /dev/null +++ b/build/torch29-cxx11-cu126-aarch64-linux/flash_attn3/_flash_attn3_847092b_dirty.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:17179eb1daba5483276f8536733febb2623ef14c002b5315859f7eed3f73fa81 +size 822395648 diff --git a/build/torch29-cxx11-cu126-aarch64-linux/flash_attn3/_ops.py b/build/torch29-cxx11-cu126-aarch64-linux/flash_attn3/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..b7a3a5f4470eb321053647e1d601c8448d21490a --- /dev/null +++ b/build/torch29-cxx11-cu126-aarch64-linux/flash_attn3/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _flash_attn3_847092b_dirty +ops = torch.ops._flash_attn3_847092b_dirty + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_flash_attn3_847092b_dirty::{op_name}" \ No newline at end of file diff --git a/build/torch29-cxx11-cu126-aarch64-linux/flash_attn3/flash_attn_interface.py b/build/torch29-cxx11-cu126-aarch64-linux/flash_attn3/flash_attn_interface.py new file mode 100644 index 0000000000000000000000000000000000000000..7572297784a73fdf891af6aafee88428120b72f2 --- /dev/null +++ b/build/torch29-cxx11-cu126-aarch64-linux/flash_attn3/flash_attn_interface.py @@ -0,0 +1,828 @@ +# Copyright (c) 2023, Tri Dao. + +from typing import Optional, Union + +import torch +import torch.nn as nn + +from ._ops import ops as flash_attn_3_cuda + +def maybe_contiguous(x): + return x.contiguous() if x is not None and x.stride(-1) != 1 else x + + +def _flash_attn_forward( + q, + k, + v, + k_new, + v_new, + qv, + out, + cu_seqlens_q, + cu_seqlens_k, + cu_seqlens_k_new, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + page_table, + kv_batch_idx, + leftpad_k, + rotary_cos, + rotary_sin, + seqlens_rotary, + q_descale, + k_descale, + v_descale, + softmax_scale, + causal, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + rotary_interleaved=True, + scheduler_metadata=None, + num_splits=1, + pack_gqa=None, + sm_margin=0): + q, k, k_new, v_new = [maybe_contiguous(x) for x in (q, k, k_new, v_new)] + v = v.contiguous() if v.stride(-1) != 1 and v.stride(-3) != 1 else v + cu_seqlens_q, cu_seqlens_k, cu_seqlens_k_new = [ + maybe_contiguous(x) for x in (cu_seqlens_q, cu_seqlens_k, cu_seqlens_k_new) + ] + seqused_q, seqused_k = [maybe_contiguous(x) for x in (seqused_q, seqused_k)] + page_table, kv_batch_idx, leftpad_k = [ + maybe_contiguous(x) for x in (page_table, kv_batch_idx, leftpad_k) + ] + rotary_cos, rotary_sin = [maybe_contiguous(x) for x in (rotary_cos, rotary_sin)] + seqlens_rotary = maybe_contiguous(seqlens_rotary) + out, softmax_lse, *rest = flash_attn_3_cuda.fwd( + q, + k, + v, + k_new, + v_new, + qv, + out, + cu_seqlens_q, + cu_seqlens_k, + cu_seqlens_k_new, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + page_table, + kv_batch_idx, + leftpad_k, + rotary_cos, + rotary_sin, + seqlens_rotary, + q_descale, + k_descale, + v_descale, + softmax_scale, + causal, + window_size[0], + window_size[1], + attention_chunk, + softcap, + rotary_interleaved, + scheduler_metadata, + num_splits, + pack_gqa, + sm_margin, + ) + return out, softmax_lse, *rest + + +def _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + cu_seqlens_q, + cu_seqlens_k, + sequed_q, + sequed_k, + max_seqlen_q, + max_seqlen_k, + dq, + dk, + dv, + softmax_scale, + causal, + window_size=(-1, -1), + softcap=0.0, + deterministic=False, + sm_margin=0, +): + # dq, dk, dv are allocated by us so they should already be contiguous + dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] + dq, dk, dv, softmax_d, *rest = flash_attn_3_cuda.bwd( + dout, + q, + k, + v, + out, + softmax_lse, + dq, + dk, + dv, + cu_seqlens_q, + cu_seqlens_k, + sequed_q, + sequed_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + window_size[0], + window_size[1], + softcap, + deterministic, + sm_margin, + ) + return dq, dk, dv, softmax_d + + +class FlashAttnQKVPackedFunc(torch.autograd.Function): + @staticmethod + def forward( + ctx, + qkv, + softmax_scale, + causal, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + deterministic=False, + num_heads_q=None, + sm_margin=0, + ): + if softmax_scale is None: + softmax_scale = qkv.shape[-1] ** (-0.5) + if qkv.dim() == 5: + assert qkv.shape[-3] == 3 + q, k, v = qkv.unbind(dim=-3) + else: + assert qkv.dim() == 4 + assert num_heads_q is not None + num_heads_k = (qkv.shape[2] - num_heads_q) // 2 + assert num_heads_k * 2 + num_heads_q == qkv.shape[2] + q, k, v = qkv.split([num_heads_q, num_heads_k, num_heads_k], dim=-2) + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + None, # qv + None, # out + None, None, None, # cu_seqlens_q/k/k_new + None, None, # seqused_q/k + None, None, # max_seqlen_q/k + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size=window_size, + attention_chunk=attention_chunk, + softcap=softcap, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse) + ctx.save_for_backward(q, k, v, out, softmax_lse) + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.ndim = qkv.dim() + ctx.sm_margin = sm_margin + # return out, softmax_lse + return out + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + if ctx.ndim == 5: + qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) + dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) + dq, dk, dv = dqkv.unbind(dim=-3) + else: + num_heads_q = q.shape[2] + num_heads_k = k.shape[2] + qkv_shape = q.shape[:-2] + (num_heads_q + num_heads_k * 2, *q.shape[-1:]) + dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) + dq, dk, dv = dqkv.split([num_heads_q, num_heads_k, num_heads_k], dim=-2) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + None, None, # cu_seqlens_q, cu_seqlens_k, + None, None, # sequed_q, sequed_k, + None, None, # max_seqlen_q, max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size, + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dqkv = dqkv[..., : dout.shape[-1]] # We could have padded the head dimension + return dqkv, None, None, None, None, None, None, None, None, None, None, None + + +class FlashAttnFunc(torch.autograd.Function): + + @staticmethod + def forward( + ctx, + q, + k, + v, + softmax_scale, + causal, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + ): + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + # out, q, k, v, out_padded, softmax_lse = _flash_attn_forward( + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + qv, # qv + None, # out + None, None, None, # cu_seqlens_q/k/k_new + None, None, # seqused_q/k + None, None, # max_seqlen_q/k + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size=window_size, + attention_chunk=attention_chunk, + softcap=softcap, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse) + ctx.save_for_backward(q, k, v, out, softmax_lse) + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.sm_margin = sm_margin + return out, softmax_lse + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + None, None, # cu_seqlens_q, cu_seqlens_k, + None, None, # sequed_q, sequed_k, + None, None, # max_seqlen_q, max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size, + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dq = dq[..., : q.shape[-1]] # We could have padded the head dimension + dk = dk[..., : k.shape[-1]] + dv = dv[..., : v.shape[-1]] + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None + + +class FlashAttnVarlenFunc(torch.autograd.Function): + + @staticmethod + def forward( + ctx, + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + ): + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + # out, q, k, v, out_padded, softmax_lse = _flash_attn_varlen_forward( + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + qv, # qv + None, # out + cu_seqlens_q, + cu_seqlens_k, + None, # cu_seqlens_k_new + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size=window_size, + attention_chunk=attention_chunk, + softcap=softcap, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k) + ctx.save_for_backward(q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k) + ctx.max_seqlen_q = max_seqlen_q + ctx.max_seqlen_k = max_seqlen_k + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.sm_margin = sm_margin + return out, softmax_lse + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + ctx.max_seqlen_q, + ctx.max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size, + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dq = dq[..., : q.shape[-1]] # We could have padded the head dimension + dk = dk[..., : k.shape[-1]] + dv = dv[..., : v.shape[-1]] + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None + + +def flash_attn_qkvpacked_func( + qkv, + softmax_scale=None, + causal=False, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + deterministic=False, + num_heads_q=None, + sm_margin=0, +): + """dropout_p should be set to 0.0 during evaluation + If Q, K, V are already stacked into 1 tensor, this function will be faster than + calling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation + of the gradients of Q, K, V. + For multi-query and grouped-query attention (MQA/GQA), please see + flash_attn_kvpacked_func and flash_attn_func. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between [i - window_size[0], i + window_size[1]] inclusive. + + Arguments: + qkv: (batch_size, seqlen, 3, nheads, headdim) + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + softcap: float. Anything > 0 activates softcapping attention. + alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of (-alibi_slope * |i - j|) is added to + the attention score of query i and key j. + deterministic: bool. Whether to use the deterministic implementation of the backward pass, + which is slightly slower and uses more memory. The forward pass is always deterministic. + return_attn_probs: bool. Whether to return the attention probabilities. This option is for + testing only. The returned probabilities are not guaranteed to be correct + (they might not have the right scaling). + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen). + The output of softmax (possibly with different scaling). It also encodes the dropout + pattern (negative means that location was dropped, nonnegative means it was kept). + """ + return FlashAttnQKVPackedFunc.apply( + qkv, + softmax_scale, + causal, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + deterministic, + num_heads_q, + sm_margin, + ) + + +def flash_attn_func( + q, + k, + v, + softmax_scale=None, + causal=False, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, +): + """dropout_p should be set to 0.0 during evaluation + Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads + than Q. Note that the number of heads in Q must be divisible by the number of heads in KV. + For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head + 0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V. + + If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix. + For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is: + 1 1 1 1 0 + 1 1 1 1 1 + If seqlen_q = 5 and seqlen_k = 2, the causal mask is: + 0 0 + 0 0 + 0 0 + 1 0 + 1 1 + If the row of the mask is all zero, the output will be zero. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between + [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive. + + Arguments: + q: (batch_size, seqlen, nheads, headdim) + k: (batch_size, seqlen, nheads_k, headdim) + v: (batch_size, seqlen, nheads_k, headdim) + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of + (-alibi_slope * |i + seqlen_k - seqlen_q - j|) + is added to the attention score of query i and key j. + deterministic: bool. Whether to use the deterministic implementation of the backward pass, + which is slightly slower and uses more memory. The forward pass is always deterministic. + return_attn_probs: bool. Whether to return the attention probabilities. This option is for + testing only. The returned probabilities are not guaranteed to be correct + (they might not have the right scaling). + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + """ + return FlashAttnFunc.apply( + q, + k, + v, + softmax_scale, + causal, + qv, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + num_splits, + pack_gqa, + deterministic, + sm_margin, + ) + + +def flash_attn_varlen_func( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + seqused_q=None, + seqused_k=None, + softmax_scale=None, + causal=False, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, +): + return FlashAttnVarlenFunc.apply( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + qv, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + num_splits, + pack_gqa, + deterministic, + sm_margin, + ) + + +def flash_attn_combine(out_partial, lse_partial, out=None, out_dtype=None): + return flash_attn_3_cuda.fwd_combine(out_partial, lse_partial, out, out_dtype) + + +def flash_attn_with_kvcache( + q, + k_cache, + v_cache, + k=None, + v=None, + qv=None, + rotary_cos=None, + rotary_sin=None, + cache_seqlens: Optional[Union[(int, torch.Tensor)]] = None, + cache_batch_idx: Optional[torch.Tensor] = None, + cache_leftpad: Optional[torch.Tensor] = None, + page_table: Optional[torch.Tensor] = None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + rotary_seqlens: Optional[torch.Tensor] = None, + q_descale: Optional[torch.Tensor] = None, + k_descale: Optional[torch.Tensor] = None, + v_descale: Optional[torch.Tensor] = None, + softmax_scale=None, + causal=False, + window_size=(-1, -1), # -1 means infinite context window + attention_chunk=0, + softcap=0.0, # 0.0 means deactivated + rotary_interleaved=True, + scheduler_metadata=None, + num_splits=0, # Can be tuned for speed + pack_gqa=None, # Can be tuned for speed + sm_margin=0, # Can be tuned if some SMs are used for communication + return_softmax_lse=False, +): + """ + If k and v are not None, k_cache and v_cache will be updated *inplace* with the new values from + k and v. This is useful for incremental decoding: you can pass in the cached keys/values from + the previous step, and update them with the new keys/values from the current step, and do + attention with the updated cache, all in 1 kernel. + + If you pass in k / v, you must make sure that the cache is large enough to hold the new values. + For example, the KV cache could be pre-allocated with the max sequence length, and you can use + cache_seqlens to keep track of the current sequence lengths of each sequence in the batch. + + Also apply rotary embedding if rotary_cos and rotary_sin are passed in. The key @k will be + rotated by rotary_cos and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc. + If causal or local (i.e., window_size != (-1, -1)), the query @q will be rotated by rotary_cos + and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc. + If not causal and not local, the query @q will be rotated by rotary_cos and rotary_sin at + indices cache_seqlens only (i.e. we consider all tokens in @q to be at position cache_seqlens). + + See tests/test_flash_attn.py::test_flash_attn_kvcache for examples of how to use this function. + + Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads + than Q. Note that the number of heads in Q must be divisible by the number of heads in KV. + For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head + 0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V. + + If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix. + For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is: + 1 1 1 1 0 + 1 1 1 1 1 + If seqlen_q = 5 and seqlen_k = 2, the causal mask is: + 0 0 + 0 0 + 0 0 + 1 0 + 1 1 + If the row of the mask is all zero, the output will be zero. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between + [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive. + + Note: Does not support backward pass. + + Arguments: + q: (batch_size, seqlen, nheads, headdim) + k_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim) if there's no page_table, + or (num_blocks, page_block_size, nheads_k, headdim) if there's a page_table (i.e. paged KV cache) + page_block_size must be a multiple of 256. + v_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim_v) if there's no page_table, + or (num_blocks, page_block_size, nheads_k, headdim_v) if there's a page_table (i.e. paged KV cache) + k [optional]: (batch_size, seqlen_new, nheads_k, headdim). If not None, we concatenate + k with k_cache, starting at the indices specified by cache_seqlens. + v [optional]: (batch_size, seqlen_new, nheads_k, headdim_v). Similar to k. + qv [optional]: (batch_size, seqlen, nheads, headdim_v) + rotary_cos [optional]: (seqlen_ro, rotary_dim / 2). If not None, we apply rotary embedding + to k and q. Only applicable if k and v are passed in. rotary_dim must be divisible by 16. + rotary_sin [optional]: (seqlen_ro, rotary_dim / 2). Similar to rotary_cos. + cache_seqlens: int, or (batch_size,), dtype torch.int32. The sequence lengths of the + KV cache. + cache_batch_idx: (batch_size,), dtype torch.int32. The indices used to index into the KV cache. + If None, we assume that the batch indices are [0, 1, 2, ..., batch_size - 1]. + If the indices are not distinct, and k and v are provided, the values updated in the cache + might come from any of the duplicate indices. + cache_leftpad: (batch_size,), dtype torch.int32. The index that the KV cache starts. If None, assume 0. + page_table [optional]: (batch_size, max_num_blocks_per_seq), dtype torch.int32. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + softcap: float. Anything > 0 activates softcapping attention. + rotary_interleaved: bool. Only applicable if rotary_cos and rotary_sin are passed in. + If True, rotary embedding will combine dimensions 0 & 1, 2 & 3, etc. If False, + rotary embedding will combine dimensions 0 & rotary_dim / 2, 1 & rotary_dim / 2 + 1 + (i.e. GPT-NeoX style). + num_splits: int. If > 1, split the key/value into this many chunks along the sequence. + If num_splits == 1, we don't split the key/value. If num_splits == 0, we use a heuristic + to automatically determine the number of splits. + Don't change this unless you know what you are doing. + return_softmax_lse: bool. Whether to return the logsumexp of the attention scores. + + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_softmax_lse=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + """ + assert k_cache.stride(-1) == 1, "k_cache must have contiguous last dimension" + assert v_cache.stride(-1) == 1, "v_cache must have contiguous last dimension" + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + if cache_seqlens is not None and isinstance(cache_seqlens, int): + cache_seqlens = torch.full( + (k_cache.shape[0],), cache_seqlens, dtype=torch.int32, device=k_cache.device + ) + cache_seqlens = maybe_contiguous(cache_seqlens) + out, softmax_lse, *rest = _flash_attn_forward( + q, + k_cache, + v_cache, + k, + v, + qv, + None, # out + cu_seqlens_q, + None, # cu_seqlens_k + cu_seqlens_k_new, + None, # seqused_q + cache_seqlens, + max_seqlen_q, + None, # max_seqlen_k + page_table, + cache_batch_idx, + cache_leftpad, + rotary_cos, + rotary_sin, + rotary_seqlens, + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size=window_size, + attention_chunk=attention_chunk, + softcap=softcap, + rotary_interleaved=rotary_interleaved, + scheduler_metadata=scheduler_metadata, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # return (out, softmax_lse) if return_softmax_lse else out + return (out, softmax_lse, *rest) if return_softmax_lse else out + + +def get_scheduler_metadata( + batch_size, max_seqlen_q, max_seqlen_k, num_heads_q, num_heads_kv, headdim, + cache_seqlens: torch.Tensor, + qkv_dtype=torch.bfloat16, + headdim_v=None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + cache_leftpad: Optional[torch.Tensor] = None, + page_size: Optional[int] = None, + max_seqlen_k_new=0, + causal=False, + window_size=(-1, -1), # -1 means infinite context window + attention_chunk=0, + has_softcap=False, + num_splits=0, # Can be tuned for speed + pack_gqa=None, # Can be tuned for speed + sm_margin=0, # Can be tuned if some SMs are used for communication +): + cache_seqlens = maybe_contiguous(cache_seqlens) + if headdim_v is None: + headdim_v = headdim + scheduler_metadata = flash_attn_3_cuda.get_scheduler_metadata( + batch_size, max_seqlen_q, max_seqlen_k, num_heads_q, num_heads_kv, headdim, headdim_v, + qkv_dtype, + cache_seqlens, + cu_seqlens_q, + None, # cu_seqlens_k + cu_seqlens_k_new, + None, # seqused_q + cache_leftpad, + page_size, + max_seqlen_k_new, + causal, + window_size[0], window_size[1], + attention_chunk, + has_softcap, + num_splits, + pack_gqa, + sm_margin, + ) + return scheduler_metadata diff --git a/build/torch29-cxx11-cu126-x86_64-linux/flash_attn3/__init__.py b/build/torch29-cxx11-cu126-x86_64-linux/flash_attn3/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3d12edd447e84d6e65c4fa9ab91ae3c3192884c1 --- /dev/null +++ b/build/torch29-cxx11-cu126-x86_64-linux/flash_attn3/__init__.py @@ -0,0 +1,17 @@ +from .flash_attn_interface import ( + flash_attn_combine, + flash_attn_func, + flash_attn_qkvpacked_func, + flash_attn_varlen_func, + flash_attn_with_kvcache, + get_scheduler_metadata, +) + +__all__ = [ + "flash_attn_combine", + "flash_attn_func", + "flash_attn_qkvpacked_func", + "flash_attn_varlen_func", + "flash_attn_with_kvcache", + "get_scheduler_metadata", +] diff --git a/build/torch29-cxx11-cu126-x86_64-linux/flash_attn3/__pycache__/__init__.cpython-313.pyc b/build/torch29-cxx11-cu126-x86_64-linux/flash_attn3/__pycache__/__init__.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7b97519dd99aca24b9e1c0189ed7d32825c19e45 Binary files /dev/null and b/build/torch29-cxx11-cu126-x86_64-linux/flash_attn3/__pycache__/__init__.cpython-313.pyc differ diff --git a/build/torch29-cxx11-cu126-x86_64-linux/flash_attn3/__pycache__/_ops.cpython-313.pyc b/build/torch29-cxx11-cu126-x86_64-linux/flash_attn3/__pycache__/_ops.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3afd1ee288e04cf21c4b4eb5383ab21dbfbf4acf Binary files /dev/null and b/build/torch29-cxx11-cu126-x86_64-linux/flash_attn3/__pycache__/_ops.cpython-313.pyc differ diff --git a/build/torch29-cxx11-cu126-x86_64-linux/flash_attn3/__pycache__/flash_attn_interface.cpython-313.pyc b/build/torch29-cxx11-cu126-x86_64-linux/flash_attn3/__pycache__/flash_attn_interface.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f6ad3174124a9f7166b4a5324ba7dd98605f51a8 Binary files /dev/null and b/build/torch29-cxx11-cu126-x86_64-linux/flash_attn3/__pycache__/flash_attn_interface.cpython-313.pyc differ diff --git a/build/torch29-cxx11-cu126-x86_64-linux/flash_attn3/_flash_attn3_7cf630c.abi3.so b/build/torch29-cxx11-cu126-x86_64-linux/flash_attn3/_flash_attn3_7cf630c.abi3.so new file mode 100755 index 0000000000000000000000000000000000000000..158541a34b9ccbc53906cd62482ef392030d3fb6 --- /dev/null +++ b/build/torch29-cxx11-cu126-x86_64-linux/flash_attn3/_flash_attn3_7cf630c.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ff6a253fa28a6e6578287f8e4e9ad03db9bb809a3392a858f61dc9fb9f904d6d +size 838540624 diff --git a/build/torch29-cxx11-cu126-x86_64-linux/flash_attn3/_ops.py b/build/torch29-cxx11-cu126-x86_64-linux/flash_attn3/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..92dc978d85ff1e8e4dd42e8192adb63ce3ca1a22 --- /dev/null +++ b/build/torch29-cxx11-cu126-x86_64-linux/flash_attn3/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _flash_attn3_7cf630c +ops = torch.ops._flash_attn3_7cf630c + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_flash_attn3_7cf630c::{op_name}" \ No newline at end of file diff --git a/build/torch29-cxx11-cu126-x86_64-linux/flash_attn3/flash_attn_interface.py b/build/torch29-cxx11-cu126-x86_64-linux/flash_attn3/flash_attn_interface.py new file mode 100644 index 0000000000000000000000000000000000000000..7572297784a73fdf891af6aafee88428120b72f2 --- /dev/null +++ b/build/torch29-cxx11-cu126-x86_64-linux/flash_attn3/flash_attn_interface.py @@ -0,0 +1,828 @@ +# Copyright (c) 2023, Tri Dao. + +from typing import Optional, Union + +import torch +import torch.nn as nn + +from ._ops import ops as flash_attn_3_cuda + +def maybe_contiguous(x): + return x.contiguous() if x is not None and x.stride(-1) != 1 else x + + +def _flash_attn_forward( + q, + k, + v, + k_new, + v_new, + qv, + out, + cu_seqlens_q, + cu_seqlens_k, + cu_seqlens_k_new, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + page_table, + kv_batch_idx, + leftpad_k, + rotary_cos, + rotary_sin, + seqlens_rotary, + q_descale, + k_descale, + v_descale, + softmax_scale, + causal, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + rotary_interleaved=True, + scheduler_metadata=None, + num_splits=1, + pack_gqa=None, + sm_margin=0): + q, k, k_new, v_new = [maybe_contiguous(x) for x in (q, k, k_new, v_new)] + v = v.contiguous() if v.stride(-1) != 1 and v.stride(-3) != 1 else v + cu_seqlens_q, cu_seqlens_k, cu_seqlens_k_new = [ + maybe_contiguous(x) for x in (cu_seqlens_q, cu_seqlens_k, cu_seqlens_k_new) + ] + seqused_q, seqused_k = [maybe_contiguous(x) for x in (seqused_q, seqused_k)] + page_table, kv_batch_idx, leftpad_k = [ + maybe_contiguous(x) for x in (page_table, kv_batch_idx, leftpad_k) + ] + rotary_cos, rotary_sin = [maybe_contiguous(x) for x in (rotary_cos, rotary_sin)] + seqlens_rotary = maybe_contiguous(seqlens_rotary) + out, softmax_lse, *rest = flash_attn_3_cuda.fwd( + q, + k, + v, + k_new, + v_new, + qv, + out, + cu_seqlens_q, + cu_seqlens_k, + cu_seqlens_k_new, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + page_table, + kv_batch_idx, + leftpad_k, + rotary_cos, + rotary_sin, + seqlens_rotary, + q_descale, + k_descale, + v_descale, + softmax_scale, + causal, + window_size[0], + window_size[1], + attention_chunk, + softcap, + rotary_interleaved, + scheduler_metadata, + num_splits, + pack_gqa, + sm_margin, + ) + return out, softmax_lse, *rest + + +def _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + cu_seqlens_q, + cu_seqlens_k, + sequed_q, + sequed_k, + max_seqlen_q, + max_seqlen_k, + dq, + dk, + dv, + softmax_scale, + causal, + window_size=(-1, -1), + softcap=0.0, + deterministic=False, + sm_margin=0, +): + # dq, dk, dv are allocated by us so they should already be contiguous + dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] + dq, dk, dv, softmax_d, *rest = flash_attn_3_cuda.bwd( + dout, + q, + k, + v, + out, + softmax_lse, + dq, + dk, + dv, + cu_seqlens_q, + cu_seqlens_k, + sequed_q, + sequed_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + window_size[0], + window_size[1], + softcap, + deterministic, + sm_margin, + ) + return dq, dk, dv, softmax_d + + +class FlashAttnQKVPackedFunc(torch.autograd.Function): + @staticmethod + def forward( + ctx, + qkv, + softmax_scale, + causal, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + deterministic=False, + num_heads_q=None, + sm_margin=0, + ): + if softmax_scale is None: + softmax_scale = qkv.shape[-1] ** (-0.5) + if qkv.dim() == 5: + assert qkv.shape[-3] == 3 + q, k, v = qkv.unbind(dim=-3) + else: + assert qkv.dim() == 4 + assert num_heads_q is not None + num_heads_k = (qkv.shape[2] - num_heads_q) // 2 + assert num_heads_k * 2 + num_heads_q == qkv.shape[2] + q, k, v = qkv.split([num_heads_q, num_heads_k, num_heads_k], dim=-2) + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + None, # qv + None, # out + None, None, None, # cu_seqlens_q/k/k_new + None, None, # seqused_q/k + None, None, # max_seqlen_q/k + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size=window_size, + attention_chunk=attention_chunk, + softcap=softcap, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse) + ctx.save_for_backward(q, k, v, out, softmax_lse) + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.ndim = qkv.dim() + ctx.sm_margin = sm_margin + # return out, softmax_lse + return out + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + if ctx.ndim == 5: + qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) + dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) + dq, dk, dv = dqkv.unbind(dim=-3) + else: + num_heads_q = q.shape[2] + num_heads_k = k.shape[2] + qkv_shape = q.shape[:-2] + (num_heads_q + num_heads_k * 2, *q.shape[-1:]) + dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) + dq, dk, dv = dqkv.split([num_heads_q, num_heads_k, num_heads_k], dim=-2) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + None, None, # cu_seqlens_q, cu_seqlens_k, + None, None, # sequed_q, sequed_k, + None, None, # max_seqlen_q, max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size, + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dqkv = dqkv[..., : dout.shape[-1]] # We could have padded the head dimension + return dqkv, None, None, None, None, None, None, None, None, None, None, None + + +class FlashAttnFunc(torch.autograd.Function): + + @staticmethod + def forward( + ctx, + q, + k, + v, + softmax_scale, + causal, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + ): + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + # out, q, k, v, out_padded, softmax_lse = _flash_attn_forward( + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + qv, # qv + None, # out + None, None, None, # cu_seqlens_q/k/k_new + None, None, # seqused_q/k + None, None, # max_seqlen_q/k + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size=window_size, + attention_chunk=attention_chunk, + softcap=softcap, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse) + ctx.save_for_backward(q, k, v, out, softmax_lse) + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.sm_margin = sm_margin + return out, softmax_lse + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + None, None, # cu_seqlens_q, cu_seqlens_k, + None, None, # sequed_q, sequed_k, + None, None, # max_seqlen_q, max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size, + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dq = dq[..., : q.shape[-1]] # We could have padded the head dimension + dk = dk[..., : k.shape[-1]] + dv = dv[..., : v.shape[-1]] + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None + + +class FlashAttnVarlenFunc(torch.autograd.Function): + + @staticmethod + def forward( + ctx, + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + ): + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + # out, q, k, v, out_padded, softmax_lse = _flash_attn_varlen_forward( + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + qv, # qv + None, # out + cu_seqlens_q, + cu_seqlens_k, + None, # cu_seqlens_k_new + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size=window_size, + attention_chunk=attention_chunk, + softcap=softcap, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k) + ctx.save_for_backward(q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k) + ctx.max_seqlen_q = max_seqlen_q + ctx.max_seqlen_k = max_seqlen_k + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.sm_margin = sm_margin + return out, softmax_lse + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + ctx.max_seqlen_q, + ctx.max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size, + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dq = dq[..., : q.shape[-1]] # We could have padded the head dimension + dk = dk[..., : k.shape[-1]] + dv = dv[..., : v.shape[-1]] + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None + + +def flash_attn_qkvpacked_func( + qkv, + softmax_scale=None, + causal=False, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + deterministic=False, + num_heads_q=None, + sm_margin=0, +): + """dropout_p should be set to 0.0 during evaluation + If Q, K, V are already stacked into 1 tensor, this function will be faster than + calling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation + of the gradients of Q, K, V. + For multi-query and grouped-query attention (MQA/GQA), please see + flash_attn_kvpacked_func and flash_attn_func. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between [i - window_size[0], i + window_size[1]] inclusive. + + Arguments: + qkv: (batch_size, seqlen, 3, nheads, headdim) + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + softcap: float. Anything > 0 activates softcapping attention. + alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of (-alibi_slope * |i - j|) is added to + the attention score of query i and key j. + deterministic: bool. Whether to use the deterministic implementation of the backward pass, + which is slightly slower and uses more memory. The forward pass is always deterministic. + return_attn_probs: bool. Whether to return the attention probabilities. This option is for + testing only. The returned probabilities are not guaranteed to be correct + (they might not have the right scaling). + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen). + The output of softmax (possibly with different scaling). It also encodes the dropout + pattern (negative means that location was dropped, nonnegative means it was kept). + """ + return FlashAttnQKVPackedFunc.apply( + qkv, + softmax_scale, + causal, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + deterministic, + num_heads_q, + sm_margin, + ) + + +def flash_attn_func( + q, + k, + v, + softmax_scale=None, + causal=False, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, +): + """dropout_p should be set to 0.0 during evaluation + Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads + than Q. Note that the number of heads in Q must be divisible by the number of heads in KV. + For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head + 0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V. + + If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix. + For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is: + 1 1 1 1 0 + 1 1 1 1 1 + If seqlen_q = 5 and seqlen_k = 2, the causal mask is: + 0 0 + 0 0 + 0 0 + 1 0 + 1 1 + If the row of the mask is all zero, the output will be zero. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between + [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive. + + Arguments: + q: (batch_size, seqlen, nheads, headdim) + k: (batch_size, seqlen, nheads_k, headdim) + v: (batch_size, seqlen, nheads_k, headdim) + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of + (-alibi_slope * |i + seqlen_k - seqlen_q - j|) + is added to the attention score of query i and key j. + deterministic: bool. Whether to use the deterministic implementation of the backward pass, + which is slightly slower and uses more memory. The forward pass is always deterministic. + return_attn_probs: bool. Whether to return the attention probabilities. This option is for + testing only. The returned probabilities are not guaranteed to be correct + (they might not have the right scaling). + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + """ + return FlashAttnFunc.apply( + q, + k, + v, + softmax_scale, + causal, + qv, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + num_splits, + pack_gqa, + deterministic, + sm_margin, + ) + + +def flash_attn_varlen_func( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + seqused_q=None, + seqused_k=None, + softmax_scale=None, + causal=False, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, +): + return FlashAttnVarlenFunc.apply( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + qv, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + num_splits, + pack_gqa, + deterministic, + sm_margin, + ) + + +def flash_attn_combine(out_partial, lse_partial, out=None, out_dtype=None): + return flash_attn_3_cuda.fwd_combine(out_partial, lse_partial, out, out_dtype) + + +def flash_attn_with_kvcache( + q, + k_cache, + v_cache, + k=None, + v=None, + qv=None, + rotary_cos=None, + rotary_sin=None, + cache_seqlens: Optional[Union[(int, torch.Tensor)]] = None, + cache_batch_idx: Optional[torch.Tensor] = None, + cache_leftpad: Optional[torch.Tensor] = None, + page_table: Optional[torch.Tensor] = None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + rotary_seqlens: Optional[torch.Tensor] = None, + q_descale: Optional[torch.Tensor] = None, + k_descale: Optional[torch.Tensor] = None, + v_descale: Optional[torch.Tensor] = None, + softmax_scale=None, + causal=False, + window_size=(-1, -1), # -1 means infinite context window + attention_chunk=0, + softcap=0.0, # 0.0 means deactivated + rotary_interleaved=True, + scheduler_metadata=None, + num_splits=0, # Can be tuned for speed + pack_gqa=None, # Can be tuned for speed + sm_margin=0, # Can be tuned if some SMs are used for communication + return_softmax_lse=False, +): + """ + If k and v are not None, k_cache and v_cache will be updated *inplace* with the new values from + k and v. This is useful for incremental decoding: you can pass in the cached keys/values from + the previous step, and update them with the new keys/values from the current step, and do + attention with the updated cache, all in 1 kernel. + + If you pass in k / v, you must make sure that the cache is large enough to hold the new values. + For example, the KV cache could be pre-allocated with the max sequence length, and you can use + cache_seqlens to keep track of the current sequence lengths of each sequence in the batch. + + Also apply rotary embedding if rotary_cos and rotary_sin are passed in. The key @k will be + rotated by rotary_cos and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc. + If causal or local (i.e., window_size != (-1, -1)), the query @q will be rotated by rotary_cos + and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc. + If not causal and not local, the query @q will be rotated by rotary_cos and rotary_sin at + indices cache_seqlens only (i.e. we consider all tokens in @q to be at position cache_seqlens). + + See tests/test_flash_attn.py::test_flash_attn_kvcache for examples of how to use this function. + + Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads + than Q. Note that the number of heads in Q must be divisible by the number of heads in KV. + For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head + 0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V. + + If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix. + For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is: + 1 1 1 1 0 + 1 1 1 1 1 + If seqlen_q = 5 and seqlen_k = 2, the causal mask is: + 0 0 + 0 0 + 0 0 + 1 0 + 1 1 + If the row of the mask is all zero, the output will be zero. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between + [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive. + + Note: Does not support backward pass. + + Arguments: + q: (batch_size, seqlen, nheads, headdim) + k_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim) if there's no page_table, + or (num_blocks, page_block_size, nheads_k, headdim) if there's a page_table (i.e. paged KV cache) + page_block_size must be a multiple of 256. + v_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim_v) if there's no page_table, + or (num_blocks, page_block_size, nheads_k, headdim_v) if there's a page_table (i.e. paged KV cache) + k [optional]: (batch_size, seqlen_new, nheads_k, headdim). If not None, we concatenate + k with k_cache, starting at the indices specified by cache_seqlens. + v [optional]: (batch_size, seqlen_new, nheads_k, headdim_v). Similar to k. + qv [optional]: (batch_size, seqlen, nheads, headdim_v) + rotary_cos [optional]: (seqlen_ro, rotary_dim / 2). If not None, we apply rotary embedding + to k and q. Only applicable if k and v are passed in. rotary_dim must be divisible by 16. + rotary_sin [optional]: (seqlen_ro, rotary_dim / 2). Similar to rotary_cos. + cache_seqlens: int, or (batch_size,), dtype torch.int32. The sequence lengths of the + KV cache. + cache_batch_idx: (batch_size,), dtype torch.int32. The indices used to index into the KV cache. + If None, we assume that the batch indices are [0, 1, 2, ..., batch_size - 1]. + If the indices are not distinct, and k and v are provided, the values updated in the cache + might come from any of the duplicate indices. + cache_leftpad: (batch_size,), dtype torch.int32. The index that the KV cache starts. If None, assume 0. + page_table [optional]: (batch_size, max_num_blocks_per_seq), dtype torch.int32. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + softcap: float. Anything > 0 activates softcapping attention. + rotary_interleaved: bool. Only applicable if rotary_cos and rotary_sin are passed in. + If True, rotary embedding will combine dimensions 0 & 1, 2 & 3, etc. If False, + rotary embedding will combine dimensions 0 & rotary_dim / 2, 1 & rotary_dim / 2 + 1 + (i.e. GPT-NeoX style). + num_splits: int. If > 1, split the key/value into this many chunks along the sequence. + If num_splits == 1, we don't split the key/value. If num_splits == 0, we use a heuristic + to automatically determine the number of splits. + Don't change this unless you know what you are doing. + return_softmax_lse: bool. Whether to return the logsumexp of the attention scores. + + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_softmax_lse=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + """ + assert k_cache.stride(-1) == 1, "k_cache must have contiguous last dimension" + assert v_cache.stride(-1) == 1, "v_cache must have contiguous last dimension" + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + if cache_seqlens is not None and isinstance(cache_seqlens, int): + cache_seqlens = torch.full( + (k_cache.shape[0],), cache_seqlens, dtype=torch.int32, device=k_cache.device + ) + cache_seqlens = maybe_contiguous(cache_seqlens) + out, softmax_lse, *rest = _flash_attn_forward( + q, + k_cache, + v_cache, + k, + v, + qv, + None, # out + cu_seqlens_q, + None, # cu_seqlens_k + cu_seqlens_k_new, + None, # seqused_q + cache_seqlens, + max_seqlen_q, + None, # max_seqlen_k + page_table, + cache_batch_idx, + cache_leftpad, + rotary_cos, + rotary_sin, + rotary_seqlens, + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size=window_size, + attention_chunk=attention_chunk, + softcap=softcap, + rotary_interleaved=rotary_interleaved, + scheduler_metadata=scheduler_metadata, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # return (out, softmax_lse) if return_softmax_lse else out + return (out, softmax_lse, *rest) if return_softmax_lse else out + + +def get_scheduler_metadata( + batch_size, max_seqlen_q, max_seqlen_k, num_heads_q, num_heads_kv, headdim, + cache_seqlens: torch.Tensor, + qkv_dtype=torch.bfloat16, + headdim_v=None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + cache_leftpad: Optional[torch.Tensor] = None, + page_size: Optional[int] = None, + max_seqlen_k_new=0, + causal=False, + window_size=(-1, -1), # -1 means infinite context window + attention_chunk=0, + has_softcap=False, + num_splits=0, # Can be tuned for speed + pack_gqa=None, # Can be tuned for speed + sm_margin=0, # Can be tuned if some SMs are used for communication +): + cache_seqlens = maybe_contiguous(cache_seqlens) + if headdim_v is None: + headdim_v = headdim + scheduler_metadata = flash_attn_3_cuda.get_scheduler_metadata( + batch_size, max_seqlen_q, max_seqlen_k, num_heads_q, num_heads_kv, headdim, headdim_v, + qkv_dtype, + cache_seqlens, + cu_seqlens_q, + None, # cu_seqlens_k + cu_seqlens_k_new, + None, # seqused_q + cache_leftpad, + page_size, + max_seqlen_k_new, + causal, + window_size[0], window_size[1], + attention_chunk, + has_softcap, + num_splits, + pack_gqa, + sm_margin, + ) + return scheduler_metadata diff --git a/build/torch29-cxx11-cu126-x86_64-linux/flash_attn_config.py b/build/torch29-cxx11-cu126-x86_64-linux/flash_attn_config.py new file mode 100644 index 0000000000000000000000000000000000000000..c4b7130324be2c4cabee7ada7000e170a2a81865 --- /dev/null +++ b/build/torch29-cxx11-cu126-x86_64-linux/flash_attn_config.py @@ -0,0 +1,7 @@ +# Auto-generated by flash attention 3 setup.py +CONFIG = {'build_flags': {'FLASHATTENTION_DISABLE_BACKWARD': False, 'FLASHATTENTION_DISABLE_SPLIT': False, 'FLASHATTENTION_DISABLE_PAGEDKV': False, 'FLASHATTENTION_DISABLE_APPENDKV': False, 'FLASHATTENTION_DISABLE_LOCAL': False, 'FLASHATTENTION_DISABLE_SOFTCAP': False, 'FLASHATTENTION_DISABLE_PACKGQA': False, 'FLASHATTENTION_DISABLE_FP16': False, 'FLASHATTENTION_DISABLE_FP8': False, 'FLASHATTENTION_DISABLE_VARLEN': False, 'FLASHATTENTION_DISABLE_CLUSTER': False, 'FLASHATTENTION_DISABLE_HDIM64': False, 'FLASHATTENTION_DISABLE_HDIM96': False, 'FLASHATTENTION_DISABLE_HDIM128': False, 'FLASHATTENTION_DISABLE_HDIM192': False, 'FLASHATTENTION_DISABLE_HDIM256': False, 'FLASHATTENTION_DISABLE_SM8x': False, 'FLASHATTENTION_ENABLE_VCOLMAJOR': False, 'FLASH_ATTENTION_DISABLE_HDIMDIFF64': False, 'FLASH_ATTENTION_DISABLE_HDIMDIFF192': False}} + +def show(): + from pprint import pprint + pprint(CONFIG) + diff --git a/build/torch29-cxx11-cu128-aarch64-linux/flash_attn3/__init__.py b/build/torch29-cxx11-cu128-aarch64-linux/flash_attn3/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3d12edd447e84d6e65c4fa9ab91ae3c3192884c1 --- /dev/null +++ b/build/torch29-cxx11-cu128-aarch64-linux/flash_attn3/__init__.py @@ -0,0 +1,17 @@ +from .flash_attn_interface import ( + flash_attn_combine, + flash_attn_func, + flash_attn_qkvpacked_func, + flash_attn_varlen_func, + flash_attn_with_kvcache, + get_scheduler_metadata, +) + +__all__ = [ + "flash_attn_combine", + "flash_attn_func", + "flash_attn_qkvpacked_func", + "flash_attn_varlen_func", + "flash_attn_with_kvcache", + "get_scheduler_metadata", +] diff --git a/build/torch29-cxx11-cu128-aarch64-linux/flash_attn3/__pycache__/__init__.cpython-313.pyc b/build/torch29-cxx11-cu128-aarch64-linux/flash_attn3/__pycache__/__init__.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..857fa9ac00bb1e42aa50ce5e49b975e27147ecc5 Binary files /dev/null and b/build/torch29-cxx11-cu128-aarch64-linux/flash_attn3/__pycache__/__init__.cpython-313.pyc differ diff --git a/build/torch29-cxx11-cu128-aarch64-linux/flash_attn3/__pycache__/_ops.cpython-313.pyc b/build/torch29-cxx11-cu128-aarch64-linux/flash_attn3/__pycache__/_ops.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a1b03ee6f27dc2017aa2aade427aeab18dbf8486 Binary files /dev/null and b/build/torch29-cxx11-cu128-aarch64-linux/flash_attn3/__pycache__/_ops.cpython-313.pyc differ diff --git a/build/torch29-cxx11-cu128-aarch64-linux/flash_attn3/__pycache__/flash_attn_interface.cpython-313.pyc b/build/torch29-cxx11-cu128-aarch64-linux/flash_attn3/__pycache__/flash_attn_interface.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..485b2c2591cb28b809904f1769de106ca9e97481 Binary files /dev/null and b/build/torch29-cxx11-cu128-aarch64-linux/flash_attn3/__pycache__/flash_attn_interface.cpython-313.pyc differ diff --git a/build/torch29-cxx11-cu128-aarch64-linux/flash_attn3/_flash_attn3_847092b_dirty.abi3.so b/build/torch29-cxx11-cu128-aarch64-linux/flash_attn3/_flash_attn3_847092b_dirty.abi3.so new file mode 100755 index 0000000000000000000000000000000000000000..dec7879b4b5bca58ed4c772d780c368462632625 --- /dev/null +++ b/build/torch29-cxx11-cu128-aarch64-linux/flash_attn3/_flash_attn3_847092b_dirty.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:17179eb1daba5483276f8536733febb2623ef14c002b5315859f7eed3f73fa81 +size 822395648 diff --git a/build/torch29-cxx11-cu128-aarch64-linux/flash_attn3/_ops.py b/build/torch29-cxx11-cu128-aarch64-linux/flash_attn3/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..b7a3a5f4470eb321053647e1d601c8448d21490a --- /dev/null +++ b/build/torch29-cxx11-cu128-aarch64-linux/flash_attn3/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _flash_attn3_847092b_dirty +ops = torch.ops._flash_attn3_847092b_dirty + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_flash_attn3_847092b_dirty::{op_name}" \ No newline at end of file diff --git a/build/torch29-cxx11-cu128-aarch64-linux/flash_attn3/flash_attn_interface.py b/build/torch29-cxx11-cu128-aarch64-linux/flash_attn3/flash_attn_interface.py new file mode 100644 index 0000000000000000000000000000000000000000..7572297784a73fdf891af6aafee88428120b72f2 --- /dev/null +++ b/build/torch29-cxx11-cu128-aarch64-linux/flash_attn3/flash_attn_interface.py @@ -0,0 +1,828 @@ +# Copyright (c) 2023, Tri Dao. + +from typing import Optional, Union + +import torch +import torch.nn as nn + +from ._ops import ops as flash_attn_3_cuda + +def maybe_contiguous(x): + return x.contiguous() if x is not None and x.stride(-1) != 1 else x + + +def _flash_attn_forward( + q, + k, + v, + k_new, + v_new, + qv, + out, + cu_seqlens_q, + cu_seqlens_k, + cu_seqlens_k_new, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + page_table, + kv_batch_idx, + leftpad_k, + rotary_cos, + rotary_sin, + seqlens_rotary, + q_descale, + k_descale, + v_descale, + softmax_scale, + causal, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + rotary_interleaved=True, + scheduler_metadata=None, + num_splits=1, + pack_gqa=None, + sm_margin=0): + q, k, k_new, v_new = [maybe_contiguous(x) for x in (q, k, k_new, v_new)] + v = v.contiguous() if v.stride(-1) != 1 and v.stride(-3) != 1 else v + cu_seqlens_q, cu_seqlens_k, cu_seqlens_k_new = [ + maybe_contiguous(x) for x in (cu_seqlens_q, cu_seqlens_k, cu_seqlens_k_new) + ] + seqused_q, seqused_k = [maybe_contiguous(x) for x in (seqused_q, seqused_k)] + page_table, kv_batch_idx, leftpad_k = [ + maybe_contiguous(x) for x in (page_table, kv_batch_idx, leftpad_k) + ] + rotary_cos, rotary_sin = [maybe_contiguous(x) for x in (rotary_cos, rotary_sin)] + seqlens_rotary = maybe_contiguous(seqlens_rotary) + out, softmax_lse, *rest = flash_attn_3_cuda.fwd( + q, + k, + v, + k_new, + v_new, + qv, + out, + cu_seqlens_q, + cu_seqlens_k, + cu_seqlens_k_new, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + page_table, + kv_batch_idx, + leftpad_k, + rotary_cos, + rotary_sin, + seqlens_rotary, + q_descale, + k_descale, + v_descale, + softmax_scale, + causal, + window_size[0], + window_size[1], + attention_chunk, + softcap, + rotary_interleaved, + scheduler_metadata, + num_splits, + pack_gqa, + sm_margin, + ) + return out, softmax_lse, *rest + + +def _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + cu_seqlens_q, + cu_seqlens_k, + sequed_q, + sequed_k, + max_seqlen_q, + max_seqlen_k, + dq, + dk, + dv, + softmax_scale, + causal, + window_size=(-1, -1), + softcap=0.0, + deterministic=False, + sm_margin=0, +): + # dq, dk, dv are allocated by us so they should already be contiguous + dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] + dq, dk, dv, softmax_d, *rest = flash_attn_3_cuda.bwd( + dout, + q, + k, + v, + out, + softmax_lse, + dq, + dk, + dv, + cu_seqlens_q, + cu_seqlens_k, + sequed_q, + sequed_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + window_size[0], + window_size[1], + softcap, + deterministic, + sm_margin, + ) + return dq, dk, dv, softmax_d + + +class FlashAttnQKVPackedFunc(torch.autograd.Function): + @staticmethod + def forward( + ctx, + qkv, + softmax_scale, + causal, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + deterministic=False, + num_heads_q=None, + sm_margin=0, + ): + if softmax_scale is None: + softmax_scale = qkv.shape[-1] ** (-0.5) + if qkv.dim() == 5: + assert qkv.shape[-3] == 3 + q, k, v = qkv.unbind(dim=-3) + else: + assert qkv.dim() == 4 + assert num_heads_q is not None + num_heads_k = (qkv.shape[2] - num_heads_q) // 2 + assert num_heads_k * 2 + num_heads_q == qkv.shape[2] + q, k, v = qkv.split([num_heads_q, num_heads_k, num_heads_k], dim=-2) + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + None, # qv + None, # out + None, None, None, # cu_seqlens_q/k/k_new + None, None, # seqused_q/k + None, None, # max_seqlen_q/k + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size=window_size, + attention_chunk=attention_chunk, + softcap=softcap, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse) + ctx.save_for_backward(q, k, v, out, softmax_lse) + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.ndim = qkv.dim() + ctx.sm_margin = sm_margin + # return out, softmax_lse + return out + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + if ctx.ndim == 5: + qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) + dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) + dq, dk, dv = dqkv.unbind(dim=-3) + else: + num_heads_q = q.shape[2] + num_heads_k = k.shape[2] + qkv_shape = q.shape[:-2] + (num_heads_q + num_heads_k * 2, *q.shape[-1:]) + dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) + dq, dk, dv = dqkv.split([num_heads_q, num_heads_k, num_heads_k], dim=-2) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + None, None, # cu_seqlens_q, cu_seqlens_k, + None, None, # sequed_q, sequed_k, + None, None, # max_seqlen_q, max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size, + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dqkv = dqkv[..., : dout.shape[-1]] # We could have padded the head dimension + return dqkv, None, None, None, None, None, None, None, None, None, None, None + + +class FlashAttnFunc(torch.autograd.Function): + + @staticmethod + def forward( + ctx, + q, + k, + v, + softmax_scale, + causal, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + ): + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + # out, q, k, v, out_padded, softmax_lse = _flash_attn_forward( + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + qv, # qv + None, # out + None, None, None, # cu_seqlens_q/k/k_new + None, None, # seqused_q/k + None, None, # max_seqlen_q/k + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size=window_size, + attention_chunk=attention_chunk, + softcap=softcap, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse) + ctx.save_for_backward(q, k, v, out, softmax_lse) + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.sm_margin = sm_margin + return out, softmax_lse + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + None, None, # cu_seqlens_q, cu_seqlens_k, + None, None, # sequed_q, sequed_k, + None, None, # max_seqlen_q, max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size, + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dq = dq[..., : q.shape[-1]] # We could have padded the head dimension + dk = dk[..., : k.shape[-1]] + dv = dv[..., : v.shape[-1]] + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None + + +class FlashAttnVarlenFunc(torch.autograd.Function): + + @staticmethod + def forward( + ctx, + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + ): + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + # out, q, k, v, out_padded, softmax_lse = _flash_attn_varlen_forward( + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + qv, # qv + None, # out + cu_seqlens_q, + cu_seqlens_k, + None, # cu_seqlens_k_new + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size=window_size, + attention_chunk=attention_chunk, + softcap=softcap, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k) + ctx.save_for_backward(q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k) + ctx.max_seqlen_q = max_seqlen_q + ctx.max_seqlen_k = max_seqlen_k + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.sm_margin = sm_margin + return out, softmax_lse + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + ctx.max_seqlen_q, + ctx.max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size, + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dq = dq[..., : q.shape[-1]] # We could have padded the head dimension + dk = dk[..., : k.shape[-1]] + dv = dv[..., : v.shape[-1]] + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None + + +def flash_attn_qkvpacked_func( + qkv, + softmax_scale=None, + causal=False, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + deterministic=False, + num_heads_q=None, + sm_margin=0, +): + """dropout_p should be set to 0.0 during evaluation + If Q, K, V are already stacked into 1 tensor, this function will be faster than + calling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation + of the gradients of Q, K, V. + For multi-query and grouped-query attention (MQA/GQA), please see + flash_attn_kvpacked_func and flash_attn_func. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between [i - window_size[0], i + window_size[1]] inclusive. + + Arguments: + qkv: (batch_size, seqlen, 3, nheads, headdim) + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + softcap: float. Anything > 0 activates softcapping attention. + alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of (-alibi_slope * |i - j|) is added to + the attention score of query i and key j. + deterministic: bool. Whether to use the deterministic implementation of the backward pass, + which is slightly slower and uses more memory. The forward pass is always deterministic. + return_attn_probs: bool. Whether to return the attention probabilities. This option is for + testing only. The returned probabilities are not guaranteed to be correct + (they might not have the right scaling). + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen). + The output of softmax (possibly with different scaling). It also encodes the dropout + pattern (negative means that location was dropped, nonnegative means it was kept). + """ + return FlashAttnQKVPackedFunc.apply( + qkv, + softmax_scale, + causal, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + deterministic, + num_heads_q, + sm_margin, + ) + + +def flash_attn_func( + q, + k, + v, + softmax_scale=None, + causal=False, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, +): + """dropout_p should be set to 0.0 during evaluation + Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads + than Q. Note that the number of heads in Q must be divisible by the number of heads in KV. + For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head + 0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V. + + If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix. + For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is: + 1 1 1 1 0 + 1 1 1 1 1 + If seqlen_q = 5 and seqlen_k = 2, the causal mask is: + 0 0 + 0 0 + 0 0 + 1 0 + 1 1 + If the row of the mask is all zero, the output will be zero. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between + [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive. + + Arguments: + q: (batch_size, seqlen, nheads, headdim) + k: (batch_size, seqlen, nheads_k, headdim) + v: (batch_size, seqlen, nheads_k, headdim) + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of + (-alibi_slope * |i + seqlen_k - seqlen_q - j|) + is added to the attention score of query i and key j. + deterministic: bool. Whether to use the deterministic implementation of the backward pass, + which is slightly slower and uses more memory. The forward pass is always deterministic. + return_attn_probs: bool. Whether to return the attention probabilities. This option is for + testing only. The returned probabilities are not guaranteed to be correct + (they might not have the right scaling). + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + """ + return FlashAttnFunc.apply( + q, + k, + v, + softmax_scale, + causal, + qv, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + num_splits, + pack_gqa, + deterministic, + sm_margin, + ) + + +def flash_attn_varlen_func( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + seqused_q=None, + seqused_k=None, + softmax_scale=None, + causal=False, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, +): + return FlashAttnVarlenFunc.apply( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + qv, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + num_splits, + pack_gqa, + deterministic, + sm_margin, + ) + + +def flash_attn_combine(out_partial, lse_partial, out=None, out_dtype=None): + return flash_attn_3_cuda.fwd_combine(out_partial, lse_partial, out, out_dtype) + + +def flash_attn_with_kvcache( + q, + k_cache, + v_cache, + k=None, + v=None, + qv=None, + rotary_cos=None, + rotary_sin=None, + cache_seqlens: Optional[Union[(int, torch.Tensor)]] = None, + cache_batch_idx: Optional[torch.Tensor] = None, + cache_leftpad: Optional[torch.Tensor] = None, + page_table: Optional[torch.Tensor] = None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + rotary_seqlens: Optional[torch.Tensor] = None, + q_descale: Optional[torch.Tensor] = None, + k_descale: Optional[torch.Tensor] = None, + v_descale: Optional[torch.Tensor] = None, + softmax_scale=None, + causal=False, + window_size=(-1, -1), # -1 means infinite context window + attention_chunk=0, + softcap=0.0, # 0.0 means deactivated + rotary_interleaved=True, + scheduler_metadata=None, + num_splits=0, # Can be tuned for speed + pack_gqa=None, # Can be tuned for speed + sm_margin=0, # Can be tuned if some SMs are used for communication + return_softmax_lse=False, +): + """ + If k and v are not None, k_cache and v_cache will be updated *inplace* with the new values from + k and v. This is useful for incremental decoding: you can pass in the cached keys/values from + the previous step, and update them with the new keys/values from the current step, and do + attention with the updated cache, all in 1 kernel. + + If you pass in k / v, you must make sure that the cache is large enough to hold the new values. + For example, the KV cache could be pre-allocated with the max sequence length, and you can use + cache_seqlens to keep track of the current sequence lengths of each sequence in the batch. + + Also apply rotary embedding if rotary_cos and rotary_sin are passed in. The key @k will be + rotated by rotary_cos and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc. + If causal or local (i.e., window_size != (-1, -1)), the query @q will be rotated by rotary_cos + and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc. + If not causal and not local, the query @q will be rotated by rotary_cos and rotary_sin at + indices cache_seqlens only (i.e. we consider all tokens in @q to be at position cache_seqlens). + + See tests/test_flash_attn.py::test_flash_attn_kvcache for examples of how to use this function. + + Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads + than Q. Note that the number of heads in Q must be divisible by the number of heads in KV. + For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head + 0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V. + + If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix. + For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is: + 1 1 1 1 0 + 1 1 1 1 1 + If seqlen_q = 5 and seqlen_k = 2, the causal mask is: + 0 0 + 0 0 + 0 0 + 1 0 + 1 1 + If the row of the mask is all zero, the output will be zero. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between + [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive. + + Note: Does not support backward pass. + + Arguments: + q: (batch_size, seqlen, nheads, headdim) + k_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim) if there's no page_table, + or (num_blocks, page_block_size, nheads_k, headdim) if there's a page_table (i.e. paged KV cache) + page_block_size must be a multiple of 256. + v_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim_v) if there's no page_table, + or (num_blocks, page_block_size, nheads_k, headdim_v) if there's a page_table (i.e. paged KV cache) + k [optional]: (batch_size, seqlen_new, nheads_k, headdim). If not None, we concatenate + k with k_cache, starting at the indices specified by cache_seqlens. + v [optional]: (batch_size, seqlen_new, nheads_k, headdim_v). Similar to k. + qv [optional]: (batch_size, seqlen, nheads, headdim_v) + rotary_cos [optional]: (seqlen_ro, rotary_dim / 2). If not None, we apply rotary embedding + to k and q. Only applicable if k and v are passed in. rotary_dim must be divisible by 16. + rotary_sin [optional]: (seqlen_ro, rotary_dim / 2). Similar to rotary_cos. + cache_seqlens: int, or (batch_size,), dtype torch.int32. The sequence lengths of the + KV cache. + cache_batch_idx: (batch_size,), dtype torch.int32. The indices used to index into the KV cache. + If None, we assume that the batch indices are [0, 1, 2, ..., batch_size - 1]. + If the indices are not distinct, and k and v are provided, the values updated in the cache + might come from any of the duplicate indices. + cache_leftpad: (batch_size,), dtype torch.int32. The index that the KV cache starts. If None, assume 0. + page_table [optional]: (batch_size, max_num_blocks_per_seq), dtype torch.int32. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + softcap: float. Anything > 0 activates softcapping attention. + rotary_interleaved: bool. Only applicable if rotary_cos and rotary_sin are passed in. + If True, rotary embedding will combine dimensions 0 & 1, 2 & 3, etc. If False, + rotary embedding will combine dimensions 0 & rotary_dim / 2, 1 & rotary_dim / 2 + 1 + (i.e. GPT-NeoX style). + num_splits: int. If > 1, split the key/value into this many chunks along the sequence. + If num_splits == 1, we don't split the key/value. If num_splits == 0, we use a heuristic + to automatically determine the number of splits. + Don't change this unless you know what you are doing. + return_softmax_lse: bool. Whether to return the logsumexp of the attention scores. + + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_softmax_lse=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + """ + assert k_cache.stride(-1) == 1, "k_cache must have contiguous last dimension" + assert v_cache.stride(-1) == 1, "v_cache must have contiguous last dimension" + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + if cache_seqlens is not None and isinstance(cache_seqlens, int): + cache_seqlens = torch.full( + (k_cache.shape[0],), cache_seqlens, dtype=torch.int32, device=k_cache.device + ) + cache_seqlens = maybe_contiguous(cache_seqlens) + out, softmax_lse, *rest = _flash_attn_forward( + q, + k_cache, + v_cache, + k, + v, + qv, + None, # out + cu_seqlens_q, + None, # cu_seqlens_k + cu_seqlens_k_new, + None, # seqused_q + cache_seqlens, + max_seqlen_q, + None, # max_seqlen_k + page_table, + cache_batch_idx, + cache_leftpad, + rotary_cos, + rotary_sin, + rotary_seqlens, + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size=window_size, + attention_chunk=attention_chunk, + softcap=softcap, + rotary_interleaved=rotary_interleaved, + scheduler_metadata=scheduler_metadata, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # return (out, softmax_lse) if return_softmax_lse else out + return (out, softmax_lse, *rest) if return_softmax_lse else out + + +def get_scheduler_metadata( + batch_size, max_seqlen_q, max_seqlen_k, num_heads_q, num_heads_kv, headdim, + cache_seqlens: torch.Tensor, + qkv_dtype=torch.bfloat16, + headdim_v=None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + cache_leftpad: Optional[torch.Tensor] = None, + page_size: Optional[int] = None, + max_seqlen_k_new=0, + causal=False, + window_size=(-1, -1), # -1 means infinite context window + attention_chunk=0, + has_softcap=False, + num_splits=0, # Can be tuned for speed + pack_gqa=None, # Can be tuned for speed + sm_margin=0, # Can be tuned if some SMs are used for communication +): + cache_seqlens = maybe_contiguous(cache_seqlens) + if headdim_v is None: + headdim_v = headdim + scheduler_metadata = flash_attn_3_cuda.get_scheduler_metadata( + batch_size, max_seqlen_q, max_seqlen_k, num_heads_q, num_heads_kv, headdim, headdim_v, + qkv_dtype, + cache_seqlens, + cu_seqlens_q, + None, # cu_seqlens_k + cu_seqlens_k_new, + None, # seqused_q + cache_leftpad, + page_size, + max_seqlen_k_new, + causal, + window_size[0], window_size[1], + attention_chunk, + has_softcap, + num_splits, + pack_gqa, + sm_margin, + ) + return scheduler_metadata diff --git a/build/torch29-cxx11-cu128-x86_64-linux/__init__.py b/build/torch29-cxx11-cu128-x86_64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3d12edd447e84d6e65c4fa9ab91ae3c3192884c1 --- /dev/null +++ b/build/torch29-cxx11-cu128-x86_64-linux/__init__.py @@ -0,0 +1,17 @@ +from .flash_attn_interface import ( + flash_attn_combine, + flash_attn_func, + flash_attn_qkvpacked_func, + flash_attn_varlen_func, + flash_attn_with_kvcache, + get_scheduler_metadata, +) + +__all__ = [ + "flash_attn_combine", + "flash_attn_func", + "flash_attn_qkvpacked_func", + "flash_attn_varlen_func", + "flash_attn_with_kvcache", + "get_scheduler_metadata", +] diff --git a/build/torch29-cxx11-cu128-x86_64-linux/_flash_attn3_cuda_2d990fe.abi3.so b/build/torch29-cxx11-cu128-x86_64-linux/_flash_attn3_cuda_2d990fe.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..0b6dbe3ced41484bd5dc1229b4a0012063640dd1 --- /dev/null +++ b/build/torch29-cxx11-cu128-x86_64-linux/_flash_attn3_cuda_2d990fe.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:385c5960f738f6c2f217ca93a46fb9be003e7076d60a5711cef8bebdf19eee20 +size 804184680 diff --git a/build/torch29-cxx11-cu128-x86_64-linux/_ops.py b/build/torch29-cxx11-cu128-x86_64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..1d9ac17ca94ff06dfccb0a48aeefbdf75e28043c --- /dev/null +++ b/build/torch29-cxx11-cu128-x86_64-linux/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _flash_attn3_cuda_2d990fe +ops = torch.ops._flash_attn3_cuda_2d990fe + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_flash_attn3_cuda_2d990fe::{op_name}" diff --git a/build/torch29-cxx11-cu128-x86_64-linux/flash_attn3/__init__.py b/build/torch29-cxx11-cu128-x86_64-linux/flash_attn3/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..03dbc1afe1cf156661a2b1b22003cd5f599a0309 --- /dev/null +++ b/build/torch29-cxx11-cu128-x86_64-linux/flash_attn3/__init__.py @@ -0,0 +1,26 @@ +import ctypes +import sys + +import importlib +from pathlib import Path +from types import ModuleType + +def _import_from_path(file_path: Path) -> ModuleType: + # We cannot use the module name as-is, after adding it to `sys.modules`, + # it would also be used for other imports. So, we make a module name that + # depends on the path for it to be unique using the hex-encoded hash of + # the path. + path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value) + module_name = path_hash + spec = importlib.util.spec_from_file_location(module_name, file_path) + if spec is None: + raise ImportError(f"Cannot load spec for {module_name} from {file_path}") + module = importlib.util.module_from_spec(spec) + if module is None: + raise ImportError(f"Cannot load module {module_name} from spec") + sys.modules[module_name] = module + spec.loader.exec_module(module) # type: ignore + return module + + +globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py"))) diff --git a/build/torch29-cxx11-cu128-x86_64-linux/flash_attn_config.py b/build/torch29-cxx11-cu128-x86_64-linux/flash_attn_config.py new file mode 100644 index 0000000000000000000000000000000000000000..c4b7130324be2c4cabee7ada7000e170a2a81865 --- /dev/null +++ b/build/torch29-cxx11-cu128-x86_64-linux/flash_attn_config.py @@ -0,0 +1,7 @@ +# Auto-generated by flash attention 3 setup.py +CONFIG = {'build_flags': {'FLASHATTENTION_DISABLE_BACKWARD': False, 'FLASHATTENTION_DISABLE_SPLIT': False, 'FLASHATTENTION_DISABLE_PAGEDKV': False, 'FLASHATTENTION_DISABLE_APPENDKV': False, 'FLASHATTENTION_DISABLE_LOCAL': False, 'FLASHATTENTION_DISABLE_SOFTCAP': False, 'FLASHATTENTION_DISABLE_PACKGQA': False, 'FLASHATTENTION_DISABLE_FP16': False, 'FLASHATTENTION_DISABLE_FP8': False, 'FLASHATTENTION_DISABLE_VARLEN': False, 'FLASHATTENTION_DISABLE_CLUSTER': False, 'FLASHATTENTION_DISABLE_HDIM64': False, 'FLASHATTENTION_DISABLE_HDIM96': False, 'FLASHATTENTION_DISABLE_HDIM128': False, 'FLASHATTENTION_DISABLE_HDIM192': False, 'FLASHATTENTION_DISABLE_HDIM256': False, 'FLASHATTENTION_DISABLE_SM8x': False, 'FLASHATTENTION_ENABLE_VCOLMAJOR': False, 'FLASH_ATTENTION_DISABLE_HDIMDIFF64': False, 'FLASH_ATTENTION_DISABLE_HDIMDIFF192': False}} + +def show(): + from pprint import pprint + pprint(CONFIG) + diff --git a/build/torch29-cxx11-cu128-x86_64-linux/flash_attn_interface.py b/build/torch29-cxx11-cu128-x86_64-linux/flash_attn_interface.py new file mode 100644 index 0000000000000000000000000000000000000000..2e2a3cdd3e89cbc59bf6c956454e2776c305e0e5 --- /dev/null +++ b/build/torch29-cxx11-cu128-x86_64-linux/flash_attn_interface.py @@ -0,0 +1,1127 @@ +# Copyright (c) 2023, Tri Dao. + +from typing import Optional, Union, List, Tuple + +import torch +import torch.nn as nn + +from ._ops import ops as flash_attn_3_cuda +from ._ops import add_op_namespace_prefix + +def maybe_contiguous(x): + return x.contiguous() if x is not None and x.stride(-1) != 1 else x + + +def round_multiple(x, m): + return (x + m - 1) // m * m + + +def round_up_headdim(head_size: int) -> int: + from .flash_attn_config import CONFIG + + if not CONFIG["build_flags"]["FLASHATTENTION_DISABLE_HDIM64"]: + if head_size <= 64: + return 64 + if not CONFIG["build_flags"]["FLASHATTENTION_DISABLE_HDIM96"]: + if head_size <= 96: + return 96 + if not CONFIG["build_flags"]["FLASHATTENTION_DISABLE_HDIM128"]: + if head_size <= 128: + return 128 + if not CONFIG["build_flags"]["FLASHATTENTION_DISABLE_HDIM192"]: + if head_size <= 192: + return 192 + if not CONFIG["build_flags"]["FLASHATTENTION_DISABLE_HDIM256"]: + if head_size <= 256: + return 256 + return 256 + + +@torch.library.custom_op(add_op_namespace_prefix("_flash_attn_forward"), mutates_args=(), device_types="cuda") +def _flash_attn_forward( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_new: Optional[torch.Tensor] = None, + v_new: Optional[torch.Tensor] = None, + qv: Optional[torch.Tensor] = None, + out_: Optional[torch.Tensor] = None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + seqused_q: Optional[torch.Tensor] = None, + seqused_k: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_k: Optional[int] = None, + page_table: Optional[torch.Tensor] = None, + kv_batch_idx: Optional[torch.Tensor] = None, + leftpad_k: Optional[torch.Tensor] = None, + rotary_cos: Optional[torch.Tensor] = None, + rotary_sin: Optional[torch.Tensor] = None, + seqlens_rotary: Optional[torch.Tensor] = None, + q_descale: Optional[torch.Tensor] = None, + k_descale: Optional[torch.Tensor] = None, + v_descale: Optional[torch.Tensor] = None, + softmax_scale: Optional[float] = None, + causal: bool = False, + window_size_left: int = -1, + window_size_right: int = -1, + attention_chunk: int = 0, + softcap: float = 0.0, + rotary_interleaved: bool = True, + scheduler_metadata: Optional[torch.Tensor] = None, + num_splits: int = 1, + pack_gqa: Optional[bool] = None, + sm_margin: int = 0, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + q, k, k_new, v_new = [maybe_contiguous(x) for x in (q, k, k_new, v_new)] + v = v.contiguous() if v.stride(-1) != 1 and v.stride(-3) != 1 else v + cu_seqlens_q, cu_seqlens_k, cu_seqlens_k_new = [ + maybe_contiguous(x) for x in (cu_seqlens_q, cu_seqlens_k, cu_seqlens_k_new) + ] + seqused_q, seqused_k = [maybe_contiguous(x) for x in (seqused_q, seqused_k)] + page_table, kv_batch_idx, leftpad_k = [ + maybe_contiguous(x) for x in (page_table, kv_batch_idx, leftpad_k) + ] + rotary_cos, rotary_sin = [maybe_contiguous(x) for x in (rotary_cos, rotary_sin)] + seqlens_rotary = maybe_contiguous(seqlens_rotary) + out, softmax_lse, out_accum, softmax_lse_accum = flash_attn_3_cuda.fwd( + q, + k, + v, + k_new, + v_new, + qv, + out_, + cu_seqlens_q, + cu_seqlens_k, + cu_seqlens_k_new, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + page_table, + kv_batch_idx, + leftpad_k, + rotary_cos, + rotary_sin, + seqlens_rotary, + q_descale, + k_descale, + v_descale, + softmax_scale, + causal, + window_size_left, + window_size_right, + attention_chunk, + softcap, + rotary_interleaved, + scheduler_metadata, + num_splits, + pack_gqa, + sm_margin, + ) + + if out_accum is None: + out_accum = torch.tensor([], device=out.device) + + if softmax_lse_accum is None: + softmax_lse_accum = torch.tensor([], device=out.device) + + return out, softmax_lse, out_accum, softmax_lse_accum + + +@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_forward")) +def _flash_attn_forward_fake( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_new: Optional[torch.Tensor] = None, + v_new: Optional[torch.Tensor] = None, + qv: Optional[torch.Tensor] = None, + out_: Optional[torch.Tensor] = None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + seqused_q: Optional[torch.Tensor] = None, + seqused_k: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_k: Optional[int] = None, + page_table: Optional[torch.Tensor] = None, + kv_batch_idx: Optional[torch.Tensor] = None, + leftpad_k: Optional[torch.Tensor] = None, + rotary_cos: Optional[torch.Tensor] = None, + rotary_sin: Optional[torch.Tensor] = None, + seqlens_rotary: Optional[torch.Tensor] = None, + q_descale: Optional[torch.Tensor] = None, + k_descale: Optional[torch.Tensor] = None, + v_descale: Optional[torch.Tensor] = None, + softmax_scale: Optional[float] = None, + causal: bool = False, + window_size_left: int = -1, + window_size_right: int = -1, + attention_chunk: int = 0, + softcap: float = 0.0, + rotary_interleaved: bool = True, + scheduler_metadata: Optional[torch.Tensor] = None, + num_splits: int = 1, + pack_gqa: Optional[bool] = None, + sm_margin: int = 0, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Symbolic fake implementation of flash attention forward. + Returns tensors with the correct shapes and dtypes without actual computation. + """ + + # Determine if we're in varlen mode + is_varlen_q = cu_seqlens_q is not None + + # Get dimensions from query tensor + if is_varlen_q: + # varlen mode: q is (total_q, num_heads, head_size) + total_q, num_heads, head_size = q.shape + batch_size = cu_seqlens_q.shape[0] - 1 + + if max_seqlen_q is None: + raise ValueError("max_seqlen_q must be provided if cu_seqlens_q is provided") + seqlen_q = max_seqlen_q + else: + # batch mode: q is (batch_size, seqlen_q, num_heads, head_size) + batch_size, seqlen_q, num_heads, head_size = q.shape + total_q = batch_size * q.shape[1] + # Get value head dimension + head_size_v = v.shape[-1] + + # Determine output dtype (FP8 inputs produce BF16 outputs) + q_type = q.dtype + if q_type == torch.float8_e4m3fn: + out_dtype = torch.bfloat16 + else: + out_dtype = q_type + + # Create output tensor + if out_ is not None: + # If out_ is provided, _flash_attn_forward becomes non-functional + raise TypeError("Tracing (torch.compile/torch.export) with pre-allocated output tensor is not supported.") + + if is_varlen_q: + out = torch.empty((total_q, num_heads, head_size_v), dtype=out_dtype, device=q.device) + else: + out = torch.empty((batch_size, seqlen_q, num_heads, head_size_v), dtype=out_dtype, device=q.device) + + # Create softmax_lse tensor + if is_varlen_q: + softmax_lse = torch.empty((num_heads, total_q), dtype=torch.float32, device=q.device) + else: + softmax_lse = torch.empty((batch_size, num_heads, seqlen_q), dtype=torch.float32, device=q.device) + + # TODO(guilhermeleobas): Implement "get_num_splits" + # There's an heuristic to compute num_splits when "num_splits <= 0" + # assert that num_splits is > 0 for now + if num_splits <= 0: + raise ValueError(f"tracing (torch.compile/torch.export) with num_splits <= 0 not supported. Got {num_splits=}") + + if num_splits > 1: + if is_varlen_q: + out_accum = torch.empty((num_splits, num_heads, total_q, head_size_v), dtype=torch.float32, device=q.device) + softmax_lse_accum = torch.empty((num_splits, num_heads, total_q), dtype=torch.float32, device=q.device) + else: + out_accum = torch.empty((num_splits, batch_size, num_heads, seqlen_q, head_size_v), dtype=torch.float32, device=q.device) + softmax_lse_accum = torch.empty((num_splits, batch_size, num_heads, seqlen_q), dtype=torch.float32, device=q.device) + else: + # Tensors are not set when num_splits < 1 + out_accum = torch.tensor([], device=out.device) + softmax_lse_accum = torch.tensor([], device=out.device) + + return out, softmax_lse, out_accum, softmax_lse_accum + + +@torch.library.custom_op(add_op_namespace_prefix("_flash_attn_backward"), mutates_args=("dq", "dk", "dv"), device_types="cuda") +def _flash_attn_backward( + dout: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + out: torch.Tensor, + softmax_lse: torch.Tensor, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k: Optional[torch.Tensor] = None, + sequed_q: Optional[torch.Tensor] = None, + sequed_k: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_k: Optional[int] = None, + dq: Optional[torch.Tensor] = None, + dk: Optional[torch.Tensor] = None, + dv: Optional[torch.Tensor] = None, + softmax_scale: Optional[float] = None, + is_causal: bool = False, + window_size_left: int = -1, + window_size_right: int = -1, + softcap: float = 0.0, + deterministic: bool = False, + sm_margin: int = 0, +) -> torch.Tensor: + # dq, dk, dv are allocated by us so they should already be contiguous + dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] + softmax_d, *rest = flash_attn_3_cuda.bwd( + dout, + q, + k, + v, + out, + softmax_lse, + dq, + dk, + dv, + cu_seqlens_q, + cu_seqlens_k, + sequed_q, + sequed_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + is_causal, + window_size_left, + window_size_right, + softcap, + deterministic, + sm_margin, + ) + return softmax_d + + +@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_backward")) +def _flash_attn_backward_fake( + dout: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + out: torch.Tensor, + softmax_lse: torch.Tensor, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k: Optional[torch.Tensor] = None, + sequed_q: Optional[torch.Tensor] = None, + sequed_k: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_k: Optional[int] = None, + dq: Optional[torch.Tensor] = None, + dk: Optional[torch.Tensor] = None, + dv: Optional[torch.Tensor] = None, + softmax_scale: Optional[float] = None, + is_causal: bool = False, + window_size_left: int = -1, + window_size_right: int = -1, + softcap: float = 0.0, + deterministic: bool = False, + sm_margin: int = 0, +) -> torch.Tensor: + + is_varlen_q = cu_seqlens_q is not None + is_varlen_k = cu_seqlens_q is not None + is_varlen = is_varlen_q or is_varlen_k or sequed_q is not None or sequed_k is not None + + if not is_varlen_q: + batch_size = q.size(0) + seqlen_q = q.size(1) + seqlen_k = k.size(1) + total_q = batch_size * q.size(1) + else: + batch_size = cu_seqlens_q.size(0) - 1 + total_q = q.size(0) + seqlen_q = max_seqlen_q + seqlen_k = max_seqlen_k + + if window_size_left >= seqlen_k - 1: + window_size_left = -1 + + if window_size_right >= seqlen_q - 1: + window_size_right = -1 + + if is_causal: + window_size_right = 0 + + is_causal = window_size_left < 0 and window_size_right == 0 + + head_size = q.size(-1) + head_size_v = v.size(-1) + head_size_rounded = round_up_headdim(max(head_size, head_size_v)) + + # Hopper gpus uses cuda compute capabilities 9.0 + cap = torch.cuda.get_device_capability(q.device) + arch = cap[0] * 10 + cap[1] + + is_local = (window_size_left >= 0 or window_size_right >= 0) and not is_causal + + if head_size_rounded <= 64: + kBlockM_sm90 = 96 if (is_causal and softcap > 0.0) else 128 + elif head_size_rounded <= 96: + kBlockM_sm90 = 64 + elif head_size_rounded <= 128: + kBlockM_sm90 = 64 if (is_causal or is_local or softcap > 0.0) else 80 + else: + kBlockM_sm90 = 64 + + kBlockM_sm80 = 128 if head_size_rounded <= 64 else 64 + kBlockM_sm86 = 64 if head_size_rounded <= 192 else 32 + + if arch >= 90: + kBlockM = kBlockM_sm90 + elif arch == 86 or arch == 89: + kBlockM = kBlockM_sm86 + else: + kBlockM = kBlockM_sm80 + + num_heads = q.shape[-2] + seqlen_q_rounded = round_multiple(seqlen_q, kBlockM) + + total_q_padded_rounded = round_multiple(total_q + batch_size * kBlockM, kBlockM) + + dq = torch.empty_like(q) if dq is None else dq + dk = torch.empty_like(k) if dk is None else dk + dv = torch.empty_like(v) if dv is None else dv + + if not is_varlen: + softmax_d = torch.empty((batch_size, num_heads, seqlen_q_rounded), dtype=torch.float32, device=q.device) + else: + softmax_d = torch.empty((num_heads, total_q_padded_rounded), dtype=torch.float32, device=q.device) + + return softmax_d + + +def setup_context(ctx, inputs, output): + q, k, v = inputs[:3] + out, softmax_lse, _, _ = output + ctx.save_for_backward(q, k, v, out, softmax_lse) + ctx.softmax_scale = inputs[-11] + ctx.causal = inputs[-10] + ctx.window_size = [inputs[-9], inputs[-8]] + ctx.attention_chunk = inputs[-7] + ctx.softcap = inputs[-6] + ctx.sm_margin = inputs[-1] + + +def _backward(ctx, dout, *grads): + q, k, v, out, softmax_lse = ctx.saved_tensors + dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + None, None, # cu_seqlens_q, cu_seqlens_k, + None, None, # sequed_q, sequed_k, + None, None, # max_seqlen_q, max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size[0], + ctx.window_size[1], + ctx.softcap, + False, # deterministic + ctx.sm_margin, + ) + return dq, dk, dv, *((None,) * 21) + + +_flash_attn_forward.register_autograd(_backward, setup_context=setup_context) + + + +class FlashAttnQKVPackedFunc(torch.autograd.Function): + @staticmethod + def forward( + ctx, + qkv, + softmax_scale, + causal, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + deterministic=False, + num_heads_q=None, + sm_margin=0, + return_softmax=False, + ): + if softmax_scale is None: + softmax_scale = qkv.shape[-1] ** (-0.5) + if qkv.dim() == 5: + assert qkv.shape[-3] == 3 + q, k, v = qkv.unbind(dim=-3) + else: + assert qkv.dim() == 4 + assert num_heads_q is not None + num_heads_k = (qkv.shape[2] - num_heads_q) // 2 + assert num_heads_k * 2 + num_heads_q == qkv.shape[2] + q, k, v = qkv.split([num_heads_q, num_heads_k, num_heads_k], dim=-2) + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + None, # qv + None, # out + None, None, None, # cu_seqlens_q/k/k_new + None, None, # seqused_q/k + None, None, # max_seqlen_q/k + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size_left=window_size[0], + window_size_right=window_size[1], + attention_chunk=attention_chunk, + softcap=softcap, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse) + ctx.save_for_backward(q, k, v, out, softmax_lse) + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.ndim = qkv.dim() + ctx.sm_margin = sm_margin + return (out, softmax_lse) if return_softmax else out + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + if ctx.ndim == 5: + qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) + dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) + dq, dk, dv = dqkv.unbind(dim=-3) + else: + num_heads_q = q.shape[2] + num_heads_k = k.shape[2] + qkv_shape = q.shape[:-2] + (num_heads_q + num_heads_k * 2, *q.shape[-1:]) + dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) + dq, dk, dv = dqkv.split([num_heads_q, num_heads_k, num_heads_k], dim=-2) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + None, None, # cu_seqlens_q, cu_seqlens_k, + None, None, # sequed_q, sequed_k, + None, None, # max_seqlen_q, max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size[0], + ctx.window_size[1], + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dqkv = dqkv[..., : dout.shape[-1]] # We could have padded the head dimension + return dqkv, None, None, None, None, None, None, None, None, None, None, None, None + + +class FlashAttnFunc(torch.autograd.Function): + + @staticmethod + def forward( + ctx, + q, + k, + v, + softmax_scale, + causal, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + return_softmax=False, + ): + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + # out, q, k, v, out_padded, softmax_lse = _flash_attn_forward( + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + qv, # qv + None, # out + None, None, None, # cu_seqlens_q/k/k_new + None, None, # seqused_q/k + None, None, # max_seqlen_q/k + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size_left=window_size[0], + window_size_right=window_size[1], + attention_chunk=attention_chunk, + softcap=softcap, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse) + ctx.save_for_backward(q, k, v, out, softmax_lse) + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.sm_margin = sm_margin + return (out, softmax_lse) if return_softmax else out + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + None, None, # cu_seqlens_q, cu_seqlens_k, + None, None, # sequed_q, sequed_k, + None, None, # max_seqlen_q, max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size[0], + ctx.window_size[1], + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dq = dq[..., : q.shape[-1]] # We could have padded the head dimension + dk = dk[..., : k.shape[-1]] + dv = dv[..., : v.shape[-1]] + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None + + +class FlashAttnVarlenFunc(torch.autograd.Function): + + @staticmethod + def forward( + ctx, + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + return_softmax=False, + ): + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + # out, q, k, v, out_padded, softmax_lse = _flash_attn_varlen_forward( + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + qv, # qv + None, # out + cu_seqlens_q, + cu_seqlens_k, + None, # cu_seqlens_k_new + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size_left=window_size[0], + window_size_right=window_size[1], + attention_chunk=attention_chunk, + softcap=softcap, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k) + ctx.save_for_backward(q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k) + ctx.max_seqlen_q = max_seqlen_q + ctx.max_seqlen_k = max_seqlen_k + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.sm_margin = sm_margin + return (out, softmax_lse) if return_softmax else out + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + ctx.max_seqlen_q, + ctx.max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size[0], + ctx.window_size[1], + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dq = dq[..., : q.shape[-1]] # We could have padded the head dimension + dk = dk[..., : k.shape[-1]] + dv = dv[..., : v.shape[-1]] + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None + + +def flash_attn_qkvpacked_func( + qkv, + softmax_scale=None, + causal=False, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + deterministic=False, + num_heads_q=None, + sm_margin=0, + return_attn_probs=False, +): + """dropout_p should be set to 0.0 during evaluation + If Q, K, V are already stacked into 1 tensor, this function will be faster than + calling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation + of the gradients of Q, K, V. + For multi-query and grouped-query attention (MQA/GQA), please see + flash_attn_kvpacked_func and flash_attn_func. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between [i - window_size[0], i + window_size[1]] inclusive. + + Arguments: + qkv: (batch_size, seqlen, 3, nheads, headdim) + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + softcap: float. Anything > 0 activates softcapping attention. + alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of (-alibi_slope * |i - j|) is added to + the attention score of query i and key j. + deterministic: bool. Whether to use the deterministic implementation of the backward pass, + which is slightly slower and uses more memory. The forward pass is always deterministic. + return_attn_probs: bool. Whether to return the attention probabilities. This option is for + testing only. The returned probabilities are not guaranteed to be correct + (they might not have the right scaling). + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen). + The output of softmax (possibly with different scaling). It also encodes the dropout + pattern (negative means that location was dropped, nonnegative means it was kept). + """ + return FlashAttnQKVPackedFunc.apply( + qkv, + softmax_scale, + causal, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + deterministic, + num_heads_q, + sm_margin, + return_attn_probs, + ) + + +def flash_attn_func( + q, + k, + v, + softmax_scale=None, + causal=False, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + return_attn_probs=False, +): + """dropout_p should be set to 0.0 during evaluation + Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads + than Q. Note that the number of heads in Q must be divisible by the number of heads in KV. + For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head + 0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V. + + If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix. + For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is: + 1 1 1 1 0 + 1 1 1 1 1 + If seqlen_q = 5 and seqlen_k = 2, the causal mask is: + 0 0 + 0 0 + 0 0 + 1 0 + 1 1 + If the row of the mask is all zero, the output will be zero. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between + [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive. + + Arguments: + q: (batch_size, seqlen, nheads, headdim) + k: (batch_size, seqlen, nheads_k, headdim) + v: (batch_size, seqlen, nheads_k, headdim) + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of + (-alibi_slope * |i + seqlen_k - seqlen_q - j|) + is added to the attention score of query i and key j. + deterministic: bool. Whether to use the deterministic implementation of the backward pass, + which is slightly slower and uses more memory. The forward pass is always deterministic. + return_attn_probs: bool. Whether to return the attention probabilities. This option is for + testing only. The returned probabilities are not guaranteed to be correct + (they might not have the right scaling). + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + """ + return FlashAttnFunc.apply( + q, + k, + v, + softmax_scale, + causal, + qv, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + num_splits, + pack_gqa, + deterministic, + sm_margin, + return_attn_probs, + ) + + +def flash_attn_varlen_func( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + seqused_q=None, + seqused_k=None, + softmax_scale=None, + causal=False, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + return_attn_probs=False, +): + return FlashAttnVarlenFunc.apply( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + qv, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + num_splits, + pack_gqa, + deterministic, + sm_margin, + return_attn_probs, + ) + + +def flash_attn_combine(out_partial, lse_partial, out=None, out_dtype=None): + return flash_attn_3_cuda.fwd_combine(out_partial, lse_partial, out, out_dtype) + + +def flash_attn_with_kvcache( + q, + k_cache, + v_cache, + k=None, + v=None, + qv=None, + rotary_cos=None, + rotary_sin=None, + cache_seqlens: Optional[Union[(int, torch.Tensor)]] = None, + cache_batch_idx: Optional[torch.Tensor] = None, + cache_leftpad: Optional[torch.Tensor] = None, + page_table: Optional[torch.Tensor] = None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + rotary_seqlens: Optional[torch.Tensor] = None, + q_descale: Optional[torch.Tensor] = None, + k_descale: Optional[torch.Tensor] = None, + v_descale: Optional[torch.Tensor] = None, + softmax_scale=None, + causal=False, + window_size=(-1, -1), # -1 means infinite context window + attention_chunk=0, + softcap=0.0, # 0.0 means deactivated + rotary_interleaved=True, + scheduler_metadata=None, + num_splits=0, # Can be tuned for speed + pack_gqa=None, # Can be tuned for speed + sm_margin=0, # Can be tuned if some SMs are used for communication + return_softmax_lse=False, +): + """ + If k and v are not None, k_cache and v_cache will be updated *inplace* with the new values from + k and v. This is useful for incremental decoding: you can pass in the cached keys/values from + the previous step, and update them with the new keys/values from the current step, and do + attention with the updated cache, all in 1 kernel. + + If you pass in k / v, you must make sure that the cache is large enough to hold the new values. + For example, the KV cache could be pre-allocated with the max sequence length, and you can use + cache_seqlens to keep track of the current sequence lengths of each sequence in the batch. + + Also apply rotary embedding if rotary_cos and rotary_sin are passed in. The key @k will be + rotated by rotary_cos and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc. + If causal or local (i.e., window_size != (-1, -1)), the query @q will be rotated by rotary_cos + and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc. + If not causal and not local, the query @q will be rotated by rotary_cos and rotary_sin at + indices cache_seqlens only (i.e. we consider all tokens in @q to be at position cache_seqlens). + + See tests/test_flash_attn.py::test_flash_attn_kvcache for examples of how to use this function. + + Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads + than Q. Note that the number of heads in Q must be divisible by the number of heads in KV. + For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head + 0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V. + + If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix. + For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is: + 1 1 1 1 0 + 1 1 1 1 1 + If seqlen_q = 5 and seqlen_k = 2, the causal mask is: + 0 0 + 0 0 + 0 0 + 1 0 + 1 1 + If the row of the mask is all zero, the output will be zero. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between + [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive. + + Note: Does not support backward pass. + + Arguments: + q: (batch_size, seqlen, nheads, headdim) + k_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim) if there's no page_table, + or (num_blocks, page_block_size, nheads_k, headdim) if there's a page_table (i.e. paged KV cache) + page_block_size can be arbitrary (e.g, 1, 2, 3, 64, etc.). + v_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim_v) if there's no page_table, + or (num_blocks, page_block_size, nheads_k, headdim_v) if there's a page_table (i.e. paged KV cache) + k [optional]: (batch_size, seqlen_new, nheads_k, headdim). If not None, we concatenate + k with k_cache, starting at the indices specified by cache_seqlens. + v [optional]: (batch_size, seqlen_new, nheads_k, headdim_v). Similar to k. + qv [optional]: (batch_size, seqlen, nheads, headdim_v) + rotary_cos [optional]: (seqlen_ro, rotary_dim / 2). If not None, we apply rotary embedding + to k and q. Only applicable if k and v are passed in. rotary_dim must be divisible by 16. + rotary_sin [optional]: (seqlen_ro, rotary_dim / 2). Similar to rotary_cos. + cache_seqlens: int, or (batch_size,), dtype torch.int32. The sequence lengths of the + KV cache. + cache_batch_idx: (batch_size,), dtype torch.int32. The indices used to index into the KV cache. + If None, we assume that the batch indices are [0, 1, 2, ..., batch_size - 1]. + If the indices are not distinct, and k and v are provided, the values updated in the cache + might come from any of the duplicate indices. + cache_leftpad: (batch_size,), dtype torch.int32. The index that the KV cache starts. If None, assume 0. + page_table [optional]: (batch_size, max_num_blocks_per_seq), dtype torch.int32. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + softcap: float. Anything > 0 activates softcapping attention. + rotary_interleaved: bool. Only applicable if rotary_cos and rotary_sin are passed in. + If True, rotary embedding will combine dimensions 0 & 1, 2 & 3, etc. If False, + rotary embedding will combine dimensions 0 & rotary_dim / 2, 1 & rotary_dim / 2 + 1 + (i.e. GPT-NeoX style). + num_splits: int. If > 1, split the key/value into this many chunks along the sequence. + If num_splits == 1, we don't split the key/value. If num_splits == 0, we use a heuristic + to automatically determine the number of splits. + Don't change this unless you know what you are doing. + return_softmax_lse: bool. Whether to return the logsumexp of the attention scores. + + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_softmax_lse=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + """ + assert k_cache.stride(-1) == 1, "k_cache must have contiguous last dimension" + assert v_cache.stride(-1) == 1, "v_cache must have contiguous last dimension" + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + if cache_seqlens is not None and isinstance(cache_seqlens, int): + cache_seqlens = torch.full( + (q.shape[0],), cache_seqlens, dtype=torch.int32, device=k_cache.device + ) + cache_seqlens = maybe_contiguous(cache_seqlens) + out, softmax_lse, *rest = _flash_attn_forward( + q, + k_cache, + v_cache, + k, + v, + qv, + None, # out + cu_seqlens_q, + None, # cu_seqlens_k + cu_seqlens_k_new, + None, # seqused_q + cache_seqlens, + max_seqlen_q, + None, # max_seqlen_k + page_table, + cache_batch_idx, + cache_leftpad, + rotary_cos, + rotary_sin, + rotary_seqlens, + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size_left=window_size[0], + window_size_right=window_size[1], + attention_chunk=attention_chunk, + softcap=softcap, + rotary_interleaved=rotary_interleaved, + scheduler_metadata=scheduler_metadata, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # return (out, softmax_lse) if return_softmax_lse else out + return (out, softmax_lse, *rest) if return_softmax_lse else out + + +def get_scheduler_metadata( + batch_size, max_seqlen_q, max_seqlen_k, num_heads_q, num_heads_kv, headdim, + cache_seqlens: torch.Tensor, + qkv_dtype=torch.bfloat16, + headdim_v=None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + cache_leftpad: Optional[torch.Tensor] = None, + page_size: Optional[int] = None, + max_seqlen_k_new=0, + causal=False, + window_size=(-1, -1), # -1 means infinite context window + attention_chunk=0, + has_softcap=False, + num_splits=0, # Can be tuned for speed + pack_gqa=None, # Can be tuned for speed + sm_margin=0, # Can be tuned if some SMs are used for communication +): + cache_seqlens = maybe_contiguous(cache_seqlens) + if headdim_v is None: + headdim_v = headdim + scheduler_metadata = flash_attn_3_cuda.get_scheduler_metadata( + batch_size, max_seqlen_q, max_seqlen_k, num_heads_q, num_heads_kv, headdim, headdim_v, + qkv_dtype, + cache_seqlens, + cu_seqlens_q, + None, # cu_seqlens_k + cu_seqlens_k_new, + None, # seqused_q + cache_leftpad, + page_size, + max_seqlen_k_new, + causal, + window_size[0], window_size[1], + attention_chunk, + has_softcap, + num_splits, + pack_gqa, + sm_margin, + ) + return scheduler_metadata diff --git a/build/torch29-cxx11-cu128-x86_64-linux/metadata.json b/build/torch29-cxx11-cu128-x86_64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..e0cc0c4f14abd10162c22499be39c75e6d922f94 --- /dev/null +++ b/build/torch29-cxx11-cu128-x86_64-linux/metadata.json @@ -0,0 +1,12 @@ +{ + "version": 1, + "license": "BSD-3-Clause", + "python-depends": [], + "backend": { + "type": "cuda", + "archs": [ + "8.0", + "9.0a" + ] + } +} diff --git a/build/torch29-cxx11-cu129-x86_64-linux/__init__.py b/build/torch29-cxx11-cu129-x86_64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3d12edd447e84d6e65c4fa9ab91ae3c3192884c1 --- /dev/null +++ b/build/torch29-cxx11-cu129-x86_64-linux/__init__.py @@ -0,0 +1,17 @@ +from .flash_attn_interface import ( + flash_attn_combine, + flash_attn_func, + flash_attn_qkvpacked_func, + flash_attn_varlen_func, + flash_attn_with_kvcache, + get_scheduler_metadata, +) + +__all__ = [ + "flash_attn_combine", + "flash_attn_func", + "flash_attn_qkvpacked_func", + "flash_attn_varlen_func", + "flash_attn_with_kvcache", + "get_scheduler_metadata", +] diff --git a/build/torch29-cxx11-cu129-x86_64-linux/_flash_attn3_cuda_e1d5be2.abi3.so b/build/torch29-cxx11-cu129-x86_64-linux/_flash_attn3_cuda_e1d5be2.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..fa7922b5c457759add8903b25bb189a745c73be6 --- /dev/null +++ b/build/torch29-cxx11-cu129-x86_64-linux/_flash_attn3_cuda_e1d5be2.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:4ce0c933aa2a1e19b7d439dbe25049bf32eadd53636ba53b9a55370720cced97 +size 804425064 diff --git a/build/torch29-cxx11-cu129-x86_64-linux/_ops.py b/build/torch29-cxx11-cu129-x86_64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..b845f529a7b7c424bbe0db01a32c2cd6db45bb1a --- /dev/null +++ b/build/torch29-cxx11-cu129-x86_64-linux/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _flash_attn3_cuda_e1d5be2 +ops = torch.ops._flash_attn3_cuda_e1d5be2 + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_flash_attn3_cuda_e1d5be2::{op_name}" diff --git a/build/torch29-cxx11-cu129-x86_64-linux/flash_attn3/__init__.py b/build/torch29-cxx11-cu129-x86_64-linux/flash_attn3/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a9b2672c1cd85b74c1b3ded0fc0b2100e1aeac23 --- /dev/null +++ b/build/torch29-cxx11-cu129-x86_64-linux/flash_attn3/__init__.py @@ -0,0 +1,26 @@ +import ctypes +import importlib.util +import sys +from pathlib import Path +from types import ModuleType + + +def _import_from_path(file_path: Path) -> ModuleType: + # We cannot use the module name as-is, after adding it to `sys.modules`, + # it would also be used for other imports. So, we make a module name that + # depends on the path for it to be unique using the hex-encoded hash of + # the path. + path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value) + module_name = path_hash + spec = importlib.util.spec_from_file_location(module_name, file_path) + if spec is None: + raise ImportError(f"Cannot load spec for {module_name} from {file_path}") + module = importlib.util.module_from_spec(spec) + if module is None: + raise ImportError(f"Cannot load module {module_name} from spec") + sys.modules[module_name] = module + spec.loader.exec_module(module) # type: ignore + return module + + +globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py"))) diff --git a/build/torch29-cxx11-cu129-x86_64-linux/flash_attn_config.py b/build/torch29-cxx11-cu129-x86_64-linux/flash_attn_config.py new file mode 100644 index 0000000000000000000000000000000000000000..c4b7130324be2c4cabee7ada7000e170a2a81865 --- /dev/null +++ b/build/torch29-cxx11-cu129-x86_64-linux/flash_attn_config.py @@ -0,0 +1,7 @@ +# Auto-generated by flash attention 3 setup.py +CONFIG = {'build_flags': {'FLASHATTENTION_DISABLE_BACKWARD': False, 'FLASHATTENTION_DISABLE_SPLIT': False, 'FLASHATTENTION_DISABLE_PAGEDKV': False, 'FLASHATTENTION_DISABLE_APPENDKV': False, 'FLASHATTENTION_DISABLE_LOCAL': False, 'FLASHATTENTION_DISABLE_SOFTCAP': False, 'FLASHATTENTION_DISABLE_PACKGQA': False, 'FLASHATTENTION_DISABLE_FP16': False, 'FLASHATTENTION_DISABLE_FP8': False, 'FLASHATTENTION_DISABLE_VARLEN': False, 'FLASHATTENTION_DISABLE_CLUSTER': False, 'FLASHATTENTION_DISABLE_HDIM64': False, 'FLASHATTENTION_DISABLE_HDIM96': False, 'FLASHATTENTION_DISABLE_HDIM128': False, 'FLASHATTENTION_DISABLE_HDIM192': False, 'FLASHATTENTION_DISABLE_HDIM256': False, 'FLASHATTENTION_DISABLE_SM8x': False, 'FLASHATTENTION_ENABLE_VCOLMAJOR': False, 'FLASH_ATTENTION_DISABLE_HDIMDIFF64': False, 'FLASH_ATTENTION_DISABLE_HDIMDIFF192': False}} + +def show(): + from pprint import pprint + pprint(CONFIG) + diff --git a/build/torch29-cxx11-cu129-x86_64-linux/flash_attn_interface.py b/build/torch29-cxx11-cu129-x86_64-linux/flash_attn_interface.py new file mode 100644 index 0000000000000000000000000000000000000000..2e2a3cdd3e89cbc59bf6c956454e2776c305e0e5 --- /dev/null +++ b/build/torch29-cxx11-cu129-x86_64-linux/flash_attn_interface.py @@ -0,0 +1,1127 @@ +# Copyright (c) 2023, Tri Dao. + +from typing import Optional, Union, List, Tuple + +import torch +import torch.nn as nn + +from ._ops import ops as flash_attn_3_cuda +from ._ops import add_op_namespace_prefix + +def maybe_contiguous(x): + return x.contiguous() if x is not None and x.stride(-1) != 1 else x + + +def round_multiple(x, m): + return (x + m - 1) // m * m + + +def round_up_headdim(head_size: int) -> int: + from .flash_attn_config import CONFIG + + if not CONFIG["build_flags"]["FLASHATTENTION_DISABLE_HDIM64"]: + if head_size <= 64: + return 64 + if not CONFIG["build_flags"]["FLASHATTENTION_DISABLE_HDIM96"]: + if head_size <= 96: + return 96 + if not CONFIG["build_flags"]["FLASHATTENTION_DISABLE_HDIM128"]: + if head_size <= 128: + return 128 + if not CONFIG["build_flags"]["FLASHATTENTION_DISABLE_HDIM192"]: + if head_size <= 192: + return 192 + if not CONFIG["build_flags"]["FLASHATTENTION_DISABLE_HDIM256"]: + if head_size <= 256: + return 256 + return 256 + + +@torch.library.custom_op(add_op_namespace_prefix("_flash_attn_forward"), mutates_args=(), device_types="cuda") +def _flash_attn_forward( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_new: Optional[torch.Tensor] = None, + v_new: Optional[torch.Tensor] = None, + qv: Optional[torch.Tensor] = None, + out_: Optional[torch.Tensor] = None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + seqused_q: Optional[torch.Tensor] = None, + seqused_k: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_k: Optional[int] = None, + page_table: Optional[torch.Tensor] = None, + kv_batch_idx: Optional[torch.Tensor] = None, + leftpad_k: Optional[torch.Tensor] = None, + rotary_cos: Optional[torch.Tensor] = None, + rotary_sin: Optional[torch.Tensor] = None, + seqlens_rotary: Optional[torch.Tensor] = None, + q_descale: Optional[torch.Tensor] = None, + k_descale: Optional[torch.Tensor] = None, + v_descale: Optional[torch.Tensor] = None, + softmax_scale: Optional[float] = None, + causal: bool = False, + window_size_left: int = -1, + window_size_right: int = -1, + attention_chunk: int = 0, + softcap: float = 0.0, + rotary_interleaved: bool = True, + scheduler_metadata: Optional[torch.Tensor] = None, + num_splits: int = 1, + pack_gqa: Optional[bool] = None, + sm_margin: int = 0, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + q, k, k_new, v_new = [maybe_contiguous(x) for x in (q, k, k_new, v_new)] + v = v.contiguous() if v.stride(-1) != 1 and v.stride(-3) != 1 else v + cu_seqlens_q, cu_seqlens_k, cu_seqlens_k_new = [ + maybe_contiguous(x) for x in (cu_seqlens_q, cu_seqlens_k, cu_seqlens_k_new) + ] + seqused_q, seqused_k = [maybe_contiguous(x) for x in (seqused_q, seqused_k)] + page_table, kv_batch_idx, leftpad_k = [ + maybe_contiguous(x) for x in (page_table, kv_batch_idx, leftpad_k) + ] + rotary_cos, rotary_sin = [maybe_contiguous(x) for x in (rotary_cos, rotary_sin)] + seqlens_rotary = maybe_contiguous(seqlens_rotary) + out, softmax_lse, out_accum, softmax_lse_accum = flash_attn_3_cuda.fwd( + q, + k, + v, + k_new, + v_new, + qv, + out_, + cu_seqlens_q, + cu_seqlens_k, + cu_seqlens_k_new, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + page_table, + kv_batch_idx, + leftpad_k, + rotary_cos, + rotary_sin, + seqlens_rotary, + q_descale, + k_descale, + v_descale, + softmax_scale, + causal, + window_size_left, + window_size_right, + attention_chunk, + softcap, + rotary_interleaved, + scheduler_metadata, + num_splits, + pack_gqa, + sm_margin, + ) + + if out_accum is None: + out_accum = torch.tensor([], device=out.device) + + if softmax_lse_accum is None: + softmax_lse_accum = torch.tensor([], device=out.device) + + return out, softmax_lse, out_accum, softmax_lse_accum + + +@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_forward")) +def _flash_attn_forward_fake( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_new: Optional[torch.Tensor] = None, + v_new: Optional[torch.Tensor] = None, + qv: Optional[torch.Tensor] = None, + out_: Optional[torch.Tensor] = None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + seqused_q: Optional[torch.Tensor] = None, + seqused_k: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_k: Optional[int] = None, + page_table: Optional[torch.Tensor] = None, + kv_batch_idx: Optional[torch.Tensor] = None, + leftpad_k: Optional[torch.Tensor] = None, + rotary_cos: Optional[torch.Tensor] = None, + rotary_sin: Optional[torch.Tensor] = None, + seqlens_rotary: Optional[torch.Tensor] = None, + q_descale: Optional[torch.Tensor] = None, + k_descale: Optional[torch.Tensor] = None, + v_descale: Optional[torch.Tensor] = None, + softmax_scale: Optional[float] = None, + causal: bool = False, + window_size_left: int = -1, + window_size_right: int = -1, + attention_chunk: int = 0, + softcap: float = 0.0, + rotary_interleaved: bool = True, + scheduler_metadata: Optional[torch.Tensor] = None, + num_splits: int = 1, + pack_gqa: Optional[bool] = None, + sm_margin: int = 0, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Symbolic fake implementation of flash attention forward. + Returns tensors with the correct shapes and dtypes without actual computation. + """ + + # Determine if we're in varlen mode + is_varlen_q = cu_seqlens_q is not None + + # Get dimensions from query tensor + if is_varlen_q: + # varlen mode: q is (total_q, num_heads, head_size) + total_q, num_heads, head_size = q.shape + batch_size = cu_seqlens_q.shape[0] - 1 + + if max_seqlen_q is None: + raise ValueError("max_seqlen_q must be provided if cu_seqlens_q is provided") + seqlen_q = max_seqlen_q + else: + # batch mode: q is (batch_size, seqlen_q, num_heads, head_size) + batch_size, seqlen_q, num_heads, head_size = q.shape + total_q = batch_size * q.shape[1] + # Get value head dimension + head_size_v = v.shape[-1] + + # Determine output dtype (FP8 inputs produce BF16 outputs) + q_type = q.dtype + if q_type == torch.float8_e4m3fn: + out_dtype = torch.bfloat16 + else: + out_dtype = q_type + + # Create output tensor + if out_ is not None: + # If out_ is provided, _flash_attn_forward becomes non-functional + raise TypeError("Tracing (torch.compile/torch.export) with pre-allocated output tensor is not supported.") + + if is_varlen_q: + out = torch.empty((total_q, num_heads, head_size_v), dtype=out_dtype, device=q.device) + else: + out = torch.empty((batch_size, seqlen_q, num_heads, head_size_v), dtype=out_dtype, device=q.device) + + # Create softmax_lse tensor + if is_varlen_q: + softmax_lse = torch.empty((num_heads, total_q), dtype=torch.float32, device=q.device) + else: + softmax_lse = torch.empty((batch_size, num_heads, seqlen_q), dtype=torch.float32, device=q.device) + + # TODO(guilhermeleobas): Implement "get_num_splits" + # There's an heuristic to compute num_splits when "num_splits <= 0" + # assert that num_splits is > 0 for now + if num_splits <= 0: + raise ValueError(f"tracing (torch.compile/torch.export) with num_splits <= 0 not supported. Got {num_splits=}") + + if num_splits > 1: + if is_varlen_q: + out_accum = torch.empty((num_splits, num_heads, total_q, head_size_v), dtype=torch.float32, device=q.device) + softmax_lse_accum = torch.empty((num_splits, num_heads, total_q), dtype=torch.float32, device=q.device) + else: + out_accum = torch.empty((num_splits, batch_size, num_heads, seqlen_q, head_size_v), dtype=torch.float32, device=q.device) + softmax_lse_accum = torch.empty((num_splits, batch_size, num_heads, seqlen_q), dtype=torch.float32, device=q.device) + else: + # Tensors are not set when num_splits < 1 + out_accum = torch.tensor([], device=out.device) + softmax_lse_accum = torch.tensor([], device=out.device) + + return out, softmax_lse, out_accum, softmax_lse_accum + + +@torch.library.custom_op(add_op_namespace_prefix("_flash_attn_backward"), mutates_args=("dq", "dk", "dv"), device_types="cuda") +def _flash_attn_backward( + dout: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + out: torch.Tensor, + softmax_lse: torch.Tensor, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k: Optional[torch.Tensor] = None, + sequed_q: Optional[torch.Tensor] = None, + sequed_k: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_k: Optional[int] = None, + dq: Optional[torch.Tensor] = None, + dk: Optional[torch.Tensor] = None, + dv: Optional[torch.Tensor] = None, + softmax_scale: Optional[float] = None, + is_causal: bool = False, + window_size_left: int = -1, + window_size_right: int = -1, + softcap: float = 0.0, + deterministic: bool = False, + sm_margin: int = 0, +) -> torch.Tensor: + # dq, dk, dv are allocated by us so they should already be contiguous + dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] + softmax_d, *rest = flash_attn_3_cuda.bwd( + dout, + q, + k, + v, + out, + softmax_lse, + dq, + dk, + dv, + cu_seqlens_q, + cu_seqlens_k, + sequed_q, + sequed_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + is_causal, + window_size_left, + window_size_right, + softcap, + deterministic, + sm_margin, + ) + return softmax_d + + +@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_backward")) +def _flash_attn_backward_fake( + dout: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + out: torch.Tensor, + softmax_lse: torch.Tensor, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k: Optional[torch.Tensor] = None, + sequed_q: Optional[torch.Tensor] = None, + sequed_k: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_k: Optional[int] = None, + dq: Optional[torch.Tensor] = None, + dk: Optional[torch.Tensor] = None, + dv: Optional[torch.Tensor] = None, + softmax_scale: Optional[float] = None, + is_causal: bool = False, + window_size_left: int = -1, + window_size_right: int = -1, + softcap: float = 0.0, + deterministic: bool = False, + sm_margin: int = 0, +) -> torch.Tensor: + + is_varlen_q = cu_seqlens_q is not None + is_varlen_k = cu_seqlens_q is not None + is_varlen = is_varlen_q or is_varlen_k or sequed_q is not None or sequed_k is not None + + if not is_varlen_q: + batch_size = q.size(0) + seqlen_q = q.size(1) + seqlen_k = k.size(1) + total_q = batch_size * q.size(1) + else: + batch_size = cu_seqlens_q.size(0) - 1 + total_q = q.size(0) + seqlen_q = max_seqlen_q + seqlen_k = max_seqlen_k + + if window_size_left >= seqlen_k - 1: + window_size_left = -1 + + if window_size_right >= seqlen_q - 1: + window_size_right = -1 + + if is_causal: + window_size_right = 0 + + is_causal = window_size_left < 0 and window_size_right == 0 + + head_size = q.size(-1) + head_size_v = v.size(-1) + head_size_rounded = round_up_headdim(max(head_size, head_size_v)) + + # Hopper gpus uses cuda compute capabilities 9.0 + cap = torch.cuda.get_device_capability(q.device) + arch = cap[0] * 10 + cap[1] + + is_local = (window_size_left >= 0 or window_size_right >= 0) and not is_causal + + if head_size_rounded <= 64: + kBlockM_sm90 = 96 if (is_causal and softcap > 0.0) else 128 + elif head_size_rounded <= 96: + kBlockM_sm90 = 64 + elif head_size_rounded <= 128: + kBlockM_sm90 = 64 if (is_causal or is_local or softcap > 0.0) else 80 + else: + kBlockM_sm90 = 64 + + kBlockM_sm80 = 128 if head_size_rounded <= 64 else 64 + kBlockM_sm86 = 64 if head_size_rounded <= 192 else 32 + + if arch >= 90: + kBlockM = kBlockM_sm90 + elif arch == 86 or arch == 89: + kBlockM = kBlockM_sm86 + else: + kBlockM = kBlockM_sm80 + + num_heads = q.shape[-2] + seqlen_q_rounded = round_multiple(seqlen_q, kBlockM) + + total_q_padded_rounded = round_multiple(total_q + batch_size * kBlockM, kBlockM) + + dq = torch.empty_like(q) if dq is None else dq + dk = torch.empty_like(k) if dk is None else dk + dv = torch.empty_like(v) if dv is None else dv + + if not is_varlen: + softmax_d = torch.empty((batch_size, num_heads, seqlen_q_rounded), dtype=torch.float32, device=q.device) + else: + softmax_d = torch.empty((num_heads, total_q_padded_rounded), dtype=torch.float32, device=q.device) + + return softmax_d + + +def setup_context(ctx, inputs, output): + q, k, v = inputs[:3] + out, softmax_lse, _, _ = output + ctx.save_for_backward(q, k, v, out, softmax_lse) + ctx.softmax_scale = inputs[-11] + ctx.causal = inputs[-10] + ctx.window_size = [inputs[-9], inputs[-8]] + ctx.attention_chunk = inputs[-7] + ctx.softcap = inputs[-6] + ctx.sm_margin = inputs[-1] + + +def _backward(ctx, dout, *grads): + q, k, v, out, softmax_lse = ctx.saved_tensors + dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + None, None, # cu_seqlens_q, cu_seqlens_k, + None, None, # sequed_q, sequed_k, + None, None, # max_seqlen_q, max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size[0], + ctx.window_size[1], + ctx.softcap, + False, # deterministic + ctx.sm_margin, + ) + return dq, dk, dv, *((None,) * 21) + + +_flash_attn_forward.register_autograd(_backward, setup_context=setup_context) + + + +class FlashAttnQKVPackedFunc(torch.autograd.Function): + @staticmethod + def forward( + ctx, + qkv, + softmax_scale, + causal, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + deterministic=False, + num_heads_q=None, + sm_margin=0, + return_softmax=False, + ): + if softmax_scale is None: + softmax_scale = qkv.shape[-1] ** (-0.5) + if qkv.dim() == 5: + assert qkv.shape[-3] == 3 + q, k, v = qkv.unbind(dim=-3) + else: + assert qkv.dim() == 4 + assert num_heads_q is not None + num_heads_k = (qkv.shape[2] - num_heads_q) // 2 + assert num_heads_k * 2 + num_heads_q == qkv.shape[2] + q, k, v = qkv.split([num_heads_q, num_heads_k, num_heads_k], dim=-2) + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + None, # qv + None, # out + None, None, None, # cu_seqlens_q/k/k_new + None, None, # seqused_q/k + None, None, # max_seqlen_q/k + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size_left=window_size[0], + window_size_right=window_size[1], + attention_chunk=attention_chunk, + softcap=softcap, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse) + ctx.save_for_backward(q, k, v, out, softmax_lse) + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.ndim = qkv.dim() + ctx.sm_margin = sm_margin + return (out, softmax_lse) if return_softmax else out + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + if ctx.ndim == 5: + qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) + dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) + dq, dk, dv = dqkv.unbind(dim=-3) + else: + num_heads_q = q.shape[2] + num_heads_k = k.shape[2] + qkv_shape = q.shape[:-2] + (num_heads_q + num_heads_k * 2, *q.shape[-1:]) + dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) + dq, dk, dv = dqkv.split([num_heads_q, num_heads_k, num_heads_k], dim=-2) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + None, None, # cu_seqlens_q, cu_seqlens_k, + None, None, # sequed_q, sequed_k, + None, None, # max_seqlen_q, max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size[0], + ctx.window_size[1], + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dqkv = dqkv[..., : dout.shape[-1]] # We could have padded the head dimension + return dqkv, None, None, None, None, None, None, None, None, None, None, None, None + + +class FlashAttnFunc(torch.autograd.Function): + + @staticmethod + def forward( + ctx, + q, + k, + v, + softmax_scale, + causal, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + return_softmax=False, + ): + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + # out, q, k, v, out_padded, softmax_lse = _flash_attn_forward( + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + qv, # qv + None, # out + None, None, None, # cu_seqlens_q/k/k_new + None, None, # seqused_q/k + None, None, # max_seqlen_q/k + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size_left=window_size[0], + window_size_right=window_size[1], + attention_chunk=attention_chunk, + softcap=softcap, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse) + ctx.save_for_backward(q, k, v, out, softmax_lse) + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.sm_margin = sm_margin + return (out, softmax_lse) if return_softmax else out + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + None, None, # cu_seqlens_q, cu_seqlens_k, + None, None, # sequed_q, sequed_k, + None, None, # max_seqlen_q, max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size[0], + ctx.window_size[1], + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dq = dq[..., : q.shape[-1]] # We could have padded the head dimension + dk = dk[..., : k.shape[-1]] + dv = dv[..., : v.shape[-1]] + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None + + +class FlashAttnVarlenFunc(torch.autograd.Function): + + @staticmethod + def forward( + ctx, + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + return_softmax=False, + ): + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + # out, q, k, v, out_padded, softmax_lse = _flash_attn_varlen_forward( + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + qv, # qv + None, # out + cu_seqlens_q, + cu_seqlens_k, + None, # cu_seqlens_k_new + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size_left=window_size[0], + window_size_right=window_size[1], + attention_chunk=attention_chunk, + softcap=softcap, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k) + ctx.save_for_backward(q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k) + ctx.max_seqlen_q = max_seqlen_q + ctx.max_seqlen_k = max_seqlen_k + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.sm_margin = sm_margin + return (out, softmax_lse) if return_softmax else out + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + ctx.max_seqlen_q, + ctx.max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size[0], + ctx.window_size[1], + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dq = dq[..., : q.shape[-1]] # We could have padded the head dimension + dk = dk[..., : k.shape[-1]] + dv = dv[..., : v.shape[-1]] + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None + + +def flash_attn_qkvpacked_func( + qkv, + softmax_scale=None, + causal=False, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + deterministic=False, + num_heads_q=None, + sm_margin=0, + return_attn_probs=False, +): + """dropout_p should be set to 0.0 during evaluation + If Q, K, V are already stacked into 1 tensor, this function will be faster than + calling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation + of the gradients of Q, K, V. + For multi-query and grouped-query attention (MQA/GQA), please see + flash_attn_kvpacked_func and flash_attn_func. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between [i - window_size[0], i + window_size[1]] inclusive. + + Arguments: + qkv: (batch_size, seqlen, 3, nheads, headdim) + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + softcap: float. Anything > 0 activates softcapping attention. + alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of (-alibi_slope * |i - j|) is added to + the attention score of query i and key j. + deterministic: bool. Whether to use the deterministic implementation of the backward pass, + which is slightly slower and uses more memory. The forward pass is always deterministic. + return_attn_probs: bool. Whether to return the attention probabilities. This option is for + testing only. The returned probabilities are not guaranteed to be correct + (they might not have the right scaling). + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen). + The output of softmax (possibly with different scaling). It also encodes the dropout + pattern (negative means that location was dropped, nonnegative means it was kept). + """ + return FlashAttnQKVPackedFunc.apply( + qkv, + softmax_scale, + causal, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + deterministic, + num_heads_q, + sm_margin, + return_attn_probs, + ) + + +def flash_attn_func( + q, + k, + v, + softmax_scale=None, + causal=False, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + return_attn_probs=False, +): + """dropout_p should be set to 0.0 during evaluation + Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads + than Q. Note that the number of heads in Q must be divisible by the number of heads in KV. + For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head + 0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V. + + If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix. + For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is: + 1 1 1 1 0 + 1 1 1 1 1 + If seqlen_q = 5 and seqlen_k = 2, the causal mask is: + 0 0 + 0 0 + 0 0 + 1 0 + 1 1 + If the row of the mask is all zero, the output will be zero. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between + [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive. + + Arguments: + q: (batch_size, seqlen, nheads, headdim) + k: (batch_size, seqlen, nheads_k, headdim) + v: (batch_size, seqlen, nheads_k, headdim) + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of + (-alibi_slope * |i + seqlen_k - seqlen_q - j|) + is added to the attention score of query i and key j. + deterministic: bool. Whether to use the deterministic implementation of the backward pass, + which is slightly slower and uses more memory. The forward pass is always deterministic. + return_attn_probs: bool. Whether to return the attention probabilities. This option is for + testing only. The returned probabilities are not guaranteed to be correct + (they might not have the right scaling). + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + """ + return FlashAttnFunc.apply( + q, + k, + v, + softmax_scale, + causal, + qv, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + num_splits, + pack_gqa, + deterministic, + sm_margin, + return_attn_probs, + ) + + +def flash_attn_varlen_func( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + seqused_q=None, + seqused_k=None, + softmax_scale=None, + causal=False, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + return_attn_probs=False, +): + return FlashAttnVarlenFunc.apply( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + qv, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + num_splits, + pack_gqa, + deterministic, + sm_margin, + return_attn_probs, + ) + + +def flash_attn_combine(out_partial, lse_partial, out=None, out_dtype=None): + return flash_attn_3_cuda.fwd_combine(out_partial, lse_partial, out, out_dtype) + + +def flash_attn_with_kvcache( + q, + k_cache, + v_cache, + k=None, + v=None, + qv=None, + rotary_cos=None, + rotary_sin=None, + cache_seqlens: Optional[Union[(int, torch.Tensor)]] = None, + cache_batch_idx: Optional[torch.Tensor] = None, + cache_leftpad: Optional[torch.Tensor] = None, + page_table: Optional[torch.Tensor] = None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + rotary_seqlens: Optional[torch.Tensor] = None, + q_descale: Optional[torch.Tensor] = None, + k_descale: Optional[torch.Tensor] = None, + v_descale: Optional[torch.Tensor] = None, + softmax_scale=None, + causal=False, + window_size=(-1, -1), # -1 means infinite context window + attention_chunk=0, + softcap=0.0, # 0.0 means deactivated + rotary_interleaved=True, + scheduler_metadata=None, + num_splits=0, # Can be tuned for speed + pack_gqa=None, # Can be tuned for speed + sm_margin=0, # Can be tuned if some SMs are used for communication + return_softmax_lse=False, +): + """ + If k and v are not None, k_cache and v_cache will be updated *inplace* with the new values from + k and v. This is useful for incremental decoding: you can pass in the cached keys/values from + the previous step, and update them with the new keys/values from the current step, and do + attention with the updated cache, all in 1 kernel. + + If you pass in k / v, you must make sure that the cache is large enough to hold the new values. + For example, the KV cache could be pre-allocated with the max sequence length, and you can use + cache_seqlens to keep track of the current sequence lengths of each sequence in the batch. + + Also apply rotary embedding if rotary_cos and rotary_sin are passed in. The key @k will be + rotated by rotary_cos and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc. + If causal or local (i.e., window_size != (-1, -1)), the query @q will be rotated by rotary_cos + and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc. + If not causal and not local, the query @q will be rotated by rotary_cos and rotary_sin at + indices cache_seqlens only (i.e. we consider all tokens in @q to be at position cache_seqlens). + + See tests/test_flash_attn.py::test_flash_attn_kvcache for examples of how to use this function. + + Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads + than Q. Note that the number of heads in Q must be divisible by the number of heads in KV. + For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head + 0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V. + + If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix. + For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is: + 1 1 1 1 0 + 1 1 1 1 1 + If seqlen_q = 5 and seqlen_k = 2, the causal mask is: + 0 0 + 0 0 + 0 0 + 1 0 + 1 1 + If the row of the mask is all zero, the output will be zero. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between + [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive. + + Note: Does not support backward pass. + + Arguments: + q: (batch_size, seqlen, nheads, headdim) + k_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim) if there's no page_table, + or (num_blocks, page_block_size, nheads_k, headdim) if there's a page_table (i.e. paged KV cache) + page_block_size can be arbitrary (e.g, 1, 2, 3, 64, etc.). + v_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim_v) if there's no page_table, + or (num_blocks, page_block_size, nheads_k, headdim_v) if there's a page_table (i.e. paged KV cache) + k [optional]: (batch_size, seqlen_new, nheads_k, headdim). If not None, we concatenate + k with k_cache, starting at the indices specified by cache_seqlens. + v [optional]: (batch_size, seqlen_new, nheads_k, headdim_v). Similar to k. + qv [optional]: (batch_size, seqlen, nheads, headdim_v) + rotary_cos [optional]: (seqlen_ro, rotary_dim / 2). If not None, we apply rotary embedding + to k and q. Only applicable if k and v are passed in. rotary_dim must be divisible by 16. + rotary_sin [optional]: (seqlen_ro, rotary_dim / 2). Similar to rotary_cos. + cache_seqlens: int, or (batch_size,), dtype torch.int32. The sequence lengths of the + KV cache. + cache_batch_idx: (batch_size,), dtype torch.int32. The indices used to index into the KV cache. + If None, we assume that the batch indices are [0, 1, 2, ..., batch_size - 1]. + If the indices are not distinct, and k and v are provided, the values updated in the cache + might come from any of the duplicate indices. + cache_leftpad: (batch_size,), dtype torch.int32. The index that the KV cache starts. If None, assume 0. + page_table [optional]: (batch_size, max_num_blocks_per_seq), dtype torch.int32. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + softcap: float. Anything > 0 activates softcapping attention. + rotary_interleaved: bool. Only applicable if rotary_cos and rotary_sin are passed in. + If True, rotary embedding will combine dimensions 0 & 1, 2 & 3, etc. If False, + rotary embedding will combine dimensions 0 & rotary_dim / 2, 1 & rotary_dim / 2 + 1 + (i.e. GPT-NeoX style). + num_splits: int. If > 1, split the key/value into this many chunks along the sequence. + If num_splits == 1, we don't split the key/value. If num_splits == 0, we use a heuristic + to automatically determine the number of splits. + Don't change this unless you know what you are doing. + return_softmax_lse: bool. Whether to return the logsumexp of the attention scores. + + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_softmax_lse=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + """ + assert k_cache.stride(-1) == 1, "k_cache must have contiguous last dimension" + assert v_cache.stride(-1) == 1, "v_cache must have contiguous last dimension" + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + if cache_seqlens is not None and isinstance(cache_seqlens, int): + cache_seqlens = torch.full( + (q.shape[0],), cache_seqlens, dtype=torch.int32, device=k_cache.device + ) + cache_seqlens = maybe_contiguous(cache_seqlens) + out, softmax_lse, *rest = _flash_attn_forward( + q, + k_cache, + v_cache, + k, + v, + qv, + None, # out + cu_seqlens_q, + None, # cu_seqlens_k + cu_seqlens_k_new, + None, # seqused_q + cache_seqlens, + max_seqlen_q, + None, # max_seqlen_k + page_table, + cache_batch_idx, + cache_leftpad, + rotary_cos, + rotary_sin, + rotary_seqlens, + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size_left=window_size[0], + window_size_right=window_size[1], + attention_chunk=attention_chunk, + softcap=softcap, + rotary_interleaved=rotary_interleaved, + scheduler_metadata=scheduler_metadata, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # return (out, softmax_lse) if return_softmax_lse else out + return (out, softmax_lse, *rest) if return_softmax_lse else out + + +def get_scheduler_metadata( + batch_size, max_seqlen_q, max_seqlen_k, num_heads_q, num_heads_kv, headdim, + cache_seqlens: torch.Tensor, + qkv_dtype=torch.bfloat16, + headdim_v=None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + cache_leftpad: Optional[torch.Tensor] = None, + page_size: Optional[int] = None, + max_seqlen_k_new=0, + causal=False, + window_size=(-1, -1), # -1 means infinite context window + attention_chunk=0, + has_softcap=False, + num_splits=0, # Can be tuned for speed + pack_gqa=None, # Can be tuned for speed + sm_margin=0, # Can be tuned if some SMs are used for communication +): + cache_seqlens = maybe_contiguous(cache_seqlens) + if headdim_v is None: + headdim_v = headdim + scheduler_metadata = flash_attn_3_cuda.get_scheduler_metadata( + batch_size, max_seqlen_q, max_seqlen_k, num_heads_q, num_heads_kv, headdim, headdim_v, + qkv_dtype, + cache_seqlens, + cu_seqlens_q, + None, # cu_seqlens_k + cu_seqlens_k_new, + None, # seqused_q + cache_leftpad, + page_size, + max_seqlen_k_new, + causal, + window_size[0], window_size[1], + attention_chunk, + has_softcap, + num_splits, + pack_gqa, + sm_margin, + ) + return scheduler_metadata diff --git a/build/torch29-cxx11-cu129-x86_64-linux/metadata.json b/build/torch29-cxx11-cu129-x86_64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..e0cc0c4f14abd10162c22499be39c75e6d922f94 --- /dev/null +++ b/build/torch29-cxx11-cu129-x86_64-linux/metadata.json @@ -0,0 +1,12 @@ +{ + "version": 1, + "license": "BSD-3-Clause", + "python-depends": [], + "backend": { + "type": "cuda", + "archs": [ + "8.0", + "9.0a" + ] + } +} diff --git a/build/torch29-cxx11-cu130-x86_64-linux/__init__.py b/build/torch29-cxx11-cu130-x86_64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3d12edd447e84d6e65c4fa9ab91ae3c3192884c1 --- /dev/null +++ b/build/torch29-cxx11-cu130-x86_64-linux/__init__.py @@ -0,0 +1,17 @@ +from .flash_attn_interface import ( + flash_attn_combine, + flash_attn_func, + flash_attn_qkvpacked_func, + flash_attn_varlen_func, + flash_attn_with_kvcache, + get_scheduler_metadata, +) + +__all__ = [ + "flash_attn_combine", + "flash_attn_func", + "flash_attn_qkvpacked_func", + "flash_attn_varlen_func", + "flash_attn_with_kvcache", + "get_scheduler_metadata", +] diff --git a/build/torch29-cxx11-cu130-x86_64-linux/_flash_attn3_cuda_2d990fe.abi3.so b/build/torch29-cxx11-cu130-x86_64-linux/_flash_attn3_cuda_2d990fe.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..d0ea0faf0dbd8ee8024b05533fc95732bd86173f --- /dev/null +++ b/build/torch29-cxx11-cu130-x86_64-linux/_flash_attn3_cuda_2d990fe.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:40dbed6283e41653bd8a7b050636856fa9c5c03ddb1ba24a5ea20ea7aadc08d5 +size 823692904 diff --git a/build/torch29-cxx11-cu130-x86_64-linux/_ops.py b/build/torch29-cxx11-cu130-x86_64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..1d9ac17ca94ff06dfccb0a48aeefbdf75e28043c --- /dev/null +++ b/build/torch29-cxx11-cu130-x86_64-linux/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _flash_attn3_cuda_2d990fe +ops = torch.ops._flash_attn3_cuda_2d990fe + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_flash_attn3_cuda_2d990fe::{op_name}" diff --git a/build/torch29-cxx11-cu130-x86_64-linux/flash_attn3/__init__.py b/build/torch29-cxx11-cu130-x86_64-linux/flash_attn3/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..03dbc1afe1cf156661a2b1b22003cd5f599a0309 --- /dev/null +++ b/build/torch29-cxx11-cu130-x86_64-linux/flash_attn3/__init__.py @@ -0,0 +1,26 @@ +import ctypes +import sys + +import importlib +from pathlib import Path +from types import ModuleType + +def _import_from_path(file_path: Path) -> ModuleType: + # We cannot use the module name as-is, after adding it to `sys.modules`, + # it would also be used for other imports. So, we make a module name that + # depends on the path for it to be unique using the hex-encoded hash of + # the path. + path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value) + module_name = path_hash + spec = importlib.util.spec_from_file_location(module_name, file_path) + if spec is None: + raise ImportError(f"Cannot load spec for {module_name} from {file_path}") + module = importlib.util.module_from_spec(spec) + if module is None: + raise ImportError(f"Cannot load module {module_name} from spec") + sys.modules[module_name] = module + spec.loader.exec_module(module) # type: ignore + return module + + +globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py"))) diff --git a/build/torch29-cxx11-cu130-x86_64-linux/flash_attn_config.py b/build/torch29-cxx11-cu130-x86_64-linux/flash_attn_config.py new file mode 100644 index 0000000000000000000000000000000000000000..c4b7130324be2c4cabee7ada7000e170a2a81865 --- /dev/null +++ b/build/torch29-cxx11-cu130-x86_64-linux/flash_attn_config.py @@ -0,0 +1,7 @@ +# Auto-generated by flash attention 3 setup.py +CONFIG = {'build_flags': {'FLASHATTENTION_DISABLE_BACKWARD': False, 'FLASHATTENTION_DISABLE_SPLIT': False, 'FLASHATTENTION_DISABLE_PAGEDKV': False, 'FLASHATTENTION_DISABLE_APPENDKV': False, 'FLASHATTENTION_DISABLE_LOCAL': False, 'FLASHATTENTION_DISABLE_SOFTCAP': False, 'FLASHATTENTION_DISABLE_PACKGQA': False, 'FLASHATTENTION_DISABLE_FP16': False, 'FLASHATTENTION_DISABLE_FP8': False, 'FLASHATTENTION_DISABLE_VARLEN': False, 'FLASHATTENTION_DISABLE_CLUSTER': False, 'FLASHATTENTION_DISABLE_HDIM64': False, 'FLASHATTENTION_DISABLE_HDIM96': False, 'FLASHATTENTION_DISABLE_HDIM128': False, 'FLASHATTENTION_DISABLE_HDIM192': False, 'FLASHATTENTION_DISABLE_HDIM256': False, 'FLASHATTENTION_DISABLE_SM8x': False, 'FLASHATTENTION_ENABLE_VCOLMAJOR': False, 'FLASH_ATTENTION_DISABLE_HDIMDIFF64': False, 'FLASH_ATTENTION_DISABLE_HDIMDIFF192': False}} + +def show(): + from pprint import pprint + pprint(CONFIG) + diff --git a/build/torch29-cxx11-cu130-x86_64-linux/flash_attn_interface.py b/build/torch29-cxx11-cu130-x86_64-linux/flash_attn_interface.py new file mode 100644 index 0000000000000000000000000000000000000000..2e2a3cdd3e89cbc59bf6c956454e2776c305e0e5 --- /dev/null +++ b/build/torch29-cxx11-cu130-x86_64-linux/flash_attn_interface.py @@ -0,0 +1,1127 @@ +# Copyright (c) 2023, Tri Dao. + +from typing import Optional, Union, List, Tuple + +import torch +import torch.nn as nn + +from ._ops import ops as flash_attn_3_cuda +from ._ops import add_op_namespace_prefix + +def maybe_contiguous(x): + return x.contiguous() if x is not None and x.stride(-1) != 1 else x + + +def round_multiple(x, m): + return (x + m - 1) // m * m + + +def round_up_headdim(head_size: int) -> int: + from .flash_attn_config import CONFIG + + if not CONFIG["build_flags"]["FLASHATTENTION_DISABLE_HDIM64"]: + if head_size <= 64: + return 64 + if not CONFIG["build_flags"]["FLASHATTENTION_DISABLE_HDIM96"]: + if head_size <= 96: + return 96 + if not CONFIG["build_flags"]["FLASHATTENTION_DISABLE_HDIM128"]: + if head_size <= 128: + return 128 + if not CONFIG["build_flags"]["FLASHATTENTION_DISABLE_HDIM192"]: + if head_size <= 192: + return 192 + if not CONFIG["build_flags"]["FLASHATTENTION_DISABLE_HDIM256"]: + if head_size <= 256: + return 256 + return 256 + + +@torch.library.custom_op(add_op_namespace_prefix("_flash_attn_forward"), mutates_args=(), device_types="cuda") +def _flash_attn_forward( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_new: Optional[torch.Tensor] = None, + v_new: Optional[torch.Tensor] = None, + qv: Optional[torch.Tensor] = None, + out_: Optional[torch.Tensor] = None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + seqused_q: Optional[torch.Tensor] = None, + seqused_k: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_k: Optional[int] = None, + page_table: Optional[torch.Tensor] = None, + kv_batch_idx: Optional[torch.Tensor] = None, + leftpad_k: Optional[torch.Tensor] = None, + rotary_cos: Optional[torch.Tensor] = None, + rotary_sin: Optional[torch.Tensor] = None, + seqlens_rotary: Optional[torch.Tensor] = None, + q_descale: Optional[torch.Tensor] = None, + k_descale: Optional[torch.Tensor] = None, + v_descale: Optional[torch.Tensor] = None, + softmax_scale: Optional[float] = None, + causal: bool = False, + window_size_left: int = -1, + window_size_right: int = -1, + attention_chunk: int = 0, + softcap: float = 0.0, + rotary_interleaved: bool = True, + scheduler_metadata: Optional[torch.Tensor] = None, + num_splits: int = 1, + pack_gqa: Optional[bool] = None, + sm_margin: int = 0, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + q, k, k_new, v_new = [maybe_contiguous(x) for x in (q, k, k_new, v_new)] + v = v.contiguous() if v.stride(-1) != 1 and v.stride(-3) != 1 else v + cu_seqlens_q, cu_seqlens_k, cu_seqlens_k_new = [ + maybe_contiguous(x) for x in (cu_seqlens_q, cu_seqlens_k, cu_seqlens_k_new) + ] + seqused_q, seqused_k = [maybe_contiguous(x) for x in (seqused_q, seqused_k)] + page_table, kv_batch_idx, leftpad_k = [ + maybe_contiguous(x) for x in (page_table, kv_batch_idx, leftpad_k) + ] + rotary_cos, rotary_sin = [maybe_contiguous(x) for x in (rotary_cos, rotary_sin)] + seqlens_rotary = maybe_contiguous(seqlens_rotary) + out, softmax_lse, out_accum, softmax_lse_accum = flash_attn_3_cuda.fwd( + q, + k, + v, + k_new, + v_new, + qv, + out_, + cu_seqlens_q, + cu_seqlens_k, + cu_seqlens_k_new, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + page_table, + kv_batch_idx, + leftpad_k, + rotary_cos, + rotary_sin, + seqlens_rotary, + q_descale, + k_descale, + v_descale, + softmax_scale, + causal, + window_size_left, + window_size_right, + attention_chunk, + softcap, + rotary_interleaved, + scheduler_metadata, + num_splits, + pack_gqa, + sm_margin, + ) + + if out_accum is None: + out_accum = torch.tensor([], device=out.device) + + if softmax_lse_accum is None: + softmax_lse_accum = torch.tensor([], device=out.device) + + return out, softmax_lse, out_accum, softmax_lse_accum + + +@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_forward")) +def _flash_attn_forward_fake( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_new: Optional[torch.Tensor] = None, + v_new: Optional[torch.Tensor] = None, + qv: Optional[torch.Tensor] = None, + out_: Optional[torch.Tensor] = None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + seqused_q: Optional[torch.Tensor] = None, + seqused_k: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_k: Optional[int] = None, + page_table: Optional[torch.Tensor] = None, + kv_batch_idx: Optional[torch.Tensor] = None, + leftpad_k: Optional[torch.Tensor] = None, + rotary_cos: Optional[torch.Tensor] = None, + rotary_sin: Optional[torch.Tensor] = None, + seqlens_rotary: Optional[torch.Tensor] = None, + q_descale: Optional[torch.Tensor] = None, + k_descale: Optional[torch.Tensor] = None, + v_descale: Optional[torch.Tensor] = None, + softmax_scale: Optional[float] = None, + causal: bool = False, + window_size_left: int = -1, + window_size_right: int = -1, + attention_chunk: int = 0, + softcap: float = 0.0, + rotary_interleaved: bool = True, + scheduler_metadata: Optional[torch.Tensor] = None, + num_splits: int = 1, + pack_gqa: Optional[bool] = None, + sm_margin: int = 0, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Symbolic fake implementation of flash attention forward. + Returns tensors with the correct shapes and dtypes without actual computation. + """ + + # Determine if we're in varlen mode + is_varlen_q = cu_seqlens_q is not None + + # Get dimensions from query tensor + if is_varlen_q: + # varlen mode: q is (total_q, num_heads, head_size) + total_q, num_heads, head_size = q.shape + batch_size = cu_seqlens_q.shape[0] - 1 + + if max_seqlen_q is None: + raise ValueError("max_seqlen_q must be provided if cu_seqlens_q is provided") + seqlen_q = max_seqlen_q + else: + # batch mode: q is (batch_size, seqlen_q, num_heads, head_size) + batch_size, seqlen_q, num_heads, head_size = q.shape + total_q = batch_size * q.shape[1] + # Get value head dimension + head_size_v = v.shape[-1] + + # Determine output dtype (FP8 inputs produce BF16 outputs) + q_type = q.dtype + if q_type == torch.float8_e4m3fn: + out_dtype = torch.bfloat16 + else: + out_dtype = q_type + + # Create output tensor + if out_ is not None: + # If out_ is provided, _flash_attn_forward becomes non-functional + raise TypeError("Tracing (torch.compile/torch.export) with pre-allocated output tensor is not supported.") + + if is_varlen_q: + out = torch.empty((total_q, num_heads, head_size_v), dtype=out_dtype, device=q.device) + else: + out = torch.empty((batch_size, seqlen_q, num_heads, head_size_v), dtype=out_dtype, device=q.device) + + # Create softmax_lse tensor + if is_varlen_q: + softmax_lse = torch.empty((num_heads, total_q), dtype=torch.float32, device=q.device) + else: + softmax_lse = torch.empty((batch_size, num_heads, seqlen_q), dtype=torch.float32, device=q.device) + + # TODO(guilhermeleobas): Implement "get_num_splits" + # There's an heuristic to compute num_splits when "num_splits <= 0" + # assert that num_splits is > 0 for now + if num_splits <= 0: + raise ValueError(f"tracing (torch.compile/torch.export) with num_splits <= 0 not supported. Got {num_splits=}") + + if num_splits > 1: + if is_varlen_q: + out_accum = torch.empty((num_splits, num_heads, total_q, head_size_v), dtype=torch.float32, device=q.device) + softmax_lse_accum = torch.empty((num_splits, num_heads, total_q), dtype=torch.float32, device=q.device) + else: + out_accum = torch.empty((num_splits, batch_size, num_heads, seqlen_q, head_size_v), dtype=torch.float32, device=q.device) + softmax_lse_accum = torch.empty((num_splits, batch_size, num_heads, seqlen_q), dtype=torch.float32, device=q.device) + else: + # Tensors are not set when num_splits < 1 + out_accum = torch.tensor([], device=out.device) + softmax_lse_accum = torch.tensor([], device=out.device) + + return out, softmax_lse, out_accum, softmax_lse_accum + + +@torch.library.custom_op(add_op_namespace_prefix("_flash_attn_backward"), mutates_args=("dq", "dk", "dv"), device_types="cuda") +def _flash_attn_backward( + dout: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + out: torch.Tensor, + softmax_lse: torch.Tensor, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k: Optional[torch.Tensor] = None, + sequed_q: Optional[torch.Tensor] = None, + sequed_k: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_k: Optional[int] = None, + dq: Optional[torch.Tensor] = None, + dk: Optional[torch.Tensor] = None, + dv: Optional[torch.Tensor] = None, + softmax_scale: Optional[float] = None, + is_causal: bool = False, + window_size_left: int = -1, + window_size_right: int = -1, + softcap: float = 0.0, + deterministic: bool = False, + sm_margin: int = 0, +) -> torch.Tensor: + # dq, dk, dv are allocated by us so they should already be contiguous + dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] + softmax_d, *rest = flash_attn_3_cuda.bwd( + dout, + q, + k, + v, + out, + softmax_lse, + dq, + dk, + dv, + cu_seqlens_q, + cu_seqlens_k, + sequed_q, + sequed_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + is_causal, + window_size_left, + window_size_right, + softcap, + deterministic, + sm_margin, + ) + return softmax_d + + +@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_backward")) +def _flash_attn_backward_fake( + dout: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + out: torch.Tensor, + softmax_lse: torch.Tensor, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k: Optional[torch.Tensor] = None, + sequed_q: Optional[torch.Tensor] = None, + sequed_k: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_k: Optional[int] = None, + dq: Optional[torch.Tensor] = None, + dk: Optional[torch.Tensor] = None, + dv: Optional[torch.Tensor] = None, + softmax_scale: Optional[float] = None, + is_causal: bool = False, + window_size_left: int = -1, + window_size_right: int = -1, + softcap: float = 0.0, + deterministic: bool = False, + sm_margin: int = 0, +) -> torch.Tensor: + + is_varlen_q = cu_seqlens_q is not None + is_varlen_k = cu_seqlens_q is not None + is_varlen = is_varlen_q or is_varlen_k or sequed_q is not None or sequed_k is not None + + if not is_varlen_q: + batch_size = q.size(0) + seqlen_q = q.size(1) + seqlen_k = k.size(1) + total_q = batch_size * q.size(1) + else: + batch_size = cu_seqlens_q.size(0) - 1 + total_q = q.size(0) + seqlen_q = max_seqlen_q + seqlen_k = max_seqlen_k + + if window_size_left >= seqlen_k - 1: + window_size_left = -1 + + if window_size_right >= seqlen_q - 1: + window_size_right = -1 + + if is_causal: + window_size_right = 0 + + is_causal = window_size_left < 0 and window_size_right == 0 + + head_size = q.size(-1) + head_size_v = v.size(-1) + head_size_rounded = round_up_headdim(max(head_size, head_size_v)) + + # Hopper gpus uses cuda compute capabilities 9.0 + cap = torch.cuda.get_device_capability(q.device) + arch = cap[0] * 10 + cap[1] + + is_local = (window_size_left >= 0 or window_size_right >= 0) and not is_causal + + if head_size_rounded <= 64: + kBlockM_sm90 = 96 if (is_causal and softcap > 0.0) else 128 + elif head_size_rounded <= 96: + kBlockM_sm90 = 64 + elif head_size_rounded <= 128: + kBlockM_sm90 = 64 if (is_causal or is_local or softcap > 0.0) else 80 + else: + kBlockM_sm90 = 64 + + kBlockM_sm80 = 128 if head_size_rounded <= 64 else 64 + kBlockM_sm86 = 64 if head_size_rounded <= 192 else 32 + + if arch >= 90: + kBlockM = kBlockM_sm90 + elif arch == 86 or arch == 89: + kBlockM = kBlockM_sm86 + else: + kBlockM = kBlockM_sm80 + + num_heads = q.shape[-2] + seqlen_q_rounded = round_multiple(seqlen_q, kBlockM) + + total_q_padded_rounded = round_multiple(total_q + batch_size * kBlockM, kBlockM) + + dq = torch.empty_like(q) if dq is None else dq + dk = torch.empty_like(k) if dk is None else dk + dv = torch.empty_like(v) if dv is None else dv + + if not is_varlen: + softmax_d = torch.empty((batch_size, num_heads, seqlen_q_rounded), dtype=torch.float32, device=q.device) + else: + softmax_d = torch.empty((num_heads, total_q_padded_rounded), dtype=torch.float32, device=q.device) + + return softmax_d + + +def setup_context(ctx, inputs, output): + q, k, v = inputs[:3] + out, softmax_lse, _, _ = output + ctx.save_for_backward(q, k, v, out, softmax_lse) + ctx.softmax_scale = inputs[-11] + ctx.causal = inputs[-10] + ctx.window_size = [inputs[-9], inputs[-8]] + ctx.attention_chunk = inputs[-7] + ctx.softcap = inputs[-6] + ctx.sm_margin = inputs[-1] + + +def _backward(ctx, dout, *grads): + q, k, v, out, softmax_lse = ctx.saved_tensors + dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + None, None, # cu_seqlens_q, cu_seqlens_k, + None, None, # sequed_q, sequed_k, + None, None, # max_seqlen_q, max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size[0], + ctx.window_size[1], + ctx.softcap, + False, # deterministic + ctx.sm_margin, + ) + return dq, dk, dv, *((None,) * 21) + + +_flash_attn_forward.register_autograd(_backward, setup_context=setup_context) + + + +class FlashAttnQKVPackedFunc(torch.autograd.Function): + @staticmethod + def forward( + ctx, + qkv, + softmax_scale, + causal, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + deterministic=False, + num_heads_q=None, + sm_margin=0, + return_softmax=False, + ): + if softmax_scale is None: + softmax_scale = qkv.shape[-1] ** (-0.5) + if qkv.dim() == 5: + assert qkv.shape[-3] == 3 + q, k, v = qkv.unbind(dim=-3) + else: + assert qkv.dim() == 4 + assert num_heads_q is not None + num_heads_k = (qkv.shape[2] - num_heads_q) // 2 + assert num_heads_k * 2 + num_heads_q == qkv.shape[2] + q, k, v = qkv.split([num_heads_q, num_heads_k, num_heads_k], dim=-2) + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + None, # qv + None, # out + None, None, None, # cu_seqlens_q/k/k_new + None, None, # seqused_q/k + None, None, # max_seqlen_q/k + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size_left=window_size[0], + window_size_right=window_size[1], + attention_chunk=attention_chunk, + softcap=softcap, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse) + ctx.save_for_backward(q, k, v, out, softmax_lse) + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.ndim = qkv.dim() + ctx.sm_margin = sm_margin + return (out, softmax_lse) if return_softmax else out + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + if ctx.ndim == 5: + qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) + dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) + dq, dk, dv = dqkv.unbind(dim=-3) + else: + num_heads_q = q.shape[2] + num_heads_k = k.shape[2] + qkv_shape = q.shape[:-2] + (num_heads_q + num_heads_k * 2, *q.shape[-1:]) + dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) + dq, dk, dv = dqkv.split([num_heads_q, num_heads_k, num_heads_k], dim=-2) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + None, None, # cu_seqlens_q, cu_seqlens_k, + None, None, # sequed_q, sequed_k, + None, None, # max_seqlen_q, max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size[0], + ctx.window_size[1], + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dqkv = dqkv[..., : dout.shape[-1]] # We could have padded the head dimension + return dqkv, None, None, None, None, None, None, None, None, None, None, None, None + + +class FlashAttnFunc(torch.autograd.Function): + + @staticmethod + def forward( + ctx, + q, + k, + v, + softmax_scale, + causal, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + return_softmax=False, + ): + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + # out, q, k, v, out_padded, softmax_lse = _flash_attn_forward( + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + qv, # qv + None, # out + None, None, None, # cu_seqlens_q/k/k_new + None, None, # seqused_q/k + None, None, # max_seqlen_q/k + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size_left=window_size[0], + window_size_right=window_size[1], + attention_chunk=attention_chunk, + softcap=softcap, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse) + ctx.save_for_backward(q, k, v, out, softmax_lse) + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.sm_margin = sm_margin + return (out, softmax_lse) if return_softmax else out + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + None, None, # cu_seqlens_q, cu_seqlens_k, + None, None, # sequed_q, sequed_k, + None, None, # max_seqlen_q, max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size[0], + ctx.window_size[1], + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dq = dq[..., : q.shape[-1]] # We could have padded the head dimension + dk = dk[..., : k.shape[-1]] + dv = dv[..., : v.shape[-1]] + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None + + +class FlashAttnVarlenFunc(torch.autograd.Function): + + @staticmethod + def forward( + ctx, + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + return_softmax=False, + ): + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + # out, q, k, v, out_padded, softmax_lse = _flash_attn_varlen_forward( + out, softmax_lse, *rest = _flash_attn_forward( + q, + k, + v, + None, None, # k_new, v_new + qv, # qv + None, # out + cu_seqlens_q, + cu_seqlens_k, + None, # cu_seqlens_k_new + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + None, None, None, # page_table, kv_batch_idx, leftpad_k, + None, None, None, # rotary_cos/sin, seqlens_rotary + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size_left=window_size[0], + window_size_right=window_size[1], + attention_chunk=attention_chunk, + softcap=softcap, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k) + ctx.save_for_backward(q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k) + ctx.max_seqlen_q = max_seqlen_q + ctx.max_seqlen_k = max_seqlen_k + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.attention_chunk = attention_chunk + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.sm_margin = sm_margin + return (out, softmax_lse) if return_softmax else out + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k = ctx.saved_tensors + assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" + dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + _flash_attn_backward( + dout, + q, + k, + v, + out, + softmax_lse, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + ctx.max_seqlen_q, + ctx.max_seqlen_k, + dq, + dk, + dv, + ctx.softmax_scale, + ctx.causal, + ctx.window_size[0], + ctx.window_size[1], + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + dq = dq[..., : q.shape[-1]] # We could have padded the head dimension + dk = dk[..., : k.shape[-1]] + dv = dv[..., : v.shape[-1]] + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None + + +def flash_attn_qkvpacked_func( + qkv, + softmax_scale=None, + causal=False, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + deterministic=False, + num_heads_q=None, + sm_margin=0, + return_attn_probs=False, +): + """dropout_p should be set to 0.0 during evaluation + If Q, K, V are already stacked into 1 tensor, this function will be faster than + calling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation + of the gradients of Q, K, V. + For multi-query and grouped-query attention (MQA/GQA), please see + flash_attn_kvpacked_func and flash_attn_func. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between [i - window_size[0], i + window_size[1]] inclusive. + + Arguments: + qkv: (batch_size, seqlen, 3, nheads, headdim) + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + softcap: float. Anything > 0 activates softcapping attention. + alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of (-alibi_slope * |i - j|) is added to + the attention score of query i and key j. + deterministic: bool. Whether to use the deterministic implementation of the backward pass, + which is slightly slower and uses more memory. The forward pass is always deterministic. + return_attn_probs: bool. Whether to return the attention probabilities. This option is for + testing only. The returned probabilities are not guaranteed to be correct + (they might not have the right scaling). + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen). + The output of softmax (possibly with different scaling). It also encodes the dropout + pattern (negative means that location was dropped, nonnegative means it was kept). + """ + return FlashAttnQKVPackedFunc.apply( + qkv, + softmax_scale, + causal, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + deterministic, + num_heads_q, + sm_margin, + return_attn_probs, + ) + + +def flash_attn_func( + q, + k, + v, + softmax_scale=None, + causal=False, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + return_attn_probs=False, +): + """dropout_p should be set to 0.0 during evaluation + Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads + than Q. Note that the number of heads in Q must be divisible by the number of heads in KV. + For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head + 0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V. + + If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix. + For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is: + 1 1 1 1 0 + 1 1 1 1 1 + If seqlen_q = 5 and seqlen_k = 2, the causal mask is: + 0 0 + 0 0 + 0 0 + 1 0 + 1 1 + If the row of the mask is all zero, the output will be zero. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between + [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive. + + Arguments: + q: (batch_size, seqlen, nheads, headdim) + k: (batch_size, seqlen, nheads_k, headdim) + v: (batch_size, seqlen, nheads_k, headdim) + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of + (-alibi_slope * |i + seqlen_k - seqlen_q - j|) + is added to the attention score of query i and key j. + deterministic: bool. Whether to use the deterministic implementation of the backward pass, + which is slightly slower and uses more memory. The forward pass is always deterministic. + return_attn_probs: bool. Whether to return the attention probabilities. This option is for + testing only. The returned probabilities are not guaranteed to be correct + (they might not have the right scaling). + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + """ + return FlashAttnFunc.apply( + q, + k, + v, + softmax_scale, + causal, + qv, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + num_splits, + pack_gqa, + deterministic, + sm_margin, + return_attn_probs, + ) + + +def flash_attn_varlen_func( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + seqused_q=None, + seqused_k=None, + softmax_scale=None, + causal=False, + qv=None, + q_descale=None, k_descale=None, v_descale=None, + window_size=(-1, -1), + attention_chunk=0, + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + return_attn_probs=False, +): + return FlashAttnVarlenFunc.apply( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + qv, + q_descale, k_descale, v_descale, + window_size, + attention_chunk, + softcap, + num_splits, + pack_gqa, + deterministic, + sm_margin, + return_attn_probs, + ) + + +def flash_attn_combine(out_partial, lse_partial, out=None, out_dtype=None): + return flash_attn_3_cuda.fwd_combine(out_partial, lse_partial, out, out_dtype) + + +def flash_attn_with_kvcache( + q, + k_cache, + v_cache, + k=None, + v=None, + qv=None, + rotary_cos=None, + rotary_sin=None, + cache_seqlens: Optional[Union[(int, torch.Tensor)]] = None, + cache_batch_idx: Optional[torch.Tensor] = None, + cache_leftpad: Optional[torch.Tensor] = None, + page_table: Optional[torch.Tensor] = None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + rotary_seqlens: Optional[torch.Tensor] = None, + q_descale: Optional[torch.Tensor] = None, + k_descale: Optional[torch.Tensor] = None, + v_descale: Optional[torch.Tensor] = None, + softmax_scale=None, + causal=False, + window_size=(-1, -1), # -1 means infinite context window + attention_chunk=0, + softcap=0.0, # 0.0 means deactivated + rotary_interleaved=True, + scheduler_metadata=None, + num_splits=0, # Can be tuned for speed + pack_gqa=None, # Can be tuned for speed + sm_margin=0, # Can be tuned if some SMs are used for communication + return_softmax_lse=False, +): + """ + If k and v are not None, k_cache and v_cache will be updated *inplace* with the new values from + k and v. This is useful for incremental decoding: you can pass in the cached keys/values from + the previous step, and update them with the new keys/values from the current step, and do + attention with the updated cache, all in 1 kernel. + + If you pass in k / v, you must make sure that the cache is large enough to hold the new values. + For example, the KV cache could be pre-allocated with the max sequence length, and you can use + cache_seqlens to keep track of the current sequence lengths of each sequence in the batch. + + Also apply rotary embedding if rotary_cos and rotary_sin are passed in. The key @k will be + rotated by rotary_cos and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc. + If causal or local (i.e., window_size != (-1, -1)), the query @q will be rotated by rotary_cos + and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc. + If not causal and not local, the query @q will be rotated by rotary_cos and rotary_sin at + indices cache_seqlens only (i.e. we consider all tokens in @q to be at position cache_seqlens). + + See tests/test_flash_attn.py::test_flash_attn_kvcache for examples of how to use this function. + + Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads + than Q. Note that the number of heads in Q must be divisible by the number of heads in KV. + For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head + 0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V. + + If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix. + For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is: + 1 1 1 1 0 + 1 1 1 1 1 + If seqlen_q = 5 and seqlen_k = 2, the causal mask is: + 0 0 + 0 0 + 0 0 + 1 0 + 1 1 + If the row of the mask is all zero, the output will be zero. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between + [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive. + + Note: Does not support backward pass. + + Arguments: + q: (batch_size, seqlen, nheads, headdim) + k_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim) if there's no page_table, + or (num_blocks, page_block_size, nheads_k, headdim) if there's a page_table (i.e. paged KV cache) + page_block_size can be arbitrary (e.g, 1, 2, 3, 64, etc.). + v_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim_v) if there's no page_table, + or (num_blocks, page_block_size, nheads_k, headdim_v) if there's a page_table (i.e. paged KV cache) + k [optional]: (batch_size, seqlen_new, nheads_k, headdim). If not None, we concatenate + k with k_cache, starting at the indices specified by cache_seqlens. + v [optional]: (batch_size, seqlen_new, nheads_k, headdim_v). Similar to k. + qv [optional]: (batch_size, seqlen, nheads, headdim_v) + rotary_cos [optional]: (seqlen_ro, rotary_dim / 2). If not None, we apply rotary embedding + to k and q. Only applicable if k and v are passed in. rotary_dim must be divisible by 16. + rotary_sin [optional]: (seqlen_ro, rotary_dim / 2). Similar to rotary_cos. + cache_seqlens: int, or (batch_size,), dtype torch.int32. The sequence lengths of the + KV cache. + cache_batch_idx: (batch_size,), dtype torch.int32. The indices used to index into the KV cache. + If None, we assume that the batch indices are [0, 1, 2, ..., batch_size - 1]. + If the indices are not distinct, and k and v are provided, the values updated in the cache + might come from any of the duplicate indices. + cache_leftpad: (batch_size,), dtype torch.int32. The index that the KV cache starts. If None, assume 0. + page_table [optional]: (batch_size, max_num_blocks_per_seq), dtype torch.int32. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + softcap: float. Anything > 0 activates softcapping attention. + rotary_interleaved: bool. Only applicable if rotary_cos and rotary_sin are passed in. + If True, rotary embedding will combine dimensions 0 & 1, 2 & 3, etc. If False, + rotary embedding will combine dimensions 0 & rotary_dim / 2, 1 & rotary_dim / 2 + 1 + (i.e. GPT-NeoX style). + num_splits: int. If > 1, split the key/value into this many chunks along the sequence. + If num_splits == 1, we don't split the key/value. If num_splits == 0, we use a heuristic + to automatically determine the number of splits. + Don't change this unless you know what you are doing. + return_softmax_lse: bool. Whether to return the logsumexp of the attention scores. + + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_softmax_lse=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + """ + assert k_cache.stride(-1) == 1, "k_cache must have contiguous last dimension" + assert v_cache.stride(-1) == 1, "v_cache must have contiguous last dimension" + if softmax_scale is None: + softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) + if cache_seqlens is not None and isinstance(cache_seqlens, int): + cache_seqlens = torch.full( + (q.shape[0],), cache_seqlens, dtype=torch.int32, device=k_cache.device + ) + cache_seqlens = maybe_contiguous(cache_seqlens) + out, softmax_lse, *rest = _flash_attn_forward( + q, + k_cache, + v_cache, + k, + v, + qv, + None, # out + cu_seqlens_q, + None, # cu_seqlens_k + cu_seqlens_k_new, + None, # seqused_q + cache_seqlens, + max_seqlen_q, + None, # max_seqlen_k + page_table, + cache_batch_idx, + cache_leftpad, + rotary_cos, + rotary_sin, + rotary_seqlens, + q_descale, k_descale, v_descale, + softmax_scale, + causal=causal, + window_size_left=window_size[0], + window_size_right=window_size[1], + attention_chunk=attention_chunk, + softcap=softcap, + rotary_interleaved=rotary_interleaved, + scheduler_metadata=scheduler_metadata, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + # return (out, softmax_lse) if return_softmax_lse else out + return (out, softmax_lse, *rest) if return_softmax_lse else out + + +def get_scheduler_metadata( + batch_size, max_seqlen_q, max_seqlen_k, num_heads_q, num_heads_kv, headdim, + cache_seqlens: torch.Tensor, + qkv_dtype=torch.bfloat16, + headdim_v=None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + cache_leftpad: Optional[torch.Tensor] = None, + page_size: Optional[int] = None, + max_seqlen_k_new=0, + causal=False, + window_size=(-1, -1), # -1 means infinite context window + attention_chunk=0, + has_softcap=False, + num_splits=0, # Can be tuned for speed + pack_gqa=None, # Can be tuned for speed + sm_margin=0, # Can be tuned if some SMs are used for communication +): + cache_seqlens = maybe_contiguous(cache_seqlens) + if headdim_v is None: + headdim_v = headdim + scheduler_metadata = flash_attn_3_cuda.get_scheduler_metadata( + batch_size, max_seqlen_q, max_seqlen_k, num_heads_q, num_heads_kv, headdim, headdim_v, + qkv_dtype, + cache_seqlens, + cu_seqlens_q, + None, # cu_seqlens_k + cu_seqlens_k_new, + None, # seqused_q + cache_leftpad, + page_size, + max_seqlen_k_new, + causal, + window_size[0], window_size[1], + attention_chunk, + has_softcap, + num_splits, + pack_gqa, + sm_margin, + ) + return scheduler_metadata diff --git a/build/torch29-cxx11-cu130-x86_64-linux/metadata.json b/build/torch29-cxx11-cu130-x86_64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..e0cc0c4f14abd10162c22499be39c75e6d922f94 --- /dev/null +++ b/build/torch29-cxx11-cu130-x86_64-linux/metadata.json @@ -0,0 +1,12 @@ +{ + "version": 1, + "license": "BSD-3-Clause", + "python-depends": [], + "backend": { + "type": "cuda", + "archs": [ + "8.0", + "9.0a" + ] + } +}