diff --git a/.gitattributes b/.gitattributes index 0f4540bf6658540f38e1526b997c53996f38f2b5..0cd58331b2a989b68be4ec5676383437fca8687b 100644 --- a/.gitattributes +++ b/.gitattributes @@ -34,5 +34,3 @@ saved_model/**/* 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 -build/torch210-cu130-x86_64-windows/_flash_attn2_d6b301e.pyd filter=lfs diff=lfs merge=lfs -text -build/torch210-cu128-x86_64-windows/_flash_attn2_d6b301e.pyd filter=lfs diff=lfs merge=lfs -text diff --git a/README.md b/README.md index 89ffa2a473172f67cb90e85c155339d64b704af6..c7eb8a7dfa3aa0ddd0343aebb252c5d5dc918463 100644 --- a/README.md +++ b/README.md @@ -1,36 +1,149 @@ --- -library_name: kernels license: bsd-3-clause +tags: + - kernels --- -This is the repository card of kernels-community/flash-attn2 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 +# Flash Attention +Flash Attention is a fast and memory-efficient implementation of the attention mechanism, designed to work with large models and long sequences. This is a Hugging Face compliant kernel build of Flash Attention. + +Original code here [https://github.com/Dao-AILab/flash-attention](https://github.com/Dao-AILab/flash-attention). + + +[`scripts/readme_example.py`](scripts/readme_example.py) provides a simple example of how to use the Flash Attention kernel in PyTorch. It demonstrates standard attention, causal attention, and variable-length sequences. ```python -# make sure `kernels` is installed: `pip install -U kernels` +# /// script +# dependencies = [ +# "numpy", +# "torch", +# "kernels" +# ] +# /// +import torch from kernels import get_kernel -kernel_module = get_kernel("kernels-community/flash-attn2", version=2) -flash_attn_func = kernel_module.flash_attn_func +# Setup +torch.manual_seed(42) +flash_attn = get_kernel("kernels-community/flash-attn2") +device = torch.device("cuda") + +# Create test tensors +B, S, H, D = 2, 5, 4, 8 # batch, seq_len, heads, head_dim +q = k = v = torch.randn(B, S, H, D, device=device, dtype=torch.float16) + +# Reference implementation using PyTorch SDPA +def reference_attention(query, key, value, causal=False): + query, key, value = (x.transpose(1, 2).contiguous() for x in (query, key, value)) + with torch.nn.attention.sdpa_kernel(torch.nn.attention.SDPBackend.MATH): + out = torch.nn.functional.scaled_dot_product_attention(query, key, value, is_causal=causal) + return out.transpose(1, 2).contiguous() + +# 1. Standard attention +print("\n1. Standard attention:") +out_ref = reference_attention(q, k, v) +out_flash = flash_attn.fwd( + q=q, + k=k, + v=v, + is_causal=False, +)[0] +print(f"Reference output: {out_ref.shape}") +print(f"Flash output: {out_flash.shape}") +print(f"Outputs close: {torch.allclose(out_flash, out_ref, atol=1e-2, rtol=1e-3)}") + +# 2. Causal attention (for autoregressive models) +print("\n2. Causal attention:") + +out_ref_causal = reference_attention(q, k, v, causal=True) +out_causal = flash_attn.fwd( + q=q, + k=k, + v=v, + is_causal=True, +)[0] +print(f"Reference causal output: {out_ref_causal.shape}") +print(f"Flash causal output: {out_causal.shape}") +print(f"Outputs close: {torch.allclose(out_causal, out_ref_causal, atol=1e-2, rtol=1e-3)}") + +def var_reference_attention(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, causal=False): + batch_size = cu_seqlens_q.shape[0] - 1 + # Return output in packed format (same as flash attention) + total_tokens_q = q.shape[0] + out = torch.zeros((total_tokens_q, q.shape[1], q.shape[2]), device=q.device, dtype=q.dtype) + + for b in range(batch_size): + start_q, end_q = cu_seqlens_q[b], cu_seqlens_q[b + 1] + start_k, end_k = cu_seqlens_k[b], cu_seqlens_k[b + 1] + + # Extract slices for this batch + q_slice = q[start_q:end_q] # Shape: (seq_len_q, H, D) + k_slice = k[start_k:end_k] # Shape: (seq_len_k, H, D) + v_slice = v[start_k:end_k] # Shape: (seq_len_k, H, D) + + # Add batch dimension for reference_attention + q_slice = q_slice.unsqueeze(0) # Shape: (1, seq_len_q, H, D) + k_slice = k_slice.unsqueeze(0) # Shape: (1, seq_len_k, H, D) + v_slice = v_slice.unsqueeze(0) # Shape: (1, seq_len_k, H, D) + + # Compute attention and remove batch dimension + attn_out = reference_attention(q_slice, k_slice, v_slice, causal=causal) + attn_out = attn_out.squeeze(0) # Shape: (seq_len_q, H, D) + + # Place result in output tensor (packed format) + out[start_q:end_q] = attn_out + + return out + +# 3. Variable length sequences (packed format) +print("\n3. Variable length sequences:") +# Pack sequences of lengths [3,4,3] for q and [4,5,3] for k into single tensors +q_var = torch.randn(10, H, D, device=device, dtype=torch.float16) # total_q=10 +k_var = v_var = torch.randn(12, H, D, device=device, dtype=torch.float16) # total_k=12 +cu_q = torch.tensor([0, 3, 7, 10], device=device, dtype=torch.int32) # cumulative sequence lengths +cu_k = torch.tensor([0, 4, 9, 12], device=device, dtype=torch.int32) + +out_var_ref = var_reference_attention(q_var, k_var, v_var, cu_q, cu_k, max_seqlen_q=4, max_seqlen_k=5, causal=False) +# Custom function to handle variable +out_var = flash_attn.varlen_fwd( + q=q_var, + k=k_var, + v=v_var, + cu_seqlens_q=cu_q, + cu_seqlens_k=cu_k, + max_seqlen_q=4, + max_seqlen_k=5, +)[0] +print(f"Variable length output: {out_var.shape}") +print(f"Reference variable length output: {out_var_ref.shape}") +print(f"Outputs close: {torch.allclose(out_var, out_var_ref, atol=1e-2, rtol=1e-3)}") +``` + +run it using the following command: + +```bash +uv run scripts/readme_example.py +``` + +```txt +Reading inline script metadata from `scripts/readme_example.py` +Fetching 20 files: 100%|██████████████████████████████████████████████████| 20/20 [00:00<00:00, 16371.21it/s] + +1. Standard attention: +Reference output: torch.Size([2, 5, 4, 8]) +Flash output: torch.Size([2, 5, 4, 8]) +Outputs close: True + +2. Causal attention: +Reference causal output: torch.Size([2, 5, 4, 8]) +Flash causal output: torch.Size([2, 5, 4, 8]) +Outputs close: True -flash_attn_func(...) +3. Variable length sequences: +Variable length output: torch.Size([10, 4, 8]) +Reference variable length output: torch.Size([10, 4, 8]) +Outputs close: True ``` -## Available functions -- `flash_attn_func` -- `flash_attn_kvpacked_func` -- `flash_attn_qkvpacked_func` -- `flash_attn_varlen_func` -- `flash_attn_varlen_kvpacked_func` -- `flash_attn_varlen_qkvpacked_func` -- `flash_attn_with_kvcache` -- `fwd` -- `varlen_fwd` -- `bwd` -- `varlen_bwd` -- `fwd_kvcache` - -## Benchmarks - -Benchmarking script is available for this kernel. Run `kernels benchmark kernels-community/flash-attn2 --version 2`. diff --git a/build/torch210-cu128-x86_64-windows/__init__.py b/build/torch210-cu128-x86_64-windows/__init__.py deleted file mode 100644 index 4b3892ab7794a39e058f2eeb27e6c97cc6bbe6b4..0000000000000000000000000000000000000000 --- a/build/torch210-cu128-x86_64-windows/__init__.py +++ /dev/null @@ -1,393 +0,0 @@ -from typing import Optional, List -import torch -from ._ops import ops as flash_attn_ops -from .flash_attn_interface import ( - flash_attn_func, - flash_attn_kvpacked_func, - flash_attn_qkvpacked_func, - flash_attn_varlen_func, - flash_attn_varlen_kvpacked_func, - flash_attn_varlen_qkvpacked_func, - flash_attn_with_kvcache, -) - - -def fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - return_softmax: bool = False, - gen: Optional[torch.Generator] = None, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention. - - Args: - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Optional output tensor, same shape as q - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - return_softmax: Whether to return softmax weights - gen: Optional random number generator - - Returns: - List of tensors: [output, softmax_lse, (softmax if return_softmax)] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.fwd( - q, - k, - v, - out, - alibi_slopes, - p_dropout, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - gen, - ) - - -def varlen_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - out: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - block_table: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - max_seqlen_q: int = 0, - max_seqlen_k: int = 0, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - zero_tensors: bool = False, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - return_softmax: bool = False, - gen: Optional[torch.Generator] = None, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention with variable sequence lengths. - - Args: - q: Query tensor of shape [total_q, num_heads, head_size] - k: Key tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - v: Value tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1] - cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1] - out: Optional output tensor of shape [total_q, num_heads, head_size] - seqused_k: Optional tensor specifying how many keys to use per batch element [batch_size] - leftpad_k: Optional left padding for keys of shape [batch_size] - block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq] - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - max_seqlen_q: Maximum sequence length for queries - max_seqlen_k: Maximum sequence length for keys - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - zero_tensors: Whether to zero tensors before computation - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - return_softmax: Whether to return softmax weights - gen: Optional random number generator - - Returns: - List of tensors: [output, softmax_lse, (softmax if return_softmax)] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.varlen_fwd( - q, - k, - v, - out, - cu_seqlens_q, - cu_seqlens_k, - seqused_k, - leftpad_k, - block_table, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - p_dropout, - softmax_scale, - zero_tensors, - is_causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - gen, - ) - - -def bwd( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - p_dropout: float = 0.0, - 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, - gen: Optional[torch.Generator] = None, - rng_state: Optional[torch.Tensor] = None, -) -> List[torch.Tensor]: - """ - Backward pass for multi-head attention. - - Args: - dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size] - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size] - softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q] - dq: Optional gradient tensor for queries, same shape as q - dk: Optional gradient tensor for keys, same shape as k - dv: Optional gradient tensor for values, same shape as v - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - deterministic: Whether to use deterministic algorithms - gen: Optional random number generator - rng_state: Optional RNG state from forward pass - - Returns: - List of tensors: [dq, dk, dv] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - alibi_slopes, - p_dropout, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - deterministic, - gen, - rng_state, - ) - - -def varlen_bwd( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - max_seqlen_q: int = 0, - max_seqlen_k: int = 0, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - zero_tensors: bool = False, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - deterministic: bool = False, - gen: Optional[torch.Generator] = None, - rng_state: Optional[torch.Tensor] = None, -) -> List[torch.Tensor]: - """ - Backward pass for multi-head attention with variable sequence lengths. - - Args: - dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size] - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size] - softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q] - cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1] - cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1] - dq: Optional gradient tensor for queries, same shape as q - dk: Optional gradient tensor for keys, same shape as k - dv: Optional gradient tensor for values, same shape as v - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - max_seqlen_q: Maximum sequence length for queries - max_seqlen_k: Maximum sequence length for keys - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - zero_tensors: Whether to zero tensors before computation - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - deterministic: Whether to use deterministic algorithms - gen: Optional random number generator - rng_state: Optional RNG state from forward pass - - Returns: - List of tensors: [dq, dk, dv] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.varlen_bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - p_dropout, - softmax_scale, - zero_tensors, - is_causal, - window_size_left, - window_size_right, - softcap, - deterministic, - gen, - rng_state, - ) - - -def fwd_kvcache( - q: torch.Tensor, - kcache: torch.Tensor, - vcache: torch.Tensor, - k: Optional[torch.Tensor] = None, - v: Optional[torch.Tensor] = None, - seqlens_k: Optional[torch.Tensor] = None, - rotary_cos: Optional[torch.Tensor] = None, - rotary_sin: Optional[torch.Tensor] = None, - cache_batch_idx: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - block_table: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - out: 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, - is_rotary_interleaved: bool = False, - num_splits: int = 1, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention with KV cache. - - Args: - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - kcache: Key cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - vcache: Value cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - k: Optional new keys tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size] - v: Optional new values tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size] - seqlens_k: Optional sequence lengths for keys of shape [batch_size] - rotary_cos: Optional rotary cosine tensor of shape [seqlen_ro, rotary_dim/2] - rotary_sin: Optional rotary sine tensor of shape [seqlen_ro, rotary_dim/2] - cache_batch_idx: Optional indices to index into the KV cache - leftpad_k: Optional left padding for keys of shape [batch_size] - block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq] - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - out: Optional output tensor, same shape as q - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - is_rotary_interleaved: Whether rotary embeddings are interleaved - num_splits: Number of splits for computation - - Returns: - List of tensors: [output, softmax_lse] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.fwd_kvcache( - q, - kcache, - vcache, - k, - v, - seqlens_k, - rotary_cos, - rotary_sin, - cache_batch_idx, - leftpad_k, - block_table, - alibi_slopes, - out, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - is_rotary_interleaved, - num_splits, - ) diff --git a/build/torch210-cu128-x86_64-windows/_flash_attn2_d6b301e.exp b/build/torch210-cu128-x86_64-windows/_flash_attn2_d6b301e.exp deleted file mode 100644 index cb8556dd2a3a4ef1756b4b2b53e1359f91389084..0000000000000000000000000000000000000000 Binary files a/build/torch210-cu128-x86_64-windows/_flash_attn2_d6b301e.exp and /dev/null differ diff --git a/build/torch210-cu128-x86_64-windows/_flash_attn2_d6b301e.lib b/build/torch210-cu128-x86_64-windows/_flash_attn2_d6b301e.lib deleted file mode 100644 index 42c2c3013bdb2214fc5ab0ad8f64aa517b2e1362..0000000000000000000000000000000000000000 Binary files a/build/torch210-cu128-x86_64-windows/_flash_attn2_d6b301e.lib and /dev/null differ diff --git a/build/torch210-cu128-x86_64-windows/_flash_attn2_d6b301e.pyd b/build/torch210-cu128-x86_64-windows/_flash_attn2_d6b301e.pyd deleted file mode 100644 index 388fc92e10a1fd0aa2af03c53723f7941488d4f1..0000000000000000000000000000000000000000 --- a/build/torch210-cu128-x86_64-windows/_flash_attn2_d6b301e.pyd +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:d6a62d14386a4a156301af046a57bb718dc8ef7f103dd01b9529fc7c276b5e8a -size 1011416064 diff --git a/build/torch210-cu128-x86_64-windows/_ops.py b/build/torch210-cu128-x86_64-windows/_ops.py deleted file mode 100644 index 504c7db9ac15fb981ffd2b45e3fa950162393430..0000000000000000000000000000000000000000 --- a/build/torch210-cu128-x86_64-windows/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _flash_attn2_d6b301e -ops = torch.ops._flash_attn2_d6b301e - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_flash_attn2_d6b301e::{op_name}" \ No newline at end of file diff --git a/build/torch210-cu128-x86_64-windows/bert_padding.py b/build/torch210-cu128-x86_64-windows/bert_padding.py deleted file mode 100644 index 9717d840687facafd159739df68575a20514b09d..0000000000000000000000000000000000000000 --- a/build/torch210-cu128-x86_64-windows/bert_padding.py +++ /dev/null @@ -1,218 +0,0 @@ -# Adapted from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/padding.py - -import torch -import torch.nn.functional as F -from einops import rearrange, repeat - - -class IndexFirstAxis(torch.autograd.Function): - @staticmethod - def forward(ctx, input, indices): - ctx.save_for_backward(indices) - assert input.ndim >= 2 - ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] - second_dim = other_shape.numel() - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - # return input[indices] - return torch.gather( - rearrange(input, "b ... -> b (...)"), 0, repeat(indices, "z -> z d", d=second_dim) - ).reshape(-1, *other_shape) - - @staticmethod - def backward(ctx, grad_output): - (indices,) = ctx.saved_tensors - assert grad_output.ndim >= 2 - other_shape = grad_output.shape[1:] - grad_output = rearrange(grad_output, "b ... -> b (...)") - grad_input = torch.zeros( - [ctx.first_axis_dim, grad_output.shape[1]], - device=grad_output.device, - dtype=grad_output.dtype, - ) - # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. - # grad_input[indices] = grad_output - grad_input.scatter_(0, repeat(indices, "z -> z d", d=grad_output.shape[1]), grad_output) - return grad_input.reshape(ctx.first_axis_dim, *other_shape), None - - -index_first_axis = IndexFirstAxis.apply - - -class IndexPutFirstAxis(torch.autograd.Function): - @staticmethod - def forward(ctx, values, indices, first_axis_dim): - ctx.save_for_backward(indices) - assert indices.ndim == 1 - assert values.ndim >= 2 - output = torch.zeros( - first_axis_dim, *values.shape[1:], device=values.device, dtype=values.dtype - ) - # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. - output[indices] = values - # output.scatter_(0, repeat(indices, 'z -> z d', d=values.shape[1]), values) - return output - - @staticmethod - def backward(ctx, grad_output): - (indices,) = ctx.saved_tensors - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - grad_values = grad_output[indices] - # grad_values = torch.gather(grad_output, 0, repeat(indices, 'z -> z d', d=grad_output.shape[1])) - return grad_values, None, None - - -index_put_first_axis = IndexPutFirstAxis.apply - - -class IndexFirstAxisResidual(torch.autograd.Function): - @staticmethod - def forward(ctx, input, indices): - ctx.save_for_backward(indices) - assert input.ndim >= 2 - ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] - second_dim = other_shape.numel() - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - output = input[indices] - # We don't want to reshape input (b ... -> b (...)) since it could change the channel_last - # memory format to channel_first. In other words, input might not be contiguous. - # If we don't detach, Pytorch complains about output being a view and is being modified inplace - return output, input.detach() - - @staticmethod - def backward(ctx, grad_output, grad_residual): - (indices,) = ctx.saved_tensors - assert grad_output.ndim >= 2 - other_shape = grad_output.shape[1:] - assert grad_residual.shape[1:] == other_shape - grad_input = grad_residual - # grad_input[indices] += grad_output - indices = indices.reshape(indices.shape[0], *((1,) * (grad_output.ndim - 1))) - indices = indices.expand_as(grad_output) - grad_input.scatter_add_(0, indices, grad_output) - return grad_input.reshape(ctx.first_axis_dim, *other_shape), None - - -index_first_axis_residual = IndexFirstAxisResidual.apply - - -def unpad_input(hidden_states, attention_mask, unused_mask=None): - """ - Arguments: - hidden_states: (batch, seqlen, ...) - attention_mask: (batch, seqlen), bool / int, 1 means valid and 0 means not valid. - unused_mask: (batch, seqlen), bool / int, 1 means the element is allocated but unused. - Return: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens selected in attention_mask + unused_mask. - indices: (total_nnz), the indices of masked tokens from the flattened input sequence. - cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. - max_seqlen_in_batch: int - seqused: (batch), returns the number of tokens selected in attention_mask + unused_mask. - """ - all_masks = (attention_mask + unused_mask) if unused_mask is not None else attention_mask - seqlens_in_batch = all_masks.sum(dim=-1, dtype=torch.int32) - used_seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) - indices = torch.nonzero(all_masks.flatten(), as_tuple=False).flatten() - max_seqlen_in_batch = seqlens_in_batch.max().item() - cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) - # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the - # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim - # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to - # index with integer indices. Moreover, torch's index is a bit slower than it needs to be, - # so we write custom forward and backward to make it a bit faster. - return ( - index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices), - indices, - cu_seqlens, - max_seqlen_in_batch, - used_seqlens_in_batch, - ) - - -def unpad_input_for_concatenated_sequences(hidden_states, attention_mask_in_length): - """ - Supports concatenating short samples in one sequence. The attention_mask_in_length is utilized to mask other short samples. It helps efficient training of variant lengths-based samples (e.g., the supervised fine-tuning task in large language model). - The motivation for this function is explained [here](https://github.com/Dao-AILab/flash-attention/issues/432#issuecomment-1668822286). - - For example, if batch = 3 and seqlen = 6, the attention_mask_in_length is: - ``` - [ - [2, 3, 0, 0, 0, 0], - [3, 2, 0, 0, 0, 0], - [6, 0, 0, 0, 0, 0] - ] - ``` - , which refers to the 3D-attention mask: - ``` - [ - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [0, 0, 1, 0, 0, 0], - [0, 0, 1, 1, 0, 0], - [0, 0, 1, 1, 1, 0], - [0, 0, 0, 0, 0, 1] - ], - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [0, 0, 0, 1, 0, 0], - [0, 0, 0, 1, 1, 0], - [0, 0, 0, 0, 0, 1] - ], - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [1, 1, 1, 1, 0, 0], - [1, 1, 1, 1, 1, 0], - [1, 1, 1, 1, 1, 1] - ] - ] - ```. - - Arguments: - hidden_states: (batch, seqlen, ...) - attention_mask_in_length: (batch, seqlen), int, a nonzero number (e.g., 1, 2, 3, etc.) means length of concatenated sequence in b-th batch, and 0 means none. - Return: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. - indices: (total_nnz), the indices of non-masked tokens from the flattened input sequence. - cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. - max_seqlen_in_batch: int - """ - length = attention_mask_in_length.sum(dim=-1) - seqlen = attention_mask_in_length.size(-1) - attention_mask_2d = torch.arange(seqlen, device=length.device, dtype=length.dtype).expand(len(length), seqlen) < length.unsqueeze(1) - real_indices_idx = torch.nonzero(attention_mask_in_length.flatten(), as_tuple=False).flatten() - seqlens_in_batch = attention_mask_in_length.flatten()[real_indices_idx] - indices = torch.nonzero(attention_mask_2d.flatten(), as_tuple=False).flatten() - max_seqlen_in_batch = seqlens_in_batch.max().item() - cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) - # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the - # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim - # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to - # index with integer indices. Moreover, torch's index is a bit slower than it needs to be, - # so we write custom forward and backward to make it a bit faster. - return ( - index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices), - indices, - cu_seqlens, - max_seqlen_in_batch, - ) - - -def pad_input(hidden_states, indices, batch, seqlen): - """ - Arguments: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. - indices: (total_nnz), the indices that represent the non-masked tokens of the original padded input sequence. - batch: int, batch size for the padded sequence. - seqlen: int, maximum sequence length for the padded sequence. - Return: - hidden_states: (batch, seqlen, ...) - """ - dim = hidden_states.shape[-1] - # output = torch.zeros((batch * seqlen), dim, device=hidden_states.device, dtype=hidden_states.dtype) - # output[indices] = hidden_states - output = index_put_first_axis(hidden_states, indices, batch * seqlen) - return rearrange(output, "(b s) ... -> b s ...", b=batch) diff --git a/build/torch210-cu128-x86_64-windows/flash-attn2/__init__.py b/build/torch210-cu128-x86_64-windows/flash-attn2/__init__.py deleted file mode 100644 index 03dbc1afe1cf156661a2b1b22003cd5f599a0309..0000000000000000000000000000000000000000 --- a/build/torch210-cu128-x86_64-windows/flash-attn2/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -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/torch210-cu128-x86_64-windows/flash_attn_interface.py b/build/torch210-cu128-x86_64-windows/flash_attn_interface.py deleted file mode 100644 index 73ab1f72743f182449cd5b3197108b738a84c9fb..0000000000000000000000000000000000000000 --- a/build/torch210-cu128-x86_64-windows/flash_attn_interface.py +++ /dev/null @@ -1,1620 +0,0 @@ -# Copyright (c) 2023, Tri Dao. - -from typing import Optional, Sequence, Tuple, Union - -import torch -import torch.nn as nn -import os - -# # isort: off -# # We need to import the CUDA kernels after importing torch -# USE_TRITON_ROCM = os.getenv("FLASH_ATTENTION_TRITON_AMD_ENABLE", "FALSE") == "TRUE" -# if USE_TRITON_ROCM: -# from .flash_attn_triton_amd import interface_fa as flash_attn -# else: -# import flash_attn_2_cuda as flash_attn - - -from ._ops import ops as flash_attn - -# # isort: on - -def maybe_contiguous(x): - return x.contiguous() if x is not None and x.stride(-1) != 1 else x - - -def _get_device(): - if torch.xpu.is_available(): - return "xpu" - elif torch.cuda.is_available(): - return "cuda" - else: - return "cpu" - -_XPU_AVAILABLE = torch.xpu.is_available() if hasattr(torch, "xpu") else False # TODO remove hasattr check when bwd is supported on XPU - - -def _get_block_size_n(device, head_dim, is_dropout, is_causal): - # This should match the block sizes in the CUDA kernel - assert head_dim <= 256 - major, minor = torch.cuda.get_device_capability(device) - is_sm8x = major == 8 and minor > 0 # Only include sm86 and sm89, exclude sm80 (A100) - is_sm80 = major == 8 and minor == 0 - is_sm90 = major == 9 and minor == 0 - if head_dim <= 32: - return 128 - if head_dim <= 64: - return 128 if not is_dropout else 64 - elif head_dim <= 96: - return 64 - elif head_dim <= 128: - if is_sm8x: - return 64 if (not is_dropout and is_causal) else 32 - else: - return 64 if not is_dropout else 32 - elif head_dim <= 192: - return 64 - elif head_dim <= 224: - return 64 - elif head_dim <= 256: - return 64 - - -def round_multiple(x, m): - return (x + m - 1) // m * m - - -# torch.compile() support is only enabled for pytorch >= 2.4 -# The reason for this is that we are using the new custom_op and register_fake -# APIs, which support inplace modification of inputs in the function itself -if torch.__version__ >= "2.4.0": - _torch_custom_op_wrapper = torch.library.custom_op - _torch_register_fake_wrapper = torch.library.register_fake -else: - def noop_custom_op_wrapper(name, fn=None, /, *, mutates_args, device_types=None, schema=None): - def wrap(func): - return func - if fn is None: - return wrap - return fn - def noop_register_fake_wrapper(op, fn=None, /, *, lib=None, _stacklevel=1): - def wrap(func): - return func - if fn is None: - return wrap - return fn - _torch_custom_op_wrapper = noop_custom_op_wrapper - _torch_register_fake_wrapper = noop_register_fake_wrapper - - -@_torch_custom_op_wrapper("flash_attn::_flash_attn_forward", mutates_args=(), device_types=_get_device()) -def _flash_attn_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - return_softmax: bool -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - out, softmax_lse, S_dmask, rng_state = flash_attn.fwd( - q, - k, - v, - None, - alibi_slopes, - dropout_p, - softmax_scale, - causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - None, - ) - return out, softmax_lse, S_dmask, rng_state - - -@_torch_register_fake_wrapper("flash_attn::_flash_attn_forward") -def _flash_attn_forward_fake( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - return_softmax: bool -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - batch_size, seqlen_q, num_heads, head_size = q.shape - seqlen_k = k.shape[1] - out = torch.empty_like(q) - softmax_lse = torch.empty((batch_size, num_heads, seqlen_q), dtype=torch.float32, device=q.device, layout=q.layout) - p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) - if return_softmax: - p = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128), round_multiple(seqlen_k, 128)), dtype=q.dtype, device=q.device, layout=q.layout) - rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) - - return out, softmax_lse, p, rng_state - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_forward = torch.ops.flash_attn._flash_attn_forward -else: - _wrapped_flash_attn_forward = _flash_attn_forward - - -@_torch_custom_op_wrapper("flash_attn::_flash_attn_varlen_forward", mutates_args=(), device_types=_get_device()) -def _flash_attn_varlen_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - alibi_slopes: Optional[torch.Tensor] = None, - return_softmax: bool = False, - block_table: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - out, softmax_lse, S_dmask, rng_state = flash_attn.varlen_fwd( - q, - k, - v, - None, - cu_seqlens_q, - cu_seqlens_k, - seqused_k, - leftpad_k, - block_table, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - zero_tensors, - causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - None, - ) - # if out.isnan().any() or softmax_lse.isnan().any(): - # breakpoint() - return out, softmax_lse, S_dmask, rng_state - - -@_torch_register_fake_wrapper("flash_attn::_flash_attn_varlen_forward") -def _flash_attn_varlen_forward_fake( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - alibi_slopes: Optional[torch.Tensor] = None, - return_softmax: bool = False, - block_table: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - paged_kv = block_table is not None - batch_size = cu_seqlens_q.numel() - 1 - total_q, num_heads, _ = q.shape - - out = torch.empty_like(q) - softmax_lse = torch.empty((num_heads, total_q), dtype=torch.float32, device=q.device, layout=q.layout) - p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) - seqlen_q_rounded = round_multiple(max_seqlen_q, 128) - seqlen_k_rounded = round_multiple(max_seqlen_k, 128) - if return_softmax: - p = torch.empty((batch_size, num_heads, seqlen_q_rounded, seqlen_k_rounded), dtype=q.dtype, device=q.device, layout=q.layout) - rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) - return out, softmax_lse, p, rng_state - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_varlen_forward = torch.ops.flash_attn._flash_attn_varlen_forward -else: - _wrapped_flash_attn_varlen_forward = _flash_attn_varlen_forward - - -@_torch_custom_op_wrapper("flash_attn::_flash_attn_backward", mutates_args=("dq", "dk", "dv"), device_types=_get_device()) -def _flash_attn_backward( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, -) -> 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)] - ( - dq, - dk, - dv, - softmax_d, - ) = flash_attn.bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - alibi_slopes, - dropout_p, - softmax_scale, - causal, - window_size_left, - window_size_right, - softcap, - deterministic, - None, - rng_state, - ) - return softmax_d - - -@_torch_register_fake_wrapper("flash_attn::_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, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, -) -> torch.Tensor: - dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] - if dq is None: - dq = torch.empty_like(q) - if dk is None: - dk = torch.empty_like(k) - if dv is None: - dv = torch.empty_like(v) - batch_size, seqlen_q, num_heads, _ = q.shape - softmax_d = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128)), device=q.device, dtype=torch.float32) - - return softmax_d - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_backward = torch.ops.flash_attn._flash_attn_backward -else: - _wrapped_flash_attn_backward = _flash_attn_backward - - -@_torch_custom_op_wrapper("flash_attn::_flash_attn_varlen_backward", mutates_args=("dq", "dk", "dv"), device_types=_get_device()) -def _flash_attn_varlen_backward( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> 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)] - ( - dq, - dk, - dv, - softmax_d, - ) = flash_attn.varlen_bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - zero_tensors, - causal, - window_size_left, - window_size_right, - softcap, - deterministic, - None, - rng_state, - ) - # if dk.isnan().any() or dk.isnan().any() or dv.isnan().any() or softmax_d.isnan().any(): - # breakpoint() - return softmax_d - - -@_torch_register_fake_wrapper("flash_attn::_flash_attn_varlen_backward") -def _flash_attn_varlen_backward_fake( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> torch.Tensor: - dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] - batch_size = cu_seqlens_q.numel() - 1 - total_q, num_heads, _ = q.shape - - if dq is None: - dq = torch.empty_like(q) - if dk is None: - dk = torch.empty_like(k) - if dv is None: - dv = torch.empty_like(v) - softmax_d = torch.empty((num_heads, total_q + 128 * batch_size), device=q.device, dtype=torch.float32) - - return softmax_d - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_varlen_backward = torch.ops.flash_attn._flash_attn_varlen_backward -else: - _wrapped_flash_attn_varlen_backward = _flash_attn_varlen_backward - - -class FlashAttnQKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and qkv.requires_grad - if softmax_scale is None: - softmax_scale = qkv.shape[-1] ** (-0.5) - q, k, v = qkv[:, :, 0].detach(), qkv[:, :, 1].detach(), qkv[:, :, 2].detach() - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) - dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dqkv[:, :, 0], - dqkv[:, :, 1], - dqkv[:, :, 2], - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dqkv = dqkv[..., : dout.shape[-1]] # We could have padded the head dimension - return dqkv, None, None, None, None, None, None, None, None, None - - -class FlashAttnVarlenQKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - cu_seqlens, - max_seqlen, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and qkv.requires_grad - if softmax_scale is None: - softmax_scale = qkv.shape[-1] ** (-0.5) - q, k, v = qkv[:, 0].detach(), qkv[:, 1].detach(), qkv[:, 2].detach() - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens, - cu_seqlens, - max_seqlen, - max_seqlen, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=None, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens, rng_state) - ctx.dropout_p = dropout_p - ctx.max_seqlen = max_seqlen - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens, rng_state = ctx.saved_tensors - qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) - dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dqkv[:, 0], - dqkv[:, 1], - dqkv[:, 2], - cu_seqlens, - cu_seqlens, - ctx.max_seqlen, - ctx.max_seqlen, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - 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 FlashAttnKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - kv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, kv] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - k, v = kv[:, :, 0].detach(), kv[:, :, 1].detach() - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - dq = torch.empty_like(q) - kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) - dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dkv[:, :, 0], - dkv[:, :, 1], - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dkv = dkv[..., : dout.shape[-1]] - return dq, dkv, None, None, None, None, None, None, None, None, None - - -class FlashAttnVarlenKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, kv] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - k, v = kv[:, 0].detach(), kv[:, 1].detach() - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=None, - ) - if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state - ) - ctx.dropout_p = dropout_p - 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.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors - dq = torch.empty_like(q) - kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) - dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dkv[:, 0], - dkv[:, 1], - cu_seqlens_q, - cu_seqlens_k, - ctx.max_seqlen_q, - ctx.max_seqlen_k, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dkv = dkv[..., : dout.shape[-1]] - return dq, dkv, None, None, None, None, None, None, None, None, None, None, None, None, None - - -class FlashAttnFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - k, - v, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, k, v] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dk = dk[..., : dout.shape[-1]] - dv = dv[..., : dout.shape[-1]] - return dq, dk, dv, 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, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - block_table, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, k, v] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=block_table, - ) - if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state - ) - ctx.dropout_p = dropout_p - 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.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors - dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - ctx.max_seqlen_q, - ctx.max_seqlen_k, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dk = dk[..., : dout.shape[-1]] - dv = dv[..., : dout.shape[-1]] - return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None - - -def flash_attn_qkvpacked_func( - qkv, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # <=0.0 means deactivate - alibi_slopes=None, - deterministic=False, - 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, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - False if _XPU_AVAILABLE else torch.is_grad_enabled(), - ) - - -def flash_attn_kvpacked_func( - q, - kv, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, -): - """dropout_p should be set to 0.0 during evaluation - If 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 K, V. - 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) - kv: (batch_size, seqlen, 2, 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. - softcap: float. Anything > 0 activates softcapping 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). - 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 FlashAttnKVPackedFunc.apply( - q, - kv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - False if _XPU_AVAILABLE else torch.is_grad_enabled(), - ) - - -def flash_attn_func( - q, - k, - v, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - 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). - 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 FlashAttnFunc.apply( - q, - k, - v, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - False if _XPU_AVAILABLE else torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_qkvpacked_func( - qkv, - cu_seqlens, - max_seqlen, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - 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_varlen_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_varlen_kvpacked_func and flash_attn_varlen_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: (total, 3, nheads, headdim), where total = total number of tokens in the batch. - cu_seqlens: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into qkv. - max_seqlen: int. Maximum sequence length in the batch. - 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenQKVPackedFunc.apply( - qkv, - cu_seqlens, - max_seqlen, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - False if _XPU_AVAILABLE else torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_kvpacked_func( - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, -): - """dropout_p should be set to 0.0 during evaluation - If 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 K, V. - 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: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch. - kv: (total_k, 2, nheads_k, headdim), where total_k = total number of key tokens in the batch. - cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into q. - cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into kv. - max_seqlen_q: int. Maximum query sequence length in the batch. - max_seqlen_k: int. Maximum key sequence length in the batch. - 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 + 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenKVPackedFunc.apply( - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - False if _XPU_AVAILABLE else torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_func( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, - block_table=None, -): - """dropout_p should be set to 0.0 during evaluation - Supports multi-query and grouped-query attention (MQA/GQA) by passing in K, V 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: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch. - k: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. - v: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. - cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into q. - cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into kv. - max_seqlen_q: int. Maximum query sequence length in the batch. - max_seqlen_k: int. Maximum key sequence length in the batch. - 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 + 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenFunc.apply( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - block_table, - False if _XPU_AVAILABLE or q.device.type == "cpu" else torch.is_grad_enabled(), - ) - - -def flash_attn_with_kvcache( - q, - k_cache, - v_cache, - k=None, - v=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, - block_table: Optional[torch.Tensor] = None, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - rotary_interleaved=True, - alibi_slopes=None, - num_splits=0, - 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 block_table, - or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_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) if there's no block_table, - or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_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). Similar to k. - 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. - block_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). - 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. - 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" - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-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) - cache_batch_idx = maybe_contiguous(cache_batch_idx) - block_table = maybe_contiguous(block_table) - out, softmax_lse = flash_attn.fwd_kvcache( - q, - k_cache, - v_cache, - k, - v, - cache_seqlens, - rotary_cos, - rotary_sin, - cache_batch_idx, - cache_leftpad, - block_table, - alibi_slopes, - None, - softmax_scale, - causal, - window_size[0], - window_size[1], - softcap, - rotary_interleaved, - num_splits, - ) - return (out, softmax_lse) if return_softmax_lse else out diff --git a/build/torch210-cu128-x86_64-windows/layers/patch_embed.py b/build/torch210-cu128-x86_64-windows/layers/patch_embed.py deleted file mode 100644 index 70c64fd30d36c3fa04a09278053065f2699fd4a5..0000000000000000000000000000000000000000 --- a/build/torch210-cu128-x86_64-windows/layers/patch_embed.py +++ /dev/null @@ -1,67 +0,0 @@ -# We use the same API as https://github.com/rwightman/pytorch-image-models/blob/v0.6.11/timm/models/layers/patch_embed.py -# But we use nn.Linear instead of Conv2d and it's about 8x faster. - -from functools import partial - -import torch.nn as nn -from einops import rearrange -from torch import _assert -from torch.nn.modules.utils import _pair - -try: - from flash_attn.ops.fused_dense import FusedDense -except ImportError: - FusedDense = None - - -class PatchEmbed(nn.Module): - """2D Image to Patch Embedding""" - - def __init__( - self, - img_size=224, - patch_size=16, - in_chans=3, - embed_dim=768, - norm_layer=None, - flatten=True, - bias=True, - fused_bias_fc=False, - ): - super().__init__() - img_size = _pair(img_size) - patch_size = _pair(patch_size) - self.img_size = img_size - self.patch_size = patch_size - self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1]) - self.num_patches = self.grid_size[0] * self.grid_size[1] - self.flatten = flatten - if fused_bias_fc and FusedDense is None: - raise ImportError("fused_dense is not installed") - - linear_cls = nn.Linear if not fused_bias_fc or not bias else FusedDense - self.proj = linear_cls(in_chans * patch_size[0] * patch_size[1], embed_dim, bias=bias) - self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() - - def forward(self, x): - _, _, H, W = x.shape - _assert( - H == self.img_size[0], - f"Input image height ({H}) doesn't match model ({self.img_size[0]}).", - ) - _assert( - W == self.img_size[1], - f"Input image width ({W}) doesn't match model ({self.img_size[1]}).", - ) - x = self.proj( - rearrange( - x, - "b c (h p1) (w p2) -> b h w (c p1 p2)", - p1=self.patch_size[0], - p2=self.patch_size[1], - ) - ) - if self.flatten: - x = rearrange(x, "b h w c -> b (h w) c") - x = self.norm(x) - return x diff --git a/build/torch210-cu128-x86_64-windows/layers/rotary.py b/build/torch210-cu128-x86_64-windows/layers/rotary.py deleted file mode 100644 index 1b67f847eeae87f8612f54f1f46f242e379117bf..0000000000000000000000000000000000000000 --- a/build/torch210-cu128-x86_64-windows/layers/rotary.py +++ /dev/null @@ -1,483 +0,0 @@ -# Copyright (c) 2025, Tri Dao - -import math -from functools import partial -from typing import Optional, Tuple, Union - -import torch -from torch import Tensor - -from einops import rearrange, repeat -# from flash_attn.ops.triton.rotary import apply_rotary -from ..ops.triton.rotary import apply_rotary - - -def rotate_half(x, interleaved=False): - if not interleaved: - x1, x2 = x.chunk(2, dim=-1) - return torch.cat((-x2, x1), dim=-1) - else: - x1, x2 = x[..., ::2], x[..., 1::2] - return rearrange(torch.stack((-x2, x1), dim=-1), "... d two -> ... (d two)", two=2) - - -def apply_rotary_emb_torch(x, cos, sin, interleaved=False): - """ - x: (batch_size, seqlen, nheads, headdim) - cos, sin: (seqlen, rotary_dim / 2) or (batch_size, seqlen, rotary_dim / 2) - """ - ro_dim = cos.shape[-1] * 2 - assert ro_dim <= x.shape[-1] - cos = repeat(cos, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)") - sin = repeat(sin, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)") - return torch.cat( - [x[..., :ro_dim] * cos + rotate_half(x[..., :ro_dim], interleaved) * sin, x[..., ro_dim:]], - dim=-1, - ) - - -class ApplyRotaryEmb(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x, - cos, - sin, - interleaved=False, - inplace=False, - seqlen_offsets: Union[int, Tensor] = 0, - cu_seqlens: Optional[Tensor] = None, - max_seqlen: Optional[int] = None, - ): - out = apply_rotary( - x, - cos, - sin, - seqlen_offsets=seqlen_offsets, - cu_seqlens=cu_seqlens, - max_seqlen=max_seqlen, - interleaved=interleaved, - inplace=inplace, - ) - if isinstance(seqlen_offsets, int): - ctx.save_for_backward(cos, sin, cu_seqlens) # Can't save int with save_for_backward - ctx.seqlen_offsets = seqlen_offsets - else: - ctx.save_for_backward(cos, sin, cu_seqlens, seqlen_offsets) - ctx.seqlen_offsets = None - ctx.interleaved = interleaved - ctx.inplace = inplace - ctx.max_seqlen = max_seqlen - return out if not inplace else x - - @staticmethod - def backward(ctx, do): - seqlen_offsets = ctx.seqlen_offsets - if seqlen_offsets is None: - cos, sin, cu_seqlens, seqlen_offsets = ctx.saved_tensors - else: - cos, sin, cu_seqlens = ctx.saved_tensors - dx = apply_rotary( - do, - cos, - sin, - seqlen_offsets=seqlen_offsets, - cu_seqlens=cu_seqlens, - max_seqlen=ctx.max_seqlen, - interleaved=ctx.interleaved, - inplace=ctx.inplace, - conjugate=True, - ) - return dx, None, None, None, None, None, None, None - - -def apply_rotary_emb( - x, - cos, - sin, - interleaved=False, - inplace=False, - seqlen_offsets: Union[int, Tensor] = 0, - cu_seqlens: Optional[Tensor] = None, - max_seqlen: Optional[int] = None, -): - """ - Arguments: - x: (batch_size, seqlen, nheads, headdim) if cu_seqlens is None - else (total_seqlen, nheads, headdim) - cos, sin: (seqlen_rotary, rotary_dim / 2) - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead - of 1st half and 2nd half (GPT-NeoX style). - inplace: if True, apply rotary embedding in-place. - seqlen_offsets: (batch_size,) or int. Each sequence in x is shifted by this amount. - Most commonly used in inference when we have KV cache. - cu_seqlens: (batch + 1,) or None - max_seqlen: int - Return: - out: (batch_size, seqlen, nheads, headdim) if cu_seqlens is None - else (total_seqlen, nheads, headdim) - rotary_dim must be <= headdim - Apply rotary embedding to the first rotary_dim of x. - """ - return ApplyRotaryEmb.apply( - x, cos, sin, interleaved, inplace, seqlen_offsets, cu_seqlens, max_seqlen - ) - - -# For backward compatibility -apply_rotary_emb_func = apply_rotary_emb - - -def _apply_rotary_emb_qkv( - qkv, - cos, - sin, - cos_k=None, - sin_k=None, - interleaved=False, - inplace=False, - conjugate=False, - seqlen_offsets: Union[int, Tensor] = 0, - num_heads_q: Optional[int] = None, -): - apply_rotary_fn = partial( - apply_rotary, - interleaved=interleaved, - inplace=inplace, - conjugate=conjugate, - seqlen_offsets=seqlen_offsets - ) - if cos_k is None and sin_k is None and qkv.is_contiguous(): - # Call 1 kernel instead of 2 kernels - # We need qkv to be contiguous so that when we reshape to combine (3, nheads) - # dimensions, we get the same tensor - if qkv.dim() == 5: - batch, seqlen, three, nheads, headdim = qkv.shape - assert three == 3 - # qk = rearrange(qkv[:, :, :2], "b s t h d -> b s (t h) d") - qk = qkv[:, :, :2].reshape(batch, seqlen, -1, headdim) - qk = apply_rotary_fn(qk, cos, sin) - else: - assert qkv.dim() == 4 - assert num_heads_q is not None - num_heads_k = (qkv.shape[2] - num_heads_q) // 2 - assert qkv.shape[2] == num_heads_q + 2 * num_heads_k - qk = qkv[:, :, :num_heads_q + num_heads_k] - qk = apply_rotary_fn(qk, cos, sin) - if not inplace: - if qkv.dim() == 5: - qkv = torch.cat([rearrange(qk, "b s (t h) d -> b s t h d", t=2), qkv[:, :, 2:]], dim=2) - else: - qkv = torch.cat([qk, qkv[:, :, num_heads_q + num_heads_k :]], dim=2) - else: - cos_k = cos if cos_k is None else cos_k - sin_k = sin if sin_k is None else sin_k - if qkv.dim() == 5: - batch, seqlen, three, nheads, headdim = qkv.shape - assert three == 3 - q, k = qkv[:, :, 0], qkv[:, :, 1] - else: - assert qkv.dim() == 4 - assert num_heads_q is not None - num_heads_k = (qkv.shape[2] - num_heads_q) // 2 - assert qkv.shape[2] == num_heads_q + 2 * num_heads_k - q, k = qkv[:, :, :num_heads_q], qkv[:, :, num_heads_q : num_heads_q + num_heads_k] - q = apply_rotary_fn(q, cos, sin) - k = apply_rotary_fn(k, cos_k, sin_k) - if not inplace: - if qkv.dim() == 5: - qkv = torch.stack([q, k, qkv[:, :, 2]], dim=2) - else: - qkv = torch.cat([q, k, qkv[:, :, num_heads_q + num_heads_k:]], dim=2) - return qkv - - -class ApplyRotaryEmbQKV_(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - cos, - sin, - cos_k=None, - sin_k=None, - interleaved=False, - seqlen_offsets: Union[int, torch.Tensor] = 0, - num_heads_q: Optional[int] = None, - ): - # apply_rotary_emb_qkv_inplace( - qkv = _apply_rotary_emb_qkv( - qkv, cos, sin, cos_k, sin_k, interleaved=interleaved, inplace=True, - seqlen_offsets=seqlen_offsets, num_heads_q=num_heads_q, - ) - if isinstance(seqlen_offsets, int): - ctx.save_for_backward(cos, sin, cos_k, sin_k) - ctx.seqlen_offsets = seqlen_offsets - else: - ctx.save_for_backward(cos, sin, cos_k, sin_k, seqlen_offsets) - ctx.seqlen_offsets = None - ctx.interleaved = interleaved - ctx.num_heads_q = num_heads_q - return qkv - - @staticmethod - def backward(ctx, dqkv): - seqlen_offsets = ctx.seqlen_offsets - if seqlen_offsets is None: - cos, sin, cos_k, sin_k, seqlen_offsets = ctx.saved_tensors - else: - cos, sin, cos_k, sin_k = ctx.saved_tensors - dqkv = _apply_rotary_emb_qkv( - dqkv, cos, sin, cos_k, sin_k, interleaved=ctx.interleaved, inplace=True, - seqlen_offsets=seqlen_offsets, num_heads_q=ctx.num_heads_q, conjugate=True, - ) - return dqkv, None, None, None, None, None, None, None - - -def apply_rotary_emb_qkv_( - qkv, - cos, - sin, - cos_k=None, - sin_k=None, - interleaved=False, - seqlen_offsets: Union[int, torch.Tensor] = 0, - num_heads_q: Optional[int] = None, -): - """ - Arguments: - qkv: (batch_size, seqlen, 3, nheads, headdim) or (batch_size, seqlen, num_heads_q + 2 * num_heads_k, headdim). - If qkv has shape (batch_size, seqlen, num_heads_q + 2 * num_heads_k, headdim) (e.g. MQA / GQA), - then num_heads_q must be provided. - cos, sin: (seqlen, rotary_dim / 2) - cos_k, sin_k: (seqlen, rotary_dim / 2), optional - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead of - 1st half and 2nd half (GPT-NeoX style). - seqlen_offsets: (batch_size,) or int. Each sequence in Q and K is shifted by this amount. - Most commonly used in inference when we have KV cache. - Return: - qkv: (batch_size, seqlen, 3, nheads, headdim) or (batch_size, seqlen, num_heads_q + 2 * num_heads_k, headdim) - rotary_dim must be <= headdim - Apply rotary embedding *inplace* to the first rotary_dim of Q and K. - """ - return ApplyRotaryEmbQKV_.apply( - qkv, cos, sin, cos_k, sin_k, interleaved, seqlen_offsets, num_heads_q - ) - - -class ApplyRotaryEmbKV_(torch.autograd.Function): - - @staticmethod - def forward(ctx, kv, cos, sin, interleaved=False, seqlen_offsets: Union[int, torch.Tensor] = 0): - batch, seqlen, two, nheads, headdim = kv.shape - assert two == 2 - k = kv[:, :, 0] - apply_rotary( - k, cos, sin, seqlen_offsets=seqlen_offsets, interleaved=interleaved, inplace=True - ) - if isinstance(seqlen_offsets, int): - ctx.save_for_backward(cos, sin) # Can't save int with save_for_backward - ctx.seqlen_offsets = seqlen_offsets - else: - ctx.save_for_backward(cos, sin, seqlen_offsets) - ctx.seqlen_offsets = None - ctx.interleaved = interleaved - return kv - - @staticmethod - def backward(ctx, dkv): - seqlen_offsets = ctx.seqlen_offsets - if seqlen_offsets is None: - cos, sin, seqlen_offsets = ctx.saved_tensors - else: - cos, sin = ctx.saved_tensors - apply_rotary( - dkv[:, :, 0], - cos, - sin, - seqlen_offsets=seqlen_offsets, - interleaved=ctx.interleaved, - inplace=True, - conjugate=True, - ) - return dkv, None, None, None, None - - -apply_rotary_emb_kv_ = ApplyRotaryEmbKV_.apply - - -def apply_rotary_emb_kv_( - kv, - cos, - sin, - interleaved=False, - seqlen_offsets: Union[int, torch.Tensor] = 0, -): - """ - Arguments: - kv: (batch_size, seqlen, 2, nheads, headdim) - cos, sin: (seqlen, rotary_dim / 2) - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead of - 1st half and 2nd half (GPT-NeoX style). - seqlen_offsets: (batch_size,) or int. Each sequence in Q and K is shifted by this amount. - Most commonly used in inference when we have KV cache. - Return: - kv: (batch_size, seqlen, 2, nheads, headdim) - rotary_dim must be <= headdim - Apply rotary embedding *inplace* to the first rotary_dim of K. - """ - return ApplyRotaryEmbKV_.apply(kv, cos, sin, interleaved, seqlen_offsets) - - -class RotaryEmbedding(torch.nn.Module): - """ - The rotary position embeddings from RoFormer_ (Su et. al). - A crucial insight from the method is that the query and keys are - transformed by rotation matrices which depend on the relative positions. - - Other implementations are available in the Rotary Transformer repo_ and in - GPT-NeoX_, GPT-NeoX was an inspiration - - .. _RoFormer: https://arxiv.org/abs/2104.09864 - .. _repo: https://github.com/ZhuiyiTechnology/roformer - .. _GPT-NeoX: https://github.com/EleutherAI/gpt-neox - - If scale_base is not None, this implements XPos (Sun et al., https://arxiv.org/abs/2212.10554). - A recommended value for scale_base is 512: https://github.com/HazyResearch/flash-attention/issues/96 - Reference: https://github.com/sunyt32/torchscale/blob/main/torchscale/component/xpos_relative_position.py - """ - - def __init__( - self, - dim: int, - base=10000.0, - interleaved=False, - scale_base=None, - device=None, - ): - """ - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead - of 1st half and 2nd half (GPT-NeoX style). - """ - super().__init__() - self.dim = dim - self.base = float(base) - # Generate and save the inverse frequency buffer (non trainable) - inv_freq = self._compute_inv_freq(device) - self.register_buffer("inv_freq", inv_freq, persistent=False) - self.interleaved = interleaved - self.scale_base = scale_base - scale = ( - (torch.arange(0, dim, 2, device=device, dtype=torch.float32) + 0.4 * dim) / (1.4 * dim) - if scale_base is not None - else None - ) - self.register_buffer("scale", scale, persistent=False) - - self._seq_len_cached = 0 - self._cos_cached = None - self._sin_cached = None - self._cos_k_cached = None - self._sin_k_cached = None - - def _compute_inv_freq(self, device=None): - return 1.0 / ( - self.base - ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim) - ) - - def _update_cos_sin_cache(self, seqlen, device=None, dtype=None): - # Reset the tables if the sequence length has changed, - # if we're on a new device (possibly due to tracing for instance), - # or if we're switching from inference mode to training - if ( - seqlen > self._seq_len_cached - or self._cos_cached is None - or self._cos_cached.device != device - or self._cos_cached.dtype != dtype - or (self.training and self._cos_cached.is_inference()) - ): - self._seq_len_cached = seqlen - # We want fp32 here, not self.inv_freq.dtype, since the model could be loaded in bf16 - # And the output of arange can be quite large, so bf16 would lose a lot of precision. - t = torch.arange(seqlen, device=device, dtype=torch.float32) - # We want fp32 here as well since inv_freq will be multiplied with t, and the output - # will be large. Having it in bf16 will lose a lot of precision and cause the - # cos & sin output to change significantly. - # We want to recompute self.inv_freq if it was not loaded in fp32 - if self.inv_freq.dtype != torch.float32: - inv_freq = self._compute_inv_freq(device=device) - else: - inv_freq = self.inv_freq - # Don't do einsum, it converts fp32 to bf16 under AMP - # freqs = torch.einsum("i,j->ij", t, self.inv_freq) - freqs = torch.outer(t, inv_freq) - if self.scale is None: - self._cos_cached = torch.cos(freqs).to(dtype) - self._sin_cached = torch.sin(freqs).to(dtype) - else: - power = ( - torch.arange(seqlen, dtype=self.scale.dtype, device=self.scale.device) - - seqlen // 2 - ) / self.scale_base - scale = self.scale.to(device=power.device) ** rearrange(power, "s -> s 1") - # We want the multiplication by scale to happen in fp32 - self._cos_cached = (torch.cos(freqs) * scale).to(dtype) - self._sin_cached = (torch.sin(freqs) * scale).to(dtype) - self._cos_k_cached = (torch.cos(freqs) / scale).to(dtype) - self._sin_k_cached = (torch.sin(freqs) / scale).to(dtype) - - def forward( - self, - qkv: torch.Tensor, - kv: Optional[torch.Tensor] = None, - seqlen_offset: Union[int, torch.Tensor] = 0, - max_seqlen: Optional[int] = None, - num_heads_q: Optional[int] = None, - ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: - """ - qkv: (batch, seqlen, 3, nheads, headdim) or (batch, seqlen, num_heads_q + 2 * num_heads_k, headdim) - if kv is none, else it's just q of shape (batch, seqlen, nheads, headdim). - If qkv has shape (batch, seqlen, num_heads_q + 2 * num_heads_k, headdim) (e.g. MQA / GQA), - then num_heads_q must be provided. - kv: (batch, seqlen, 2, nheads, headdim) - seqlen_offset: (batch_size,) or int. Each sequence in x is shifted by this amount. - Most commonly used in inference when we have KV cache. - If it's a tensor of shape (batch_size,), then to update the cos / sin cache, one - should pass in max_seqlen, which will update the cos / sin cache up to that length. - Apply rotary embedding *inplace* to qkv and / or kv. - """ - seqlen = qkv.shape[1] - if max_seqlen is not None: - self._update_cos_sin_cache(max_seqlen, device=qkv.device, dtype=qkv.dtype) - elif isinstance(seqlen_offset, int): - self._update_cos_sin_cache(seqlen + seqlen_offset, device=qkv.device, dtype=qkv.dtype) - if kv is None: - return apply_rotary_emb_qkv_( - qkv, - self._cos_cached, - self._sin_cached, - self._cos_k_cached if self.scale is not None else None, - self._sin_k_cached if self.scale is not None else None, - interleaved=self.interleaved, - seqlen_offsets=seqlen_offset, - num_heads_q=num_heads_q, - ) - else: - q = qkv - q = apply_rotary_emb_func( - q, - self._cos_cached, - self._sin_cached, - interleaved=self.interleaved, - inplace=True, - seqlen_offsets=seqlen_offset, - ) - kv = apply_rotary_emb_kv_( - kv, - self._cos_cached if self.scale is None else self._cos_k_cached, - self._sin_cached if self.scale is None else self._sin_k_cached, - interleaved=self.interleaved, - seqlen_offsets=seqlen_offset, - ) - return q, kv diff --git a/build/torch210-cu128-x86_64-windows/metadata.json b/build/torch210-cu128-x86_64-windows/metadata.json deleted file mode 100644 index 9cf5deed9898dce769f4cc73913d3530b92a0bd8..0000000000000000000000000000000000000000 --- a/build/torch210-cu128-x86_64-windows/metadata.json +++ /dev/null @@ -1,4 +0,0 @@ -{ - "version": 1, - "python-depends": [] -} \ No newline at end of file diff --git a/build/torch210-cu128-x86_64-windows/ops/activations.py b/build/torch210-cu128-x86_64-windows/ops/activations.py deleted file mode 100644 index d407510070e2b369b0fd8855ed8969cc80b6bc2f..0000000000000000000000000000000000000000 --- a/build/torch210-cu128-x86_64-windows/ops/activations.py +++ /dev/null @@ -1,135 +0,0 @@ -# Copied from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/model/layers/activations.py -import math - -import torch -import torch.nn as nn -import torch.nn.functional as F - -# 1/sqrt(2*pi)-> 0.3989423 -# 1/sqrt(2) -> 0.70710678 -# sqrt(2/pi) -> 0.79788456 - -# this function is tanh approximation of gelu -# actual gelu is: -# x * 0.5 * (1.0 + torch.erf(x * 0.70710678)) -@torch.jit.script -def bias_gelu(y, bias): - x = bias + y - return (x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))).to(dtype=y.dtype) - - -# gradient of tanh approximation of gelu -# gradient of actual gelu is: -# 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x) -@torch.jit.script -def bias_gelu_back(g, y, bias): - """Assume that y has shape (B, D) and bias has shape (D)""" - x = bias + y - tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)) - # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243 - ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( - 1 + tanh_out - ) - grad_y = ff * g - return grad_y.to(dtype=y.dtype), grad_y.sum(dim=(0), dtype=bias.dtype) - - -class GeLUFunction(torch.autograd.Function): - @staticmethod - # bias is an optional argument - def forward(ctx, input, bias): - ctx.save_for_backward(input, bias) - return bias_gelu(input, bias) - - @staticmethod - def backward(ctx, grad_output): - input, bias = ctx.saved_tensors - tmp = bias_gelu_back(grad_output, input, bias) - return tmp, tmp - - -bias_gelu_impl = GeLUFunction.apply - -# this function is tanh approximation of gelu -# actual gelu is: -# x * 0.5 * (1.0 + torch.erf(x * 0.70710678)) -@torch.jit.script -def gelu_fwd(x): - return (x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))).to(dtype=x.dtype) - - -# gradient of tanh approximation of gelu -# gradient of actual gelu is: -# 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x) -@torch.jit.script -def gelu_bwd(g, x): - tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)) - # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243 - ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( - 1 + tanh_out - ) - return (ff * g).to(dtype=x.dtype) - - -class FastGeLUFunction(torch.autograd.Function): - @staticmethod - # bias is an optional argument - def forward(ctx, input): - ctx.save_for_backward(input) - return gelu_fwd(input) - - @staticmethod - def backward(ctx, grad_output): - (input,) = ctx.saved_tensors - tmp = gelu_bwd(grad_output, input) - return tmp - - -fast_gelu_impl = FastGeLUFunction.apply - - -@torch.jit.script -def relu_bwd(g, x): - return torch.where(x >= 0, g, 0.0).to(dtype=x.dtype) - - -@torch.jit.script -def sqrelu_fwd(x): - r = F.relu(x) - return (r * r).to(dtype=x.dtype) - - -@torch.jit.script -def sqrelu_bwd(g, x): - return (2.0 * g * F.relu(x)).to(dtype=x.dtype) - - -swiglu_fwd_codestring = """ -template T swiglu_fwd(T x, T y) { - return float(x) * float(y) / (1.0f + ::exp(-float(x))); -} -""" -swiglu_bwd_codestring = """ -template void swiglu_bwd(T x, T y, T g, T& dx, T& dy) { - float x_sigmoid = 1.0f / (1.0f + ::exp(-float(x))); - dx = x_sigmoid * (1 + float(x) * (1.0f - x_sigmoid)) * float(g) * float(y); - dy = float(x) * x_sigmoid * float(g); -} -""" -swiglu_fwd = torch.cuda.jiterator._create_jit_fn(swiglu_fwd_codestring) -swiglu_bwd = torch.cuda.jiterator._create_multi_output_jit_fn(swiglu_bwd_codestring, num_outputs=2) - - -class SwiGLUFunction(torch.autograd.Function): - - @staticmethod - def forward(ctx, x, y): - ctx.save_for_backward(x, y) - return swiglu_fwd(x, y) - - @staticmethod - def backward(ctx, dout): - x, y = ctx.saved_tensors - return swiglu_bwd(x, y, dout) - -swiglu = SwiGLUFunction.apply diff --git a/build/torch210-cu128-x86_64-windows/ops/fused_dense.py b/build/torch210-cu128-x86_64-windows/ops/fused_dense.py deleted file mode 100644 index 61ac40b4aaf97101f63ad411364fcc9f78938596..0000000000000000000000000000000000000000 --- a/build/torch210-cu128-x86_64-windows/ops/fused_dense.py +++ /dev/null @@ -1,688 +0,0 @@ -# Copyright (c) 2023, Tri Dao. -# Inspired by https://github.com/NVIDIA/apex/blob/master/apex/fused_dense/fused_dense.py -# We make it work with pytorch amp and with bfloat16. -# The TensorParallel linear modules are inspired by https://github.com/NVIDIA/apex/blob/master/apex/transformer/tensor_parallel/layers.py -from functools import partial -from typing import Optional - -# import fused_dense_cuda # from apex -import fused_dense_lib as fused_dense_cuda -import torch -import torch.nn as nn -import torch.nn.functional as F -from torch import Tensor -from torch.distributed import ProcessGroup - -from flash_attn.utils.torch import custom_fwd, custom_bwd -from flash_attn.ops.activations import gelu_bwd, relu_bwd, sqrelu_bwd, sqrelu_fwd -from flash_attn.utils.distributed import ( - all_gather_raw, - all_reduce, - all_reduce_raw, - reduce_scatter, - reduce_scatter_raw, -) - - -class FusedDenseFunc(torch.autograd.Function): - @staticmethod - @custom_fwd - def forward( - ctx, x, weight, bias, return_residual=False, process_group=None, sequence_parallel=True - ): - """ - If process_group is not None and sequence_parallel=True, we're doing Tensor Parallel - with sequence parallelism: we do an all_gather_raw of x before doing the matmul. - """ - ctx.compute_weight_gradient = weight.requires_grad - ctx.return_residual = return_residual - ctx.process_group = process_group - ctx.sequence_parallel = sequence_parallel - - if torch.is_autocast_enabled(): - x = x.to(dtype=torch.get_autocast_gpu_dtype()) - x = x.contiguous() - if process_group is not None and sequence_parallel: - # We want to kick off the all_gather early, before weight dtype conversion - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - else: - total_x = x - - if torch.is_autocast_enabled(): - weight = weight.to(dtype=torch.get_autocast_gpu_dtype()) - bias = bias.to(dtype=torch.get_autocast_gpu_dtype()) if bias is not None else None - weight = weight.contiguous() - if process_group is not None and sequence_parallel: - handle_x.wait() - batch_shape, n = total_x.shape[:-1], total_x.shape[-1] - batch_dim = batch_shape.numel() - # https://github.com/pytorch/pytorch/blob/5b51849b48a7dbccd297286cc0110def4706f9e7/aten/src/ATen/native/cuda/Blas.cpp#L174 - if min(batch_dim, n, *weight.shape) > 65535 * 32: - raise RuntimeError("fused_dense only supports matrix dims <= 2M") - output = F.linear(total_x, weight, bias) - if ctx.compute_weight_gradient: - ctx.save_for_backward(x, weight) - else: - ctx.save_for_backward(weight) - return output if not return_residual else (output, x) - - @staticmethod - @custom_bwd - def backward(ctx, grad_output, *args): - grad_output = grad_output.contiguous() - if ctx.return_residual: - (grad_input,) = args - grad_input = grad_input.contiguous() - process_group = ctx.process_group - sequence_parallel = ctx.sequence_parallel - if ctx.compute_weight_gradient: - x, weight = ctx.saved_tensors - if process_group is not None and sequence_parallel: - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - else: - total_x = x - else: - (weight,) = ctx.saved_tensors - total_x = None - batch_shape = grad_output.shape[:-1] - batch_dim = batch_shape.numel() - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - if ctx.needs_input_grad[0]: - if not ctx.return_residual: - grad_input = F.linear(grad_output, weight.t()) - else: - grad_input = torch.addmm( - grad_input.reshape(batch_dim, grad_input.shape[-1]), grad_output, weight - ) - grad_input = grad_input.reshape(*batch_shape, grad_input.shape[-1]) - if process_group is not None: - reduce_fn = reduce_scatter_raw if sequence_parallel else all_reduce_raw - grad_input, handle_grad_input = reduce_fn(grad_input, process_group, async_op=True) - else: - grad_input = None - if ctx.needs_input_grad[1]: - assert ctx.compute_weight_gradient - if process_group is not None and sequence_parallel: - handle_x.wait() - grad_weight, grad_bias = fused_dense_cuda.linear_bias_wgrad( - total_x.reshape(batch_dim, total_x.shape[-1]), grad_output, ctx.needs_input_grad[2] - ) - else: - grad_weight = None - grad_bias = grad_output if ctx.needs_input_grad[2] else None - if process_group is not None and ctx.needs_input_grad[0]: - handle_grad_input.wait() - return grad_input, grad_weight, grad_bias, None, None, None - - -def fused_dense_func( - x: Tensor, - weight: Tensor, - bias: Optional[Tensor] = None, - return_residual: bool = False, - process_group: Optional[ProcessGroup] = None, - sequence_parallel: bool = True, -): - dtype_eligible = x.dtype in [torch.float16, torch.bfloat16] or ( - x.dtype == torch.float32 and torch.is_autocast_enabled() - ) - if x.is_cuda and weight.is_cuda and (bias is None or bias.is_cuda) and dtype_eligible: - return FusedDenseFunc.apply( - x, weight, bias, return_residual, process_group, sequence_parallel - ) - else: - assert process_group is None - out = F.linear(x, weight, bias) - return out if not return_residual else (out, x) - - -class FusedDense(nn.Linear): - def __init__( - self, - in_features: int, - out_features: int, - bias: bool = True, - return_residual: bool = False, - device=None, - dtype=None, - ) -> None: - super().__init__(in_features, out_features, bias=bias, device=device, dtype=dtype) - self.return_residual = return_residual - - def forward(self, x, process_group=None): - """ - If process_group is not None, we're doing Tensor Parallel with sequence parallelism: - we do an all_gather of x before doing the matmul. - """ - return fused_dense_func( - x, - self.weight, - self.bias, - return_residual=self.return_residual, - process_group=process_group, - ) - - -class ColumnParallelLinear(nn.Linear): - def __init__( - self, - in_features: int, - out_features: int, - process_group: ProcessGroup, - bias: bool = True, - sequence_parallel=True, - multiple_of=1, - device=None, - dtype=None, - ) -> None: - world_size = torch.distributed.get_world_size(process_group) - if out_features % multiple_of: - raise ValueError(f"out_features ({out_features}) must be a multiple of {multiple_of}") - multiple = out_features // multiple_of - # We want to split @multiple across world_size, but it could be an uneven split - div = multiple // world_size - mod = multiple % world_size - # The first @mod ranks get @div + 1 copies, the rest get @div copies - local_multiple = div + int(torch.distributed.get_rank(process_group) < mod) - super().__init__( - in_features, local_multiple * multiple_of, bias=bias, device=device, dtype=dtype - ) - self.process_group = process_group - self.sequence_parallel = sequence_parallel - - def forward(self, x): - # If self.sequence_parallel is True, we're doing Tensor Parallel with sequence parallelism: - # we do an all_gather of x before doing the matmul. - # If not, then the input is already gathered. - return fused_dense_func( - x, - self.weight, - self.bias, - process_group=self.process_group, - sequence_parallel=self.sequence_parallel, - ) - - -class RowParallelLinear(nn.Linear): - def __init__( - self, - in_features: int, - out_features: int, - process_group: ProcessGroup, - bias: bool = True, - sequence_parallel=True, - multiple_of=1, - device=None, - dtype=None, - ) -> None: - world_size = torch.distributed.get_world_size(process_group) - rank = torch.distributed.get_rank(process_group) - if in_features % multiple_of: - raise ValueError(f"in_features ({in_features}) must be a multiple of {multiple_of}") - multiple = in_features // multiple_of - # We want to split @multiple across world_size, but it could be an uneven split - div = multiple // world_size - mod = multiple % world_size - # The first @mod ranks get @div + 1 copies, the rest get @div copies - local_multiple = div + int(torch.distributed.get_rank(process_group) < mod) - # Only rank 0 will have bias - super().__init__( - local_multiple * multiple_of, - out_features, - bias=bias and rank == 0, - device=device, - dtype=dtype, - ) - self.process_group = process_group - self.sequence_parallel = sequence_parallel - - def forward(self, x): - """ - We're doing Tensor Parallel with sequence parallelism: we do the matmul and then - a reduce_scatter of the result. - """ - out = fused_dense_func(x, self.weight, self.bias) - reduce_fn = reduce_scatter if self.sequence_parallel else all_reduce - return reduce_fn(out, self.process_group) - - -class FusedMLPFunc(torch.autograd.Function): - @staticmethod - @custom_fwd - def forward( - ctx, - x, - weight1, - bias1, - weight2, - bias2, - activation="gelu_approx", - save_pre_act=True, - return_residual=False, - checkpoint_lvl=0, - heuristic=0, - process_group=None, - sequence_parallel=True, - ): - """ - If process_group is not None and sequence_parallel=True, we're doing Tensor Parallel - with sequence parallelism: we do an all_gather of x before doing the matmul. - If sequence_parallel=False, then the input is already gathered. - - checkpoint_lvl: - 0: no recomputation in the bwd - 1: recompute gelu_out / relu_out in the bwd - 2: recompute pre_act and gelu_out / relu_out in the bwd - """ - assert -1 <= heuristic <= 4 - assert activation in ["gelu_approx", "relu", "sqrelu"] - if activation == "sqrelu": - assert heuristic == -1 - if not save_pre_act: - checkpoint_lvl = 2 - assert checkpoint_lvl in [0, 1, 2] - ctx.return_residual = return_residual - ctx.process_group = process_group - ctx.sequence_parallel = sequence_parallel - ctx.checkpoint_lvl = checkpoint_lvl - ctx.activation = activation - ctx.heuristic = heuristic - - if torch.is_autocast_enabled(): - x = x.to(dtype=torch.get_autocast_gpu_dtype()) - x = x.contiguous() - if process_group is not None and sequence_parallel: - # We want to kick off the all_gather early, before weight dtype conversion - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - else: - total_x = x - - if torch.is_autocast_enabled(): - dtype = torch.get_autocast_gpu_dtype() - weight1, weight2 = [a.to(dtype=dtype) for a in [weight1, weight2]] - bias1 = bias1.to(dtype=dtype) if bias1 is not None else None - bias2 = bias2.to(dtype=dtype) if bias2 is not None else None - weight1 = weight1.contiguous() - bias1 = bias1.contiguous() if bias1 is not None else None - weight2 = weight2.contiguous() - bias2 = bias2.contiguous() if bias2 is not None else None - if process_group is not None and sequence_parallel: - handle_x.wait() - batch_shape, n = total_x.shape[:-1], total_x.shape[-1] - batch_dim = batch_shape.numel() - # https://github.com/pytorch/pytorch/blob/5b51849b48a7dbccd297286cc0110def4706f9e7/aten/src/ATen/native/cuda/Blas.cpp#L174 - if min(batch_dim, n, *weight1.shape, *weight2.shape) > 65535 * 32: - raise RuntimeError("fused_dense only supports matrix dims <= 2M") - if heuristic == -1: - pre_act = F.linear(total_x, weight1, bias1) - activation_fn = ( - partial(F.gelu, approximate="tanh") - if activation == "gelu_approx" - else (sqrelu_fwd if activation == "sqrelu" else F.relu) - ) - with torch.jit.fuser("fuser2"): - output1 = activation_fn(pre_act) - # This is before adding bias1 - # pre_act = F.linear(total_x.reshape(batch_dim, n), weight1) - # with torch.jit.fuser('fuser2'): - # output1 = bias_gelu(pre_act, bias1) - else: - is_gelu = activation == "gelu_approx" - output1, *rest = fused_dense_cuda.linear_act_forward( - total_x.reshape(batch_dim, n), weight1, bias1, is_gelu, save_pre_act, heuristic - ) - if save_pre_act: - pre_act = rest[0] - output2 = F.linear(output1, weight2, bias2) - if checkpoint_lvl == 0 or (checkpoint_lvl == 1 and activation == "relu"): - # For RELU the pre_act is very small (just a bit-mask) so we just save it - ctx.save_for_backward(x, weight1, weight2, pre_act, output1) - elif checkpoint_lvl == 1: - ctx.save_for_backward(x, weight1, weight2, pre_act) - elif checkpoint_lvl == 2: - ctx.save_for_backward(x, weight1, weight2, bias1) - output2 = output2.reshape(*batch_shape, output2.shape[-1]) - return output2 if not return_residual else (output2, x) - - @staticmethod - @custom_bwd - def backward(ctx, grad_output, *args): - grad_output = grad_output.contiguous() - checkpoint_lvl = ctx.checkpoint_lvl - activation = ctx.activation - activation_fn = ( - partial(F.gelu, approximate="tanh") - if activation == "gelu_approx" - else (sqrelu_fwd if activation == "sqrelu" else F.relu) - ) - if ctx.return_residual: - (grad_input,) = args - grad_input = grad_input.contiguous() - process_group = ctx.process_group - sequence_parallel = ctx.sequence_parallel - x, weight1, weight2, *rest = ctx.saved_tensors - if process_group is None or not sequence_parallel: - total_x = x - batch_shape = grad_output.shape[:-1] - batch_dim = batch_shape.numel() - if checkpoint_lvl in [0, 1]: - if process_group is not None and sequence_parallel: - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - if checkpoint_lvl == 0 or (checkpoint_lvl == 1 and activation == "relu"): - pre_act, output1 = rest - elif checkpoint_lvl == 1: - (pre_act,) = rest - with torch.jit.fuser("fuser2"): - output1 = activation_fn(pre_act) - elif checkpoint_lvl == 2: - (bias1,) = rest - if process_group is not None and sequence_parallel: - total_x, _ = all_gather_raw(x, process_group) - if ctx.heuristic == -1: - pre_act = F.linear(total_x, weight1, bias1) - with torch.jit.fuser("fuser2"): - output1 = activation_fn(pre_act) - else: - output1, pre_act = fused_dense_cuda.linear_act_forward( - total_x.reshape(batch_dim, total_x.shape[-1]), - weight1, - bias1, - activation == "gelu_approx", - True, - ctx.heuristic, - ) - - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - output1 = output1.reshape(batch_dim, output1.shape[-1]) - pre_act = pre_act.reshape(batch_dim, pre_act.shape[-1]) - if ctx.needs_input_grad[3]: - grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad( - output1, grad_output, ctx.needs_input_grad[4] - ) - else: - grad_weight2 = None - grad_bias2 = grad_output if ctx.needs_input_grad[4] else None - if ctx.heuristic == -1: - # grad_pre_act = matmul_dgelu(grad_output, weight2, pre_act) - grad_output1 = F.linear(grad_output, weight2.t()) - activation_grad_fn = ( - gelu_bwd - if activation == "gelu_approx" - else (sqrelu_bwd if activation == "sqrelu" else relu_bwd) - ) - with torch.jit.fuser("fuser2"): - grad_pre_act = activation_grad_fn(grad_output1, pre_act) - else: - # The cublasLt epilogue has to compute both gelu/relu grad and bias grad, we can't - # just compute gelu/relu grad - grad_pre_act, grad_bias1 = fused_dense_cuda.bias_act_linear_dgrad_bgrad( - weight2, grad_output, pre_act, activation == "gelu_approx", ctx.heuristic - ) - if not ctx.needs_input_grad[2]: - grad_bias1 = None - if ctx.needs_input_grad[0]: - if not ctx.return_residual: - grad_input = F.linear(grad_pre_act, weight1.t()) - else: - grad_input = torch.addmm( - grad_input.reshape(batch_dim, grad_input.shape[-1]), grad_pre_act, weight1 - ) - grad_input = grad_input.reshape(*batch_shape, grad_input.shape[-1]) - if process_group is not None: - reduce_fn = reduce_scatter_raw if sequence_parallel else all_reduce_raw - grad_input, handle_grad_input = reduce_fn(grad_input, process_group, async_op=True) - else: - grad_input = None - if ctx.heuristic == -1: - if ctx.needs_input_grad[1]: - if process_group is not None and sequence_parallel and checkpoint_lvl != 2: - handle_x.wait() - grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_wgrad( - total_x.reshape(batch_dim, total_x.shape[-1]), - grad_pre_act, - ctx.needs_input_grad[2], - ) - else: - grad_weight1 = None - grad_bias1 = grad_pre_act if ctx.needs_input_grad[2] else None - else: - if ctx.needs_input_grad[1]: - if process_group is not None and sequence_parallel and checkpoint_lvl != 2: - handle_x.wait() - grad_weight1 = F.linear( - grad_pre_act.t(), total_x.reshape(batch_dim, total_x.shape[-1]).t() - ) - else: - grad_weight1 = None - if process_group is not None and ctx.needs_input_grad[0]: - handle_grad_input.wait() - return ( - grad_input, - grad_weight1, - grad_bias1, - grad_weight2, - grad_bias2, - None, - None, - None, - None, - None, - None, - None, - ) - - -def fused_mlp_func( - x: Tensor, - weight1: Tensor, - weight2: Tensor, - bias1: Optional[Tensor] = None, - bias2: Optional[Tensor] = None, - activation: str = "gelu_approx", - save_pre_act: bool = True, - return_residual: bool = False, - checkpoint_lvl: int = 0, - heuristic: int = 0, - process_group: Optional[ProcessGroup] = None, - sequence_parallel: bool = True, -): - assert activation in ["gelu_approx", "relu", "sqrelu"] - dtype_eligible = x.dtype in [torch.float16, torch.bfloat16] or ( - x.dtype == torch.float32 and torch.is_autocast_enabled() - ) - # If we save pre-activation, dimension must be divisible by 128 (relu) or 8 (gelu) - dim_eligible = not save_pre_act or (x.shape[-1] % (128 if activation == "relu" else 8) == 0) - if ( - x.is_cuda - and weight1.is_cuda - and weight2.is_cuda - and (bias1 is None or bias1.is_cuda) - and (bias2 is None or bias2.is_cuda) - and dtype_eligible - and dim_eligible - ): - return FusedMLPFunc.apply( - x, - weight1, - bias1, - weight2, - bias2, - activation, - save_pre_act, - return_residual, - checkpoint_lvl, - heuristic, - process_group, - sequence_parallel, - ) - else: - assert process_group is None - pre_act = F.linear(x, weight1, bias1) - activation_fn = ( - partial(F.gelu, approximate="tanh") - if activation == "gelu_approx" - else partial(F.relu, inplace=True) - ) - output1 = activation_fn(pre_act) - output2 = F.linear(output1, weight2, bias2) - return output2 if not return_residual else (output2, x) - - -class FusedMLP(nn.Module): - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - bias1=True, - bias2=True, - activation="gelu_approx", - return_residual=False, - checkpoint_lvl=0, - heuristic="auto", - device=None, - dtype=None, - ): - """ - If process_group is not None, we're doing Tensor Parallel with sequence parallelism: - we do an all_gather of x before doing the matmul, gelu, then matmul. - Finally we do a reduce_scatter of the output. - - checkpoint_lvl (increasing lvl means slower but more memory saving): - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute pre_act and gelu_out in the bwd - heuristic: - -1: don't fuse gemm + gelu (separate kernel) - 0..4: use this heuristic for the algo section in the fused gemm + gelu - 'auto': heuristic will be picked automatically: - For CUDA >= 11.8, we set heuristic=0 for both fp16 and bf16 for best perf. - For CUDA <= 11.7, we set heuristic=1 for fp16 and heuristic=-1 for bf16. - For H100, we set heuristic=-1 for both fp16 and bf16 as the fused cuBlasLt implementation - is slower than the unfused version. - return_residual: whether to return the input x along with the output. This is for - performance reason: for post-norm architecture, returning the input allows us - to fuse the backward of nn.Linear with the residual connection. - """ - assert checkpoint_lvl in [0, 1, 2] - assert activation in ["gelu_approx", "relu", "sqrelu"] - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features * 4 - self.activation = activation - self.return_residual = return_residual - self.checkpoint_lvl = checkpoint_lvl - self.heuristic = heuristic if activation != "sqrelu" else -1 - self.fc1 = nn.Linear(in_features, hidden_features, bias=bias1, **factory_kwargs) - self.fc2 = nn.Linear(hidden_features, out_features, bias=bias2, **factory_kwargs) - - def forward(self, x, process_group=None): - dtype = x.dtype if not torch.is_autocast_enabled() else torch.get_autocast_gpu_dtype() - if self.heuristic == "auto": - if self.activation == "gelu_approx": - if torch.cuda.get_device_capability("cuda") == (9, 0): - heuristic = -1 - else: - cuda_ver = tuple(map(int, torch.version.cuda.split("."))) - heuristic = 0 if cuda_ver >= (11, 8) else (1 if dtype == torch.float16 else -1) - else: - heuristic = 0 - else: - heuristic = self.heuristic - out = fused_mlp_func( - x, - self.fc1.weight, - self.fc2.weight, - self.fc1.bias, - self.fc2.bias, - activation=self.activation, - save_pre_act=self.training, - return_residual=self.return_residual, - checkpoint_lvl=self.checkpoint_lvl, - heuristic=heuristic, - process_group=process_group, - ) - if self.return_residual: - out, x = out - if process_group is not None: - out = reduce_scatter(out, process_group) - return out if not self.return_residual else (out, x) - - -class ParallelFusedMLP(nn.Module): - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - activation="gelu_approx", - process_group: ProcessGroup = None, - bias1=True, - bias2=True, - sequence_parallel=True, - checkpoint_lvl=0, - heuristic="auto", - device=None, - dtype=None, - ): - """ - process_group is required. We're doing Tensor Parallel with sequence parallelism: - we do an all_gather of x before doing the matmul, gelu, then matmul. - Finally we do a reduce_scatter of the output. - - checkpoint_lvl (increasing lvl means slower but more memory saving): - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute pre_act and gelu_out in the bwd - heuristic: - -1: don't fuse gemm + gelu (separate kernel) - 0..4: use this heuristic for the algo section in the fused gemm + gelu - 'auto': heuristic will be picked automatically: - For CUDA >= 11.8, we set heuristic=0 for both fp16 and bf16 for best perf. - For CUDA <= 11.7, we set heuristic=1 for fp16 and heuristic=-1 for bf16. - """ - assert checkpoint_lvl in [0, 1, 2] - assert activation in ["gelu_approx", "relu", "sqrelu"] - assert process_group is not None - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features * 4 - self.activation = activation - self.process_group = process_group - self.sequence_parallel = sequence_parallel - self.checkpoint_lvl = checkpoint_lvl - self.heuristic = heuristic if activation != "sqrelu" else -1 - self.fc1 = ColumnParallelLinear( - in_features, hidden_features, process_group, bias=bias1, **factory_kwargs - ) - self.fc2 = RowParallelLinear( - hidden_features, out_features, process_group, bias=bias2, **factory_kwargs - ) - - def forward(self, x): - dtype = x.dtype if not torch.is_autocast_enabled() else torch.get_autocast_gpu_dtype() - if self.heuristic == "auto": - if self.activation == "gelu_approx": - cuda_ver = tuple(map(int, torch.version.cuda.split("."))) - heuristic = 0 if cuda_ver >= (11, 8) else (1 if dtype == torch.float16 else -1) - else: - heuristic = 0 - else: - heuristic = self.heuristic - out = fused_mlp_func( - x, - self.fc1.weight, - self.fc2.weight, - self.fc1.bias, - self.fc2.bias, - activation=self.activation, - save_pre_act=self.training, - checkpoint_lvl=self.checkpoint_lvl, - heuristic=heuristic, - process_group=self.process_group, - sequence_parallel=self.sequence_parallel, - ) - reduce_fn = reduce_scatter if self.sequence_parallel else all_reduce - return reduce_fn(out, self.process_group) diff --git a/build/torch210-cu128-x86_64-windows/ops/layer_norm.py b/build/torch210-cu128-x86_64-windows/ops/layer_norm.py deleted file mode 100644 index c8bcae3b080845f56464dceca508278f76cde68f..0000000000000000000000000000000000000000 --- a/build/torch210-cu128-x86_64-windows/ops/layer_norm.py +++ /dev/null @@ -1,800 +0,0 @@ -# Copyright (c) 2022, Tri Dao. -# Adapted from https://github.com/NVIDIA/apex/blob/master/apex/contrib/layer_norm/layer_norm.py - -import dropout_layer_norm -import torch -from torch.nn import init - - -def maybe_align(x, alignment_in_bytes=16): - """Assume that x already has last dim divisible by alignment_in_bytes""" - # TD [2023-07-04] I'm not 100% sure that clone will align the memory - # https://discuss.pytorch.org/t/how-to-ensure-that-tensor-data-ptr-is-aligned-to-16-bytes/183440 - return x if x.data_ptr() % alignment_in_bytes == 0 else x.clone() - - -def _dropout_add_layer_norm_forward( - x0, - residual, - gamma, - beta, - rowscale, - colscale, - dropout_p, - epsilon, - residual_in_fp32=False, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes""" - hidden_size = gamma.numel() - x0mat = x0.view((-1, hidden_size)) - residualmat = residual.view((-1, hidden_size)) if residual is not None else None - rowscale = rowscale.view(-1) if rowscale is not None else None - zmat, xmat, dmask, mu, rsigma = dropout_layer_norm.dropout_add_ln_fwd( - x0mat, - residualmat, - gamma, - beta, - rowscale, - colscale, - None, - None, - dropout_p, - epsilon, - 1.0, - 0, - None, - residual_in_fp32, - is_rms_norm, - ) - # dmask is None if dropout_p == 0.0 - # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype - return zmat, xmat if xmat is not None else x0mat, dmask, mu, rsigma - - -def _dropout_add_layer_norm_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - rowscale, - colscale, - dropout_p, - has_residual, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes - dx == None means that it was a post-norm architecture - (x = drop(x0) + residual was not returned in the fwd). - x0 must not be None if we have colscale. - """ - hidden_size = gamma.numel() - xmat = x.view((-1, hidden_size)) - dzmat = dz.view(xmat.shape) - dxmat = dx.view(xmat.shape) if dx is not None else None - x0mat = x0.view((-1, hidden_size)) if x0 is not None else None - rowscale = rowscale.view(-1) if rowscale is not None else None - if colscale is not None: - assert x0 is not None, "x0 is required to compute the gradient of colscale" - dx0mat, dresidualmat, dgamma, dbeta, _, _, *rest = dropout_layer_norm.dropout_add_ln_bwd( - dzmat, - dxmat, - xmat, - x0mat, - dmask, - mu, - rsigma, - gamma, - rowscale, - colscale, - None, - None, - dropout_p, - 1.0, - 0, - has_residual, - is_rms_norm, - ) - # dresidualmat is None if not has_residual - if colscale is None: - return dx0mat, dresidualmat, dgamma, dbeta - else: - dcolscale = rest[0] - return dx0mat, dresidualmat, dgamma, dbeta, dcolscale - - -def _dropout_add_layer_norm_subset_forward( - x0, - residual, - gamma, - beta, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32=False, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes""" - hidden_size = gamma.numel() - x0mat = x0.view((-1, hidden_size)) - residualmat = residual.view((-1, hidden_size)) if residual is not None else None - x0_subset = x0_subset.view(-1) if x0_subset is not None else None - out_subset = out_subset.view(-1) if out_subset is not None else None - zmat, xmat, dmask, mu, rsigma = dropout_layer_norm.dropout_add_ln_fwd( - x0mat, - residualmat, - gamma, - beta, - None, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - None, - residual_in_fp32, - is_rms_norm, - ) - # dmask is None if dropout_p == 0.0 - # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype - return zmat, xmat if xmat is not None else x0mat, dmask, mu, rsigma - - -def _dropout_add_layer_norm_subset_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - colscale, - x0_subset, - out_subset, - dropout_p, - rowscale_const, - x0_numrows, - has_residual, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes - dx == None means that it was a post-norm architecture - (x = drop(x0) + residual was not returned in the fwd). - x0 must not be None if we have colscale. - """ - hidden_size = gamma.numel() - xmat = x.view((-1, hidden_size)) - dzmat = dz.view(-1, hidden_size) - dxmat = dx.view(xmat.shape) if dx is not None else None - x0mat = x0.view((-1, hidden_size)) if x0 is not None else None - x0_subset = x0_subset.view(-1) if x0_subset is not None else None - out_subset = out_subset.view(-1) if out_subset is not None else None - if colscale is not None: - assert x0 is not None, "x0 is required to compute the gradient of colscale" - dx0mat, dresidualmat, dgamma, dbeta, _, _, *rest = dropout_layer_norm.dropout_add_ln_bwd( - dzmat, - dxmat, - xmat, - x0mat, - dmask, - mu, - rsigma, - gamma, - None, - colscale, - x0_subset, - out_subset, - dropout_p, - rowscale_const, - x0_numrows, - has_residual, - is_rms_norm, - ) - # dresidualmat is None if not has_residual - if colscale is None: - return dx0mat, dresidualmat, dgamma, dbeta - else: - dcolscale = rest[0] - return dx0mat, dresidualmat, dgamma, dbeta, dcolscale - - -def _dropout_add_layer_norm_parallel_residual_forward( - x0, - x1, - residual, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - residual_in_fp32=False, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes""" - hidden_size = gamma0.numel() - x0mat = x0.view((-1, hidden_size)) - x1mat = x1.view((-1, hidden_size)) if x1 is not None else None - residualmat = residual.view((-1, hidden_size)) if residual is not None else None - ( - z0mat, - z1mat, - xmat, - dmask0, - dmask1, - mu, - rsigma, - ) = dropout_layer_norm.dropout_add_ln_parallel_residual_fwd( - x0mat, - x1mat, - residualmat, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - None, - residual_in_fp32, - is_rms_norm, - ) - # dmask0 and dmask1 are None if dropout_p == 0.0 - # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype - return z0mat, z1mat, xmat if xmat is not None else x0mat, dmask0, dmask1, mu, rsigma - - -def _dropout_add_layer_norm_parallel_residual_backward( - dz0, - dz1, - dx, - x, - dmask0, - dmask1, - mu, - rsigma, - gamma0, - gamma1, - dropout_p, - has_x1, - has_residual, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes - dx == None means that it was a post-norm architecture - (x = drop(x0) + residual was not returned in the fwd). - """ - hidden_size = gamma0.numel() - xmat = x.view((-1, hidden_size)) - dz0mat = dz0.view(xmat.shape) - dz1mat = dz1.view(xmat.shape) if dz1 is not None else None - dxmat = dx.view(xmat.shape) if dx is not None else None - ( - dx0mat, - dx1mat, - dresidualmat, - dgamma0, - dbeta0, - dgamma1, - dbeta1, - *rest, - ) = dropout_layer_norm.dropout_add_ln_parallel_residual_bwd( - dz0mat, - dz1mat, - dxmat, - xmat, - dmask0, - dmask1, - mu, - rsigma, - gamma0, - gamma1, - dropout_p, - has_x1, - has_residual, - is_rms_norm, - ) - # dresidualmat is None if not has_residual - return dx0mat, dx1mat, dresidualmat, dgamma0, dbeta0, dgamma1, dbeta1 - - -class DropoutAddLayerNormFn(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x0, - residual, - gamma, - beta, - rowscale, - colscale, - dropout_p, - epsilon, - residual_in_fp32=False, - prenorm=False, - is_rms_norm=False, - return_dmask=False, - ): - x0 = maybe_align(x0.contiguous(), 16) - residual = maybe_align(residual.contiguous(), 16) if residual is not None else None - gamma = maybe_align(gamma.contiguous(), 16) - beta = maybe_align(beta.contiguous(), 16) if beta is not None else None - rowscale = maybe_align(rowscale.contiguous(), 16) if rowscale is not None else None - colscale = maybe_align(colscale.contiguous(), 16) if colscale is not None else None - zmat, xmat, dmask, mu, rsigma = _dropout_add_layer_norm_forward( - x0, - residual, - gamma, - beta, - rowscale, - colscale, - dropout_p, - epsilon, - residual_in_fp32, - is_rms_norm, - ) - # Only need to save x0 if we need to compute gradient wrt colscale - x0_saved = x0 if colscale is not None else None - ctx.save_for_backward( - xmat.view(x0.shape), x0_saved, dmask, gamma, mu, rsigma, rowscale, colscale - ) - ctx.prenorm = prenorm - ctx.dropout_p = dropout_p - ctx.has_residual = residual is not None - ctx.is_rms_norm = is_rms_norm - ctx.has_beta = beta is not None - if not return_dmask: - return ( - zmat.view(x0.shape) if not prenorm else (zmat.view(x0.shape), xmat.view(x0.shape)) - ) - else: - dmask = ( - dmask.view(x0.shape) - if dropout_p > 0.0 - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - ctx.mark_non_differentiable(dmask) - return ( - (zmat.view(x0.shape), dmask) - if not prenorm - else (zmat.view(x0.shape), xmat.view(x0.shape), dmask) - ) - - @staticmethod - def backward(ctx, dz, *args): - # assert dz.is_contiguous() - dz = maybe_align(dz.contiguous(), 16) # this happens! - dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None - x, x0, dmask, gamma, mu, rsigma, rowscale, colscale = ctx.saved_tensors - # x0 is None if colscale is None - dropout_p = ctx.dropout_p - has_residual = ctx.has_residual - dx0mat, dresidualmat, dgamma, dbeta, *rest = _dropout_add_layer_norm_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - rowscale, - colscale, - dropout_p, - has_residual, - ctx.is_rms_norm, - ) - dx0 = dx0mat.view(x.shape) - dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None - dcolscale = rest[0] if colscale is not None else None - return ( - dx0, - dresidual, - dgamma, - dbeta if ctx.has_beta else None, - None, - dcolscale, - None, - None, - None, - None, - None, - None, - ) - - -class DropoutAddLayerNormSubsetFn(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x0, - residual, - gamma, - beta, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32=False, - prenorm=False, - is_rms_norm=False, - return_dmask=False, - ): - x0 = maybe_align(x0.contiguous(), 16) - residual = maybe_align(residual.contiguous(), 16) if residual is not None else None - gamma = maybe_align(gamma.contiguous(), 16) - beta = maybe_align(beta.contiguous(), 16) if beta is not None else None - colscale = maybe_align(colscale.contiguous(), 16) if colscale is not None else None - zmat, xmat, dmask, mu, rsigma = _dropout_add_layer_norm_subset_forward( - x0, - residual, - gamma, - beta, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32, - is_rms_norm, - ) - # Only need to save x0 if we need to compute gradient wrt colscale - x0_saved = x0 if colscale is not None else None - x_shape = (-1, *x0.shape[1:]) - ctx.save_for_backward( - xmat.view(x_shape), x0_saved, dmask, gamma, mu, rsigma, colscale, x0_subset, out_subset - ) - ctx.prenorm = prenorm - ctx.dropout_p = dropout_p - ctx.rowscale_const = rowscale_const - ctx.x0_numrows = x0.shape[:-1].numel() - ctx.has_residual = residual is not None - ctx.is_rms_norm = is_rms_norm - ctx.has_beta = beta is not None - z_shape = (-1, *x0.shape[1:]) - if not return_dmask: - return zmat.view(z_shape) if not prenorm else (zmat.view(z_shape), xmat.view(x0.shape)) - else: - z = zmat.view(z_shape) - dmask = ( - dmask.view(x0.shape) - if dropout_p > 0.0 - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - ctx.mark_non_differentiable(dmask) - return (z, dmask) if not prenorm else (z, xmat.view(x_shape), dmask) - - @staticmethod - def backward(ctx, dz, *args): - # assert dz.is_contiguous() - dz = maybe_align(dz.contiguous(), 16) # this happens! - dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None - x, x0, dmask, gamma, mu, rsigma, colscale, x0_subset, out_subset = ctx.saved_tensors - # x0 is None if colscale is None - dropout_p = ctx.dropout_p - has_residual = ctx.has_residual - dx0mat, dresidualmat, dgamma, dbeta, *rest = _dropout_add_layer_norm_subset_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - colscale, - x0_subset, - out_subset, - dropout_p, - ctx.rowscale_const, - ctx.x0_numrows, - has_residual, - ctx.is_rms_norm, - ) - dx0 = dx0mat.view(-1, *x.shape[1:]) - dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None - dcolscale = rest[0] if colscale is not None else None - return ( - dx0, - dresidual, - dgamma, - dbeta if ctx.has_beta else None, - dcolscale, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) - - -class DropoutAddLayerNormParallelResidualFn(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x0, - x1, - residual, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - residual_in_fp32=False, - prenorm=False, - is_rms_norm=False, - return_dmask=False, - ): - x0 = maybe_align(x0.contiguous(), 16) - x1 = maybe_align(x1.contiguous(), 16) if x1 is not None else None - residual = maybe_align(residual.contiguous(), 16) if residual is not None else None - gamma0 = maybe_align(gamma0.contiguous(), 16) - beta0 = maybe_align(beta0.contiguous(), 16) if beta0 is not None else None - gamma1 = maybe_align(gamma1.contiguous(), 16) if gamma1 is not None else None - beta1 = maybe_align(beta1.contiguous(), 16) if beta1 is not None else None - ( - z0mat, - z1mat, - xmat, - dmask0, - dmask1, - mu, - rsigma, - ) = _dropout_add_layer_norm_parallel_residual_forward( - x0, - x1, - residual, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - residual_in_fp32, - is_rms_norm, - ) - ctx.save_for_backward(xmat.view(x0.shape), dmask0, dmask1, gamma0, gamma1, mu, rsigma) - ctx.prenorm = prenorm - ctx.dropout_p = dropout_p - ctx.has_x1 = x1 is not None - ctx.has_residual = residual is not None - ctx.is_rms_norm = is_rms_norm - ctx.has_beta = beta0 is not None - z = (z0mat.view(x0.shape), z1mat.view(x0.shape) if z1mat is not None else None) - if not return_dmask: - return z if not prenorm else (*z, xmat.view(x0.shape)) - else: - dmask0 = ( - dmask0.view(x0.shape) - if dropout_p > 0.0 - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - dmask1 = ( - dmask1.view(x0.shape) - if dropout_p > 0.0 and x1 is not None - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - ctx.mark_non_differentiable(dmask0) - ctx.mark_non_differentiable(dmask1) - return ( - (*z, dmask0, dmask1) if not prenorm else (*z, xmat.view(x0.shape), dmask0, dmask1) - ) - - @staticmethod - def backward(ctx, dz0, dz1, *args): - dz0 = maybe_align(dz0.contiguous(), 16) # this happens! - dz1 = maybe_align(dz1.contiguous(), 16) if dz1 is not None else None - dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None - x, dmask0, dmask1, gamma0, gamma1, mu, rsigma = ctx.saved_tensors - dropout_p = ctx.dropout_p - has_x1 = ctx.has_x1 - has_residual = ctx.has_residual - ( - dx0mat, - dx1mat, - dresidualmat, - dgamma0, - dbeta0, - dgamma1, - dbeta1, - ) = _dropout_add_layer_norm_parallel_residual_backward( - dz0, - dz1, - dx, - x, - dmask0, - dmask1, - mu, - rsigma, - gamma0, - gamma1, - dropout_p, - has_x1, - has_residual, - ctx.is_rms_norm, - ) - dx0 = dx0mat.view(x.shape) - dx1 = dx1mat.view(x.shape) if dx1mat is not None else None - dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None - return ( - dx0, - dx1, - dresidual, - dgamma0, - dbeta0 if ctx.has_beta else None, - dgamma1, - dbeta1 if ctx.has_beta else None, - None, - None, - None, - None, - None, - None, - ) - - -def layer_norm(x, weight, bias, epsilon): - return DropoutAddLayerNormFn.apply(x, None, weight, bias, None, None, 0.0, epsilon, False) - - -def dropout_add_layer_norm( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - rowscale=None, - layerscale=None, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormFn.apply( - x0, - residual, - weight, - bias, - rowscale, - layerscale, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - False, - return_dropout_mask, - ) - - -def dropout_add_layer_norm_subset( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - layerscale=None, - x0_subset=None, - out_subset=None, - rowscale_const=1.0, - out_numrows=0, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormSubsetFn.apply( - x0, - residual, - weight, - bias, - layerscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32, - prenorm, - False, - return_dropout_mask, - ) - - -def dropout_add_layer_norm_parallel_residual( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormParallelResidualFn.apply( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - False, - return_dropout_mask, - ) - - -class DropoutAddLayerNorm(torch.nn.Module): - def __init__( - self, - hidden_size, - prenorm=False, - p=0.0, - eps=1e-5, - residual_in_fp32=False, - device=None, - dtype=None, - ): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.prenorm = prenorm - self.p = p - self.eps = eps - self.residual_in_fp32 = residual_in_fp32 - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.bias = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.reset_parameters() - - def reset_parameters(self): - init.ones_(self.weight) - init.zeros_(self.bias) - - def forward(self, x0, residual=None): - return dropout_add_layer_norm( - x0, - residual, - self.weight, - self.bias, - self.p if self.training else 0.0, - self.eps, - prenorm=self.prenorm, - residual_in_fp32=self.residual_in_fp32, - ) diff --git a/build/torch210-cu128-x86_64-windows/ops/rms_norm.py b/build/torch210-cu128-x86_64-windows/ops/rms_norm.py deleted file mode 100644 index 9e286bc004a7353d81abc29759064b5383571743..0000000000000000000000000000000000000000 --- a/build/torch210-cu128-x86_64-windows/ops/rms_norm.py +++ /dev/null @@ -1,174 +0,0 @@ -# Copyright (c) 2022, Tri Dao. -# Adapted from https://github.com/NVIDIA/apex/blob/master/apex/contrib/layer_norm/layer_norm.py - -import torch -from torch.nn import init - -from flash_attn.ops.layer_norm import ( - DropoutAddLayerNormFn, - DropoutAddLayerNormParallelResidualFn, - DropoutAddLayerNormSubsetFn, -) - - -def rms_norm(x, weight, epsilon): - return DropoutAddLayerNormFn.apply( - x, None, weight, None, None, None, 0.0, epsilon, False, False, True - ) - - -def dropout_add_rms_norm( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - rowscale=None, - layerscale=None, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormFn.apply( - x0, - residual, - weight, - bias, - rowscale, - layerscale, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - True, - return_dropout_mask, - ) - - -def dropout_add_rms_norm_subset( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - layerscale=None, - x0_subset=None, - out_subset=None, - rowscale_const=1.0, - out_numrows=0, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormSubsetFn.apply( - x0, - residual, - weight, - bias, - layerscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32, - prenorm, - True, - return_dropout_mask, - ) - - -def dropout_add_rms_norm_parallel_residual( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormParallelResidualFn.apply( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - True, - return_dropout_mask, - ) - - -class RMSNorm(torch.nn.Module): - def __init__(self, hidden_size, eps=1e-5, device=None, dtype=None): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.eps = eps - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.register_parameter("bias", None) - self.reset_parameters() - - def reset_parameters(self): - init.ones_(self.weight) - - def forward(self, x): - return rms_norm(x, self.weight, self.eps) - - -class DropoutAddRMSNorm(torch.nn.Module): - def __init__( - self, - hidden_size, - prenorm=False, - p=0.0, - eps=1e-5, - residual_in_fp32=False, - device=None, - dtype=None, - ): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.prenorm = prenorm - self.p = p - self.eps = eps - self.residual_in_fp32 = residual_in_fp32 - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.register_parameter("bias", None) - self.reset_parameters() - - def reset_parameters(self): - init.ones_(self.weight) - - def forward(self, x0, residual=None): - return dropout_add_rms_norm( - x0, - residual, - self.weight, - None, - self.p if self.training else 0.0, - self.eps, - prenorm=self.prenorm, - residual_in_fp32=self.residual_in_fp32, - ) diff --git a/build/torch210-cu128-x86_64-windows/ops/triton/__init__.py b/build/torch210-cu128-x86_64-windows/ops/triton/__init__.py deleted file mode 100644 index d3f5a12faa99758192ecc4ed3fc22c9249232e86..0000000000000000000000000000000000000000 --- a/build/torch210-cu128-x86_64-windows/ops/triton/__init__.py +++ /dev/null @@ -1 +0,0 @@ - diff --git a/build/torch210-cu128-x86_64-windows/ops/triton/cross_entropy.py b/build/torch210-cu128-x86_64-windows/ops/triton/cross_entropy.py deleted file mode 100644 index f44d6bfaf87a515e3d30598f98c1dfebc0511399..0000000000000000000000000000000000000000 --- a/build/torch210-cu128-x86_64-windows/ops/triton/cross_entropy.py +++ /dev/null @@ -1,330 +0,0 @@ -# Copyright (c) 2023, Tri Dao. - -from typing import Tuple, Optional, Union - -import torch -import torch.nn.functional as F - -import triton -import triton.language as tl - -# `all_gather_into_tensor` and `reduce_scatter_tensor` are new placeholders for -# `_all_gather_base` and `_reduce_scatter_base`. They require the most recent -# version of PyTorch. The following 2 lines are for backward compatibility with -# older PyTorch. -if "all_gather_into_tensor" not in dir(torch.distributed): - torch.distributed.all_gather_into_tensor = torch.distributed._all_gather_base - - -@triton.heuristics( - { - "HAS_SMOOTHING": lambda args: args["smoothing"] > 0.0, - } -) -@triton.jit -def cross_entropy_fwd_kernel( - loss_ptr, # data ptrs - lse_ptr, - z_loss_ptr, - logits_ptr, - labels_ptr, - smoothing, - logit_scale, - lse_square_scale, - ignore_index, - total_classes, - class_start_idx, # Useful for tensor parallel when each rank only has a subset of classes - n_cols, # shapes - logits_row_stride, # strides - BLOCK_SIZE: tl.constexpr, - HAS_SMOOTHING: tl.constexpr, - # if SPLIT (e.g. tensor parallel), don't include the LSE in the loss since it's not the final LSE - SPLIT: tl.constexpr, - PRECOMPUTED_LSE: tl.constexpr, # If LSE is already computed (also no smoothing and logit_scale == 1.0) -): - row_idx = tl.program_id(0) - logits_ptr = logits_ptr + row_idx * logits_row_stride.to(tl.int64) - sum_logits = 0.0 # For smoothing - if not PRECOMPUTED_LSE: - # Statistics for online softmax - m_i = -float("inf") - l_i = 0.0 - for col_offset in range(0, n_cols, BLOCK_SIZE): - cols = col_offset + tl.arange(0, BLOCK_SIZE) - logits = tl.load(logits_ptr + cols, mask=cols < n_cols, other=-float("inf")).to( - tl.float32 - ) * logit_scale - if HAS_SMOOTHING: - sum_logits += tl.sum(tl.where(cols < n_cols, logits, 0.0)) - m_i_new = tl.maximum(m_i, tl.max(logits)) - l_i = tl.exp(m_i - m_i_new) * l_i + tl.sum(tl.exp(logits - m_i_new)) - m_i = m_i_new - lse = tl.log(l_i) + m_i - tl.store(lse_ptr + row_idx, lse) - else: - lse = tl.load(lse_ptr + row_idx) - label_idx = tl.load(labels_ptr + row_idx) - if label_idx == ignore_index: - loss = 0.0 - z_loss = 0.0 - else: - label_idx -= class_start_idx - if label_idx >= 0 and label_idx < n_cols: - logits_label = tl.load(logits_ptr + label_idx) * logit_scale - if HAS_SMOOTHING: - loss = ( - (lse if not SPLIT else 0.0) - - smoothing * sum_logits / total_classes - - (1 - smoothing) * logits_label - ) - else: - loss = (lse if not SPLIT else 0.0) - logits_label - else: - # If label is out of bounds, we set the CE loss to 0.0. But we still want the smoothing loss - if HAS_SMOOTHING: - loss = smoothing * ((lse if not SPLIT else 0.0) - sum_logits / total_classes) - else: - loss = 0.0 - if not SPLIT: - z_loss = lse_square_scale * lse * lse - loss += z_loss - else: - z_loss = 0.0 - tl.store(loss_ptr + row_idx, loss) - if not SPLIT: - tl.store(z_loss_ptr + row_idx, z_loss) - - -@triton.heuristics( - { - "HAS_SMOOTHING": lambda args: args["smoothing"] > 0.0, - } -) -@triton.jit -def cross_entropy_bwd_kernel( - dlogits_ptr, # data ptrs - dloss_ptr, - logits_ptr, - lse_ptr, - labels_ptr, - smoothing, - logit_scale, - lse_square_scale, - ignore_index, - total_classes, - class_start_idx, # Useful for tensor parallel when each rank only has a subset of classes - n_cols, # shapes - logits_row_stride, # strides - dlogits_row_stride, - dloss_row_stride, - BLOCK_SIZE: tl.constexpr, - HAS_SMOOTHING: tl.constexpr, -): - row_idx = tl.program_id(0) - col_block_idx = tl.program_id(1) - logits_ptr = logits_ptr + row_idx * logits_row_stride.to(tl.int64) - dlogits_ptr = dlogits_ptr + row_idx * dlogits_row_stride.to(tl.int64) - col_offsets = col_block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - label_idx = tl.load(labels_ptr + row_idx) - if label_idx != ignore_index: - dloss = tl.load(dloss_ptr + row_idx * dloss_row_stride) - else: - dloss = 0.0 - logits = tl.load(logits_ptr + col_offsets, mask=col_offsets < n_cols, other=-float("inf")).to( - tl.float32 - ) * logit_scale - lse = tl.load(lse_ptr + row_idx) - probs = tl.exp(logits - lse) - probs += 2.0 * lse_square_scale * lse * probs - label_idx -= class_start_idx - if HAS_SMOOTHING: - smooth_positive = 1.0 - smoothing - smooth_negative = smoothing / total_classes - probs = tl.where(col_offsets == label_idx, probs - smooth_positive, probs) - smooth_negative - else: - probs = tl.where(col_offsets == label_idx, probs - 1.0, probs) - tl.store(dlogits_ptr + col_offsets, (dloss * logit_scale) * probs, mask=col_offsets < n_cols) - - -class CrossEntropyLoss(torch.autograd.Function): - - @staticmethod - def forward( - ctx, - logits, - labels, - precomputed_lse=None, - smoothing=0.0, - logit_scale=1.0, - lse_square_scale=0.0, - ignore_index=-100, - inplace_backward=False, - process_group=None, - ): - # For some reason Triton generates wrong code when labels has dtype long and its address - # is not aligned to 16 bytes. The ld.global.b64 seems to load the wrong label index. - if labels.dtype == torch.long and labels.data_ptr() % 16 != 0: - labels = F.pad(labels, (0, 1))[..., :-1] - assert labels.data_ptr() % 16 == 0 - assert logit_scale > 0.0 - n_rows, n_cols = logits.shape - assert labels.shape == (n_rows,) - world_size = 1 if process_group is None else torch.distributed.get_world_size(process_group) - total_classes = world_size * n_cols - rank = 0 if process_group is None else torch.distributed.get_rank(process_group) - class_start_idx = rank * n_cols - use_precomputed_lse = precomputed_lse is not None and logit_scale == 1.0 and smoothing == 0.0 - - if logits.stride(-1) != 1: - logits = logits.contiguous() - MAX_BLOCK_SIZE = 16 * 1024 - BLOCK_SIZE = min(triton.next_power_of_2(n_cols), MAX_BLOCK_SIZE) - num_warps = ( - 4 - if BLOCK_SIZE < 2048 - else (8 if BLOCK_SIZE < 8192 else (16 if BLOCK_SIZE < 128 * 1024 else 32)) - ) - losses = torch.empty(n_rows, dtype=torch.float, device=logits.device) - if use_precomputed_lse: - assert precomputed_lse.shape == (n_rows,) - lse = precomputed_lse.contiguous() - else: - lse = torch.empty(n_rows, dtype=torch.float, device=logits.device) - z_losses = torch.empty(n_rows, dtype=torch.float, device=logits.device) - # Need this, otherwise Triton tries to launch from cuda:0 and we get - # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?) - with torch.cuda.device(logits.device.index): - cross_entropy_fwd_kernel[(n_rows,)]( - losses, # data ptrs - lse, - z_losses, - logits, - labels, - smoothing, - logit_scale, - lse_square_scale, - ignore_index, - total_classes, - class_start_idx, - n_cols, # shapes - logits.stride(0), # strides - BLOCK_SIZE=BLOCK_SIZE, # constants - SPLIT=world_size > 1, - PRECOMPUTED_LSE=use_precomputed_lse, - num_warps=num_warps, - ) - - if world_size > 1: - # If there's no smoothing, if labels are in the vocab of this partition, losses contains - # - predicted logit, and 0 otherwise. - # If there's smoothing=0.1, for labels in the vocab of this partition, losses contains - # -0.9 * predicted logit - 0.1 * sum logit / total_classes. - # For labels not in the vocab of this partition, losses contains - # -0.1 * sum logit / total_classes. - if world_size > 1: - lse_allgather = torch.empty(world_size, n_rows, dtype=lse.dtype, device=lse.device) - torch.distributed.all_gather_into_tensor(lse_allgather, lse, group=process_group) - handle_losses = torch.distributed.all_reduce( - losses, op=torch.distributed.ReduceOp.SUM, group=process_group, async_op=True - ) - lse = torch.logsumexp(lse_allgather, dim=0) - handle_losses.wait() - # After the allreduce, if there's no smoothing, the total losses are - predicted_logit, - # we just have to add the (global) lse. - # If there's smoothing=0.1, the total losses are - # -0.9 * predicted_logit - 0.1 * sum logit / total_classes. - # Again, we just have to add the (global) lse. - losses += lse - if lse_square_scale != 0.0: - z_losses = lse_square_scale * lse.square() - z_losses.masked_fill_(labels == ignore_index, 0.0) - losses += z_losses - else: - z_losses = torch.zeros_like(losses) - losses.masked_fill_(labels == ignore_index, 0.0) - - ctx.save_for_backward(logits, lse, labels) - ctx.mark_non_differentiable(z_losses) - ctx.smoothing = smoothing - ctx.logit_scale = logit_scale - ctx.lse_square_scale = lse_square_scale - ctx.ignore_index = ignore_index - ctx.total_classes = total_classes - ctx.class_start_idx = class_start_idx - ctx.inplace_backward = inplace_backward - return losses, z_losses - - @staticmethod - def backward(ctx, grad_losses, grad_z_losses): - del grad_z_losses # z_losses are only for logging. - - logits, lse, labels = ctx.saved_tensors - dlogits = logits if ctx.inplace_backward else torch.empty_like(logits) - n_rows, n_cols = logits.shape - BLOCK_SIZE = min(triton.next_power_of_2(n_cols), 4 * 1024) - num_warps = 4 if BLOCK_SIZE < 2048 else (8 if BLOCK_SIZE < 8192 else 16) - grid = lambda META: (n_rows, triton.cdiv(n_cols, META["BLOCK_SIZE"])) # noqa - # Need this, otherwise Triton tries to launch from cuda:0 and we get - # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?) - with torch.cuda.device(logits.device.index): - cross_entropy_bwd_kernel[grid]( - dlogits, # data ptrs - grad_losses, - logits, - lse, - labels, - ctx.smoothing, - ctx.logit_scale, - ctx.lse_square_scale, - ctx.ignore_index, - ctx.total_classes, - ctx.class_start_idx, - n_cols, # shapes - logits.stride(0), # strides - dlogits.stride(0), - grad_losses.stride(0), - BLOCK_SIZE=BLOCK_SIZE, # constants - num_warps=num_warps, - ) - return dlogits, None, None, None, None, None, None, None, None, None - - -def cross_entropy_loss( - logits: torch.Tensor, - labels: torch.Tensor, - precomputed_lse: Optional[torch.Tensor] = None, - label_smoothing: float = 0.0, - logit_scale: float = 1.0, - lse_square_scale: float = 0.0, - ignore_index=-100, - inplace_backward: bool = False, - process_group=None, -) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Arguments: - logits: (batch, vocab_size) - labels: (batch,) - label_smoothing: float - logit_scale: float. Multiply logits by this scale before calculating the loss. - lse_square_scale: float. If > 0, we add lse_square_scale * lse(logits) ^ 2 to the loss. - This is also referred to as "z-loss". - ignore_index: int. If labels == ignore_index, the loss is set to 0.0. - inplace_backward: bool. If True, we do the backward pass in-place by modifying the logits. - This saves memory. - process_group: if not None, we're doing Tensor Parallel: each process is responsible for - one part of the vocab. The loss will be aggregated across processes. - Returns: - losses: (batch,), float - z_losses: (batch,), float - """ - return CrossEntropyLoss.apply( - logits, - labels, - precomputed_lse, - label_smoothing, - logit_scale, - lse_square_scale, - ignore_index, - inplace_backward, - process_group, - ) diff --git a/build/torch210-cu128-x86_64-windows/ops/triton/k_activations.py b/build/torch210-cu128-x86_64-windows/ops/triton/k_activations.py deleted file mode 100644 index f6a3f63a3a2e699092d78b48826d01bd6a70a609..0000000000000000000000000000000000000000 --- a/build/torch210-cu128-x86_64-windows/ops/triton/k_activations.py +++ /dev/null @@ -1,162 +0,0 @@ -# Adapted from https://github.com/facebookresearch/xformers/blob/main/xformers/triton/k_activations.py -# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved. -# -# This source code is licensed under the BSD license found in the -# LICENSE file in the root directory of this source tree. - -import math -from enum import Enum -from typing import Optional - -import triton -import triton.language as tl - -_sqrt2pi = math.sqrt(2.0 / math.pi) -_sqrt1_2 = math.sqrt(1.0 / 2) -_gaussian_pdf_normalization = 1.0 / math.sqrt(2 * math.pi) - - -class Activation(str, Enum): - SquaredReLU = "squared_relu" - GeLU = "gelu" - GeLUApprox = "gelu_approx" - LeakyReLU = "leaky_relu" - ReLU = "relu" - - -def get_triton_activation_kernel(activation: Optional[Activation]): - return ( - { - Activation.ReLU: relu, - Activation.LeakyReLU: leaky_relu, - Activation.GeLU: gelu, - Activation.GeLUApprox: gelu_approx, - Activation.SquaredReLU: squared_relu, - }[activation] - if activation - else None - ) - - -def get_triton_activation_bwd_kernel(activation: Optional[Activation]): - return ( - { - Activation.ReLU: relu_grad, - Activation.LeakyReLU: leaky_relu_grad, - Activation.GeLU: gelu_grad, - Activation.GeLUApprox: gelu_approx_grad, - Activation.SquaredReLU: squared_relu_grad, - }[activation] - if activation - else None - ) - - -@triton.jit -def tanh(x): - # Tanh is just a scaled sigmoid - return 2 * tl.sigmoid(2 * x) - 1 - - -@triton.jit -def cosh(x): - exp_x = tl.exp(x) - return (exp_x + 1.0 / exp_x) * 0.5 - - -# a Triton implementation of the most used activations -# See for instance http://arxiv.org/abs/1606.08415 for an overview - -# ReLU -@triton.jit -def relu(x): - """ - ReLU_ activation function - - .. _ReLU: https://pytorch.org/docs/stable/generated/torch.nn.ReLU.html - """ - zero = 0.0 - return tl.where(x >= 0, x, zero.to(x.dtype)) - - -@triton.jit -def relu_grad(x): - # ReLU is different from other activations - # in that it does not require the input to retrospectively compute its gradient - # here the input is the downstream gradient, and we return the upstream gradient directly - zero = 0.0 - one = 1.0 - return tl.where(x >= 0, one.to(x.dtype), zero.to(x.dtype)) - - -@triton.jit -def squared_relu(x): - """ - Squared ReLU activation, as proposed in the Primer_ paper. - - .. _Primer: https://arxiv.org/abs/2109.08668 - """ - x_ = relu(x) - return (x_ * x_).to(x.dtype) - - -@triton.jit -def squared_relu_grad(x): - return tl.where(x >= 0, 2.0 * x, 0.0) - - -# Leaky ReLU -@triton.jit -def leaky_relu(x): - """ - LeakyReLU_ activation - - .. _LeakyReLU: https://pytorch.org/docs/stable/generated/torch.nn.LeakyReLU.html - """ - scale = 0.01 + 0.0 - scale = scale.to(x.dtype) - return tl.where(x >= 0, x, scale * x) - - -@triton.jit -def leaky_relu_grad(x): - min_grad = 0.01 - max_grad = 1 - - min_grad = min_grad.to(x.dtype) - max_grad = max_grad.to(x.dtype) - - return tl.where(x >= 0, max_grad, min_grad) - - -@triton.jit -def gelu(x): - """Gaussian Error Linear Unit (GELU)""" - return x * 0.5 * (1.0 + tl.libdevice.erf(x * _sqrt1_2)) - - -@triton.jit -def gelu_grad(x): - cdf = 0.5 * (1.0 + tl.libdevice.erf(x * _sqrt1_2)) - pdf = tl.exp(-0.5 * x * x) * _gaussian_pdf_normalization - return cdf + x * pdf - - -@triton.jit -def gelu_approx(x): - """ - GeLU_ activation - Gaussian error linear unit, with tanh approximation - - .. _GeLU: https://arxiv.org/pdf/1606.08415.pdf - """ - return 0.5 * x * (1.0 + tanh(_sqrt2pi * x * (1.0 + 0.044715 * x * x))) - - -@triton.jit -def gelu_approx_grad(x): - # CREDITS: Fast implementation proposed in - # https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/fused_bias_gelu.py#L30 - tanh_out = tanh(0.79788456 * x * (1 + 0.044715 * x * x)) - return 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( - 1 + tanh_out - ) diff --git a/build/torch210-cu128-x86_64-windows/ops/triton/layer_norm.py b/build/torch210-cu128-x86_64-windows/ops/triton/layer_norm.py deleted file mode 100644 index 3724dd92e91fd8af7def1c849ab74d6adee4e486..0000000000000000000000000000000000000000 --- a/build/torch210-cu128-x86_64-windows/ops/triton/layer_norm.py +++ /dev/null @@ -1,1252 +0,0 @@ -# Copyright (c) 2024, Tri Dao. -# Implement dropout + residual + layer_norm / rms_norm. - -# Based on the Triton LayerNorm tutorial: https://triton-lang.org/main/getting-started/tutorials/05-layer-norm.html -# For the backward pass, we keep weight_grad and bias_grad in registers and accumulate. -# This is faster for dimensions up to 8k, but after that it's much slower due to register spilling. -# The models we train have hidden dim up to 8k anyway (e.g. Llama 70B), so this is fine. - -import math -from typing import Optional, List - -import torch -import torch.nn.functional as F -from torch import Tensor - -import triton -import triton.language as tl - -from flash_attn.utils.torch import custom_fwd, custom_bwd -from flash_attn.utils.library import triton_op - - -def maybe_contiguous_lastdim(x): - return x.contiguous() if x is not None and x.stride(-1) != 1 else x - - -def maybe_contiguous(x): - return x.contiguous() if x is not None else None - - -def triton_autotune_configs(): - # Return configs with a valid warp count for the current device - configs = [] - # Maximum threads per block is architecture-dependent in theory, but in reality all are 1024 - max_threads_per_block = 1024 - # Default to warp size 32 if not defined by device - warp_size = getattr(torch.cuda.get_device_properties(torch.cuda.current_device()), "warp_size", 32) - # Autotune for warp counts which are powers of 2 and do not exceed thread per block limit - return [triton.Config({}, num_warps=warp_count) for warp_count in [1, 2, 4, 8, 16, 32] - if warp_count * warp_size <= max_threads_per_block] - # return [triton.Config({}, num_warps=8)] - - -def layer_norm_ref( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - zero_centered_weight=False, - dropout_mask=None, - dropout_mask1=None, - upcast=False, -): - dtype = x.dtype - if upcast: - x = x.float() - weight = weight.float() - bias = bias.float() if bias is not None else None - residual = residual.float() if residual is not None else residual - x1 = x1.float() if x1 is not None else None - weight1 = weight1.float() if weight1 is not None else None - bias1 = bias1.float() if bias1 is not None else None - if zero_centered_weight: - weight = weight + 1.0 - if weight1 is not None: - weight1 = weight1 + 1.0 - if x1 is not None: - assert rowscale is None, "rowscale is not supported with parallel LayerNorm" - if rowscale is not None: - x = x * rowscale[..., None] - if dropout_p > 0.0: - if dropout_mask is not None: - x = x.masked_fill(~dropout_mask, 0.0) / (1.0 - dropout_p) - else: - x = F.dropout(x, p=dropout_p) - if x1 is not None: - if dropout_mask1 is not None: - x1 = x1.masked_fill(~dropout_mask1, 0.0) / (1.0 - dropout_p) - else: - x1 = F.dropout(x1, p=dropout_p) - if x1 is not None: - x = x + x1 - if residual is not None: - x = (x + residual).to(x.dtype) - out = F.layer_norm(x.to(weight.dtype), x.shape[-1:], weight=weight, bias=bias, eps=eps).to( - dtype - ) - if weight1 is None: - return out if not prenorm else (out, x) - else: - out1 = F.layer_norm( - x.to(weight1.dtype), x.shape[-1:], weight=weight1, bias=bias1, eps=eps - ).to(dtype) - return (out, out1) if not prenorm else (out, out1, x) - - -def rms_norm_ref( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - zero_centered_weight=False, - dropout_mask=None, - dropout_mask1=None, - upcast=False, -): - dtype = x.dtype - if upcast: - x = x.float() - weight = weight.float() - bias = bias.float() if bias is not None else None - residual = residual.float() if residual is not None else residual - x1 = x1.float() if x1 is not None else None - weight1 = weight1.float() if weight1 is not None else None - bias1 = bias1.float() if bias1 is not None else None - if zero_centered_weight: - weight = weight + 1.0 - if weight1 is not None: - weight1 = weight1 + 1.0 - if x1 is not None: - assert rowscale is None, "rowscale is not supported with parallel LayerNorm" - if rowscale is not None: - x = x * rowscale[..., None] - if dropout_p > 0.0: - if dropout_mask is not None: - x = x.masked_fill(~dropout_mask, 0.0) / (1.0 - dropout_p) - else: - x = F.dropout(x, p=dropout_p) - if x1 is not None: - if dropout_mask1 is not None: - x1 = x1.masked_fill(~dropout_mask1, 0.0) / (1.0 - dropout_p) - else: - x1 = F.dropout(x1, p=dropout_p) - if x1 is not None: - x = x + x1 - if residual is not None: - x = (x + residual).to(x.dtype) - rstd = 1 / torch.sqrt((x.square()).mean(dim=-1, keepdim=True) + eps) - out = ((x * rstd * weight) + bias if bias is not None else (x * rstd * weight)).to(dtype) - if weight1 is None: - return out if not prenorm else (out, x) - else: - out1 = ((x * rstd * weight1) + bias1 if bias1 is not None else (x * rstd * weight1)).to( - dtype - ) - return (out, out1) if not prenorm else (out, out1, x) - - -@triton.autotune( - configs=triton_autotune_configs(), - key=["N", "HAS_RESIDUAL", "STORE_RESIDUAL_OUT", "IS_RMS_NORM", "HAS_BIAS", "HAS_X1", "HAS_W1", "HAS_B1"], -) -# torch compile doesn't like triton.heuristics, so we set these manually when calling the kernel -# @triton.heuristics({"HAS_BIAS": lambda args: args["B"] is not None}) -# @triton.heuristics({"HAS_RESIDUAL": lambda args: args["RESIDUAL"] is not None}) -# @triton.heuristics({"HAS_X1": lambda args: args["X1"] is not None}) -# @triton.heuristics({"HAS_W1": lambda args: args["W1"] is not None}) -# @triton.heuristics({"HAS_B1": lambda args: args["B1"] is not None}) -@triton.jit -def _layer_norm_fwd_1pass_kernel( - X, # pointer to the input - Y, # pointer to the output - W, # pointer to the weights - B, # pointer to the biases - RESIDUAL, # pointer to the residual - X1, - W1, - B1, - Y1, - RESIDUAL_OUT, # pointer to the residual - ROWSCALE, - SEEDS, # Dropout seeds for each row - DROPOUT_MASK, - DROPOUT_MASK1, - Mean, # pointer to the mean - Rstd, # pointer to the 1/std - stride_x_row, # how much to increase the pointer when moving by 1 row - stride_y_row, - stride_res_row, - stride_res_out_row, - stride_x1_row, - stride_y1_row, - M, # number of rows in X - N, # number of columns in X - eps, # epsilon to avoid division by zero - dropout_p, # Dropout probability - zero_centered_weight, # If true, add 1.0 to the weight - IS_RMS_NORM: tl.constexpr, - BLOCK_N: tl.constexpr, - HAS_RESIDUAL: tl.constexpr, - STORE_RESIDUAL_OUT: tl.constexpr, - HAS_BIAS: tl.constexpr, - HAS_DROPOUT: tl.constexpr, - STORE_DROPOUT_MASK: tl.constexpr, - HAS_ROWSCALE: tl.constexpr, - HAS_X1: tl.constexpr, - HAS_W1: tl.constexpr, - HAS_B1: tl.constexpr, -): - # Map the program id to the row of X and Y it should compute. - row = tl.program_id(0) - X += row * stride_x_row - Y += row * stride_y_row - if HAS_RESIDUAL: - RESIDUAL += row * stride_res_row - if STORE_RESIDUAL_OUT: - RESIDUAL_OUT += row * stride_res_out_row - if HAS_X1: - X1 += row * stride_x1_row - if HAS_W1: - Y1 += row * stride_y1_row - # Compute mean and variance - cols = tl.arange(0, BLOCK_N) - x = tl.load(X + cols, mask=cols < N, other=0.0).to(tl.float32) - if HAS_ROWSCALE: - rowscale = tl.load(ROWSCALE + row).to(tl.float32) - x *= rowscale - if HAS_DROPOUT: - # Compute dropout mask - # 7 rounds is good enough, and reduces register pressure - keep_mask = tl.rand(tl.load(SEEDS + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - x = tl.where(keep_mask, x / (1.0 - dropout_p), 0.0) - if STORE_DROPOUT_MASK: - tl.store(DROPOUT_MASK + row * N + cols, keep_mask, mask=cols < N) - if HAS_X1: - x1 = tl.load(X1 + cols, mask=cols < N, other=0.0).to(tl.float32) - if HAS_ROWSCALE: - rowscale = tl.load(ROWSCALE + M + row).to(tl.float32) - x1 *= rowscale - if HAS_DROPOUT: - # Compute dropout mask - # 7 rounds is good enough, and reduces register pressure - keep_mask = ( - tl.rand(tl.load(SEEDS + M + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - ) - x1 = tl.where(keep_mask, x1 / (1.0 - dropout_p), 0.0) - if STORE_DROPOUT_MASK: - tl.store(DROPOUT_MASK1 + row * N + cols, keep_mask, mask=cols < N) - x += x1 - if HAS_RESIDUAL: - residual = tl.load(RESIDUAL + cols, mask=cols < N, other=0.0).to(tl.float32) - x += residual - if STORE_RESIDUAL_OUT: - tl.store(RESIDUAL_OUT + cols, x, mask=cols < N) - if not IS_RMS_NORM: - mean = tl.sum(x, axis=0) / N - tl.store(Mean + row, mean) - xbar = tl.where(cols < N, x - mean, 0.0) - var = tl.sum(xbar * xbar, axis=0) / N - else: - xbar = tl.where(cols < N, x, 0.0) - var = tl.sum(xbar * xbar, axis=0) / N - rstd = 1 / tl.sqrt(var + eps) - tl.store(Rstd + row, rstd) - # Normalize and apply linear transformation - mask = cols < N - w = tl.load(W + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w += 1.0 - if HAS_BIAS: - b = tl.load(B + cols, mask=mask).to(tl.float32) - x_hat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd - y = x_hat * w + b if HAS_BIAS else x_hat * w - # Write output - tl.store(Y + cols, y, mask=mask) - if HAS_W1: - w1 = tl.load(W1 + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w1 += 1.0 - if HAS_B1: - b1 = tl.load(B1 + cols, mask=mask).to(tl.float32) - y1 = x_hat * w1 + b1 if HAS_B1 else x_hat * w1 - tl.store(Y1 + cols, y1, mask=mask) - - -def _layer_norm_fwd( - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - residual: Optional[Tensor] = None, - x1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - out_dtype: Optional[torch.dtype] = None, - residual_dtype: Optional[torch.dtype] = None, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - return_dropout_mask: bool = False, - out: Optional[Tensor] = None, - residual_out: Optional[Tensor] = None -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - # Need to wrap to handle the case where residual_out is a alias of x, which makes torch.library - # and torch.compile unhappy. Also allocate memory for out and residual_out if they are None - # so that _layer_norm_fwd_impl doesn't have to return them. - if out is None: - out = torch.empty_like(x, dtype=x.dtype if out_dtype is None else out_dtype) - if residual is not None: - residual_dtype = residual.dtype - if residual_out is None and ( - residual is not None - or (residual_dtype is not None and residual_dtype != x.dtype) - or dropout_p > 0.0 - or rowscale is not None - or x1 is not None - ): - residual_out = torch.empty_like( - x, dtype=residual_dtype if residual_dtype is not None else x.dtype - ) - else: - residual_out = None - y1, mean, rstd, seeds, dropout_mask, dropout_mask1 = _layer_norm_fwd_impl( - x, - weight, - bias, - eps, - out, - residual=residual, - x1=x1, - weight1=weight1, - bias1=bias1, - dropout_p=dropout_p, - rowscale=rowscale, - zero_centered_weight=zero_centered_weight, - is_rms_norm=is_rms_norm, - return_dropout_mask=return_dropout_mask, - residual_out=residual_out, - ) - # residual_out is None if residual is None and residual_dtype == input_dtype and dropout_p == 0.0 - if residual_out is None: - residual_out = x - return out, y1, mean, rstd, residual_out, seeds, dropout_mask, dropout_mask1 - - -# [2025-04-28] torch.library.triton_op ignores the schema argument, but here we need the schema -# since we're returning a tuple of tensors -@triton_op("flash_attn::layer_norm_fwd_impl", mutates_args={"out", "residual_out"}, - schema="(Tensor x, Tensor weight, Tensor bias, float eps, Tensor(a!) out, Tensor? residual, Tensor? x1, Tensor? weight1, Tensor? bias1, float dropout_p, Tensor? rowscale, bool zero_centered_weight, bool is_rms_norm, bool return_dropout_mask, Tensor(a!)? residual_out) -> (Tensor y1, Tensor mean, Tensor rstd, Tensor seeds, Tensor dropout_mask, Tensor dropout_mask1)") -def _layer_norm_fwd_impl( - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - out: Tensor, - residual: Optional[Tensor] = None, - x1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - return_dropout_mask: bool = False, - residual_out: Optional[Tensor] = None -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - M, N = x.shape - assert x.stride(-1) == 1 - if residual is not None: - assert residual.stride(-1) == 1 - assert residual.shape == (M, N) - assert weight.shape == (N,) - assert weight.stride(-1) == 1 - if bias is not None: - assert bias.stride(-1) == 1 - assert bias.shape == (N,) - if x1 is not None: - assert x1.shape == x.shape - assert rowscale is None - assert x1.stride(-1) == 1 - if weight1 is not None: - assert weight1.shape == (N,) - assert weight1.stride(-1) == 1 - if bias1 is not None: - assert bias1.shape == (N,) - assert bias1.stride(-1) == 1 - if rowscale is not None: - assert rowscale.is_contiguous() - assert rowscale.shape == (M,) - assert out.shape == x.shape - assert out.stride(-1) == 1 - if residual_out is not None: - assert residual_out.shape == x.shape - assert residual_out.stride(-1) == 1 - if weight1 is not None: - y1 = torch.empty_like(out) - assert y1.stride(-1) == 1 - else: - y1 = None - mean = torch.empty((M,), dtype=torch.float32, device=x.device) if not is_rms_norm else None - rstd = torch.empty((M,), dtype=torch.float32, device=x.device) - if dropout_p > 0.0: - seeds = torch.randint( - 2**32, (M if x1 is None else 2 * M,), device=x.device, dtype=torch.int64 - ) - else: - seeds = None - if return_dropout_mask and dropout_p > 0.0: - dropout_mask = torch.empty(M, N, device=x.device, dtype=torch.bool) - if x1 is not None: - dropout_mask1 = torch.empty(M, N, device=x.device, dtype=torch.bool) - else: - dropout_mask1 = None - else: - dropout_mask, dropout_mask1 = None, None - # Less than 64KB per feature: enqueue fused kernel - MAX_FUSED_SIZE = 65536 // x.element_size() - BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) - if N > BLOCK_N: - raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") - with torch.cuda.device(x.device.index): - torch.library.wrap_triton(_layer_norm_fwd_1pass_kernel)[(M,)]( - x, - out, - weight, - bias, - residual, - x1, - weight1, - bias1, - y1, - residual_out, - rowscale, - seeds, - dropout_mask, - dropout_mask1, - mean, - rstd, - x.stride(0), - out.stride(0), - residual.stride(0) if residual is not None else 0, - residual_out.stride(0) if residual_out is not None else 0, - x1.stride(0) if x1 is not None else 0, - y1.stride(0) if y1 is not None else 0, - M, - N, - eps, - dropout_p, - # Passing bool make torch inductor very unhappy since it then tries to compare to int_max - int(zero_centered_weight), - is_rms_norm, - BLOCK_N, - residual is not None, - residual_out is not None, - bias is not None, - dropout_p > 0.0, - dropout_mask is not None, - rowscale is not None, - HAS_X1=x1 is not None, - HAS_W1=weight1 is not None, - HAS_B1=bias1 is not None, - ) - return y1, mean, rstd, seeds, dropout_mask, dropout_mask1 - - -@triton.autotune( - configs=triton_autotune_configs(), - key=["N", "HAS_DRESIDUAL", "STORE_DRESIDUAL", "IS_RMS_NORM", "HAS_BIAS", "HAS_DROPOUT"], -) -# torch compile doesn't like triton.heuristics, so we set these manually when calling the kernel -# @triton.heuristics({"HAS_BIAS": lambda args: args["B"] is not None}) -# @triton.heuristics({"HAS_DRESIDUAL": lambda args: args["DRESIDUAL"] is not None}) -# @triton.heuristics({"STORE_DRESIDUAL": lambda args: args["DRESIDUAL_IN"] is not None}) -# @triton.heuristics({"HAS_ROWSCALE": lambda args: args["ROWSCALE"] is not None}) -# @triton.heuristics({"HAS_DY1": lambda args: args["DY1"] is not None}) -# @triton.heuristics({"HAS_DX1": lambda args: args["DX1"] is not None}) -# @triton.heuristics({"HAS_B1": lambda args: args["DB1"] is not None}) -# @triton.heuristics({"RECOMPUTE_OUTPUT": lambda args: args["Y"] is not None}) -@triton.jit -def _layer_norm_bwd_kernel( - X, # pointer to the input - W, # pointer to the weights - B, # pointer to the biases - Y, # pointer to the output to be recomputed - DY, # pointer to the output gradient - DX, # pointer to the input gradient - DW, # pointer to the partial sum of weights gradient - DB, # pointer to the partial sum of biases gradient - DRESIDUAL, - W1, - DY1, - DX1, - DW1, - DB1, - DRESIDUAL_IN, - ROWSCALE, - SEEDS, - Mean, # pointer to the mean - Rstd, # pointer to the 1/std - stride_x_row, # how much to increase the pointer when moving by 1 row - stride_y_row, - stride_dy_row, - stride_dx_row, - stride_dres_row, - stride_dy1_row, - stride_dx1_row, - stride_dres_in_row, - M, # number of rows in X - N, # number of columns in X - eps, # epsilon to avoid division by zero - dropout_p, - zero_centered_weight, - rows_per_program, - IS_RMS_NORM: tl.constexpr, - BLOCK_N: tl.constexpr, - HAS_DRESIDUAL: tl.constexpr, - STORE_DRESIDUAL: tl.constexpr, - HAS_BIAS: tl.constexpr, - HAS_DROPOUT: tl.constexpr, - HAS_ROWSCALE: tl.constexpr, - HAS_DY1: tl.constexpr, - HAS_DX1: tl.constexpr, - HAS_B1: tl.constexpr, - RECOMPUTE_OUTPUT: tl.constexpr, -): - # Map the program id to the elements of X, DX, and DY it should compute. - row_block_id = tl.program_id(0) - row_start = row_block_id * rows_per_program - # Do not early exit if row_start >= M, because we need to write DW and DB - cols = tl.arange(0, BLOCK_N) - mask = cols < N - X += row_start * stride_x_row - if HAS_DRESIDUAL: - DRESIDUAL += row_start * stride_dres_row - if STORE_DRESIDUAL: - DRESIDUAL_IN += row_start * stride_dres_in_row - DY += row_start * stride_dy_row - DX += row_start * stride_dx_row - if HAS_DY1: - DY1 += row_start * stride_dy1_row - if HAS_DX1: - DX1 += row_start * stride_dx1_row - if RECOMPUTE_OUTPUT: - Y += row_start * stride_y_row - w = tl.load(W + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w += 1.0 - if RECOMPUTE_OUTPUT and HAS_BIAS: - b = tl.load(B + cols, mask=mask, other=0.0).to(tl.float32) - if HAS_DY1: - w1 = tl.load(W1 + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w1 += 1.0 - dw = tl.zeros((BLOCK_N,), dtype=tl.float32) - if HAS_BIAS: - db = tl.zeros((BLOCK_N,), dtype=tl.float32) - if HAS_DY1: - dw1 = tl.zeros((BLOCK_N,), dtype=tl.float32) - if HAS_B1: - db1 = tl.zeros((BLOCK_N,), dtype=tl.float32) - row_end = min((row_block_id + 1) * rows_per_program, M) - for row in range(row_start, row_end): - # Load data to SRAM - x = tl.load(X + cols, mask=mask, other=0).to(tl.float32) - dy = tl.load(DY + cols, mask=mask, other=0).to(tl.float32) - if HAS_DY1: - dy1 = tl.load(DY1 + cols, mask=mask, other=0).to(tl.float32) - if not IS_RMS_NORM: - mean = tl.load(Mean + row) - rstd = tl.load(Rstd + row) - # Compute dx - xhat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd - xhat = tl.where(mask, xhat, 0.0) - if RECOMPUTE_OUTPUT: - y = xhat * w + b if HAS_BIAS else xhat * w - tl.store(Y + cols, y, mask=mask) - wdy = w * dy - dw += dy * xhat - if HAS_BIAS: - db += dy - if HAS_DY1: - wdy += w1 * dy1 - dw1 += dy1 * xhat - if HAS_B1: - db1 += dy1 - if not IS_RMS_NORM: - c1 = tl.sum(xhat * wdy, axis=0) / N - c2 = tl.sum(wdy, axis=0) / N - dx = (wdy - (xhat * c1 + c2)) * rstd - else: - c1 = tl.sum(xhat * wdy, axis=0) / N - dx = (wdy - xhat * c1) * rstd - if HAS_DRESIDUAL: - dres = tl.load(DRESIDUAL + cols, mask=mask, other=0).to(tl.float32) - dx += dres - # Write dx - if STORE_DRESIDUAL: - tl.store(DRESIDUAL_IN + cols, dx, mask=mask) - if HAS_DX1: - if HAS_DROPOUT: - keep_mask = ( - tl.rand(tl.load(SEEDS + M + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - ) - dx1 = tl.where(keep_mask, dx / (1.0 - dropout_p), 0.0) - else: - dx1 = dx - tl.store(DX1 + cols, dx1, mask=mask) - if HAS_DROPOUT: - keep_mask = tl.rand(tl.load(SEEDS + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - dx = tl.where(keep_mask, dx / (1.0 - dropout_p), 0.0) - if HAS_ROWSCALE: - rowscale = tl.load(ROWSCALE + row).to(tl.float32) - dx *= rowscale - tl.store(DX + cols, dx, mask=mask) - - X += stride_x_row - if HAS_DRESIDUAL: - DRESIDUAL += stride_dres_row - if STORE_DRESIDUAL: - DRESIDUAL_IN += stride_dres_in_row - if RECOMPUTE_OUTPUT: - Y += stride_y_row - DY += stride_dy_row - DX += stride_dx_row - if HAS_DY1: - DY1 += stride_dy1_row - if HAS_DX1: - DX1 += stride_dx1_row - tl.store(DW + row_block_id * N + cols, dw, mask=mask) - if HAS_BIAS: - tl.store(DB + row_block_id * N + cols, db, mask=mask) - if HAS_DY1: - tl.store(DW1 + row_block_id * N + cols, dw1, mask=mask) - if HAS_B1: - tl.store(DB1 + row_block_id * N + cols, db1, mask=mask) - - -def _layer_norm_bwd( - dy: Tensor, - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - mean: Tensor, - rstd: Tensor, - dresidual: Optional[Tensor] = None, - dy1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - seeds: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - has_residual: bool = False, - has_x1: bool = False, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - x_dtype: Optional[torch.dtype] = None, - recompute_output: bool = False, -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - # Need to wrap to handle the case where dresidual_in or dx1 are aliases of x, - # which makes torch.library unhappy - dx, dw, db, dresidual_in, dx1, dw1, db1, y = _layer_norm_bwd_impl( - dy, - x, - weight, - bias, - eps, - mean, - rstd, - dresidual, - dy1, - weight1, - bias1, - seeds, - dropout_p, - rowscale, - has_residual, - has_x1, - zero_centered_weight, - is_rms_norm, - x_dtype=x_dtype, - recompute_output=recompute_output, - ) - # Don't need to compute dresidual_in separately in this case - if has_residual and dx.dtype == x.dtype and dropout_p == 0.0 and rowscale is None: - dresidual_in = dx - if has_x1 and dropout_p == 0.0: - dx1 = dx - return dx, dw, db, dresidual_in, dx1, dw1, db1, y - - - -@triton_op("flash_attn::layer_norm_bwd_impl", mutates_args={}, - schema="(Tensor dy, Tensor x, Tensor weight, Tensor bias, float eps, Tensor mean, Tensor rstd, Tensor? dresidual, Tensor? dy1, Tensor? weight1, Tensor? bias1, Tensor? seeds, float dropout_p, Tensor? rowscale, bool has_residual, bool has_x1, bool zero_centered_weight, bool is_rms_norm, ScalarType? x_dtype, bool recompute_output) -> (Tensor dx, Tensor dw, Tensor db, Tensor dresidual_in, Tensor dx1, Tensor dw1, Tensor db1, Tensor y)", - allow_decomposition=False, # Don't let torch.compile trace inside - ) -def _layer_norm_bwd_impl( - dy: Tensor, - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - mean: Tensor, - rstd: Tensor, - dresidual: Optional[Tensor] = None, - dy1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - seeds: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - has_residual: bool = False, - has_x1: bool = False, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - x_dtype: Optional[torch.dtype] = None, - recompute_output: bool = False, -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - M, N = x.shape - assert x.stride(-1) == 1 - dy = maybe_contiguous_lastdim(dy) - assert dy.stride(-1) == 1 - assert dy.shape == (M, N) - if dresidual is not None: - dresidual = maybe_contiguous_lastdim(dresidual) - assert dresidual.stride(-1) == 1 - assert dresidual.shape == (M, N) - assert weight.shape == (N,) - assert weight.stride(-1) == 1 - if bias is not None: - assert bias.stride(-1) == 1 - assert bias.shape == (N,) - if dy1 is not None: - dy1 = maybe_contiguous_lastdim(dy1) - assert weight1 is not None - assert dy1.shape == dy.shape - assert dy1.stride(-1) == 1 - if weight1 is not None: - assert weight1.shape == (N,) - assert weight1.stride(-1) == 1 - if bias1 is not None: - assert bias1.shape == (N,) - assert bias1.stride(-1) == 1 - if seeds is not None: - assert seeds.is_contiguous() - assert seeds.shape == (M if not has_x1 else M * 2,) - if rowscale is not None: - assert rowscale.is_contiguous() - assert rowscale.shape == (M,) - # allocate output - dx = ( - torch.empty_like(x) - if x_dtype is None - else torch.empty(M, N, dtype=x_dtype, device=x.device) - ) - dresidual_in = ( - torch.empty_like(x) - if has_residual - and (dx.dtype != x.dtype or dropout_p > 0.0 or rowscale is not None or has_x1) - else None - ) - dx1 = torch.empty_like(dx) if (has_x1 and dropout_p > 0.0) else None - y = torch.empty(M, N, dtype=dy.dtype, device=dy.device) if recompute_output else None - if recompute_output: - assert weight1 is None, "recompute_output is not supported with parallel LayerNorm" - - # Less than 64KB per feature: enqueue fused kernel - MAX_FUSED_SIZE = 65536 // x.element_size() - BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) - if N > BLOCK_N: - raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") - # Increasing the multiple (e.g. 8) will allow more thread blocks to be launched and hide the - # latency of the gmem reads/writes, but will increase the time of summing up dw / db. - sm_count = torch.cuda.get_device_properties(x.device).multi_processor_count * 8 - _dw = torch.empty((sm_count, N), dtype=torch.float32, device=weight.device) - _db = ( - torch.empty((sm_count, N), dtype=torch.float32, device=bias.device) - if bias is not None - else None - ) - _dw1 = torch.empty_like(_dw) if weight1 is not None else None - _db1 = torch.empty_like(_db) if bias1 is not None else None - rows_per_program = math.ceil(M / sm_count) - grid = (sm_count,) - with torch.cuda.device(x.device.index): - torch.library.wrap_triton(_layer_norm_bwd_kernel)[grid]( - x, - weight, - bias, - y, - dy, - dx, - _dw, - _db, - dresidual, - weight1, - dy1, - dx1, - _dw1, - _db1, - dresidual_in, - rowscale, - seeds, - mean, - rstd, - x.stride(0), - 0 if not recompute_output else y.stride(0), - dy.stride(0), - dx.stride(0), - dresidual.stride(0) if dresidual is not None else 0, - dy1.stride(0) if dy1 is not None else 0, - dx1.stride(0) if dx1 is not None else 0, - dresidual_in.stride(0) if dresidual_in is not None else 0, - M, - N, - eps, - dropout_p, - # Passing bool make torch inductor very unhappy since it then tries to compare to int_max - int(zero_centered_weight), - rows_per_program, - is_rms_norm, - BLOCK_N, - dresidual is not None, - dresidual_in is not None, - bias is not None, - dropout_p > 0.0, - HAS_ROWSCALE=rowscale is not None, - HAS_DY1=dy1 is not None, - HAS_DX1=dx1 is not None, - HAS_B1=bias1 is not None, - RECOMPUTE_OUTPUT=y is not None, - ) - dw = _dw.sum(0).to(weight.dtype) - db = _db.sum(0).to(bias.dtype) if bias is not None else None - dw1 = _dw1.sum(0).to(weight1.dtype) if weight1 is not None else None - db1 = _db1.sum(0).to(bias1.dtype) if bias1 is not None else None - # dresidual_in and dx1 could be None, the wrapper will handle assigning them from dx - return dx, dw, db, dresidual_in, dx1, dw1, db1, y - - -class LayerNormFn(torch.autograd.Function): - - @staticmethod - def forward( - ctx, - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - residual_in_fp32=False, - zero_centered_weight=False, - is_rms_norm=False, - return_dropout_mask=False, - out_dtype=None, - out=None, - residual_out=None - ): - x_shape_og = x.shape - # reshape input data into 2D tensor - x = maybe_contiguous_lastdim(x.reshape(-1, x.shape[-1])) - if residual is not None: - assert residual.shape == x_shape_og - residual = maybe_contiguous_lastdim(residual.reshape(-1, residual.shape[-1])) - if x1 is not None: - assert x1.shape == x_shape_og - assert rowscale is None, "rowscale is not supported with parallel LayerNorm" - x1 = maybe_contiguous_lastdim(x1.reshape(-1, x1.shape[-1])) - weight = weight.contiguous() - bias = maybe_contiguous(bias) - weight1 = maybe_contiguous(weight1) - bias1 = maybe_contiguous(bias1) - if rowscale is not None: - rowscale = rowscale.reshape(-1).contiguous() - residual_dtype = ( - residual.dtype - if residual is not None - else (torch.float32 if residual_in_fp32 else None) - ) - if out is not None: - out = out.reshape(-1, out.shape[-1]) - if residual_out is not None: - residual_out = residual_out.reshape(-1, residual_out.shape[-1]) - y, y1, mean, rstd, residual_out, seeds, dropout_mask, dropout_mask1 = _layer_norm_fwd( - x, - weight, - bias, - eps, - residual, - x1, - weight1, - bias1, - dropout_p=dropout_p, - rowscale=rowscale, - out_dtype=out_dtype, - residual_dtype=residual_dtype, - zero_centered_weight=zero_centered_weight, - is_rms_norm=is_rms_norm, - return_dropout_mask=return_dropout_mask, - out=out, - residual_out=residual_out, - ) - ctx.save_for_backward( - residual_out, weight, bias, weight1, bias1, rowscale, seeds, mean, rstd - ) - ctx.x_shape_og = x_shape_og - ctx.eps = eps - ctx.dropout_p = dropout_p - ctx.is_rms_norm = is_rms_norm - ctx.has_residual = residual is not None - ctx.has_x1 = x1 is not None - ctx.prenorm = prenorm - ctx.x_dtype = x.dtype - ctx.zero_centered_weight = zero_centered_weight - y = y.reshape(x_shape_og) - y1 = y1.reshape(x_shape_og) if y1 is not None else None - residual_out = residual_out.reshape(x_shape_og) if residual_out is not None else None - dropout_mask = dropout_mask.reshape(x_shape_og) if dropout_mask is not None else None - dropout_mask1 = dropout_mask1.reshape(x_shape_og) if dropout_mask1 is not None else None - if not return_dropout_mask: - if weight1 is None: - return y if not prenorm else (y, residual_out) - else: - return (y, y1) if not prenorm else (y, y1, residual_out) - else: - if weight1 is None: - return ( - (y, dropout_mask, dropout_mask1) - if not prenorm - else (y, residual_out, dropout_mask, dropout_mask1) - ) - else: - return ( - (y, y1, dropout_mask, dropout_mask1) - if not prenorm - else (y, y1, residual_out, dropout_mask, dropout_mask1) - ) - - @staticmethod - def backward(ctx, dy, *args): - x, weight, bias, weight1, bias1, rowscale, seeds, mean, rstd = ctx.saved_tensors - dy = dy.reshape(-1, dy.shape[-1]) - if weight1 is not None: - dy1, args = args[0], args[1:] - dy1 = dy1.reshape(-1, dy1.shape[-1]) - assert dy1.shape == x.shape - else: - dy1 = None - if ctx.prenorm: - dresidual = args[0] - dresidual = dresidual.reshape(-1, dresidual.shape[-1]) - assert dresidual.shape == x.shape - else: - dresidual = None - dx, dw, db, dresidual_in, dx1, dw1, db1, _ = _layer_norm_bwd( - dy, - x, - weight, - bias, - ctx.eps, - mean, - rstd, - dresidual, - dy1, - weight1, - bias1, - seeds, - ctx.dropout_p, - rowscale, - ctx.has_residual, - ctx.has_x1, - ctx.zero_centered_weight, - ctx.is_rms_norm, - x_dtype=ctx.x_dtype, - recompute_output=False, - ) - return ( - dx.reshape(ctx.x_shape_og), - dw, - db, - dresidual_in.reshape(ctx.x_shape_og) if ctx.has_residual else None, - dx1.reshape(ctx.x_shape_og) if dx1 is not None else None, - dw1, - db1, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) - - -def layer_norm_fn( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - residual_in_fp32=False, - zero_centered_weight=False, - is_rms_norm=False, - return_dropout_mask=False, - out_dtype=None, - out=None, - residual_out=None -): - return LayerNormFn.apply( - x, - weight, - bias, - residual, - x1, - weight1, - bias1, - eps, - dropout_p, - rowscale, - prenorm, - residual_in_fp32, - zero_centered_weight, - is_rms_norm, - return_dropout_mask, - out_dtype, - out, - residual_out - ) - - -def rms_norm_fn( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - residual_in_fp32=False, - zero_centered_weight=False, - return_dropout_mask=False, - out_dtype=None, - out=None, - residual_out=None -): - return LayerNormFn.apply( - x, - weight, - bias, - residual, - x1, - weight1, - bias1, - eps, - dropout_p, - rowscale, - prenorm, - residual_in_fp32, - zero_centered_weight, - True, - return_dropout_mask, - out_dtype, - out, - residual_out - ) - - -class RMSNorm(torch.nn.Module): - - def __init__(self, hidden_size, eps=1e-5, dropout_p=0.0, zero_centered_weight=False, - device=None, dtype=None): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.eps = eps - if dropout_p > 0.0: - self.drop = torch.nn.Dropout(dropout_p) - else: - self.drop = None - self.zero_centered_weight = zero_centered_weight - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.register_parameter("bias", None) - self.reset_parameters() - - def reset_parameters(self): - if not self.zero_centered_weight: - torch.nn.init.ones_(self.weight) - else: - torch.nn.init.zeros_(self.weight) - - def forward(self, x, residual=None, prenorm=False, residual_in_fp32=False): - return rms_norm_fn( - x, - self.weight, - self.bias, - residual=residual, - eps=self.eps, - dropout_p=self.drop.p if self.drop is not None and self.training else 0.0, - prenorm=prenorm, - residual_in_fp32=residual_in_fp32, - zero_centered_weight=self.zero_centered_weight, - ) - - -class LayerNormLinearFn(torch.autograd.Function): - - @staticmethod - @custom_fwd - def forward( - ctx, - x, - norm_weight, - norm_bias, - linear_weight, - linear_bias, - residual=None, - eps=1e-6, - prenorm=False, - residual_in_fp32=False, - is_rms_norm=False, - ): - x_shape_og = x.shape - # reshape input data into 2D tensor - x = maybe_contiguous_lastdim(x.reshape(-1, x.shape[-1])) - if residual is not None: - assert residual.shape == x_shape_og - residual = maybe_contiguous_lastdim(residual.reshape(-1, residual.shape[-1])) - norm_weight = norm_weight.contiguous() - norm_bias = maybe_contiguous(norm_bias) - residual_dtype = ( - residual.dtype - if residual is not None - else (torch.float32 if residual_in_fp32 else None) - ) - y, _, mean, rstd, residual_out, *rest = _layer_norm_fwd( - x, - norm_weight, - norm_bias, - eps, - residual, - out_dtype=None if not torch.is_autocast_enabled() else torch.get_autocast_dtype("cuda"), - residual_dtype=residual_dtype, - is_rms_norm=is_rms_norm, - ) - y = y.reshape(x_shape_og) - dtype = torch.get_autocast_dtype("cuda") if torch.is_autocast_enabled() else y.dtype - linear_weight = linear_weight.to(dtype) - linear_bias = linear_bias.to(dtype) if linear_bias is not None else None - out = F.linear(y.to(linear_weight.dtype), linear_weight, linear_bias) - # We don't store y, will be recomputed in the backward pass to save memory - ctx.save_for_backward(residual_out, norm_weight, norm_bias, linear_weight, mean, rstd) - ctx.x_shape_og = x_shape_og - ctx.eps = eps - ctx.is_rms_norm = is_rms_norm - ctx.has_residual = residual is not None - ctx.prenorm = prenorm - ctx.x_dtype = x.dtype - ctx.linear_bias_is_none = linear_bias is None - return out if not prenorm else (out, residual_out.reshape(x_shape_og)) - - @staticmethod - @custom_bwd - def backward(ctx, dout, *args): - x, norm_weight, norm_bias, linear_weight, mean, rstd = ctx.saved_tensors - dout = dout.reshape(-1, dout.shape[-1]) - dy = F.linear(dout, linear_weight.t()) - dlinear_bias = None if ctx.linear_bias_is_none else dout.sum(0) - dy = maybe_contiguous_lastdim(dy) - assert dy.shape == x.shape - if ctx.prenorm: - dresidual = args[0] - dresidual = maybe_contiguous_lastdim(dresidual.reshape(-1, dresidual.shape[-1])) - assert dresidual.shape == x.shape - else: - dresidual = None - dx, dnorm_weight, dnorm_bias, dresidual_in, _, _, _, y = _layer_norm_bwd( - dy, - x, - norm_weight, - norm_bias, - ctx.eps, - mean, - rstd, - dresidual=dresidual, - has_residual=ctx.has_residual, - is_rms_norm=ctx.is_rms_norm, - x_dtype=ctx.x_dtype, - recompute_output=True, - ) - dlinear_weight = torch.einsum("bo,bi->oi", dout, y) - return ( - dx.reshape(ctx.x_shape_og), - dnorm_weight, - dnorm_bias, - dlinear_weight, - dlinear_bias, - dresidual_in.reshape(ctx.x_shape_og) if ctx.has_residual else None, - None, - None, - None, - None, - ) - - -def layer_norm_linear_fn( - x, - norm_weight, - norm_bias, - linear_weight, - linear_bias, - residual=None, - eps=1e-6, - prenorm=False, - residual_in_fp32=False, - is_rms_norm=False, -): - return LayerNormLinearFn.apply( - x, - norm_weight, - norm_bias, - linear_weight, - linear_bias, - residual, - eps, - prenorm, - residual_in_fp32, - is_rms_norm, - ) diff --git a/build/torch210-cu128-x86_64-windows/ops/triton/linear.py b/build/torch210-cu128-x86_64-windows/ops/triton/linear.py deleted file mode 100644 index 0d800851ebdd5a589bb03e5374b807d2c4bd148c..0000000000000000000000000000000000000000 --- a/build/torch210-cu128-x86_64-windows/ops/triton/linear.py +++ /dev/null @@ -1,594 +0,0 @@ -# Adapted from https://github.com/ELS-RD/kernl/blob/main/src/kernl/implementations/linear_layer.py -# and https://github.com/openai/triton/blob/master/python/triton/ops/matmul.py -from typing import Optional - -import torch -import triton -import triton.language as tl -from triton.ops.matmul_perf_model import early_config_prune, estimate_matmul_time - -from flash_attn.ops.triton.k_activations import ( - gelu, - gelu_approx, - gelu_approx_grad, - gelu_grad, - squared_relu, - squared_relu_grad, -) - -# CREDITS: Initially inspired by the Triton tutorial on matrix multiplications - - -def init_to_zero(name): - return lambda nargs: nargs[name].zero_() - - -def get_configs_io_bound(): - configs = [] - for num_stages in [2, 3, 4, 5, 6]: - for block_m in [16, 32]: - for block_k in [32, 64]: - for block_n in [32, 64, 128, 256]: - num_warps = 2 if block_n <= 64 else 4 - configs.append( - triton.Config( - { - "BLOCK_M": block_m, - "BLOCK_N": block_n, - "BLOCK_K": block_k, - "SPLIT_K": 1, - }, - num_stages=num_stages, - num_warps=num_warps, - ) - ) - # split_k not used - # for split_k in [2, 4, 8, 16]: - # configs.append(triton.Config( - # {'BLOCK_M': block_m, 'BLOCK_N': block_n, 'BLOCK_K': block_k, 'SPLIT_K': split_k}, - # num_stages=num_stages, num_warps=num_warps, pre_hook=init_to_zero('C'))) - return configs - - -@triton.autotune( - configs=[ - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - # good for int8 - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=4, - num_warps=4, - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - ] - + get_configs_io_bound(), - key=["CACHE_KEY_M", "CACHE_KEY_N", "CACHE_KEY_K"], - prune_configs_by={ - "early_config_prune": early_config_prune, - "perf_model": estimate_matmul_time, - "top_k": 10, - }, -) -@triton.heuristics( - { - "EVEN_K": lambda args: args["K"] % (args["BLOCK_K"] * args["SPLIT_K"]) == 0, - } -) -@triton.jit -def kernel_fwd( - C, # Pointers to matrices - ACT_INPUT, - A, - B, - bias, - # Matrix dimensions - M, - N, - K, - CACHE_KEY_M, - CACHE_KEY_N, - CACHE_KEY_K, - # The stride variables represent how much to increase the ptr by when moving by 1 - # element in a particular dimension. E.g. stride_am is how much to increase a_ptr - # by to get the element one row down (A has M rows) - stride_cm, - # stride_cn, # Assume that stride_cn == 1 - stride_am, - stride_ak, - stride_bn, - stride_bk, - # Meta-parameters - BLOCK_M: tl.constexpr, - GROUP_M: tl.constexpr, - BLOCK_N: tl.constexpr, - BLOCK_K: tl.constexpr, - # split k not used, not performant with activation, kept because early_config_prune is expecting it - SPLIT_K: tl.constexpr, - EVEN_K: tl.constexpr, - A_ROWMAJOR: tl.constexpr, - B_COLMAJOR: tl.constexpr, - BIAS: tl.constexpr, - SAVE_ACT_INPUT: tl.constexpr, - ACTIVATION: tl.constexpr, -): - - """ - Kernel for computing Out = activation(A x W + C) - - Input has shape (M, K) - - Weight has shape (K, N) - - Bias has shape (N,) - - Output has shape (M, N) - - ActInputs (optional) has shape (M, N) - 'ActInputs' optionally saves the A x W + C intermediate for backward computations - This kernel will consolidate over K - """ - - pid = tl.program_id(axis=0) - - grid_m = (M + BLOCK_M - 1) // BLOCK_M - grid_n = (N + BLOCK_N - 1) // BLOCK_N - # re-order program ID for better L2 performance - width = GROUP_M * grid_n - group_id = pid // width - group_size = min(grid_m - group_id * GROUP_M, GROUP_M) - pid_m = group_id * GROUP_M + (pid % group_size) - pid_n = (pid % width) // (group_size) - - # now compute the block that each program will go through - # rm (resp. rn) denotes a range of indices - # for rows (resp. col) of C - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - # trick to avoid masking on M and N axis - ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M) - rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N) - rk = tl.arange(0, BLOCK_K) - - if A_ROWMAJOR: - A = A + (ram[:, None] * stride_am + rk[None, :]) - else: - A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak) - if B_COLMAJOR: - B = B + (rk[:, None] + rbn[None, :] * stride_bn) - else: - B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn) - - acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) - - for k in range(K, 0, -BLOCK_K): - if EVEN_K: - a = tl.load(A) - b = tl.load(B) - else: - a = tl.load(A, mask=rk[None, :] < k, other=0.0) - b = tl.load(B, mask=rk[:, None] < k, other=0.0) - acc += tl.dot(a, b) - - if A_ROWMAJOR: - A += BLOCK_K - else: - A += BLOCK_K * stride_ak - if B_COLMAJOR: - B += BLOCK_K - else: - B += BLOCK_K * stride_bk - - # Putting bias after the matmul (instead of before) is faster, idk why - if BIAS: - bias = tl.load(bias + rn, mask=rn < N, other=0.0).to(tl.float32) - acc += bias[None, :] - - # optional: save the activation inputs - if SAVE_ACT_INPUT: - # act_in_ptrs = ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :] * stride_cn - act_in_ptrs = ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :] - tl.store(act_in_ptrs, acc) - - # optional: fused activation (while the data is in shared memory) - if ACTIVATION == "gelu": - acc = gelu(acc) - elif ACTIVATION == "gelu_approx": - acc = gelu_approx(acc) - elif ACTIVATION == "squared_relu": - acc = squared_relu(acc) - # rematerialize rm and rn to save registers - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - - # write back result - # C = C + rm[:, None] * stride_cm + rn[None, :] * stride_cn - C = C + rm[:, None] * stride_cm + rn[None, :] - mask = (rm < M)[:, None] & (rn < N)[None, :] - tl.store(C, acc) - - -def triton_linear_act( - x: torch.Tensor, - weight: torch.Tensor, - bias: Optional[torch.Tensor] = None, - activation: str = "id", - save_act_input: bool = False, -) -> torch.Tensor: - """ - Compute e = activation(x @ weight.T + bias). - This wrapper kicks the `kernel_fwd` Triton kernel - :param x: input tensor - :param weight: weight matrix - :param bias: an optional bias tensor - :param activation: Activation name. Needs to be a Triton kernel. - :param act_input: an optional tensor to save the activation inputs (for backward) - :return: result tensor - """ - # if torch.is_autocast_enabled(): - # dtype = torch.get_autocast_gpu_dtype() - # x, weight, bias = [a.to(dtype=dtype) for a in [x, weight, bias]] - - assert activation in ["id", "gelu", "gelu_approx", "squared_relu"] - - batch_shape, n = x.shape[:-1], x.shape[-1] - batch_dim = batch_shape.numel() - x_reshaped = x.reshape(batch_dim, n) - - if x_reshaped.stride(0) > 1 and x_reshaped.stride(1) > 1: - x_reshaped = x_reshaped.contiguous() - if weight.stride(0) > 1 and weight.stride(1) > 1: - weight = weight.contiguous() - bias = bias.contiguous() if bias is not None else None - - assert ( - x.dtype == weight.dtype - ), f"Input and weight must have the same dtype, got {x.dtype} and {weight.dtype}" - if bias is not None: - assert ( - x.dtype == bias.dtype - ), f"Input and bias must have the same dtype, got {x.dtype} and {bias.dtype}" - assert ( - x_reshaped.shape[1] == weight.shape[1] - ), f"Incompatible dimensions: {x_reshaped.shape} - {weight.shape}" - - assert ( - bias is None or bias.shape[0] == weight.shape[0] - ), "Incompatible dimensions in between weight and bias" - - M, K = x_reshaped.shape - N, K = weight.shape - - output = torch.empty((M, N), device=x.device, dtype=x.dtype) - act_input = torch.empty_like(output) if save_act_input else None - - # 1D launch kernel where each block gets its own program. - grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),) # noqa - - kernel_fwd[grid]( - output, - act_input, - x_reshaped, - weight, # data ptrs - bias if bias is not None else x, # auto skip bias if not present - M, # shapes - N, - K, - M // 32, # key for triton cache (limit number of compilations) - N // 32, - K // 32, - stride_cm=output.stride(0), # strides - # stride_cn=output.stride(1), - stride_am=x_reshaped.stride(0), - stride_ak=x_reshaped.stride(1), - stride_bk=weight.stride(1), - stride_bn=weight.stride(0), - BIAS=bias is not None, # optional fused bias - SAVE_ACT_INPUT=save_act_input, # optional save activation inputs - ACTIVATION=activation, # optional fused activation - A_ROWMAJOR=x_reshaped.stride(1) == 1, - B_COLMAJOR=weight.stride(1) == 1, - GROUP_M=8, # speed optimization: group the programs - ) - - if not save_act_input: - return output.reshape(*batch_shape, output.shape[-1]) - else: - return ( - output.reshape(*batch_shape, output.shape[-1]), - act_input.reshape(*batch_shape, act_input.shape[-1]), - ) - - -@triton.autotune( - configs=[ - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - # good for int8 - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=4, - num_warps=4, - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - ] - + get_configs_io_bound(), - key=["CACHE_KEY_M", "CACHE_KEY_N", "CACHE_KEY_K"], - prune_configs_by={ - "early_config_prune": early_config_prune, - "perf_model": estimate_matmul_time, - "top_k": 10, - }, -) -@triton.heuristics( - { - "EVEN_K": lambda args: args["K"] % (args["BLOCK_K"] * args["SPLIT_K"]) == 0, - } -) -@triton.jit -def kernel_bwd( - C, # Pointers to matrices - ACT_INPUT, - A, - B, - # Matrix dimensions - M, - N, - K, - CACHE_KEY_M, - CACHE_KEY_N, - CACHE_KEY_K, - # The stride variables represent how much to increase the ptr by when moving by 1 - # element in a particular dimension. E.g. stride_am is how much to increase a_ptr - # by to get the element one row down (A has M rows) - stride_cm, - # stride_cn, # Assume that stride_cn == 1 - stride_am, - stride_ak, - stride_bk, - stride_bn, - # Meta-parameters - BLOCK_M: tl.constexpr, - GROUP_M: tl.constexpr, - BLOCK_N: tl.constexpr, - BLOCK_K: tl.constexpr, - # split k not used, not performant with activation, kept because early_config_prune is expecting it - SPLIT_K: tl.constexpr, - EVEN_K: tl.constexpr, - ACTIVATION: tl.constexpr, -): - - """ - Kernel for computing Out = activation(A x W + C) - - Input has shape (M, K) - - Weight has shape (K, N) - - Output has shape (M, N) - - ActInputs (optional) has shape (M, N) - 'ActInputs' optionally saves the A x W + C intermediate for backward computations - This kernel will consolidate over K - """ - - pid = tl.program_id(axis=0) - - grid_m = (M + BLOCK_M - 1) // BLOCK_M - grid_n = (N + BLOCK_N - 1) // BLOCK_N - # re-order program ID for better L2 performance - width = GROUP_M * grid_n - group_id = pid // width - group_size = min(grid_m - group_id * GROUP_M, GROUP_M) - pid_m = group_id * GROUP_M + (pid % group_size) - pid_n = (pid % width) // (group_size) - - # now compute the block that each program will go through - # rm (resp. rn) denotes a range of indices - # for rows (resp. col) of C - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - # trick to avoid masking on M and N axis - ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M) - rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N) - rk = tl.arange(0, BLOCK_K) - - A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak) - B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn) - - acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) - - for k in range(K, 0, -BLOCK_K): - if EVEN_K: - a = tl.load(A) - b = tl.load(B) - else: - a = tl.load(A, mask=rk[None, :] < k, other=0.0) - b = tl.load(B, mask=rk[:, None] < k, other=0.0) - acc += tl.dot(a, b) - - A += BLOCK_K * stride_ak - B += BLOCK_K * stride_bk - - # optional: fused activation (while the data is in shared memory) - if ACTIVATION != "id": - act_in_ptrs = ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :] - act_input = tl.load(act_in_ptrs).to(acc.dtype) - if ACTIVATION == "gelu": - acc *= gelu_grad(act_input) - elif ACTIVATION == "gelu_approx": - acc *= gelu_approx_grad(act_input) - elif ACTIVATION == "squared_relu": - acc *= squared_relu_grad(act_input) - - # rematerialize rm and rn to save registers - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - - # write back result - C = C + rm[:, None] * stride_cm + rn[None, :] - mask = (rm < M)[:, None] & (rn < N)[None, :] - tl.store(C, acc, mask=mask) - - -def triton_dgrad_act( - grad_output: torch.Tensor, - weight: torch.Tensor, - activation: str = "id", - act_input: Optional[torch.Tensor] = None, -) -> torch.Tensor: - """ - Compute e = activation(grad_output @ weight + bias). - This wrapper kicks the `kernel_fwd` Triton kernel - :param grad_output: input tensor - :param weight: weight matrix - :param activation: Activation name. Needs to be a Triton kernel. - :param act_input: an optional tensor to save the activation inputs (for backward) - :return: result tensor - """ - assert activation in ["id", "gelu", "gelu_approx", "squared_relu"] - - batch_shape, n = grad_output.shape[:-1], grad_output.shape[-1] - batch_dim = batch_shape.numel() - grad_output_reshaped = grad_output.reshape(batch_dim, n) - - if grad_output_reshaped.stride(0) > 1 and grad_output_reshaped.stride(1) > 1: - grad_output_reshaped = grad_output_reshaped.contiguous() - if weight.stride(0) > 1 and weight.stride(1) > 1: - weight = weight.contiguous() - - assert ( - grad_output.dtype == weight.dtype - ), f"grad_output and weight must have the same dtype, got {grad_output.dtype} and {weight.dtype}" - assert ( - grad_output_reshaped.shape[1] == weight.shape[0] - ), f"Incompatible dimensions: {grad_output_reshaped.shape} - {weight.shape}" - if activation != "id": - assert act_input is not None, f"act_input is required for activation {activation}" - - # M, N, K in bwd are different from M, N, K in fwd - M, K = grad_output_reshaped.shape - K, N = weight.shape - - grad_input = torch.empty((M, N), device=grad_output.device, dtype=grad_output.dtype) - - # 1D launch kernel where each block gets its own program. - grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),) # noqa - - kernel_bwd[grid]( - grad_input, - act_input, - grad_output_reshaped, - weight, # data ptrs - M, # shapes - N, - K, - M // 32, # key for triton cache (limit number of compilations) - N // 32, - K // 32, - stride_cm=grad_input.stride(0), # strides - # stride_cn=grad_input.stride(1), - stride_am=grad_output_reshaped.stride(0), - stride_ak=grad_output_reshaped.stride(1), - stride_bk=weight.stride(0), - stride_bn=weight.stride(1), - ACTIVATION=activation, # optional fused activation - GROUP_M=8, # speed optimization: group the programs - ) - - return grad_input.reshape(*batch_shape, grad_input.shape[-1]) diff --git a/build/torch210-cu128-x86_64-windows/ops/triton/mlp.py b/build/torch210-cu128-x86_64-windows/ops/triton/mlp.py deleted file mode 100644 index ae7db36e3f1e79daca3dc8588ef891d904d79b35..0000000000000000000000000000000000000000 --- a/build/torch210-cu128-x86_64-windows/ops/triton/mlp.py +++ /dev/null @@ -1,149 +0,0 @@ -# The triton fused matmul + sqrelu is faster for fp16 but slower for bf16, compared -# to naive implementation. -import fused_dense_lib as fused_dense_cuda -import torch -import torch.nn as nn -import torch.nn.functional as F - -from flash_attn.utils.torch import custom_fwd, custom_bwd -from flash_attn.ops.activations import sqrelu_bwd, sqrelu_fwd -from flash_attn.ops.triton.linear import triton_dgrad_act, triton_linear_act - - -class FusedDenseSqreluDenseFunc(torch.autograd.Function): - @staticmethod - @custom_fwd - def forward(ctx, x, weight1, bias1, weight2, bias2, checkpoint_lvl=0): - """checkpoint_lvl: - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute act_input and gelu_out in the bwd - """ - if torch.is_autocast_enabled(): - dtype = torch.get_autocast_gpu_dtype() - x, weight1, bias1, weight2, bias2 = [ - a.to(dtype=dtype) for a in [x, weight1, bias1, weight2, bias2] - ] - is_bf16 = x.dtype == torch.bfloat16 - assert checkpoint_lvl in [0, 1, 2] - x = x.contiguous() - weight1 = weight1.contiguous() - bias1 = bias1.contiguous() - weight2 = weight2.contiguous() - bias2 = bias2.contiguous() - batch_shape, n = x.shape[:-1], x.shape[-1] - batch_dim = batch_shape.numel() - if is_bf16: - act_input = fused_dense_cuda.linear_bias_forward( - x.reshape(batch_dim, n), weight1, bias1 - ) - output1 = sqrelu_fwd(act_input) - else: - save_act_input = checkpoint_lvl != 2 - result = triton_linear_act( - x.reshape(batch_dim, n), - weight1, - bias1, - activation="squared_relu", - save_act_input=save_act_input, - ) - if save_act_input: - output1, act_input = result - else: - output1 = result - output2 = fused_dense_cuda.linear_bias_forward(output1, weight2, bias2) - ctx.checkpoint_lvl = checkpoint_lvl - if checkpoint_lvl == 0: - ctx.save_for_backward(x, weight1, bias1, weight2, act_input, output1) - elif checkpoint_lvl == 1: - ctx.save_for_backward(x, weight1, bias1, weight2, act_input) - elif checkpoint_lvl == 2: - ctx.save_for_backward(x, weight1, bias1, weight2) - return output2.reshape(*batch_shape, output2.shape[-1]) - - @staticmethod - @custom_bwd - def backward(ctx, grad_output): - grad_output = grad_output.contiguous() - checkpoint_lvl = ctx.checkpoint_lvl - x, weight1, bias1, weight2, *rest = ctx.saved_tensors - batch_shape, n = x.shape[:-1], x.shape[-1] - batch_dim = batch_shape.numel() - is_bf16 = x.dtype == torch.bfloat16 - if checkpoint_lvl == 0: - act_input, output1 = rest - elif checkpoint_lvl == 1: - (act_input,) = rest - output1 = sqrelu_fwd(act_input) - elif checkpoint_lvl == 2: - if is_bf16: - act_input = fused_dense_cuda.linear_bias_forward( - x.reshape(batch_dim, n), weight1, bias1 - ) - output1 = sqrelu_fwd(act_input) - else: - output1, act_input = triton_linear_act( - x.reshape(batch_dim, n), - weight1, - bias1, - activation="squared_relu", - save_act_input=True, - ) - - if is_bf16: - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad(output1, grad_output) - grad_output1 = grad_output @ weight2 - grad_act_input = sqrelu_bwd(grad_output1, act_input) - grad_input, grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_backward( - x.reshape(batch_dim, n), weight1, grad_act_input - ) - else: - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad(output1, grad_output) - grad_act_input = triton_dgrad_act( - grad_output, weight2, activation="squared_relu", act_input=act_input - ) - grad_input, grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_backward( - x.reshape(batch_dim, n), weight1, grad_act_input - ) - return grad_input.reshape_as(x), grad_weight1, grad_bias1, grad_weight2, grad_bias2, None - - -fused_dense_sqrelu_dense_function = FusedDenseSqreluDenseFunc.apply - - -class FusedDenseSqreluDense(nn.Module): - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - bias1=True, - bias2=True, - checkpoint_lvl=0, - device=None, - dtype=None, - ): - """ - checkpoint_lvl (increasing lvl means slower but more memory saving): - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute gelu_in and gelu_out in the bwd - """ - assert checkpoint_lvl in [0, 1, 2] - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features * 4 - assert bias1 == True, "DenseSqreluDense module without bias is currently not supported" - assert bias2 == True, "DenseSqreluDense module without bias is currently not supported" - self.checkpoint_lvl = checkpoint_lvl - self.fc1 = nn.Linear(in_features, hidden_features, bias=bias1, **factory_kwargs) - self.fc2 = nn.Linear(hidden_features, out_features, bias=bias2, **factory_kwargs) - - def forward(self, x): - assert x.is_cuda - return fused_dense_sqrelu_dense_function( - x, self.fc1.weight, self.fc1.bias, self.fc2.weight, self.fc2.bias, self.checkpoint_lvl - ) diff --git a/build/torch210-cu128-x86_64-windows/ops/triton/rotary.py b/build/torch210-cu128-x86_64-windows/ops/triton/rotary.py deleted file mode 100644 index ba1d164f4bf38d56e74169c8f1394e686d55271f..0000000000000000000000000000000000000000 --- a/build/torch210-cu128-x86_64-windows/ops/triton/rotary.py +++ /dev/null @@ -1,186 +0,0 @@ -# Copyright (c) 2025, Tri Dao. -# As of 2025-04-23, we require triton >= 3.0 - -from typing import Optional, Union - -import torch - -import triton -import triton.language as tl - - -@triton.jit -def rotary_kernel( - OUT, # Pointers to matrices - X, - COS, - SIN, - CU_SEQLENS, - SEQLEN_OFFSETS, # this could be int or a pointer - # Matrix dimensions - seqlen, - nheads, - seqlen_ro, - # strides - stride_out_batch, - stride_out_seqlen, - stride_out_nheads, - stride_out_headdim, - stride_x_batch, - stride_x_seqlen, - stride_x_nheads, - stride_x_headdim, - # Meta-parameters - # We want ROTARY_DIM to be constexpr, otherwise the triton compiler doesn't know that - # the mask is constant every 8 elements, and it will generate LDG.16 instead of LDG.128 - ROTARY_DIM: tl.constexpr, - IS_SEQLEN_OFFSETS_TENSOR: tl.constexpr, - IS_VARLEN: tl.constexpr, - INTERLEAVED: tl.constexpr, - CONJUGATE: tl.constexpr, - BLOCK_H: tl.constexpr, - BLOCK_M: tl.constexpr, -): - BLOCK_K: tl.constexpr = triton.next_power_of_2(ROTARY_DIM) - ROTARY_DIM_HALF = ROTARY_DIM // 2 - pid_head = tl.program_id(axis=0) - pid_m = tl.program_id(axis=1) - pid_batch = tl.program_id(axis=2) - - if not IS_VARLEN: - X = X + pid_batch * stride_x_batch - OUT = OUT + pid_batch * stride_out_batch - else: - start_idx = tl.load(CU_SEQLENS + pid_batch) - seqlen = tl.load(CU_SEQLENS + pid_batch + 1) - start_idx - X = X + start_idx * stride_x_seqlen - OUT = OUT + start_idx * stride_out_seqlen - - if pid_m * BLOCK_M >= seqlen: - return - - rh = pid_head * BLOCK_H + tl.arange(0, BLOCK_H) - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - if not IS_SEQLEN_OFFSETS_TENSOR: - rm_cs = rm + SEQLEN_OFFSETS - else: - rm_cs = rm + tl.load(SEQLEN_OFFSETS + pid_batch) - - rk_half = tl.arange(0, BLOCK_K // 2) - COS = COS + (rm_cs[:, None] * ROTARY_DIM_HALF + rk_half[None, :]) - SIN = SIN + (rm_cs[:, None] * ROTARY_DIM_HALF + rk_half[None, :]) - mask_cs = (rm_cs[:, None] < seqlen_ro) & (rk_half[None, :] < ROTARY_DIM_HALF) - cos = tl.load(COS, mask=mask_cs, other=1.0).to(tl.float32) - sin = tl.load(SIN, mask=mask_cs, other=0.0).to(tl.float32) - if CONJUGATE: - sin = -sin - - if not INTERLEAVED: - # Load the 1st and 2nd halves of X, do calculation, then store to 1st and 2nd halves of OUT - X = X + (rh[:, None, None] * stride_x_nheads + rm[None, :, None] * stride_x_seqlen + rk_half[None, None, :] * stride_x_headdim) - OUT = OUT + (rh[:, None, None] * stride_out_nheads + rm[None, :, None] * stride_out_seqlen + rk_half[None, None, :] * stride_out_headdim) - mask = (rh[:, None, None] < nheads) & (rm[None, :, None] < seqlen) & (rk_half[None, None, :] < ROTARY_DIM_HALF) - x0 = tl.load(X, mask=mask, other=0.0).to(tl.float32) - x1 = tl.load(X + ROTARY_DIM_HALF * stride_x_headdim, mask=mask, other=0.0,).to(tl.float32) - o0 = x0 * cos - x1 * sin - o1 = x0 * sin + x1 * cos - tl.store(OUT, o0, mask=mask) - tl.store(OUT + ROTARY_DIM_HALF * stride_out_headdim, o1, mask=mask) - else: - rk = tl.arange(0, BLOCK_K) - X = X + (rh[:, None, None] * stride_x_nheads + rm[None, :, None] * stride_x_seqlen + rk[None, None, :] * stride_x_headdim) - OUT = OUT + (rh[:, None, None] * stride_out_nheads + rm[None, :, None] * stride_out_seqlen + rk[None, None, :] * stride_out_headdim) - mask = (rh[:, None, None] < nheads) & (rm[None, :, None] < seqlen) & (rk[None, None, :] < ROTARY_DIM) - x = tl.load(X, mask=mask, other=0.0).to(tl.float32) - x0, x1 = tl.split(tl.reshape(x, [BLOCK_H, BLOCK_M, BLOCK_K // 2, 2])) - o0 = x0 * cos - x1 * sin - o1 = x0 * sin + x1 * cos - o = tl.reshape(tl.join(o0, o1), [BLOCK_H, BLOCK_M, BLOCK_K]) - tl.store(OUT, o, mask=mask) - - -def apply_rotary( - x: torch.Tensor, - cos: torch.Tensor, - sin: torch.Tensor, - seqlen_offsets: Union[int, torch.Tensor] = 0, - cu_seqlens: Optional[torch.Tensor] = None, - max_seqlen: Optional[int] = None, - interleaved=False, - inplace=False, - conjugate=False, -) -> torch.Tensor: - """ - Arguments: - x: (batch, seqlen, nheads, headdim) if cu_seqlens is None - else (total_seqlen, nheads, headdim). - cos: (seqlen_ro, rotary_dim / 2) - sin: (seqlen_ro, rotary_dim / 2) - seqlen_offsets: integer or integer tensor of size (batch,) - cu_seqlens: (batch + 1,) or None - max_seqlen: int - Returns: - y: (batch, seqlen, nheads, headdim) - """ - is_varlen = cu_seqlens is not None - if not is_varlen: - batch, seqlen, nheads, headdim = x.shape - else: - assert max_seqlen is not None, "If cu_seqlens is passed in, then max_seqlen must be passed" - total_seqlen, nheads, headdim = x.shape - batch_p_1 = cu_seqlens.shape[0] - batch = batch_p_1 - 1 - seqlen = max_seqlen - seqlen_ro, rotary_dim = cos.shape - assert sin.shape == cos.shape - rotary_dim *= 2 - assert rotary_dim <= headdim, "rotary_dim must be <= headdim" - assert headdim <= 256, "Only support headdim <= 256" - assert seqlen_ro >= seqlen, "seqlen_ro must be >= seqlen" - - cos, sin = cos.contiguous(), sin.contiguous() - if isinstance(seqlen_offsets, torch.Tensor): - assert seqlen_offsets.shape == (batch,) - assert seqlen_offsets.dtype in [torch.int32, torch.int64] - seqlen_offsets = seqlen_offsets.contiguous() - else: - assert seqlen_offsets + seqlen <= seqlen_ro - - output = torch.empty_like(x) if not inplace else x - if rotary_dim < headdim and not inplace: - output[..., rotary_dim:].copy_(x[..., rotary_dim:]) - - grid = lambda META: (triton.cdiv(nheads, META["BLOCK_H"]), triton.cdiv(seqlen, META["BLOCK_M"]), batch) # noqa - BLOCK_M = 8 if rotary_dim <= 128 else 4 - - # Need this, otherwise Triton tries to launch from cuda:0 and we get - # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?) - device_ctx = torch.cuda.device(x.device.index) if x.device.type == 'cuda' else torch.xpu.device(x.device.index) - with device_ctx: - torch.library.wrap_triton(rotary_kernel)[grid]( - output, # data ptrs - x, - cos, - sin, - cu_seqlens, - seqlen_offsets, - seqlen, # shapes - nheads, - seqlen_ro, - output.stride(0) if not is_varlen else 0, # batch_strides if not varlen else 0 - output.stride(-3), # seqlen_stride or total_seqlen_stride - output.stride(-2), # nheads_stride - output.stride(-1), # headdim_stride - x.stride(0) if not is_varlen else 0, # batch_strides if not varlen else 0 - x.stride(-3), # seqlen stride or total_seqlen_stride - x.stride(-2), # nheads stride - x.stride(-1), # headdim stride - rotary_dim, - isinstance(seqlen_offsets, torch.Tensor), - is_varlen, - interleaved, - conjugate, - BLOCK_M=BLOCK_M, - BLOCK_H=2, - ) - return output diff --git a/build/torch210-cu130-x86_64-windows/__init__.py b/build/torch210-cu130-x86_64-windows/__init__.py deleted file mode 100644 index 4b3892ab7794a39e058f2eeb27e6c97cc6bbe6b4..0000000000000000000000000000000000000000 --- a/build/torch210-cu130-x86_64-windows/__init__.py +++ /dev/null @@ -1,393 +0,0 @@ -from typing import Optional, List -import torch -from ._ops import ops as flash_attn_ops -from .flash_attn_interface import ( - flash_attn_func, - flash_attn_kvpacked_func, - flash_attn_qkvpacked_func, - flash_attn_varlen_func, - flash_attn_varlen_kvpacked_func, - flash_attn_varlen_qkvpacked_func, - flash_attn_with_kvcache, -) - - -def fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - return_softmax: bool = False, - gen: Optional[torch.Generator] = None, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention. - - Args: - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Optional output tensor, same shape as q - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - return_softmax: Whether to return softmax weights - gen: Optional random number generator - - Returns: - List of tensors: [output, softmax_lse, (softmax if return_softmax)] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.fwd( - q, - k, - v, - out, - alibi_slopes, - p_dropout, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - gen, - ) - - -def varlen_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - out: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - block_table: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - max_seqlen_q: int = 0, - max_seqlen_k: int = 0, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - zero_tensors: bool = False, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - return_softmax: bool = False, - gen: Optional[torch.Generator] = None, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention with variable sequence lengths. - - Args: - q: Query tensor of shape [total_q, num_heads, head_size] - k: Key tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - v: Value tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1] - cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1] - out: Optional output tensor of shape [total_q, num_heads, head_size] - seqused_k: Optional tensor specifying how many keys to use per batch element [batch_size] - leftpad_k: Optional left padding for keys of shape [batch_size] - block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq] - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - max_seqlen_q: Maximum sequence length for queries - max_seqlen_k: Maximum sequence length for keys - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - zero_tensors: Whether to zero tensors before computation - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - return_softmax: Whether to return softmax weights - gen: Optional random number generator - - Returns: - List of tensors: [output, softmax_lse, (softmax if return_softmax)] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.varlen_fwd( - q, - k, - v, - out, - cu_seqlens_q, - cu_seqlens_k, - seqused_k, - leftpad_k, - block_table, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - p_dropout, - softmax_scale, - zero_tensors, - is_causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - gen, - ) - - -def bwd( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - p_dropout: float = 0.0, - 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, - gen: Optional[torch.Generator] = None, - rng_state: Optional[torch.Tensor] = None, -) -> List[torch.Tensor]: - """ - Backward pass for multi-head attention. - - Args: - dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size] - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size] - softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q] - dq: Optional gradient tensor for queries, same shape as q - dk: Optional gradient tensor for keys, same shape as k - dv: Optional gradient tensor for values, same shape as v - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - deterministic: Whether to use deterministic algorithms - gen: Optional random number generator - rng_state: Optional RNG state from forward pass - - Returns: - List of tensors: [dq, dk, dv] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - alibi_slopes, - p_dropout, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - deterministic, - gen, - rng_state, - ) - - -def varlen_bwd( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - max_seqlen_q: int = 0, - max_seqlen_k: int = 0, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - zero_tensors: bool = False, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - deterministic: bool = False, - gen: Optional[torch.Generator] = None, - rng_state: Optional[torch.Tensor] = None, -) -> List[torch.Tensor]: - """ - Backward pass for multi-head attention with variable sequence lengths. - - Args: - dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size] - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size] - softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q] - cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1] - cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1] - dq: Optional gradient tensor for queries, same shape as q - dk: Optional gradient tensor for keys, same shape as k - dv: Optional gradient tensor for values, same shape as v - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - max_seqlen_q: Maximum sequence length for queries - max_seqlen_k: Maximum sequence length for keys - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - zero_tensors: Whether to zero tensors before computation - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - deterministic: Whether to use deterministic algorithms - gen: Optional random number generator - rng_state: Optional RNG state from forward pass - - Returns: - List of tensors: [dq, dk, dv] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.varlen_bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - p_dropout, - softmax_scale, - zero_tensors, - is_causal, - window_size_left, - window_size_right, - softcap, - deterministic, - gen, - rng_state, - ) - - -def fwd_kvcache( - q: torch.Tensor, - kcache: torch.Tensor, - vcache: torch.Tensor, - k: Optional[torch.Tensor] = None, - v: Optional[torch.Tensor] = None, - seqlens_k: Optional[torch.Tensor] = None, - rotary_cos: Optional[torch.Tensor] = None, - rotary_sin: Optional[torch.Tensor] = None, - cache_batch_idx: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - block_table: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - out: 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, - is_rotary_interleaved: bool = False, - num_splits: int = 1, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention with KV cache. - - Args: - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - kcache: Key cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - vcache: Value cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - k: Optional new keys tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size] - v: Optional new values tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size] - seqlens_k: Optional sequence lengths for keys of shape [batch_size] - rotary_cos: Optional rotary cosine tensor of shape [seqlen_ro, rotary_dim/2] - rotary_sin: Optional rotary sine tensor of shape [seqlen_ro, rotary_dim/2] - cache_batch_idx: Optional indices to index into the KV cache - leftpad_k: Optional left padding for keys of shape [batch_size] - block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq] - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - out: Optional output tensor, same shape as q - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - is_rotary_interleaved: Whether rotary embeddings are interleaved - num_splits: Number of splits for computation - - Returns: - List of tensors: [output, softmax_lse] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.fwd_kvcache( - q, - kcache, - vcache, - k, - v, - seqlens_k, - rotary_cos, - rotary_sin, - cache_batch_idx, - leftpad_k, - block_table, - alibi_slopes, - out, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - is_rotary_interleaved, - num_splits, - ) diff --git a/build/torch210-cu130-x86_64-windows/_flash_attn2_d6b301e.exp b/build/torch210-cu130-x86_64-windows/_flash_attn2_d6b301e.exp deleted file mode 100644 index cb8556dd2a3a4ef1756b4b2b53e1359f91389084..0000000000000000000000000000000000000000 Binary files a/build/torch210-cu130-x86_64-windows/_flash_attn2_d6b301e.exp and /dev/null differ diff --git a/build/torch210-cu130-x86_64-windows/_flash_attn2_d6b301e.lib b/build/torch210-cu130-x86_64-windows/_flash_attn2_d6b301e.lib deleted file mode 100644 index 42c2c3013bdb2214fc5ab0ad8f64aa517b2e1362..0000000000000000000000000000000000000000 Binary files a/build/torch210-cu130-x86_64-windows/_flash_attn2_d6b301e.lib and /dev/null differ diff --git a/build/torch210-cu130-x86_64-windows/_flash_attn2_d6b301e.pyd b/build/torch210-cu130-x86_64-windows/_flash_attn2_d6b301e.pyd deleted file mode 100644 index 641a522dd58041e4c061c884c0ab7ba669c2054e..0000000000000000000000000000000000000000 --- a/build/torch210-cu130-x86_64-windows/_flash_attn2_d6b301e.pyd +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:66d14115a8233fae3a597c511de0b12d172426e62678acf171d006211856e5ff -size 1011415040 diff --git a/build/torch210-cu130-x86_64-windows/_ops.py b/build/torch210-cu130-x86_64-windows/_ops.py deleted file mode 100644 index 504c7db9ac15fb981ffd2b45e3fa950162393430..0000000000000000000000000000000000000000 --- a/build/torch210-cu130-x86_64-windows/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _flash_attn2_d6b301e -ops = torch.ops._flash_attn2_d6b301e - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_flash_attn2_d6b301e::{op_name}" \ No newline at end of file diff --git a/build/torch210-cu130-x86_64-windows/bert_padding.py b/build/torch210-cu130-x86_64-windows/bert_padding.py deleted file mode 100644 index 9717d840687facafd159739df68575a20514b09d..0000000000000000000000000000000000000000 --- a/build/torch210-cu130-x86_64-windows/bert_padding.py +++ /dev/null @@ -1,218 +0,0 @@ -# Adapted from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/padding.py - -import torch -import torch.nn.functional as F -from einops import rearrange, repeat - - -class IndexFirstAxis(torch.autograd.Function): - @staticmethod - def forward(ctx, input, indices): - ctx.save_for_backward(indices) - assert input.ndim >= 2 - ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] - second_dim = other_shape.numel() - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - # return input[indices] - return torch.gather( - rearrange(input, "b ... -> b (...)"), 0, repeat(indices, "z -> z d", d=second_dim) - ).reshape(-1, *other_shape) - - @staticmethod - def backward(ctx, grad_output): - (indices,) = ctx.saved_tensors - assert grad_output.ndim >= 2 - other_shape = grad_output.shape[1:] - grad_output = rearrange(grad_output, "b ... -> b (...)") - grad_input = torch.zeros( - [ctx.first_axis_dim, grad_output.shape[1]], - device=grad_output.device, - dtype=grad_output.dtype, - ) - # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. - # grad_input[indices] = grad_output - grad_input.scatter_(0, repeat(indices, "z -> z d", d=grad_output.shape[1]), grad_output) - return grad_input.reshape(ctx.first_axis_dim, *other_shape), None - - -index_first_axis = IndexFirstAxis.apply - - -class IndexPutFirstAxis(torch.autograd.Function): - @staticmethod - def forward(ctx, values, indices, first_axis_dim): - ctx.save_for_backward(indices) - assert indices.ndim == 1 - assert values.ndim >= 2 - output = torch.zeros( - first_axis_dim, *values.shape[1:], device=values.device, dtype=values.dtype - ) - # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. - output[indices] = values - # output.scatter_(0, repeat(indices, 'z -> z d', d=values.shape[1]), values) - return output - - @staticmethod - def backward(ctx, grad_output): - (indices,) = ctx.saved_tensors - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - grad_values = grad_output[indices] - # grad_values = torch.gather(grad_output, 0, repeat(indices, 'z -> z d', d=grad_output.shape[1])) - return grad_values, None, None - - -index_put_first_axis = IndexPutFirstAxis.apply - - -class IndexFirstAxisResidual(torch.autograd.Function): - @staticmethod - def forward(ctx, input, indices): - ctx.save_for_backward(indices) - assert input.ndim >= 2 - ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] - second_dim = other_shape.numel() - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - output = input[indices] - # We don't want to reshape input (b ... -> b (...)) since it could change the channel_last - # memory format to channel_first. In other words, input might not be contiguous. - # If we don't detach, Pytorch complains about output being a view and is being modified inplace - return output, input.detach() - - @staticmethod - def backward(ctx, grad_output, grad_residual): - (indices,) = ctx.saved_tensors - assert grad_output.ndim >= 2 - other_shape = grad_output.shape[1:] - assert grad_residual.shape[1:] == other_shape - grad_input = grad_residual - # grad_input[indices] += grad_output - indices = indices.reshape(indices.shape[0], *((1,) * (grad_output.ndim - 1))) - indices = indices.expand_as(grad_output) - grad_input.scatter_add_(0, indices, grad_output) - return grad_input.reshape(ctx.first_axis_dim, *other_shape), None - - -index_first_axis_residual = IndexFirstAxisResidual.apply - - -def unpad_input(hidden_states, attention_mask, unused_mask=None): - """ - Arguments: - hidden_states: (batch, seqlen, ...) - attention_mask: (batch, seqlen), bool / int, 1 means valid and 0 means not valid. - unused_mask: (batch, seqlen), bool / int, 1 means the element is allocated but unused. - Return: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens selected in attention_mask + unused_mask. - indices: (total_nnz), the indices of masked tokens from the flattened input sequence. - cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. - max_seqlen_in_batch: int - seqused: (batch), returns the number of tokens selected in attention_mask + unused_mask. - """ - all_masks = (attention_mask + unused_mask) if unused_mask is not None else attention_mask - seqlens_in_batch = all_masks.sum(dim=-1, dtype=torch.int32) - used_seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) - indices = torch.nonzero(all_masks.flatten(), as_tuple=False).flatten() - max_seqlen_in_batch = seqlens_in_batch.max().item() - cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) - # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the - # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim - # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to - # index with integer indices. Moreover, torch's index is a bit slower than it needs to be, - # so we write custom forward and backward to make it a bit faster. - return ( - index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices), - indices, - cu_seqlens, - max_seqlen_in_batch, - used_seqlens_in_batch, - ) - - -def unpad_input_for_concatenated_sequences(hidden_states, attention_mask_in_length): - """ - Supports concatenating short samples in one sequence. The attention_mask_in_length is utilized to mask other short samples. It helps efficient training of variant lengths-based samples (e.g., the supervised fine-tuning task in large language model). - The motivation for this function is explained [here](https://github.com/Dao-AILab/flash-attention/issues/432#issuecomment-1668822286). - - For example, if batch = 3 and seqlen = 6, the attention_mask_in_length is: - ``` - [ - [2, 3, 0, 0, 0, 0], - [3, 2, 0, 0, 0, 0], - [6, 0, 0, 0, 0, 0] - ] - ``` - , which refers to the 3D-attention mask: - ``` - [ - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [0, 0, 1, 0, 0, 0], - [0, 0, 1, 1, 0, 0], - [0, 0, 1, 1, 1, 0], - [0, 0, 0, 0, 0, 1] - ], - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [0, 0, 0, 1, 0, 0], - [0, 0, 0, 1, 1, 0], - [0, 0, 0, 0, 0, 1] - ], - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [1, 1, 1, 1, 0, 0], - [1, 1, 1, 1, 1, 0], - [1, 1, 1, 1, 1, 1] - ] - ] - ```. - - Arguments: - hidden_states: (batch, seqlen, ...) - attention_mask_in_length: (batch, seqlen), int, a nonzero number (e.g., 1, 2, 3, etc.) means length of concatenated sequence in b-th batch, and 0 means none. - Return: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. - indices: (total_nnz), the indices of non-masked tokens from the flattened input sequence. - cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. - max_seqlen_in_batch: int - """ - length = attention_mask_in_length.sum(dim=-1) - seqlen = attention_mask_in_length.size(-1) - attention_mask_2d = torch.arange(seqlen, device=length.device, dtype=length.dtype).expand(len(length), seqlen) < length.unsqueeze(1) - real_indices_idx = torch.nonzero(attention_mask_in_length.flatten(), as_tuple=False).flatten() - seqlens_in_batch = attention_mask_in_length.flatten()[real_indices_idx] - indices = torch.nonzero(attention_mask_2d.flatten(), as_tuple=False).flatten() - max_seqlen_in_batch = seqlens_in_batch.max().item() - cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) - # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the - # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim - # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to - # index with integer indices. Moreover, torch's index is a bit slower than it needs to be, - # so we write custom forward and backward to make it a bit faster. - return ( - index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices), - indices, - cu_seqlens, - max_seqlen_in_batch, - ) - - -def pad_input(hidden_states, indices, batch, seqlen): - """ - Arguments: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. - indices: (total_nnz), the indices that represent the non-masked tokens of the original padded input sequence. - batch: int, batch size for the padded sequence. - seqlen: int, maximum sequence length for the padded sequence. - Return: - hidden_states: (batch, seqlen, ...) - """ - dim = hidden_states.shape[-1] - # output = torch.zeros((batch * seqlen), dim, device=hidden_states.device, dtype=hidden_states.dtype) - # output[indices] = hidden_states - output = index_put_first_axis(hidden_states, indices, batch * seqlen) - return rearrange(output, "(b s) ... -> b s ...", b=batch) diff --git a/build/torch210-cu130-x86_64-windows/flash-attn2/__init__.py b/build/torch210-cu130-x86_64-windows/flash-attn2/__init__.py deleted file mode 100644 index 03dbc1afe1cf156661a2b1b22003cd5f599a0309..0000000000000000000000000000000000000000 --- a/build/torch210-cu130-x86_64-windows/flash-attn2/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -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/torch210-cu130-x86_64-windows/flash_attn_interface.py b/build/torch210-cu130-x86_64-windows/flash_attn_interface.py deleted file mode 100644 index 73ab1f72743f182449cd5b3197108b738a84c9fb..0000000000000000000000000000000000000000 --- a/build/torch210-cu130-x86_64-windows/flash_attn_interface.py +++ /dev/null @@ -1,1620 +0,0 @@ -# Copyright (c) 2023, Tri Dao. - -from typing import Optional, Sequence, Tuple, Union - -import torch -import torch.nn as nn -import os - -# # isort: off -# # We need to import the CUDA kernels after importing torch -# USE_TRITON_ROCM = os.getenv("FLASH_ATTENTION_TRITON_AMD_ENABLE", "FALSE") == "TRUE" -# if USE_TRITON_ROCM: -# from .flash_attn_triton_amd import interface_fa as flash_attn -# else: -# import flash_attn_2_cuda as flash_attn - - -from ._ops import ops as flash_attn - -# # isort: on - -def maybe_contiguous(x): - return x.contiguous() if x is not None and x.stride(-1) != 1 else x - - -def _get_device(): - if torch.xpu.is_available(): - return "xpu" - elif torch.cuda.is_available(): - return "cuda" - else: - return "cpu" - -_XPU_AVAILABLE = torch.xpu.is_available() if hasattr(torch, "xpu") else False # TODO remove hasattr check when bwd is supported on XPU - - -def _get_block_size_n(device, head_dim, is_dropout, is_causal): - # This should match the block sizes in the CUDA kernel - assert head_dim <= 256 - major, minor = torch.cuda.get_device_capability(device) - is_sm8x = major == 8 and minor > 0 # Only include sm86 and sm89, exclude sm80 (A100) - is_sm80 = major == 8 and minor == 0 - is_sm90 = major == 9 and minor == 0 - if head_dim <= 32: - return 128 - if head_dim <= 64: - return 128 if not is_dropout else 64 - elif head_dim <= 96: - return 64 - elif head_dim <= 128: - if is_sm8x: - return 64 if (not is_dropout and is_causal) else 32 - else: - return 64 if not is_dropout else 32 - elif head_dim <= 192: - return 64 - elif head_dim <= 224: - return 64 - elif head_dim <= 256: - return 64 - - -def round_multiple(x, m): - return (x + m - 1) // m * m - - -# torch.compile() support is only enabled for pytorch >= 2.4 -# The reason for this is that we are using the new custom_op and register_fake -# APIs, which support inplace modification of inputs in the function itself -if torch.__version__ >= "2.4.0": - _torch_custom_op_wrapper = torch.library.custom_op - _torch_register_fake_wrapper = torch.library.register_fake -else: - def noop_custom_op_wrapper(name, fn=None, /, *, mutates_args, device_types=None, schema=None): - def wrap(func): - return func - if fn is None: - return wrap - return fn - def noop_register_fake_wrapper(op, fn=None, /, *, lib=None, _stacklevel=1): - def wrap(func): - return func - if fn is None: - return wrap - return fn - _torch_custom_op_wrapper = noop_custom_op_wrapper - _torch_register_fake_wrapper = noop_register_fake_wrapper - - -@_torch_custom_op_wrapper("flash_attn::_flash_attn_forward", mutates_args=(), device_types=_get_device()) -def _flash_attn_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - return_softmax: bool -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - out, softmax_lse, S_dmask, rng_state = flash_attn.fwd( - q, - k, - v, - None, - alibi_slopes, - dropout_p, - softmax_scale, - causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - None, - ) - return out, softmax_lse, S_dmask, rng_state - - -@_torch_register_fake_wrapper("flash_attn::_flash_attn_forward") -def _flash_attn_forward_fake( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - return_softmax: bool -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - batch_size, seqlen_q, num_heads, head_size = q.shape - seqlen_k = k.shape[1] - out = torch.empty_like(q) - softmax_lse = torch.empty((batch_size, num_heads, seqlen_q), dtype=torch.float32, device=q.device, layout=q.layout) - p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) - if return_softmax: - p = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128), round_multiple(seqlen_k, 128)), dtype=q.dtype, device=q.device, layout=q.layout) - rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) - - return out, softmax_lse, p, rng_state - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_forward = torch.ops.flash_attn._flash_attn_forward -else: - _wrapped_flash_attn_forward = _flash_attn_forward - - -@_torch_custom_op_wrapper("flash_attn::_flash_attn_varlen_forward", mutates_args=(), device_types=_get_device()) -def _flash_attn_varlen_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - alibi_slopes: Optional[torch.Tensor] = None, - return_softmax: bool = False, - block_table: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - out, softmax_lse, S_dmask, rng_state = flash_attn.varlen_fwd( - q, - k, - v, - None, - cu_seqlens_q, - cu_seqlens_k, - seqused_k, - leftpad_k, - block_table, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - zero_tensors, - causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - None, - ) - # if out.isnan().any() or softmax_lse.isnan().any(): - # breakpoint() - return out, softmax_lse, S_dmask, rng_state - - -@_torch_register_fake_wrapper("flash_attn::_flash_attn_varlen_forward") -def _flash_attn_varlen_forward_fake( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - alibi_slopes: Optional[torch.Tensor] = None, - return_softmax: bool = False, - block_table: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - paged_kv = block_table is not None - batch_size = cu_seqlens_q.numel() - 1 - total_q, num_heads, _ = q.shape - - out = torch.empty_like(q) - softmax_lse = torch.empty((num_heads, total_q), dtype=torch.float32, device=q.device, layout=q.layout) - p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) - seqlen_q_rounded = round_multiple(max_seqlen_q, 128) - seqlen_k_rounded = round_multiple(max_seqlen_k, 128) - if return_softmax: - p = torch.empty((batch_size, num_heads, seqlen_q_rounded, seqlen_k_rounded), dtype=q.dtype, device=q.device, layout=q.layout) - rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) - return out, softmax_lse, p, rng_state - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_varlen_forward = torch.ops.flash_attn._flash_attn_varlen_forward -else: - _wrapped_flash_attn_varlen_forward = _flash_attn_varlen_forward - - -@_torch_custom_op_wrapper("flash_attn::_flash_attn_backward", mutates_args=("dq", "dk", "dv"), device_types=_get_device()) -def _flash_attn_backward( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, -) -> 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)] - ( - dq, - dk, - dv, - softmax_d, - ) = flash_attn.bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - alibi_slopes, - dropout_p, - softmax_scale, - causal, - window_size_left, - window_size_right, - softcap, - deterministic, - None, - rng_state, - ) - return softmax_d - - -@_torch_register_fake_wrapper("flash_attn::_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, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, -) -> torch.Tensor: - dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] - if dq is None: - dq = torch.empty_like(q) - if dk is None: - dk = torch.empty_like(k) - if dv is None: - dv = torch.empty_like(v) - batch_size, seqlen_q, num_heads, _ = q.shape - softmax_d = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128)), device=q.device, dtype=torch.float32) - - return softmax_d - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_backward = torch.ops.flash_attn._flash_attn_backward -else: - _wrapped_flash_attn_backward = _flash_attn_backward - - -@_torch_custom_op_wrapper("flash_attn::_flash_attn_varlen_backward", mutates_args=("dq", "dk", "dv"), device_types=_get_device()) -def _flash_attn_varlen_backward( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> 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)] - ( - dq, - dk, - dv, - softmax_d, - ) = flash_attn.varlen_bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - zero_tensors, - causal, - window_size_left, - window_size_right, - softcap, - deterministic, - None, - rng_state, - ) - # if dk.isnan().any() or dk.isnan().any() or dv.isnan().any() or softmax_d.isnan().any(): - # breakpoint() - return softmax_d - - -@_torch_register_fake_wrapper("flash_attn::_flash_attn_varlen_backward") -def _flash_attn_varlen_backward_fake( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> torch.Tensor: - dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] - batch_size = cu_seqlens_q.numel() - 1 - total_q, num_heads, _ = q.shape - - if dq is None: - dq = torch.empty_like(q) - if dk is None: - dk = torch.empty_like(k) - if dv is None: - dv = torch.empty_like(v) - softmax_d = torch.empty((num_heads, total_q + 128 * batch_size), device=q.device, dtype=torch.float32) - - return softmax_d - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_varlen_backward = torch.ops.flash_attn._flash_attn_varlen_backward -else: - _wrapped_flash_attn_varlen_backward = _flash_attn_varlen_backward - - -class FlashAttnQKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and qkv.requires_grad - if softmax_scale is None: - softmax_scale = qkv.shape[-1] ** (-0.5) - q, k, v = qkv[:, :, 0].detach(), qkv[:, :, 1].detach(), qkv[:, :, 2].detach() - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) - dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dqkv[:, :, 0], - dqkv[:, :, 1], - dqkv[:, :, 2], - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dqkv = dqkv[..., : dout.shape[-1]] # We could have padded the head dimension - return dqkv, None, None, None, None, None, None, None, None, None - - -class FlashAttnVarlenQKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - cu_seqlens, - max_seqlen, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and qkv.requires_grad - if softmax_scale is None: - softmax_scale = qkv.shape[-1] ** (-0.5) - q, k, v = qkv[:, 0].detach(), qkv[:, 1].detach(), qkv[:, 2].detach() - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens, - cu_seqlens, - max_seqlen, - max_seqlen, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=None, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens, rng_state) - ctx.dropout_p = dropout_p - ctx.max_seqlen = max_seqlen - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens, rng_state = ctx.saved_tensors - qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) - dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dqkv[:, 0], - dqkv[:, 1], - dqkv[:, 2], - cu_seqlens, - cu_seqlens, - ctx.max_seqlen, - ctx.max_seqlen, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - 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 FlashAttnKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - kv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, kv] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - k, v = kv[:, :, 0].detach(), kv[:, :, 1].detach() - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - dq = torch.empty_like(q) - kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) - dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dkv[:, :, 0], - dkv[:, :, 1], - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dkv = dkv[..., : dout.shape[-1]] - return dq, dkv, None, None, None, None, None, None, None, None, None - - -class FlashAttnVarlenKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, kv] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - k, v = kv[:, 0].detach(), kv[:, 1].detach() - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=None, - ) - if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state - ) - ctx.dropout_p = dropout_p - 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.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors - dq = torch.empty_like(q) - kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) - dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dkv[:, 0], - dkv[:, 1], - cu_seqlens_q, - cu_seqlens_k, - ctx.max_seqlen_q, - ctx.max_seqlen_k, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dkv = dkv[..., : dout.shape[-1]] - return dq, dkv, None, None, None, None, None, None, None, None, None, None, None, None, None - - -class FlashAttnFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - k, - v, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, k, v] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dk = dk[..., : dout.shape[-1]] - dv = dv[..., : dout.shape[-1]] - return dq, dk, dv, 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, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - block_table, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, k, v] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=block_table, - ) - if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state - ) - ctx.dropout_p = dropout_p - 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.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors - dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - ctx.max_seqlen_q, - ctx.max_seqlen_k, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dk = dk[..., : dout.shape[-1]] - dv = dv[..., : dout.shape[-1]] - return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None - - -def flash_attn_qkvpacked_func( - qkv, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # <=0.0 means deactivate - alibi_slopes=None, - deterministic=False, - 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, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - False if _XPU_AVAILABLE else torch.is_grad_enabled(), - ) - - -def flash_attn_kvpacked_func( - q, - kv, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, -): - """dropout_p should be set to 0.0 during evaluation - If 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 K, V. - 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) - kv: (batch_size, seqlen, 2, 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. - softcap: float. Anything > 0 activates softcapping 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). - 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 FlashAttnKVPackedFunc.apply( - q, - kv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - False if _XPU_AVAILABLE else torch.is_grad_enabled(), - ) - - -def flash_attn_func( - q, - k, - v, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - 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). - 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 FlashAttnFunc.apply( - q, - k, - v, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - False if _XPU_AVAILABLE else torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_qkvpacked_func( - qkv, - cu_seqlens, - max_seqlen, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - 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_varlen_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_varlen_kvpacked_func and flash_attn_varlen_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: (total, 3, nheads, headdim), where total = total number of tokens in the batch. - cu_seqlens: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into qkv. - max_seqlen: int. Maximum sequence length in the batch. - 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenQKVPackedFunc.apply( - qkv, - cu_seqlens, - max_seqlen, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - False if _XPU_AVAILABLE else torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_kvpacked_func( - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, -): - """dropout_p should be set to 0.0 during evaluation - If 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 K, V. - 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: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch. - kv: (total_k, 2, nheads_k, headdim), where total_k = total number of key tokens in the batch. - cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into q. - cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into kv. - max_seqlen_q: int. Maximum query sequence length in the batch. - max_seqlen_k: int. Maximum key sequence length in the batch. - 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 + 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenKVPackedFunc.apply( - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - False if _XPU_AVAILABLE else torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_func( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, - block_table=None, -): - """dropout_p should be set to 0.0 during evaluation - Supports multi-query and grouped-query attention (MQA/GQA) by passing in K, V 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: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch. - k: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. - v: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. - cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into q. - cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into kv. - max_seqlen_q: int. Maximum query sequence length in the batch. - max_seqlen_k: int. Maximum key sequence length in the batch. - 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 + 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenFunc.apply( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - block_table, - False if _XPU_AVAILABLE or q.device.type == "cpu" else torch.is_grad_enabled(), - ) - - -def flash_attn_with_kvcache( - q, - k_cache, - v_cache, - k=None, - v=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, - block_table: Optional[torch.Tensor] = None, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - rotary_interleaved=True, - alibi_slopes=None, - num_splits=0, - 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 block_table, - or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_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) if there's no block_table, - or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_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). Similar to k. - 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. - block_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). - 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. - 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" - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-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) - cache_batch_idx = maybe_contiguous(cache_batch_idx) - block_table = maybe_contiguous(block_table) - out, softmax_lse = flash_attn.fwd_kvcache( - q, - k_cache, - v_cache, - k, - v, - cache_seqlens, - rotary_cos, - rotary_sin, - cache_batch_idx, - cache_leftpad, - block_table, - alibi_slopes, - None, - softmax_scale, - causal, - window_size[0], - window_size[1], - softcap, - rotary_interleaved, - num_splits, - ) - return (out, softmax_lse) if return_softmax_lse else out diff --git a/build/torch210-cu130-x86_64-windows/layers/patch_embed.py b/build/torch210-cu130-x86_64-windows/layers/patch_embed.py deleted file mode 100644 index 70c64fd30d36c3fa04a09278053065f2699fd4a5..0000000000000000000000000000000000000000 --- a/build/torch210-cu130-x86_64-windows/layers/patch_embed.py +++ /dev/null @@ -1,67 +0,0 @@ -# We use the same API as https://github.com/rwightman/pytorch-image-models/blob/v0.6.11/timm/models/layers/patch_embed.py -# But we use nn.Linear instead of Conv2d and it's about 8x faster. - -from functools import partial - -import torch.nn as nn -from einops import rearrange -from torch import _assert -from torch.nn.modules.utils import _pair - -try: - from flash_attn.ops.fused_dense import FusedDense -except ImportError: - FusedDense = None - - -class PatchEmbed(nn.Module): - """2D Image to Patch Embedding""" - - def __init__( - self, - img_size=224, - patch_size=16, - in_chans=3, - embed_dim=768, - norm_layer=None, - flatten=True, - bias=True, - fused_bias_fc=False, - ): - super().__init__() - img_size = _pair(img_size) - patch_size = _pair(patch_size) - self.img_size = img_size - self.patch_size = patch_size - self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1]) - self.num_patches = self.grid_size[0] * self.grid_size[1] - self.flatten = flatten - if fused_bias_fc and FusedDense is None: - raise ImportError("fused_dense is not installed") - - linear_cls = nn.Linear if not fused_bias_fc or not bias else FusedDense - self.proj = linear_cls(in_chans * patch_size[0] * patch_size[1], embed_dim, bias=bias) - self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() - - def forward(self, x): - _, _, H, W = x.shape - _assert( - H == self.img_size[0], - f"Input image height ({H}) doesn't match model ({self.img_size[0]}).", - ) - _assert( - W == self.img_size[1], - f"Input image width ({W}) doesn't match model ({self.img_size[1]}).", - ) - x = self.proj( - rearrange( - x, - "b c (h p1) (w p2) -> b h w (c p1 p2)", - p1=self.patch_size[0], - p2=self.patch_size[1], - ) - ) - if self.flatten: - x = rearrange(x, "b h w c -> b (h w) c") - x = self.norm(x) - return x diff --git a/build/torch210-cu130-x86_64-windows/layers/rotary.py b/build/torch210-cu130-x86_64-windows/layers/rotary.py deleted file mode 100644 index 1b67f847eeae87f8612f54f1f46f242e379117bf..0000000000000000000000000000000000000000 --- a/build/torch210-cu130-x86_64-windows/layers/rotary.py +++ /dev/null @@ -1,483 +0,0 @@ -# Copyright (c) 2025, Tri Dao - -import math -from functools import partial -from typing import Optional, Tuple, Union - -import torch -from torch import Tensor - -from einops import rearrange, repeat -# from flash_attn.ops.triton.rotary import apply_rotary -from ..ops.triton.rotary import apply_rotary - - -def rotate_half(x, interleaved=False): - if not interleaved: - x1, x2 = x.chunk(2, dim=-1) - return torch.cat((-x2, x1), dim=-1) - else: - x1, x2 = x[..., ::2], x[..., 1::2] - return rearrange(torch.stack((-x2, x1), dim=-1), "... d two -> ... (d two)", two=2) - - -def apply_rotary_emb_torch(x, cos, sin, interleaved=False): - """ - x: (batch_size, seqlen, nheads, headdim) - cos, sin: (seqlen, rotary_dim / 2) or (batch_size, seqlen, rotary_dim / 2) - """ - ro_dim = cos.shape[-1] * 2 - assert ro_dim <= x.shape[-1] - cos = repeat(cos, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)") - sin = repeat(sin, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)") - return torch.cat( - [x[..., :ro_dim] * cos + rotate_half(x[..., :ro_dim], interleaved) * sin, x[..., ro_dim:]], - dim=-1, - ) - - -class ApplyRotaryEmb(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x, - cos, - sin, - interleaved=False, - inplace=False, - seqlen_offsets: Union[int, Tensor] = 0, - cu_seqlens: Optional[Tensor] = None, - max_seqlen: Optional[int] = None, - ): - out = apply_rotary( - x, - cos, - sin, - seqlen_offsets=seqlen_offsets, - cu_seqlens=cu_seqlens, - max_seqlen=max_seqlen, - interleaved=interleaved, - inplace=inplace, - ) - if isinstance(seqlen_offsets, int): - ctx.save_for_backward(cos, sin, cu_seqlens) # Can't save int with save_for_backward - ctx.seqlen_offsets = seqlen_offsets - else: - ctx.save_for_backward(cos, sin, cu_seqlens, seqlen_offsets) - ctx.seqlen_offsets = None - ctx.interleaved = interleaved - ctx.inplace = inplace - ctx.max_seqlen = max_seqlen - return out if not inplace else x - - @staticmethod - def backward(ctx, do): - seqlen_offsets = ctx.seqlen_offsets - if seqlen_offsets is None: - cos, sin, cu_seqlens, seqlen_offsets = ctx.saved_tensors - else: - cos, sin, cu_seqlens = ctx.saved_tensors - dx = apply_rotary( - do, - cos, - sin, - seqlen_offsets=seqlen_offsets, - cu_seqlens=cu_seqlens, - max_seqlen=ctx.max_seqlen, - interleaved=ctx.interleaved, - inplace=ctx.inplace, - conjugate=True, - ) - return dx, None, None, None, None, None, None, None - - -def apply_rotary_emb( - x, - cos, - sin, - interleaved=False, - inplace=False, - seqlen_offsets: Union[int, Tensor] = 0, - cu_seqlens: Optional[Tensor] = None, - max_seqlen: Optional[int] = None, -): - """ - Arguments: - x: (batch_size, seqlen, nheads, headdim) if cu_seqlens is None - else (total_seqlen, nheads, headdim) - cos, sin: (seqlen_rotary, rotary_dim / 2) - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead - of 1st half and 2nd half (GPT-NeoX style). - inplace: if True, apply rotary embedding in-place. - seqlen_offsets: (batch_size,) or int. Each sequence in x is shifted by this amount. - Most commonly used in inference when we have KV cache. - cu_seqlens: (batch + 1,) or None - max_seqlen: int - Return: - out: (batch_size, seqlen, nheads, headdim) if cu_seqlens is None - else (total_seqlen, nheads, headdim) - rotary_dim must be <= headdim - Apply rotary embedding to the first rotary_dim of x. - """ - return ApplyRotaryEmb.apply( - x, cos, sin, interleaved, inplace, seqlen_offsets, cu_seqlens, max_seqlen - ) - - -# For backward compatibility -apply_rotary_emb_func = apply_rotary_emb - - -def _apply_rotary_emb_qkv( - qkv, - cos, - sin, - cos_k=None, - sin_k=None, - interleaved=False, - inplace=False, - conjugate=False, - seqlen_offsets: Union[int, Tensor] = 0, - num_heads_q: Optional[int] = None, -): - apply_rotary_fn = partial( - apply_rotary, - interleaved=interleaved, - inplace=inplace, - conjugate=conjugate, - seqlen_offsets=seqlen_offsets - ) - if cos_k is None and sin_k is None and qkv.is_contiguous(): - # Call 1 kernel instead of 2 kernels - # We need qkv to be contiguous so that when we reshape to combine (3, nheads) - # dimensions, we get the same tensor - if qkv.dim() == 5: - batch, seqlen, three, nheads, headdim = qkv.shape - assert three == 3 - # qk = rearrange(qkv[:, :, :2], "b s t h d -> b s (t h) d") - qk = qkv[:, :, :2].reshape(batch, seqlen, -1, headdim) - qk = apply_rotary_fn(qk, cos, sin) - else: - assert qkv.dim() == 4 - assert num_heads_q is not None - num_heads_k = (qkv.shape[2] - num_heads_q) // 2 - assert qkv.shape[2] == num_heads_q + 2 * num_heads_k - qk = qkv[:, :, :num_heads_q + num_heads_k] - qk = apply_rotary_fn(qk, cos, sin) - if not inplace: - if qkv.dim() == 5: - qkv = torch.cat([rearrange(qk, "b s (t h) d -> b s t h d", t=2), qkv[:, :, 2:]], dim=2) - else: - qkv = torch.cat([qk, qkv[:, :, num_heads_q + num_heads_k :]], dim=2) - else: - cos_k = cos if cos_k is None else cos_k - sin_k = sin if sin_k is None else sin_k - if qkv.dim() == 5: - batch, seqlen, three, nheads, headdim = qkv.shape - assert three == 3 - q, k = qkv[:, :, 0], qkv[:, :, 1] - else: - assert qkv.dim() == 4 - assert num_heads_q is not None - num_heads_k = (qkv.shape[2] - num_heads_q) // 2 - assert qkv.shape[2] == num_heads_q + 2 * num_heads_k - q, k = qkv[:, :, :num_heads_q], qkv[:, :, num_heads_q : num_heads_q + num_heads_k] - q = apply_rotary_fn(q, cos, sin) - k = apply_rotary_fn(k, cos_k, sin_k) - if not inplace: - if qkv.dim() == 5: - qkv = torch.stack([q, k, qkv[:, :, 2]], dim=2) - else: - qkv = torch.cat([q, k, qkv[:, :, num_heads_q + num_heads_k:]], dim=2) - return qkv - - -class ApplyRotaryEmbQKV_(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - cos, - sin, - cos_k=None, - sin_k=None, - interleaved=False, - seqlen_offsets: Union[int, torch.Tensor] = 0, - num_heads_q: Optional[int] = None, - ): - # apply_rotary_emb_qkv_inplace( - qkv = _apply_rotary_emb_qkv( - qkv, cos, sin, cos_k, sin_k, interleaved=interleaved, inplace=True, - seqlen_offsets=seqlen_offsets, num_heads_q=num_heads_q, - ) - if isinstance(seqlen_offsets, int): - ctx.save_for_backward(cos, sin, cos_k, sin_k) - ctx.seqlen_offsets = seqlen_offsets - else: - ctx.save_for_backward(cos, sin, cos_k, sin_k, seqlen_offsets) - ctx.seqlen_offsets = None - ctx.interleaved = interleaved - ctx.num_heads_q = num_heads_q - return qkv - - @staticmethod - def backward(ctx, dqkv): - seqlen_offsets = ctx.seqlen_offsets - if seqlen_offsets is None: - cos, sin, cos_k, sin_k, seqlen_offsets = ctx.saved_tensors - else: - cos, sin, cos_k, sin_k = ctx.saved_tensors - dqkv = _apply_rotary_emb_qkv( - dqkv, cos, sin, cos_k, sin_k, interleaved=ctx.interleaved, inplace=True, - seqlen_offsets=seqlen_offsets, num_heads_q=ctx.num_heads_q, conjugate=True, - ) - return dqkv, None, None, None, None, None, None, None - - -def apply_rotary_emb_qkv_( - qkv, - cos, - sin, - cos_k=None, - sin_k=None, - interleaved=False, - seqlen_offsets: Union[int, torch.Tensor] = 0, - num_heads_q: Optional[int] = None, -): - """ - Arguments: - qkv: (batch_size, seqlen, 3, nheads, headdim) or (batch_size, seqlen, num_heads_q + 2 * num_heads_k, headdim). - If qkv has shape (batch_size, seqlen, num_heads_q + 2 * num_heads_k, headdim) (e.g. MQA / GQA), - then num_heads_q must be provided. - cos, sin: (seqlen, rotary_dim / 2) - cos_k, sin_k: (seqlen, rotary_dim / 2), optional - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead of - 1st half and 2nd half (GPT-NeoX style). - seqlen_offsets: (batch_size,) or int. Each sequence in Q and K is shifted by this amount. - Most commonly used in inference when we have KV cache. - Return: - qkv: (batch_size, seqlen, 3, nheads, headdim) or (batch_size, seqlen, num_heads_q + 2 * num_heads_k, headdim) - rotary_dim must be <= headdim - Apply rotary embedding *inplace* to the first rotary_dim of Q and K. - """ - return ApplyRotaryEmbQKV_.apply( - qkv, cos, sin, cos_k, sin_k, interleaved, seqlen_offsets, num_heads_q - ) - - -class ApplyRotaryEmbKV_(torch.autograd.Function): - - @staticmethod - def forward(ctx, kv, cos, sin, interleaved=False, seqlen_offsets: Union[int, torch.Tensor] = 0): - batch, seqlen, two, nheads, headdim = kv.shape - assert two == 2 - k = kv[:, :, 0] - apply_rotary( - k, cos, sin, seqlen_offsets=seqlen_offsets, interleaved=interleaved, inplace=True - ) - if isinstance(seqlen_offsets, int): - ctx.save_for_backward(cos, sin) # Can't save int with save_for_backward - ctx.seqlen_offsets = seqlen_offsets - else: - ctx.save_for_backward(cos, sin, seqlen_offsets) - ctx.seqlen_offsets = None - ctx.interleaved = interleaved - return kv - - @staticmethod - def backward(ctx, dkv): - seqlen_offsets = ctx.seqlen_offsets - if seqlen_offsets is None: - cos, sin, seqlen_offsets = ctx.saved_tensors - else: - cos, sin = ctx.saved_tensors - apply_rotary( - dkv[:, :, 0], - cos, - sin, - seqlen_offsets=seqlen_offsets, - interleaved=ctx.interleaved, - inplace=True, - conjugate=True, - ) - return dkv, None, None, None, None - - -apply_rotary_emb_kv_ = ApplyRotaryEmbKV_.apply - - -def apply_rotary_emb_kv_( - kv, - cos, - sin, - interleaved=False, - seqlen_offsets: Union[int, torch.Tensor] = 0, -): - """ - Arguments: - kv: (batch_size, seqlen, 2, nheads, headdim) - cos, sin: (seqlen, rotary_dim / 2) - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead of - 1st half and 2nd half (GPT-NeoX style). - seqlen_offsets: (batch_size,) or int. Each sequence in Q and K is shifted by this amount. - Most commonly used in inference when we have KV cache. - Return: - kv: (batch_size, seqlen, 2, nheads, headdim) - rotary_dim must be <= headdim - Apply rotary embedding *inplace* to the first rotary_dim of K. - """ - return ApplyRotaryEmbKV_.apply(kv, cos, sin, interleaved, seqlen_offsets) - - -class RotaryEmbedding(torch.nn.Module): - """ - The rotary position embeddings from RoFormer_ (Su et. al). - A crucial insight from the method is that the query and keys are - transformed by rotation matrices which depend on the relative positions. - - Other implementations are available in the Rotary Transformer repo_ and in - GPT-NeoX_, GPT-NeoX was an inspiration - - .. _RoFormer: https://arxiv.org/abs/2104.09864 - .. _repo: https://github.com/ZhuiyiTechnology/roformer - .. _GPT-NeoX: https://github.com/EleutherAI/gpt-neox - - If scale_base is not None, this implements XPos (Sun et al., https://arxiv.org/abs/2212.10554). - A recommended value for scale_base is 512: https://github.com/HazyResearch/flash-attention/issues/96 - Reference: https://github.com/sunyt32/torchscale/blob/main/torchscale/component/xpos_relative_position.py - """ - - def __init__( - self, - dim: int, - base=10000.0, - interleaved=False, - scale_base=None, - device=None, - ): - """ - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead - of 1st half and 2nd half (GPT-NeoX style). - """ - super().__init__() - self.dim = dim - self.base = float(base) - # Generate and save the inverse frequency buffer (non trainable) - inv_freq = self._compute_inv_freq(device) - self.register_buffer("inv_freq", inv_freq, persistent=False) - self.interleaved = interleaved - self.scale_base = scale_base - scale = ( - (torch.arange(0, dim, 2, device=device, dtype=torch.float32) + 0.4 * dim) / (1.4 * dim) - if scale_base is not None - else None - ) - self.register_buffer("scale", scale, persistent=False) - - self._seq_len_cached = 0 - self._cos_cached = None - self._sin_cached = None - self._cos_k_cached = None - self._sin_k_cached = None - - def _compute_inv_freq(self, device=None): - return 1.0 / ( - self.base - ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim) - ) - - def _update_cos_sin_cache(self, seqlen, device=None, dtype=None): - # Reset the tables if the sequence length has changed, - # if we're on a new device (possibly due to tracing for instance), - # or if we're switching from inference mode to training - if ( - seqlen > self._seq_len_cached - or self._cos_cached is None - or self._cos_cached.device != device - or self._cos_cached.dtype != dtype - or (self.training and self._cos_cached.is_inference()) - ): - self._seq_len_cached = seqlen - # We want fp32 here, not self.inv_freq.dtype, since the model could be loaded in bf16 - # And the output of arange can be quite large, so bf16 would lose a lot of precision. - t = torch.arange(seqlen, device=device, dtype=torch.float32) - # We want fp32 here as well since inv_freq will be multiplied with t, and the output - # will be large. Having it in bf16 will lose a lot of precision and cause the - # cos & sin output to change significantly. - # We want to recompute self.inv_freq if it was not loaded in fp32 - if self.inv_freq.dtype != torch.float32: - inv_freq = self._compute_inv_freq(device=device) - else: - inv_freq = self.inv_freq - # Don't do einsum, it converts fp32 to bf16 under AMP - # freqs = torch.einsum("i,j->ij", t, self.inv_freq) - freqs = torch.outer(t, inv_freq) - if self.scale is None: - self._cos_cached = torch.cos(freqs).to(dtype) - self._sin_cached = torch.sin(freqs).to(dtype) - else: - power = ( - torch.arange(seqlen, dtype=self.scale.dtype, device=self.scale.device) - - seqlen // 2 - ) / self.scale_base - scale = self.scale.to(device=power.device) ** rearrange(power, "s -> s 1") - # We want the multiplication by scale to happen in fp32 - self._cos_cached = (torch.cos(freqs) * scale).to(dtype) - self._sin_cached = (torch.sin(freqs) * scale).to(dtype) - self._cos_k_cached = (torch.cos(freqs) / scale).to(dtype) - self._sin_k_cached = (torch.sin(freqs) / scale).to(dtype) - - def forward( - self, - qkv: torch.Tensor, - kv: Optional[torch.Tensor] = None, - seqlen_offset: Union[int, torch.Tensor] = 0, - max_seqlen: Optional[int] = None, - num_heads_q: Optional[int] = None, - ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: - """ - qkv: (batch, seqlen, 3, nheads, headdim) or (batch, seqlen, num_heads_q + 2 * num_heads_k, headdim) - if kv is none, else it's just q of shape (batch, seqlen, nheads, headdim). - If qkv has shape (batch, seqlen, num_heads_q + 2 * num_heads_k, headdim) (e.g. MQA / GQA), - then num_heads_q must be provided. - kv: (batch, seqlen, 2, nheads, headdim) - seqlen_offset: (batch_size,) or int. Each sequence in x is shifted by this amount. - Most commonly used in inference when we have KV cache. - If it's a tensor of shape (batch_size,), then to update the cos / sin cache, one - should pass in max_seqlen, which will update the cos / sin cache up to that length. - Apply rotary embedding *inplace* to qkv and / or kv. - """ - seqlen = qkv.shape[1] - if max_seqlen is not None: - self._update_cos_sin_cache(max_seqlen, device=qkv.device, dtype=qkv.dtype) - elif isinstance(seqlen_offset, int): - self._update_cos_sin_cache(seqlen + seqlen_offset, device=qkv.device, dtype=qkv.dtype) - if kv is None: - return apply_rotary_emb_qkv_( - qkv, - self._cos_cached, - self._sin_cached, - self._cos_k_cached if self.scale is not None else None, - self._sin_k_cached if self.scale is not None else None, - interleaved=self.interleaved, - seqlen_offsets=seqlen_offset, - num_heads_q=num_heads_q, - ) - else: - q = qkv - q = apply_rotary_emb_func( - q, - self._cos_cached, - self._sin_cached, - interleaved=self.interleaved, - inplace=True, - seqlen_offsets=seqlen_offset, - ) - kv = apply_rotary_emb_kv_( - kv, - self._cos_cached if self.scale is None else self._cos_k_cached, - self._sin_cached if self.scale is None else self._sin_k_cached, - interleaved=self.interleaved, - seqlen_offsets=seqlen_offset, - ) - return q, kv diff --git a/build/torch210-cu130-x86_64-windows/metadata.json b/build/torch210-cu130-x86_64-windows/metadata.json deleted file mode 100644 index 9cf5deed9898dce769f4cc73913d3530b92a0bd8..0000000000000000000000000000000000000000 --- a/build/torch210-cu130-x86_64-windows/metadata.json +++ /dev/null @@ -1,4 +0,0 @@ -{ - "version": 1, - "python-depends": [] -} \ No newline at end of file diff --git a/build/torch210-cu130-x86_64-windows/ops/activations.py b/build/torch210-cu130-x86_64-windows/ops/activations.py deleted file mode 100644 index d407510070e2b369b0fd8855ed8969cc80b6bc2f..0000000000000000000000000000000000000000 --- a/build/torch210-cu130-x86_64-windows/ops/activations.py +++ /dev/null @@ -1,135 +0,0 @@ -# Copied from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/model/layers/activations.py -import math - -import torch -import torch.nn as nn -import torch.nn.functional as F - -# 1/sqrt(2*pi)-> 0.3989423 -# 1/sqrt(2) -> 0.70710678 -# sqrt(2/pi) -> 0.79788456 - -# this function is tanh approximation of gelu -# actual gelu is: -# x * 0.5 * (1.0 + torch.erf(x * 0.70710678)) -@torch.jit.script -def bias_gelu(y, bias): - x = bias + y - return (x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))).to(dtype=y.dtype) - - -# gradient of tanh approximation of gelu -# gradient of actual gelu is: -# 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x) -@torch.jit.script -def bias_gelu_back(g, y, bias): - """Assume that y has shape (B, D) and bias has shape (D)""" - x = bias + y - tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)) - # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243 - ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( - 1 + tanh_out - ) - grad_y = ff * g - return grad_y.to(dtype=y.dtype), grad_y.sum(dim=(0), dtype=bias.dtype) - - -class GeLUFunction(torch.autograd.Function): - @staticmethod - # bias is an optional argument - def forward(ctx, input, bias): - ctx.save_for_backward(input, bias) - return bias_gelu(input, bias) - - @staticmethod - def backward(ctx, grad_output): - input, bias = ctx.saved_tensors - tmp = bias_gelu_back(grad_output, input, bias) - return tmp, tmp - - -bias_gelu_impl = GeLUFunction.apply - -# this function is tanh approximation of gelu -# actual gelu is: -# x * 0.5 * (1.0 + torch.erf(x * 0.70710678)) -@torch.jit.script -def gelu_fwd(x): - return (x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))).to(dtype=x.dtype) - - -# gradient of tanh approximation of gelu -# gradient of actual gelu is: -# 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x) -@torch.jit.script -def gelu_bwd(g, x): - tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)) - # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243 - ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( - 1 + tanh_out - ) - return (ff * g).to(dtype=x.dtype) - - -class FastGeLUFunction(torch.autograd.Function): - @staticmethod - # bias is an optional argument - def forward(ctx, input): - ctx.save_for_backward(input) - return gelu_fwd(input) - - @staticmethod - def backward(ctx, grad_output): - (input,) = ctx.saved_tensors - tmp = gelu_bwd(grad_output, input) - return tmp - - -fast_gelu_impl = FastGeLUFunction.apply - - -@torch.jit.script -def relu_bwd(g, x): - return torch.where(x >= 0, g, 0.0).to(dtype=x.dtype) - - -@torch.jit.script -def sqrelu_fwd(x): - r = F.relu(x) - return (r * r).to(dtype=x.dtype) - - -@torch.jit.script -def sqrelu_bwd(g, x): - return (2.0 * g * F.relu(x)).to(dtype=x.dtype) - - -swiglu_fwd_codestring = """ -template T swiglu_fwd(T x, T y) { - return float(x) * float(y) / (1.0f + ::exp(-float(x))); -} -""" -swiglu_bwd_codestring = """ -template void swiglu_bwd(T x, T y, T g, T& dx, T& dy) { - float x_sigmoid = 1.0f / (1.0f + ::exp(-float(x))); - dx = x_sigmoid * (1 + float(x) * (1.0f - x_sigmoid)) * float(g) * float(y); - dy = float(x) * x_sigmoid * float(g); -} -""" -swiglu_fwd = torch.cuda.jiterator._create_jit_fn(swiglu_fwd_codestring) -swiglu_bwd = torch.cuda.jiterator._create_multi_output_jit_fn(swiglu_bwd_codestring, num_outputs=2) - - -class SwiGLUFunction(torch.autograd.Function): - - @staticmethod - def forward(ctx, x, y): - ctx.save_for_backward(x, y) - return swiglu_fwd(x, y) - - @staticmethod - def backward(ctx, dout): - x, y = ctx.saved_tensors - return swiglu_bwd(x, y, dout) - -swiglu = SwiGLUFunction.apply diff --git a/build/torch210-cu130-x86_64-windows/ops/fused_dense.py b/build/torch210-cu130-x86_64-windows/ops/fused_dense.py deleted file mode 100644 index 61ac40b4aaf97101f63ad411364fcc9f78938596..0000000000000000000000000000000000000000 --- a/build/torch210-cu130-x86_64-windows/ops/fused_dense.py +++ /dev/null @@ -1,688 +0,0 @@ -# Copyright (c) 2023, Tri Dao. -# Inspired by https://github.com/NVIDIA/apex/blob/master/apex/fused_dense/fused_dense.py -# We make it work with pytorch amp and with bfloat16. -# The TensorParallel linear modules are inspired by https://github.com/NVIDIA/apex/blob/master/apex/transformer/tensor_parallel/layers.py -from functools import partial -from typing import Optional - -# import fused_dense_cuda # from apex -import fused_dense_lib as fused_dense_cuda -import torch -import torch.nn as nn -import torch.nn.functional as F -from torch import Tensor -from torch.distributed import ProcessGroup - -from flash_attn.utils.torch import custom_fwd, custom_bwd -from flash_attn.ops.activations import gelu_bwd, relu_bwd, sqrelu_bwd, sqrelu_fwd -from flash_attn.utils.distributed import ( - all_gather_raw, - all_reduce, - all_reduce_raw, - reduce_scatter, - reduce_scatter_raw, -) - - -class FusedDenseFunc(torch.autograd.Function): - @staticmethod - @custom_fwd - def forward( - ctx, x, weight, bias, return_residual=False, process_group=None, sequence_parallel=True - ): - """ - If process_group is not None and sequence_parallel=True, we're doing Tensor Parallel - with sequence parallelism: we do an all_gather_raw of x before doing the matmul. - """ - ctx.compute_weight_gradient = weight.requires_grad - ctx.return_residual = return_residual - ctx.process_group = process_group - ctx.sequence_parallel = sequence_parallel - - if torch.is_autocast_enabled(): - x = x.to(dtype=torch.get_autocast_gpu_dtype()) - x = x.contiguous() - if process_group is not None and sequence_parallel: - # We want to kick off the all_gather early, before weight dtype conversion - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - else: - total_x = x - - if torch.is_autocast_enabled(): - weight = weight.to(dtype=torch.get_autocast_gpu_dtype()) - bias = bias.to(dtype=torch.get_autocast_gpu_dtype()) if bias is not None else None - weight = weight.contiguous() - if process_group is not None and sequence_parallel: - handle_x.wait() - batch_shape, n = total_x.shape[:-1], total_x.shape[-1] - batch_dim = batch_shape.numel() - # https://github.com/pytorch/pytorch/blob/5b51849b48a7dbccd297286cc0110def4706f9e7/aten/src/ATen/native/cuda/Blas.cpp#L174 - if min(batch_dim, n, *weight.shape) > 65535 * 32: - raise RuntimeError("fused_dense only supports matrix dims <= 2M") - output = F.linear(total_x, weight, bias) - if ctx.compute_weight_gradient: - ctx.save_for_backward(x, weight) - else: - ctx.save_for_backward(weight) - return output if not return_residual else (output, x) - - @staticmethod - @custom_bwd - def backward(ctx, grad_output, *args): - grad_output = grad_output.contiguous() - if ctx.return_residual: - (grad_input,) = args - grad_input = grad_input.contiguous() - process_group = ctx.process_group - sequence_parallel = ctx.sequence_parallel - if ctx.compute_weight_gradient: - x, weight = ctx.saved_tensors - if process_group is not None and sequence_parallel: - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - else: - total_x = x - else: - (weight,) = ctx.saved_tensors - total_x = None - batch_shape = grad_output.shape[:-1] - batch_dim = batch_shape.numel() - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - if ctx.needs_input_grad[0]: - if not ctx.return_residual: - grad_input = F.linear(grad_output, weight.t()) - else: - grad_input = torch.addmm( - grad_input.reshape(batch_dim, grad_input.shape[-1]), grad_output, weight - ) - grad_input = grad_input.reshape(*batch_shape, grad_input.shape[-1]) - if process_group is not None: - reduce_fn = reduce_scatter_raw if sequence_parallel else all_reduce_raw - grad_input, handle_grad_input = reduce_fn(grad_input, process_group, async_op=True) - else: - grad_input = None - if ctx.needs_input_grad[1]: - assert ctx.compute_weight_gradient - if process_group is not None and sequence_parallel: - handle_x.wait() - grad_weight, grad_bias = fused_dense_cuda.linear_bias_wgrad( - total_x.reshape(batch_dim, total_x.shape[-1]), grad_output, ctx.needs_input_grad[2] - ) - else: - grad_weight = None - grad_bias = grad_output if ctx.needs_input_grad[2] else None - if process_group is not None and ctx.needs_input_grad[0]: - handle_grad_input.wait() - return grad_input, grad_weight, grad_bias, None, None, None - - -def fused_dense_func( - x: Tensor, - weight: Tensor, - bias: Optional[Tensor] = None, - return_residual: bool = False, - process_group: Optional[ProcessGroup] = None, - sequence_parallel: bool = True, -): - dtype_eligible = x.dtype in [torch.float16, torch.bfloat16] or ( - x.dtype == torch.float32 and torch.is_autocast_enabled() - ) - if x.is_cuda and weight.is_cuda and (bias is None or bias.is_cuda) and dtype_eligible: - return FusedDenseFunc.apply( - x, weight, bias, return_residual, process_group, sequence_parallel - ) - else: - assert process_group is None - out = F.linear(x, weight, bias) - return out if not return_residual else (out, x) - - -class FusedDense(nn.Linear): - def __init__( - self, - in_features: int, - out_features: int, - bias: bool = True, - return_residual: bool = False, - device=None, - dtype=None, - ) -> None: - super().__init__(in_features, out_features, bias=bias, device=device, dtype=dtype) - self.return_residual = return_residual - - def forward(self, x, process_group=None): - """ - If process_group is not None, we're doing Tensor Parallel with sequence parallelism: - we do an all_gather of x before doing the matmul. - """ - return fused_dense_func( - x, - self.weight, - self.bias, - return_residual=self.return_residual, - process_group=process_group, - ) - - -class ColumnParallelLinear(nn.Linear): - def __init__( - self, - in_features: int, - out_features: int, - process_group: ProcessGroup, - bias: bool = True, - sequence_parallel=True, - multiple_of=1, - device=None, - dtype=None, - ) -> None: - world_size = torch.distributed.get_world_size(process_group) - if out_features % multiple_of: - raise ValueError(f"out_features ({out_features}) must be a multiple of {multiple_of}") - multiple = out_features // multiple_of - # We want to split @multiple across world_size, but it could be an uneven split - div = multiple // world_size - mod = multiple % world_size - # The first @mod ranks get @div + 1 copies, the rest get @div copies - local_multiple = div + int(torch.distributed.get_rank(process_group) < mod) - super().__init__( - in_features, local_multiple * multiple_of, bias=bias, device=device, dtype=dtype - ) - self.process_group = process_group - self.sequence_parallel = sequence_parallel - - def forward(self, x): - # If self.sequence_parallel is True, we're doing Tensor Parallel with sequence parallelism: - # we do an all_gather of x before doing the matmul. - # If not, then the input is already gathered. - return fused_dense_func( - x, - self.weight, - self.bias, - process_group=self.process_group, - sequence_parallel=self.sequence_parallel, - ) - - -class RowParallelLinear(nn.Linear): - def __init__( - self, - in_features: int, - out_features: int, - process_group: ProcessGroup, - bias: bool = True, - sequence_parallel=True, - multiple_of=1, - device=None, - dtype=None, - ) -> None: - world_size = torch.distributed.get_world_size(process_group) - rank = torch.distributed.get_rank(process_group) - if in_features % multiple_of: - raise ValueError(f"in_features ({in_features}) must be a multiple of {multiple_of}") - multiple = in_features // multiple_of - # We want to split @multiple across world_size, but it could be an uneven split - div = multiple // world_size - mod = multiple % world_size - # The first @mod ranks get @div + 1 copies, the rest get @div copies - local_multiple = div + int(torch.distributed.get_rank(process_group) < mod) - # Only rank 0 will have bias - super().__init__( - local_multiple * multiple_of, - out_features, - bias=bias and rank == 0, - device=device, - dtype=dtype, - ) - self.process_group = process_group - self.sequence_parallel = sequence_parallel - - def forward(self, x): - """ - We're doing Tensor Parallel with sequence parallelism: we do the matmul and then - a reduce_scatter of the result. - """ - out = fused_dense_func(x, self.weight, self.bias) - reduce_fn = reduce_scatter if self.sequence_parallel else all_reduce - return reduce_fn(out, self.process_group) - - -class FusedMLPFunc(torch.autograd.Function): - @staticmethod - @custom_fwd - def forward( - ctx, - x, - weight1, - bias1, - weight2, - bias2, - activation="gelu_approx", - save_pre_act=True, - return_residual=False, - checkpoint_lvl=0, - heuristic=0, - process_group=None, - sequence_parallel=True, - ): - """ - If process_group is not None and sequence_parallel=True, we're doing Tensor Parallel - with sequence parallelism: we do an all_gather of x before doing the matmul. - If sequence_parallel=False, then the input is already gathered. - - checkpoint_lvl: - 0: no recomputation in the bwd - 1: recompute gelu_out / relu_out in the bwd - 2: recompute pre_act and gelu_out / relu_out in the bwd - """ - assert -1 <= heuristic <= 4 - assert activation in ["gelu_approx", "relu", "sqrelu"] - if activation == "sqrelu": - assert heuristic == -1 - if not save_pre_act: - checkpoint_lvl = 2 - assert checkpoint_lvl in [0, 1, 2] - ctx.return_residual = return_residual - ctx.process_group = process_group - ctx.sequence_parallel = sequence_parallel - ctx.checkpoint_lvl = checkpoint_lvl - ctx.activation = activation - ctx.heuristic = heuristic - - if torch.is_autocast_enabled(): - x = x.to(dtype=torch.get_autocast_gpu_dtype()) - x = x.contiguous() - if process_group is not None and sequence_parallel: - # We want to kick off the all_gather early, before weight dtype conversion - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - else: - total_x = x - - if torch.is_autocast_enabled(): - dtype = torch.get_autocast_gpu_dtype() - weight1, weight2 = [a.to(dtype=dtype) for a in [weight1, weight2]] - bias1 = bias1.to(dtype=dtype) if bias1 is not None else None - bias2 = bias2.to(dtype=dtype) if bias2 is not None else None - weight1 = weight1.contiguous() - bias1 = bias1.contiguous() if bias1 is not None else None - weight2 = weight2.contiguous() - bias2 = bias2.contiguous() if bias2 is not None else None - if process_group is not None and sequence_parallel: - handle_x.wait() - batch_shape, n = total_x.shape[:-1], total_x.shape[-1] - batch_dim = batch_shape.numel() - # https://github.com/pytorch/pytorch/blob/5b51849b48a7dbccd297286cc0110def4706f9e7/aten/src/ATen/native/cuda/Blas.cpp#L174 - if min(batch_dim, n, *weight1.shape, *weight2.shape) > 65535 * 32: - raise RuntimeError("fused_dense only supports matrix dims <= 2M") - if heuristic == -1: - pre_act = F.linear(total_x, weight1, bias1) - activation_fn = ( - partial(F.gelu, approximate="tanh") - if activation == "gelu_approx" - else (sqrelu_fwd if activation == "sqrelu" else F.relu) - ) - with torch.jit.fuser("fuser2"): - output1 = activation_fn(pre_act) - # This is before adding bias1 - # pre_act = F.linear(total_x.reshape(batch_dim, n), weight1) - # with torch.jit.fuser('fuser2'): - # output1 = bias_gelu(pre_act, bias1) - else: - is_gelu = activation == "gelu_approx" - output1, *rest = fused_dense_cuda.linear_act_forward( - total_x.reshape(batch_dim, n), weight1, bias1, is_gelu, save_pre_act, heuristic - ) - if save_pre_act: - pre_act = rest[0] - output2 = F.linear(output1, weight2, bias2) - if checkpoint_lvl == 0 or (checkpoint_lvl == 1 and activation == "relu"): - # For RELU the pre_act is very small (just a bit-mask) so we just save it - ctx.save_for_backward(x, weight1, weight2, pre_act, output1) - elif checkpoint_lvl == 1: - ctx.save_for_backward(x, weight1, weight2, pre_act) - elif checkpoint_lvl == 2: - ctx.save_for_backward(x, weight1, weight2, bias1) - output2 = output2.reshape(*batch_shape, output2.shape[-1]) - return output2 if not return_residual else (output2, x) - - @staticmethod - @custom_bwd - def backward(ctx, grad_output, *args): - grad_output = grad_output.contiguous() - checkpoint_lvl = ctx.checkpoint_lvl - activation = ctx.activation - activation_fn = ( - partial(F.gelu, approximate="tanh") - if activation == "gelu_approx" - else (sqrelu_fwd if activation == "sqrelu" else F.relu) - ) - if ctx.return_residual: - (grad_input,) = args - grad_input = grad_input.contiguous() - process_group = ctx.process_group - sequence_parallel = ctx.sequence_parallel - x, weight1, weight2, *rest = ctx.saved_tensors - if process_group is None or not sequence_parallel: - total_x = x - batch_shape = grad_output.shape[:-1] - batch_dim = batch_shape.numel() - if checkpoint_lvl in [0, 1]: - if process_group is not None and sequence_parallel: - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - if checkpoint_lvl == 0 or (checkpoint_lvl == 1 and activation == "relu"): - pre_act, output1 = rest - elif checkpoint_lvl == 1: - (pre_act,) = rest - with torch.jit.fuser("fuser2"): - output1 = activation_fn(pre_act) - elif checkpoint_lvl == 2: - (bias1,) = rest - if process_group is not None and sequence_parallel: - total_x, _ = all_gather_raw(x, process_group) - if ctx.heuristic == -1: - pre_act = F.linear(total_x, weight1, bias1) - with torch.jit.fuser("fuser2"): - output1 = activation_fn(pre_act) - else: - output1, pre_act = fused_dense_cuda.linear_act_forward( - total_x.reshape(batch_dim, total_x.shape[-1]), - weight1, - bias1, - activation == "gelu_approx", - True, - ctx.heuristic, - ) - - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - output1 = output1.reshape(batch_dim, output1.shape[-1]) - pre_act = pre_act.reshape(batch_dim, pre_act.shape[-1]) - if ctx.needs_input_grad[3]: - grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad( - output1, grad_output, ctx.needs_input_grad[4] - ) - else: - grad_weight2 = None - grad_bias2 = grad_output if ctx.needs_input_grad[4] else None - if ctx.heuristic == -1: - # grad_pre_act = matmul_dgelu(grad_output, weight2, pre_act) - grad_output1 = F.linear(grad_output, weight2.t()) - activation_grad_fn = ( - gelu_bwd - if activation == "gelu_approx" - else (sqrelu_bwd if activation == "sqrelu" else relu_bwd) - ) - with torch.jit.fuser("fuser2"): - grad_pre_act = activation_grad_fn(grad_output1, pre_act) - else: - # The cublasLt epilogue has to compute both gelu/relu grad and bias grad, we can't - # just compute gelu/relu grad - grad_pre_act, grad_bias1 = fused_dense_cuda.bias_act_linear_dgrad_bgrad( - weight2, grad_output, pre_act, activation == "gelu_approx", ctx.heuristic - ) - if not ctx.needs_input_grad[2]: - grad_bias1 = None - if ctx.needs_input_grad[0]: - if not ctx.return_residual: - grad_input = F.linear(grad_pre_act, weight1.t()) - else: - grad_input = torch.addmm( - grad_input.reshape(batch_dim, grad_input.shape[-1]), grad_pre_act, weight1 - ) - grad_input = grad_input.reshape(*batch_shape, grad_input.shape[-1]) - if process_group is not None: - reduce_fn = reduce_scatter_raw if sequence_parallel else all_reduce_raw - grad_input, handle_grad_input = reduce_fn(grad_input, process_group, async_op=True) - else: - grad_input = None - if ctx.heuristic == -1: - if ctx.needs_input_grad[1]: - if process_group is not None and sequence_parallel and checkpoint_lvl != 2: - handle_x.wait() - grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_wgrad( - total_x.reshape(batch_dim, total_x.shape[-1]), - grad_pre_act, - ctx.needs_input_grad[2], - ) - else: - grad_weight1 = None - grad_bias1 = grad_pre_act if ctx.needs_input_grad[2] else None - else: - if ctx.needs_input_grad[1]: - if process_group is not None and sequence_parallel and checkpoint_lvl != 2: - handle_x.wait() - grad_weight1 = F.linear( - grad_pre_act.t(), total_x.reshape(batch_dim, total_x.shape[-1]).t() - ) - else: - grad_weight1 = None - if process_group is not None and ctx.needs_input_grad[0]: - handle_grad_input.wait() - return ( - grad_input, - grad_weight1, - grad_bias1, - grad_weight2, - grad_bias2, - None, - None, - None, - None, - None, - None, - None, - ) - - -def fused_mlp_func( - x: Tensor, - weight1: Tensor, - weight2: Tensor, - bias1: Optional[Tensor] = None, - bias2: Optional[Tensor] = None, - activation: str = "gelu_approx", - save_pre_act: bool = True, - return_residual: bool = False, - checkpoint_lvl: int = 0, - heuristic: int = 0, - process_group: Optional[ProcessGroup] = None, - sequence_parallel: bool = True, -): - assert activation in ["gelu_approx", "relu", "sqrelu"] - dtype_eligible = x.dtype in [torch.float16, torch.bfloat16] or ( - x.dtype == torch.float32 and torch.is_autocast_enabled() - ) - # If we save pre-activation, dimension must be divisible by 128 (relu) or 8 (gelu) - dim_eligible = not save_pre_act or (x.shape[-1] % (128 if activation == "relu" else 8) == 0) - if ( - x.is_cuda - and weight1.is_cuda - and weight2.is_cuda - and (bias1 is None or bias1.is_cuda) - and (bias2 is None or bias2.is_cuda) - and dtype_eligible - and dim_eligible - ): - return FusedMLPFunc.apply( - x, - weight1, - bias1, - weight2, - bias2, - activation, - save_pre_act, - return_residual, - checkpoint_lvl, - heuristic, - process_group, - sequence_parallel, - ) - else: - assert process_group is None - pre_act = F.linear(x, weight1, bias1) - activation_fn = ( - partial(F.gelu, approximate="tanh") - if activation == "gelu_approx" - else partial(F.relu, inplace=True) - ) - output1 = activation_fn(pre_act) - output2 = F.linear(output1, weight2, bias2) - return output2 if not return_residual else (output2, x) - - -class FusedMLP(nn.Module): - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - bias1=True, - bias2=True, - activation="gelu_approx", - return_residual=False, - checkpoint_lvl=0, - heuristic="auto", - device=None, - dtype=None, - ): - """ - If process_group is not None, we're doing Tensor Parallel with sequence parallelism: - we do an all_gather of x before doing the matmul, gelu, then matmul. - Finally we do a reduce_scatter of the output. - - checkpoint_lvl (increasing lvl means slower but more memory saving): - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute pre_act and gelu_out in the bwd - heuristic: - -1: don't fuse gemm + gelu (separate kernel) - 0..4: use this heuristic for the algo section in the fused gemm + gelu - 'auto': heuristic will be picked automatically: - For CUDA >= 11.8, we set heuristic=0 for both fp16 and bf16 for best perf. - For CUDA <= 11.7, we set heuristic=1 for fp16 and heuristic=-1 for bf16. - For H100, we set heuristic=-1 for both fp16 and bf16 as the fused cuBlasLt implementation - is slower than the unfused version. - return_residual: whether to return the input x along with the output. This is for - performance reason: for post-norm architecture, returning the input allows us - to fuse the backward of nn.Linear with the residual connection. - """ - assert checkpoint_lvl in [0, 1, 2] - assert activation in ["gelu_approx", "relu", "sqrelu"] - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features * 4 - self.activation = activation - self.return_residual = return_residual - self.checkpoint_lvl = checkpoint_lvl - self.heuristic = heuristic if activation != "sqrelu" else -1 - self.fc1 = nn.Linear(in_features, hidden_features, bias=bias1, **factory_kwargs) - self.fc2 = nn.Linear(hidden_features, out_features, bias=bias2, **factory_kwargs) - - def forward(self, x, process_group=None): - dtype = x.dtype if not torch.is_autocast_enabled() else torch.get_autocast_gpu_dtype() - if self.heuristic == "auto": - if self.activation == "gelu_approx": - if torch.cuda.get_device_capability("cuda") == (9, 0): - heuristic = -1 - else: - cuda_ver = tuple(map(int, torch.version.cuda.split("."))) - heuristic = 0 if cuda_ver >= (11, 8) else (1 if dtype == torch.float16 else -1) - else: - heuristic = 0 - else: - heuristic = self.heuristic - out = fused_mlp_func( - x, - self.fc1.weight, - self.fc2.weight, - self.fc1.bias, - self.fc2.bias, - activation=self.activation, - save_pre_act=self.training, - return_residual=self.return_residual, - checkpoint_lvl=self.checkpoint_lvl, - heuristic=heuristic, - process_group=process_group, - ) - if self.return_residual: - out, x = out - if process_group is not None: - out = reduce_scatter(out, process_group) - return out if not self.return_residual else (out, x) - - -class ParallelFusedMLP(nn.Module): - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - activation="gelu_approx", - process_group: ProcessGroup = None, - bias1=True, - bias2=True, - sequence_parallel=True, - checkpoint_lvl=0, - heuristic="auto", - device=None, - dtype=None, - ): - """ - process_group is required. We're doing Tensor Parallel with sequence parallelism: - we do an all_gather of x before doing the matmul, gelu, then matmul. - Finally we do a reduce_scatter of the output. - - checkpoint_lvl (increasing lvl means slower but more memory saving): - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute pre_act and gelu_out in the bwd - heuristic: - -1: don't fuse gemm + gelu (separate kernel) - 0..4: use this heuristic for the algo section in the fused gemm + gelu - 'auto': heuristic will be picked automatically: - For CUDA >= 11.8, we set heuristic=0 for both fp16 and bf16 for best perf. - For CUDA <= 11.7, we set heuristic=1 for fp16 and heuristic=-1 for bf16. - """ - assert checkpoint_lvl in [0, 1, 2] - assert activation in ["gelu_approx", "relu", "sqrelu"] - assert process_group is not None - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features * 4 - self.activation = activation - self.process_group = process_group - self.sequence_parallel = sequence_parallel - self.checkpoint_lvl = checkpoint_lvl - self.heuristic = heuristic if activation != "sqrelu" else -1 - self.fc1 = ColumnParallelLinear( - in_features, hidden_features, process_group, bias=bias1, **factory_kwargs - ) - self.fc2 = RowParallelLinear( - hidden_features, out_features, process_group, bias=bias2, **factory_kwargs - ) - - def forward(self, x): - dtype = x.dtype if not torch.is_autocast_enabled() else torch.get_autocast_gpu_dtype() - if self.heuristic == "auto": - if self.activation == "gelu_approx": - cuda_ver = tuple(map(int, torch.version.cuda.split("."))) - heuristic = 0 if cuda_ver >= (11, 8) else (1 if dtype == torch.float16 else -1) - else: - heuristic = 0 - else: - heuristic = self.heuristic - out = fused_mlp_func( - x, - self.fc1.weight, - self.fc2.weight, - self.fc1.bias, - self.fc2.bias, - activation=self.activation, - save_pre_act=self.training, - checkpoint_lvl=self.checkpoint_lvl, - heuristic=heuristic, - process_group=self.process_group, - sequence_parallel=self.sequence_parallel, - ) - reduce_fn = reduce_scatter if self.sequence_parallel else all_reduce - return reduce_fn(out, self.process_group) diff --git a/build/torch210-cu130-x86_64-windows/ops/layer_norm.py b/build/torch210-cu130-x86_64-windows/ops/layer_norm.py deleted file mode 100644 index c8bcae3b080845f56464dceca508278f76cde68f..0000000000000000000000000000000000000000 --- a/build/torch210-cu130-x86_64-windows/ops/layer_norm.py +++ /dev/null @@ -1,800 +0,0 @@ -# Copyright (c) 2022, Tri Dao. -# Adapted from https://github.com/NVIDIA/apex/blob/master/apex/contrib/layer_norm/layer_norm.py - -import dropout_layer_norm -import torch -from torch.nn import init - - -def maybe_align(x, alignment_in_bytes=16): - """Assume that x already has last dim divisible by alignment_in_bytes""" - # TD [2023-07-04] I'm not 100% sure that clone will align the memory - # https://discuss.pytorch.org/t/how-to-ensure-that-tensor-data-ptr-is-aligned-to-16-bytes/183440 - return x if x.data_ptr() % alignment_in_bytes == 0 else x.clone() - - -def _dropout_add_layer_norm_forward( - x0, - residual, - gamma, - beta, - rowscale, - colscale, - dropout_p, - epsilon, - residual_in_fp32=False, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes""" - hidden_size = gamma.numel() - x0mat = x0.view((-1, hidden_size)) - residualmat = residual.view((-1, hidden_size)) if residual is not None else None - rowscale = rowscale.view(-1) if rowscale is not None else None - zmat, xmat, dmask, mu, rsigma = dropout_layer_norm.dropout_add_ln_fwd( - x0mat, - residualmat, - gamma, - beta, - rowscale, - colscale, - None, - None, - dropout_p, - epsilon, - 1.0, - 0, - None, - residual_in_fp32, - is_rms_norm, - ) - # dmask is None if dropout_p == 0.0 - # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype - return zmat, xmat if xmat is not None else x0mat, dmask, mu, rsigma - - -def _dropout_add_layer_norm_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - rowscale, - colscale, - dropout_p, - has_residual, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes - dx == None means that it was a post-norm architecture - (x = drop(x0) + residual was not returned in the fwd). - x0 must not be None if we have colscale. - """ - hidden_size = gamma.numel() - xmat = x.view((-1, hidden_size)) - dzmat = dz.view(xmat.shape) - dxmat = dx.view(xmat.shape) if dx is not None else None - x0mat = x0.view((-1, hidden_size)) if x0 is not None else None - rowscale = rowscale.view(-1) if rowscale is not None else None - if colscale is not None: - assert x0 is not None, "x0 is required to compute the gradient of colscale" - dx0mat, dresidualmat, dgamma, dbeta, _, _, *rest = dropout_layer_norm.dropout_add_ln_bwd( - dzmat, - dxmat, - xmat, - x0mat, - dmask, - mu, - rsigma, - gamma, - rowscale, - colscale, - None, - None, - dropout_p, - 1.0, - 0, - has_residual, - is_rms_norm, - ) - # dresidualmat is None if not has_residual - if colscale is None: - return dx0mat, dresidualmat, dgamma, dbeta - else: - dcolscale = rest[0] - return dx0mat, dresidualmat, dgamma, dbeta, dcolscale - - -def _dropout_add_layer_norm_subset_forward( - x0, - residual, - gamma, - beta, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32=False, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes""" - hidden_size = gamma.numel() - x0mat = x0.view((-1, hidden_size)) - residualmat = residual.view((-1, hidden_size)) if residual is not None else None - x0_subset = x0_subset.view(-1) if x0_subset is not None else None - out_subset = out_subset.view(-1) if out_subset is not None else None - zmat, xmat, dmask, mu, rsigma = dropout_layer_norm.dropout_add_ln_fwd( - x0mat, - residualmat, - gamma, - beta, - None, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - None, - residual_in_fp32, - is_rms_norm, - ) - # dmask is None if dropout_p == 0.0 - # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype - return zmat, xmat if xmat is not None else x0mat, dmask, mu, rsigma - - -def _dropout_add_layer_norm_subset_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - colscale, - x0_subset, - out_subset, - dropout_p, - rowscale_const, - x0_numrows, - has_residual, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes - dx == None means that it was a post-norm architecture - (x = drop(x0) + residual was not returned in the fwd). - x0 must not be None if we have colscale. - """ - hidden_size = gamma.numel() - xmat = x.view((-1, hidden_size)) - dzmat = dz.view(-1, hidden_size) - dxmat = dx.view(xmat.shape) if dx is not None else None - x0mat = x0.view((-1, hidden_size)) if x0 is not None else None - x0_subset = x0_subset.view(-1) if x0_subset is not None else None - out_subset = out_subset.view(-1) if out_subset is not None else None - if colscale is not None: - assert x0 is not None, "x0 is required to compute the gradient of colscale" - dx0mat, dresidualmat, dgamma, dbeta, _, _, *rest = dropout_layer_norm.dropout_add_ln_bwd( - dzmat, - dxmat, - xmat, - x0mat, - dmask, - mu, - rsigma, - gamma, - None, - colscale, - x0_subset, - out_subset, - dropout_p, - rowscale_const, - x0_numrows, - has_residual, - is_rms_norm, - ) - # dresidualmat is None if not has_residual - if colscale is None: - return dx0mat, dresidualmat, dgamma, dbeta - else: - dcolscale = rest[0] - return dx0mat, dresidualmat, dgamma, dbeta, dcolscale - - -def _dropout_add_layer_norm_parallel_residual_forward( - x0, - x1, - residual, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - residual_in_fp32=False, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes""" - hidden_size = gamma0.numel() - x0mat = x0.view((-1, hidden_size)) - x1mat = x1.view((-1, hidden_size)) if x1 is not None else None - residualmat = residual.view((-1, hidden_size)) if residual is not None else None - ( - z0mat, - z1mat, - xmat, - dmask0, - dmask1, - mu, - rsigma, - ) = dropout_layer_norm.dropout_add_ln_parallel_residual_fwd( - x0mat, - x1mat, - residualmat, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - None, - residual_in_fp32, - is_rms_norm, - ) - # dmask0 and dmask1 are None if dropout_p == 0.0 - # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype - return z0mat, z1mat, xmat if xmat is not None else x0mat, dmask0, dmask1, mu, rsigma - - -def _dropout_add_layer_norm_parallel_residual_backward( - dz0, - dz1, - dx, - x, - dmask0, - dmask1, - mu, - rsigma, - gamma0, - gamma1, - dropout_p, - has_x1, - has_residual, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes - dx == None means that it was a post-norm architecture - (x = drop(x0) + residual was not returned in the fwd). - """ - hidden_size = gamma0.numel() - xmat = x.view((-1, hidden_size)) - dz0mat = dz0.view(xmat.shape) - dz1mat = dz1.view(xmat.shape) if dz1 is not None else None - dxmat = dx.view(xmat.shape) if dx is not None else None - ( - dx0mat, - dx1mat, - dresidualmat, - dgamma0, - dbeta0, - dgamma1, - dbeta1, - *rest, - ) = dropout_layer_norm.dropout_add_ln_parallel_residual_bwd( - dz0mat, - dz1mat, - dxmat, - xmat, - dmask0, - dmask1, - mu, - rsigma, - gamma0, - gamma1, - dropout_p, - has_x1, - has_residual, - is_rms_norm, - ) - # dresidualmat is None if not has_residual - return dx0mat, dx1mat, dresidualmat, dgamma0, dbeta0, dgamma1, dbeta1 - - -class DropoutAddLayerNormFn(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x0, - residual, - gamma, - beta, - rowscale, - colscale, - dropout_p, - epsilon, - residual_in_fp32=False, - prenorm=False, - is_rms_norm=False, - return_dmask=False, - ): - x0 = maybe_align(x0.contiguous(), 16) - residual = maybe_align(residual.contiguous(), 16) if residual is not None else None - gamma = maybe_align(gamma.contiguous(), 16) - beta = maybe_align(beta.contiguous(), 16) if beta is not None else None - rowscale = maybe_align(rowscale.contiguous(), 16) if rowscale is not None else None - colscale = maybe_align(colscale.contiguous(), 16) if colscale is not None else None - zmat, xmat, dmask, mu, rsigma = _dropout_add_layer_norm_forward( - x0, - residual, - gamma, - beta, - rowscale, - colscale, - dropout_p, - epsilon, - residual_in_fp32, - is_rms_norm, - ) - # Only need to save x0 if we need to compute gradient wrt colscale - x0_saved = x0 if colscale is not None else None - ctx.save_for_backward( - xmat.view(x0.shape), x0_saved, dmask, gamma, mu, rsigma, rowscale, colscale - ) - ctx.prenorm = prenorm - ctx.dropout_p = dropout_p - ctx.has_residual = residual is not None - ctx.is_rms_norm = is_rms_norm - ctx.has_beta = beta is not None - if not return_dmask: - return ( - zmat.view(x0.shape) if not prenorm else (zmat.view(x0.shape), xmat.view(x0.shape)) - ) - else: - dmask = ( - dmask.view(x0.shape) - if dropout_p > 0.0 - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - ctx.mark_non_differentiable(dmask) - return ( - (zmat.view(x0.shape), dmask) - if not prenorm - else (zmat.view(x0.shape), xmat.view(x0.shape), dmask) - ) - - @staticmethod - def backward(ctx, dz, *args): - # assert dz.is_contiguous() - dz = maybe_align(dz.contiguous(), 16) # this happens! - dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None - x, x0, dmask, gamma, mu, rsigma, rowscale, colscale = ctx.saved_tensors - # x0 is None if colscale is None - dropout_p = ctx.dropout_p - has_residual = ctx.has_residual - dx0mat, dresidualmat, dgamma, dbeta, *rest = _dropout_add_layer_norm_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - rowscale, - colscale, - dropout_p, - has_residual, - ctx.is_rms_norm, - ) - dx0 = dx0mat.view(x.shape) - dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None - dcolscale = rest[0] if colscale is not None else None - return ( - dx0, - dresidual, - dgamma, - dbeta if ctx.has_beta else None, - None, - dcolscale, - None, - None, - None, - None, - None, - None, - ) - - -class DropoutAddLayerNormSubsetFn(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x0, - residual, - gamma, - beta, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32=False, - prenorm=False, - is_rms_norm=False, - return_dmask=False, - ): - x0 = maybe_align(x0.contiguous(), 16) - residual = maybe_align(residual.contiguous(), 16) if residual is not None else None - gamma = maybe_align(gamma.contiguous(), 16) - beta = maybe_align(beta.contiguous(), 16) if beta is not None else None - colscale = maybe_align(colscale.contiguous(), 16) if colscale is not None else None - zmat, xmat, dmask, mu, rsigma = _dropout_add_layer_norm_subset_forward( - x0, - residual, - gamma, - beta, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32, - is_rms_norm, - ) - # Only need to save x0 if we need to compute gradient wrt colscale - x0_saved = x0 if colscale is not None else None - x_shape = (-1, *x0.shape[1:]) - ctx.save_for_backward( - xmat.view(x_shape), x0_saved, dmask, gamma, mu, rsigma, colscale, x0_subset, out_subset - ) - ctx.prenorm = prenorm - ctx.dropout_p = dropout_p - ctx.rowscale_const = rowscale_const - ctx.x0_numrows = x0.shape[:-1].numel() - ctx.has_residual = residual is not None - ctx.is_rms_norm = is_rms_norm - ctx.has_beta = beta is not None - z_shape = (-1, *x0.shape[1:]) - if not return_dmask: - return zmat.view(z_shape) if not prenorm else (zmat.view(z_shape), xmat.view(x0.shape)) - else: - z = zmat.view(z_shape) - dmask = ( - dmask.view(x0.shape) - if dropout_p > 0.0 - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - ctx.mark_non_differentiable(dmask) - return (z, dmask) if not prenorm else (z, xmat.view(x_shape), dmask) - - @staticmethod - def backward(ctx, dz, *args): - # assert dz.is_contiguous() - dz = maybe_align(dz.contiguous(), 16) # this happens! - dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None - x, x0, dmask, gamma, mu, rsigma, colscale, x0_subset, out_subset = ctx.saved_tensors - # x0 is None if colscale is None - dropout_p = ctx.dropout_p - has_residual = ctx.has_residual - dx0mat, dresidualmat, dgamma, dbeta, *rest = _dropout_add_layer_norm_subset_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - colscale, - x0_subset, - out_subset, - dropout_p, - ctx.rowscale_const, - ctx.x0_numrows, - has_residual, - ctx.is_rms_norm, - ) - dx0 = dx0mat.view(-1, *x.shape[1:]) - dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None - dcolscale = rest[0] if colscale is not None else None - return ( - dx0, - dresidual, - dgamma, - dbeta if ctx.has_beta else None, - dcolscale, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) - - -class DropoutAddLayerNormParallelResidualFn(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x0, - x1, - residual, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - residual_in_fp32=False, - prenorm=False, - is_rms_norm=False, - return_dmask=False, - ): - x0 = maybe_align(x0.contiguous(), 16) - x1 = maybe_align(x1.contiguous(), 16) if x1 is not None else None - residual = maybe_align(residual.contiguous(), 16) if residual is not None else None - gamma0 = maybe_align(gamma0.contiguous(), 16) - beta0 = maybe_align(beta0.contiguous(), 16) if beta0 is not None else None - gamma1 = maybe_align(gamma1.contiguous(), 16) if gamma1 is not None else None - beta1 = maybe_align(beta1.contiguous(), 16) if beta1 is not None else None - ( - z0mat, - z1mat, - xmat, - dmask0, - dmask1, - mu, - rsigma, - ) = _dropout_add_layer_norm_parallel_residual_forward( - x0, - x1, - residual, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - residual_in_fp32, - is_rms_norm, - ) - ctx.save_for_backward(xmat.view(x0.shape), dmask0, dmask1, gamma0, gamma1, mu, rsigma) - ctx.prenorm = prenorm - ctx.dropout_p = dropout_p - ctx.has_x1 = x1 is not None - ctx.has_residual = residual is not None - ctx.is_rms_norm = is_rms_norm - ctx.has_beta = beta0 is not None - z = (z0mat.view(x0.shape), z1mat.view(x0.shape) if z1mat is not None else None) - if not return_dmask: - return z if not prenorm else (*z, xmat.view(x0.shape)) - else: - dmask0 = ( - dmask0.view(x0.shape) - if dropout_p > 0.0 - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - dmask1 = ( - dmask1.view(x0.shape) - if dropout_p > 0.0 and x1 is not None - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - ctx.mark_non_differentiable(dmask0) - ctx.mark_non_differentiable(dmask1) - return ( - (*z, dmask0, dmask1) if not prenorm else (*z, xmat.view(x0.shape), dmask0, dmask1) - ) - - @staticmethod - def backward(ctx, dz0, dz1, *args): - dz0 = maybe_align(dz0.contiguous(), 16) # this happens! - dz1 = maybe_align(dz1.contiguous(), 16) if dz1 is not None else None - dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None - x, dmask0, dmask1, gamma0, gamma1, mu, rsigma = ctx.saved_tensors - dropout_p = ctx.dropout_p - has_x1 = ctx.has_x1 - has_residual = ctx.has_residual - ( - dx0mat, - dx1mat, - dresidualmat, - dgamma0, - dbeta0, - dgamma1, - dbeta1, - ) = _dropout_add_layer_norm_parallel_residual_backward( - dz0, - dz1, - dx, - x, - dmask0, - dmask1, - mu, - rsigma, - gamma0, - gamma1, - dropout_p, - has_x1, - has_residual, - ctx.is_rms_norm, - ) - dx0 = dx0mat.view(x.shape) - dx1 = dx1mat.view(x.shape) if dx1mat is not None else None - dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None - return ( - dx0, - dx1, - dresidual, - dgamma0, - dbeta0 if ctx.has_beta else None, - dgamma1, - dbeta1 if ctx.has_beta else None, - None, - None, - None, - None, - None, - None, - ) - - -def layer_norm(x, weight, bias, epsilon): - return DropoutAddLayerNormFn.apply(x, None, weight, bias, None, None, 0.0, epsilon, False) - - -def dropout_add_layer_norm( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - rowscale=None, - layerscale=None, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormFn.apply( - x0, - residual, - weight, - bias, - rowscale, - layerscale, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - False, - return_dropout_mask, - ) - - -def dropout_add_layer_norm_subset( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - layerscale=None, - x0_subset=None, - out_subset=None, - rowscale_const=1.0, - out_numrows=0, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormSubsetFn.apply( - x0, - residual, - weight, - bias, - layerscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32, - prenorm, - False, - return_dropout_mask, - ) - - -def dropout_add_layer_norm_parallel_residual( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormParallelResidualFn.apply( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - False, - return_dropout_mask, - ) - - -class DropoutAddLayerNorm(torch.nn.Module): - def __init__( - self, - hidden_size, - prenorm=False, - p=0.0, - eps=1e-5, - residual_in_fp32=False, - device=None, - dtype=None, - ): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.prenorm = prenorm - self.p = p - self.eps = eps - self.residual_in_fp32 = residual_in_fp32 - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.bias = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.reset_parameters() - - def reset_parameters(self): - init.ones_(self.weight) - init.zeros_(self.bias) - - def forward(self, x0, residual=None): - return dropout_add_layer_norm( - x0, - residual, - self.weight, - self.bias, - self.p if self.training else 0.0, - self.eps, - prenorm=self.prenorm, - residual_in_fp32=self.residual_in_fp32, - ) diff --git a/build/torch210-cu130-x86_64-windows/ops/rms_norm.py b/build/torch210-cu130-x86_64-windows/ops/rms_norm.py deleted file mode 100644 index 9e286bc004a7353d81abc29759064b5383571743..0000000000000000000000000000000000000000 --- a/build/torch210-cu130-x86_64-windows/ops/rms_norm.py +++ /dev/null @@ -1,174 +0,0 @@ -# Copyright (c) 2022, Tri Dao. -# Adapted from https://github.com/NVIDIA/apex/blob/master/apex/contrib/layer_norm/layer_norm.py - -import torch -from torch.nn import init - -from flash_attn.ops.layer_norm import ( - DropoutAddLayerNormFn, - DropoutAddLayerNormParallelResidualFn, - DropoutAddLayerNormSubsetFn, -) - - -def rms_norm(x, weight, epsilon): - return DropoutAddLayerNormFn.apply( - x, None, weight, None, None, None, 0.0, epsilon, False, False, True - ) - - -def dropout_add_rms_norm( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - rowscale=None, - layerscale=None, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormFn.apply( - x0, - residual, - weight, - bias, - rowscale, - layerscale, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - True, - return_dropout_mask, - ) - - -def dropout_add_rms_norm_subset( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - layerscale=None, - x0_subset=None, - out_subset=None, - rowscale_const=1.0, - out_numrows=0, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormSubsetFn.apply( - x0, - residual, - weight, - bias, - layerscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32, - prenorm, - True, - return_dropout_mask, - ) - - -def dropout_add_rms_norm_parallel_residual( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormParallelResidualFn.apply( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - True, - return_dropout_mask, - ) - - -class RMSNorm(torch.nn.Module): - def __init__(self, hidden_size, eps=1e-5, device=None, dtype=None): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.eps = eps - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.register_parameter("bias", None) - self.reset_parameters() - - def reset_parameters(self): - init.ones_(self.weight) - - def forward(self, x): - return rms_norm(x, self.weight, self.eps) - - -class DropoutAddRMSNorm(torch.nn.Module): - def __init__( - self, - hidden_size, - prenorm=False, - p=0.0, - eps=1e-5, - residual_in_fp32=False, - device=None, - dtype=None, - ): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.prenorm = prenorm - self.p = p - self.eps = eps - self.residual_in_fp32 = residual_in_fp32 - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.register_parameter("bias", None) - self.reset_parameters() - - def reset_parameters(self): - init.ones_(self.weight) - - def forward(self, x0, residual=None): - return dropout_add_rms_norm( - x0, - residual, - self.weight, - None, - self.p if self.training else 0.0, - self.eps, - prenorm=self.prenorm, - residual_in_fp32=self.residual_in_fp32, - ) diff --git a/build/torch210-cu130-x86_64-windows/ops/triton/__init__.py b/build/torch210-cu130-x86_64-windows/ops/triton/__init__.py deleted file mode 100644 index d3f5a12faa99758192ecc4ed3fc22c9249232e86..0000000000000000000000000000000000000000 --- a/build/torch210-cu130-x86_64-windows/ops/triton/__init__.py +++ /dev/null @@ -1 +0,0 @@ - diff --git a/build/torch210-cu130-x86_64-windows/ops/triton/cross_entropy.py b/build/torch210-cu130-x86_64-windows/ops/triton/cross_entropy.py deleted file mode 100644 index f44d6bfaf87a515e3d30598f98c1dfebc0511399..0000000000000000000000000000000000000000 --- a/build/torch210-cu130-x86_64-windows/ops/triton/cross_entropy.py +++ /dev/null @@ -1,330 +0,0 @@ -# Copyright (c) 2023, Tri Dao. - -from typing import Tuple, Optional, Union - -import torch -import torch.nn.functional as F - -import triton -import triton.language as tl - -# `all_gather_into_tensor` and `reduce_scatter_tensor` are new placeholders for -# `_all_gather_base` and `_reduce_scatter_base`. They require the most recent -# version of PyTorch. The following 2 lines are for backward compatibility with -# older PyTorch. -if "all_gather_into_tensor" not in dir(torch.distributed): - torch.distributed.all_gather_into_tensor = torch.distributed._all_gather_base - - -@triton.heuristics( - { - "HAS_SMOOTHING": lambda args: args["smoothing"] > 0.0, - } -) -@triton.jit -def cross_entropy_fwd_kernel( - loss_ptr, # data ptrs - lse_ptr, - z_loss_ptr, - logits_ptr, - labels_ptr, - smoothing, - logit_scale, - lse_square_scale, - ignore_index, - total_classes, - class_start_idx, # Useful for tensor parallel when each rank only has a subset of classes - n_cols, # shapes - logits_row_stride, # strides - BLOCK_SIZE: tl.constexpr, - HAS_SMOOTHING: tl.constexpr, - # if SPLIT (e.g. tensor parallel), don't include the LSE in the loss since it's not the final LSE - SPLIT: tl.constexpr, - PRECOMPUTED_LSE: tl.constexpr, # If LSE is already computed (also no smoothing and logit_scale == 1.0) -): - row_idx = tl.program_id(0) - logits_ptr = logits_ptr + row_idx * logits_row_stride.to(tl.int64) - sum_logits = 0.0 # For smoothing - if not PRECOMPUTED_LSE: - # Statistics for online softmax - m_i = -float("inf") - l_i = 0.0 - for col_offset in range(0, n_cols, BLOCK_SIZE): - cols = col_offset + tl.arange(0, BLOCK_SIZE) - logits = tl.load(logits_ptr + cols, mask=cols < n_cols, other=-float("inf")).to( - tl.float32 - ) * logit_scale - if HAS_SMOOTHING: - sum_logits += tl.sum(tl.where(cols < n_cols, logits, 0.0)) - m_i_new = tl.maximum(m_i, tl.max(logits)) - l_i = tl.exp(m_i - m_i_new) * l_i + tl.sum(tl.exp(logits - m_i_new)) - m_i = m_i_new - lse = tl.log(l_i) + m_i - tl.store(lse_ptr + row_idx, lse) - else: - lse = tl.load(lse_ptr + row_idx) - label_idx = tl.load(labels_ptr + row_idx) - if label_idx == ignore_index: - loss = 0.0 - z_loss = 0.0 - else: - label_idx -= class_start_idx - if label_idx >= 0 and label_idx < n_cols: - logits_label = tl.load(logits_ptr + label_idx) * logit_scale - if HAS_SMOOTHING: - loss = ( - (lse if not SPLIT else 0.0) - - smoothing * sum_logits / total_classes - - (1 - smoothing) * logits_label - ) - else: - loss = (lse if not SPLIT else 0.0) - logits_label - else: - # If label is out of bounds, we set the CE loss to 0.0. But we still want the smoothing loss - if HAS_SMOOTHING: - loss = smoothing * ((lse if not SPLIT else 0.0) - sum_logits / total_classes) - else: - loss = 0.0 - if not SPLIT: - z_loss = lse_square_scale * lse * lse - loss += z_loss - else: - z_loss = 0.0 - tl.store(loss_ptr + row_idx, loss) - if not SPLIT: - tl.store(z_loss_ptr + row_idx, z_loss) - - -@triton.heuristics( - { - "HAS_SMOOTHING": lambda args: args["smoothing"] > 0.0, - } -) -@triton.jit -def cross_entropy_bwd_kernel( - dlogits_ptr, # data ptrs - dloss_ptr, - logits_ptr, - lse_ptr, - labels_ptr, - smoothing, - logit_scale, - lse_square_scale, - ignore_index, - total_classes, - class_start_idx, # Useful for tensor parallel when each rank only has a subset of classes - n_cols, # shapes - logits_row_stride, # strides - dlogits_row_stride, - dloss_row_stride, - BLOCK_SIZE: tl.constexpr, - HAS_SMOOTHING: tl.constexpr, -): - row_idx = tl.program_id(0) - col_block_idx = tl.program_id(1) - logits_ptr = logits_ptr + row_idx * logits_row_stride.to(tl.int64) - dlogits_ptr = dlogits_ptr + row_idx * dlogits_row_stride.to(tl.int64) - col_offsets = col_block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - label_idx = tl.load(labels_ptr + row_idx) - if label_idx != ignore_index: - dloss = tl.load(dloss_ptr + row_idx * dloss_row_stride) - else: - dloss = 0.0 - logits = tl.load(logits_ptr + col_offsets, mask=col_offsets < n_cols, other=-float("inf")).to( - tl.float32 - ) * logit_scale - lse = tl.load(lse_ptr + row_idx) - probs = tl.exp(logits - lse) - probs += 2.0 * lse_square_scale * lse * probs - label_idx -= class_start_idx - if HAS_SMOOTHING: - smooth_positive = 1.0 - smoothing - smooth_negative = smoothing / total_classes - probs = tl.where(col_offsets == label_idx, probs - smooth_positive, probs) - smooth_negative - else: - probs = tl.where(col_offsets == label_idx, probs - 1.0, probs) - tl.store(dlogits_ptr + col_offsets, (dloss * logit_scale) * probs, mask=col_offsets < n_cols) - - -class CrossEntropyLoss(torch.autograd.Function): - - @staticmethod - def forward( - ctx, - logits, - labels, - precomputed_lse=None, - smoothing=0.0, - logit_scale=1.0, - lse_square_scale=0.0, - ignore_index=-100, - inplace_backward=False, - process_group=None, - ): - # For some reason Triton generates wrong code when labels has dtype long and its address - # is not aligned to 16 bytes. The ld.global.b64 seems to load the wrong label index. - if labels.dtype == torch.long and labels.data_ptr() % 16 != 0: - labels = F.pad(labels, (0, 1))[..., :-1] - assert labels.data_ptr() % 16 == 0 - assert logit_scale > 0.0 - n_rows, n_cols = logits.shape - assert labels.shape == (n_rows,) - world_size = 1 if process_group is None else torch.distributed.get_world_size(process_group) - total_classes = world_size * n_cols - rank = 0 if process_group is None else torch.distributed.get_rank(process_group) - class_start_idx = rank * n_cols - use_precomputed_lse = precomputed_lse is not None and logit_scale == 1.0 and smoothing == 0.0 - - if logits.stride(-1) != 1: - logits = logits.contiguous() - MAX_BLOCK_SIZE = 16 * 1024 - BLOCK_SIZE = min(triton.next_power_of_2(n_cols), MAX_BLOCK_SIZE) - num_warps = ( - 4 - if BLOCK_SIZE < 2048 - else (8 if BLOCK_SIZE < 8192 else (16 if BLOCK_SIZE < 128 * 1024 else 32)) - ) - losses = torch.empty(n_rows, dtype=torch.float, device=logits.device) - if use_precomputed_lse: - assert precomputed_lse.shape == (n_rows,) - lse = precomputed_lse.contiguous() - else: - lse = torch.empty(n_rows, dtype=torch.float, device=logits.device) - z_losses = torch.empty(n_rows, dtype=torch.float, device=logits.device) - # Need this, otherwise Triton tries to launch from cuda:0 and we get - # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?) - with torch.cuda.device(logits.device.index): - cross_entropy_fwd_kernel[(n_rows,)]( - losses, # data ptrs - lse, - z_losses, - logits, - labels, - smoothing, - logit_scale, - lse_square_scale, - ignore_index, - total_classes, - class_start_idx, - n_cols, # shapes - logits.stride(0), # strides - BLOCK_SIZE=BLOCK_SIZE, # constants - SPLIT=world_size > 1, - PRECOMPUTED_LSE=use_precomputed_lse, - num_warps=num_warps, - ) - - if world_size > 1: - # If there's no smoothing, if labels are in the vocab of this partition, losses contains - # - predicted logit, and 0 otherwise. - # If there's smoothing=0.1, for labels in the vocab of this partition, losses contains - # -0.9 * predicted logit - 0.1 * sum logit / total_classes. - # For labels not in the vocab of this partition, losses contains - # -0.1 * sum logit / total_classes. - if world_size > 1: - lse_allgather = torch.empty(world_size, n_rows, dtype=lse.dtype, device=lse.device) - torch.distributed.all_gather_into_tensor(lse_allgather, lse, group=process_group) - handle_losses = torch.distributed.all_reduce( - losses, op=torch.distributed.ReduceOp.SUM, group=process_group, async_op=True - ) - lse = torch.logsumexp(lse_allgather, dim=0) - handle_losses.wait() - # After the allreduce, if there's no smoothing, the total losses are - predicted_logit, - # we just have to add the (global) lse. - # If there's smoothing=0.1, the total losses are - # -0.9 * predicted_logit - 0.1 * sum logit / total_classes. - # Again, we just have to add the (global) lse. - losses += lse - if lse_square_scale != 0.0: - z_losses = lse_square_scale * lse.square() - z_losses.masked_fill_(labels == ignore_index, 0.0) - losses += z_losses - else: - z_losses = torch.zeros_like(losses) - losses.masked_fill_(labels == ignore_index, 0.0) - - ctx.save_for_backward(logits, lse, labels) - ctx.mark_non_differentiable(z_losses) - ctx.smoothing = smoothing - ctx.logit_scale = logit_scale - ctx.lse_square_scale = lse_square_scale - ctx.ignore_index = ignore_index - ctx.total_classes = total_classes - ctx.class_start_idx = class_start_idx - ctx.inplace_backward = inplace_backward - return losses, z_losses - - @staticmethod - def backward(ctx, grad_losses, grad_z_losses): - del grad_z_losses # z_losses are only for logging. - - logits, lse, labels = ctx.saved_tensors - dlogits = logits if ctx.inplace_backward else torch.empty_like(logits) - n_rows, n_cols = logits.shape - BLOCK_SIZE = min(triton.next_power_of_2(n_cols), 4 * 1024) - num_warps = 4 if BLOCK_SIZE < 2048 else (8 if BLOCK_SIZE < 8192 else 16) - grid = lambda META: (n_rows, triton.cdiv(n_cols, META["BLOCK_SIZE"])) # noqa - # Need this, otherwise Triton tries to launch from cuda:0 and we get - # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?) - with torch.cuda.device(logits.device.index): - cross_entropy_bwd_kernel[grid]( - dlogits, # data ptrs - grad_losses, - logits, - lse, - labels, - ctx.smoothing, - ctx.logit_scale, - ctx.lse_square_scale, - ctx.ignore_index, - ctx.total_classes, - ctx.class_start_idx, - n_cols, # shapes - logits.stride(0), # strides - dlogits.stride(0), - grad_losses.stride(0), - BLOCK_SIZE=BLOCK_SIZE, # constants - num_warps=num_warps, - ) - return dlogits, None, None, None, None, None, None, None, None, None - - -def cross_entropy_loss( - logits: torch.Tensor, - labels: torch.Tensor, - precomputed_lse: Optional[torch.Tensor] = None, - label_smoothing: float = 0.0, - logit_scale: float = 1.0, - lse_square_scale: float = 0.0, - ignore_index=-100, - inplace_backward: bool = False, - process_group=None, -) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Arguments: - logits: (batch, vocab_size) - labels: (batch,) - label_smoothing: float - logit_scale: float. Multiply logits by this scale before calculating the loss. - lse_square_scale: float. If > 0, we add lse_square_scale * lse(logits) ^ 2 to the loss. - This is also referred to as "z-loss". - ignore_index: int. If labels == ignore_index, the loss is set to 0.0. - inplace_backward: bool. If True, we do the backward pass in-place by modifying the logits. - This saves memory. - process_group: if not None, we're doing Tensor Parallel: each process is responsible for - one part of the vocab. The loss will be aggregated across processes. - Returns: - losses: (batch,), float - z_losses: (batch,), float - """ - return CrossEntropyLoss.apply( - logits, - labels, - precomputed_lse, - label_smoothing, - logit_scale, - lse_square_scale, - ignore_index, - inplace_backward, - process_group, - ) diff --git a/build/torch210-cu130-x86_64-windows/ops/triton/k_activations.py b/build/torch210-cu130-x86_64-windows/ops/triton/k_activations.py deleted file mode 100644 index f6a3f63a3a2e699092d78b48826d01bd6a70a609..0000000000000000000000000000000000000000 --- a/build/torch210-cu130-x86_64-windows/ops/triton/k_activations.py +++ /dev/null @@ -1,162 +0,0 @@ -# Adapted from https://github.com/facebookresearch/xformers/blob/main/xformers/triton/k_activations.py -# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved. -# -# This source code is licensed under the BSD license found in the -# LICENSE file in the root directory of this source tree. - -import math -from enum import Enum -from typing import Optional - -import triton -import triton.language as tl - -_sqrt2pi = math.sqrt(2.0 / math.pi) -_sqrt1_2 = math.sqrt(1.0 / 2) -_gaussian_pdf_normalization = 1.0 / math.sqrt(2 * math.pi) - - -class Activation(str, Enum): - SquaredReLU = "squared_relu" - GeLU = "gelu" - GeLUApprox = "gelu_approx" - LeakyReLU = "leaky_relu" - ReLU = "relu" - - -def get_triton_activation_kernel(activation: Optional[Activation]): - return ( - { - Activation.ReLU: relu, - Activation.LeakyReLU: leaky_relu, - Activation.GeLU: gelu, - Activation.GeLUApprox: gelu_approx, - Activation.SquaredReLU: squared_relu, - }[activation] - if activation - else None - ) - - -def get_triton_activation_bwd_kernel(activation: Optional[Activation]): - return ( - { - Activation.ReLU: relu_grad, - Activation.LeakyReLU: leaky_relu_grad, - Activation.GeLU: gelu_grad, - Activation.GeLUApprox: gelu_approx_grad, - Activation.SquaredReLU: squared_relu_grad, - }[activation] - if activation - else None - ) - - -@triton.jit -def tanh(x): - # Tanh is just a scaled sigmoid - return 2 * tl.sigmoid(2 * x) - 1 - - -@triton.jit -def cosh(x): - exp_x = tl.exp(x) - return (exp_x + 1.0 / exp_x) * 0.5 - - -# a Triton implementation of the most used activations -# See for instance http://arxiv.org/abs/1606.08415 for an overview - -# ReLU -@triton.jit -def relu(x): - """ - ReLU_ activation function - - .. _ReLU: https://pytorch.org/docs/stable/generated/torch.nn.ReLU.html - """ - zero = 0.0 - return tl.where(x >= 0, x, zero.to(x.dtype)) - - -@triton.jit -def relu_grad(x): - # ReLU is different from other activations - # in that it does not require the input to retrospectively compute its gradient - # here the input is the downstream gradient, and we return the upstream gradient directly - zero = 0.0 - one = 1.0 - return tl.where(x >= 0, one.to(x.dtype), zero.to(x.dtype)) - - -@triton.jit -def squared_relu(x): - """ - Squared ReLU activation, as proposed in the Primer_ paper. - - .. _Primer: https://arxiv.org/abs/2109.08668 - """ - x_ = relu(x) - return (x_ * x_).to(x.dtype) - - -@triton.jit -def squared_relu_grad(x): - return tl.where(x >= 0, 2.0 * x, 0.0) - - -# Leaky ReLU -@triton.jit -def leaky_relu(x): - """ - LeakyReLU_ activation - - .. _LeakyReLU: https://pytorch.org/docs/stable/generated/torch.nn.LeakyReLU.html - """ - scale = 0.01 + 0.0 - scale = scale.to(x.dtype) - return tl.where(x >= 0, x, scale * x) - - -@triton.jit -def leaky_relu_grad(x): - min_grad = 0.01 - max_grad = 1 - - min_grad = min_grad.to(x.dtype) - max_grad = max_grad.to(x.dtype) - - return tl.where(x >= 0, max_grad, min_grad) - - -@triton.jit -def gelu(x): - """Gaussian Error Linear Unit (GELU)""" - return x * 0.5 * (1.0 + tl.libdevice.erf(x * _sqrt1_2)) - - -@triton.jit -def gelu_grad(x): - cdf = 0.5 * (1.0 + tl.libdevice.erf(x * _sqrt1_2)) - pdf = tl.exp(-0.5 * x * x) * _gaussian_pdf_normalization - return cdf + x * pdf - - -@triton.jit -def gelu_approx(x): - """ - GeLU_ activation - Gaussian error linear unit, with tanh approximation - - .. _GeLU: https://arxiv.org/pdf/1606.08415.pdf - """ - return 0.5 * x * (1.0 + tanh(_sqrt2pi * x * (1.0 + 0.044715 * x * x))) - - -@triton.jit -def gelu_approx_grad(x): - # CREDITS: Fast implementation proposed in - # https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/fused_bias_gelu.py#L30 - tanh_out = tanh(0.79788456 * x * (1 + 0.044715 * x * x)) - return 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( - 1 + tanh_out - ) diff --git a/build/torch210-cu130-x86_64-windows/ops/triton/layer_norm.py b/build/torch210-cu130-x86_64-windows/ops/triton/layer_norm.py deleted file mode 100644 index 3724dd92e91fd8af7def1c849ab74d6adee4e486..0000000000000000000000000000000000000000 --- a/build/torch210-cu130-x86_64-windows/ops/triton/layer_norm.py +++ /dev/null @@ -1,1252 +0,0 @@ -# Copyright (c) 2024, Tri Dao. -# Implement dropout + residual + layer_norm / rms_norm. - -# Based on the Triton LayerNorm tutorial: https://triton-lang.org/main/getting-started/tutorials/05-layer-norm.html -# For the backward pass, we keep weight_grad and bias_grad in registers and accumulate. -# This is faster for dimensions up to 8k, but after that it's much slower due to register spilling. -# The models we train have hidden dim up to 8k anyway (e.g. Llama 70B), so this is fine. - -import math -from typing import Optional, List - -import torch -import torch.nn.functional as F -from torch import Tensor - -import triton -import triton.language as tl - -from flash_attn.utils.torch import custom_fwd, custom_bwd -from flash_attn.utils.library import triton_op - - -def maybe_contiguous_lastdim(x): - return x.contiguous() if x is not None and x.stride(-1) != 1 else x - - -def maybe_contiguous(x): - return x.contiguous() if x is not None else None - - -def triton_autotune_configs(): - # Return configs with a valid warp count for the current device - configs = [] - # Maximum threads per block is architecture-dependent in theory, but in reality all are 1024 - max_threads_per_block = 1024 - # Default to warp size 32 if not defined by device - warp_size = getattr(torch.cuda.get_device_properties(torch.cuda.current_device()), "warp_size", 32) - # Autotune for warp counts which are powers of 2 and do not exceed thread per block limit - return [triton.Config({}, num_warps=warp_count) for warp_count in [1, 2, 4, 8, 16, 32] - if warp_count * warp_size <= max_threads_per_block] - # return [triton.Config({}, num_warps=8)] - - -def layer_norm_ref( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - zero_centered_weight=False, - dropout_mask=None, - dropout_mask1=None, - upcast=False, -): - dtype = x.dtype - if upcast: - x = x.float() - weight = weight.float() - bias = bias.float() if bias is not None else None - residual = residual.float() if residual is not None else residual - x1 = x1.float() if x1 is not None else None - weight1 = weight1.float() if weight1 is not None else None - bias1 = bias1.float() if bias1 is not None else None - if zero_centered_weight: - weight = weight + 1.0 - if weight1 is not None: - weight1 = weight1 + 1.0 - if x1 is not None: - assert rowscale is None, "rowscale is not supported with parallel LayerNorm" - if rowscale is not None: - x = x * rowscale[..., None] - if dropout_p > 0.0: - if dropout_mask is not None: - x = x.masked_fill(~dropout_mask, 0.0) / (1.0 - dropout_p) - else: - x = F.dropout(x, p=dropout_p) - if x1 is not None: - if dropout_mask1 is not None: - x1 = x1.masked_fill(~dropout_mask1, 0.0) / (1.0 - dropout_p) - else: - x1 = F.dropout(x1, p=dropout_p) - if x1 is not None: - x = x + x1 - if residual is not None: - x = (x + residual).to(x.dtype) - out = F.layer_norm(x.to(weight.dtype), x.shape[-1:], weight=weight, bias=bias, eps=eps).to( - dtype - ) - if weight1 is None: - return out if not prenorm else (out, x) - else: - out1 = F.layer_norm( - x.to(weight1.dtype), x.shape[-1:], weight=weight1, bias=bias1, eps=eps - ).to(dtype) - return (out, out1) if not prenorm else (out, out1, x) - - -def rms_norm_ref( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - zero_centered_weight=False, - dropout_mask=None, - dropout_mask1=None, - upcast=False, -): - dtype = x.dtype - if upcast: - x = x.float() - weight = weight.float() - bias = bias.float() if bias is not None else None - residual = residual.float() if residual is not None else residual - x1 = x1.float() if x1 is not None else None - weight1 = weight1.float() if weight1 is not None else None - bias1 = bias1.float() if bias1 is not None else None - if zero_centered_weight: - weight = weight + 1.0 - if weight1 is not None: - weight1 = weight1 + 1.0 - if x1 is not None: - assert rowscale is None, "rowscale is not supported with parallel LayerNorm" - if rowscale is not None: - x = x * rowscale[..., None] - if dropout_p > 0.0: - if dropout_mask is not None: - x = x.masked_fill(~dropout_mask, 0.0) / (1.0 - dropout_p) - else: - x = F.dropout(x, p=dropout_p) - if x1 is not None: - if dropout_mask1 is not None: - x1 = x1.masked_fill(~dropout_mask1, 0.0) / (1.0 - dropout_p) - else: - x1 = F.dropout(x1, p=dropout_p) - if x1 is not None: - x = x + x1 - if residual is not None: - x = (x + residual).to(x.dtype) - rstd = 1 / torch.sqrt((x.square()).mean(dim=-1, keepdim=True) + eps) - out = ((x * rstd * weight) + bias if bias is not None else (x * rstd * weight)).to(dtype) - if weight1 is None: - return out if not prenorm else (out, x) - else: - out1 = ((x * rstd * weight1) + bias1 if bias1 is not None else (x * rstd * weight1)).to( - dtype - ) - return (out, out1) if not prenorm else (out, out1, x) - - -@triton.autotune( - configs=triton_autotune_configs(), - key=["N", "HAS_RESIDUAL", "STORE_RESIDUAL_OUT", "IS_RMS_NORM", "HAS_BIAS", "HAS_X1", "HAS_W1", "HAS_B1"], -) -# torch compile doesn't like triton.heuristics, so we set these manually when calling the kernel -# @triton.heuristics({"HAS_BIAS": lambda args: args["B"] is not None}) -# @triton.heuristics({"HAS_RESIDUAL": lambda args: args["RESIDUAL"] is not None}) -# @triton.heuristics({"HAS_X1": lambda args: args["X1"] is not None}) -# @triton.heuristics({"HAS_W1": lambda args: args["W1"] is not None}) -# @triton.heuristics({"HAS_B1": lambda args: args["B1"] is not None}) -@triton.jit -def _layer_norm_fwd_1pass_kernel( - X, # pointer to the input - Y, # pointer to the output - W, # pointer to the weights - B, # pointer to the biases - RESIDUAL, # pointer to the residual - X1, - W1, - B1, - Y1, - RESIDUAL_OUT, # pointer to the residual - ROWSCALE, - SEEDS, # Dropout seeds for each row - DROPOUT_MASK, - DROPOUT_MASK1, - Mean, # pointer to the mean - Rstd, # pointer to the 1/std - stride_x_row, # how much to increase the pointer when moving by 1 row - stride_y_row, - stride_res_row, - stride_res_out_row, - stride_x1_row, - stride_y1_row, - M, # number of rows in X - N, # number of columns in X - eps, # epsilon to avoid division by zero - dropout_p, # Dropout probability - zero_centered_weight, # If true, add 1.0 to the weight - IS_RMS_NORM: tl.constexpr, - BLOCK_N: tl.constexpr, - HAS_RESIDUAL: tl.constexpr, - STORE_RESIDUAL_OUT: tl.constexpr, - HAS_BIAS: tl.constexpr, - HAS_DROPOUT: tl.constexpr, - STORE_DROPOUT_MASK: tl.constexpr, - HAS_ROWSCALE: tl.constexpr, - HAS_X1: tl.constexpr, - HAS_W1: tl.constexpr, - HAS_B1: tl.constexpr, -): - # Map the program id to the row of X and Y it should compute. - row = tl.program_id(0) - X += row * stride_x_row - Y += row * stride_y_row - if HAS_RESIDUAL: - RESIDUAL += row * stride_res_row - if STORE_RESIDUAL_OUT: - RESIDUAL_OUT += row * stride_res_out_row - if HAS_X1: - X1 += row * stride_x1_row - if HAS_W1: - Y1 += row * stride_y1_row - # Compute mean and variance - cols = tl.arange(0, BLOCK_N) - x = tl.load(X + cols, mask=cols < N, other=0.0).to(tl.float32) - if HAS_ROWSCALE: - rowscale = tl.load(ROWSCALE + row).to(tl.float32) - x *= rowscale - if HAS_DROPOUT: - # Compute dropout mask - # 7 rounds is good enough, and reduces register pressure - keep_mask = tl.rand(tl.load(SEEDS + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - x = tl.where(keep_mask, x / (1.0 - dropout_p), 0.0) - if STORE_DROPOUT_MASK: - tl.store(DROPOUT_MASK + row * N + cols, keep_mask, mask=cols < N) - if HAS_X1: - x1 = tl.load(X1 + cols, mask=cols < N, other=0.0).to(tl.float32) - if HAS_ROWSCALE: - rowscale = tl.load(ROWSCALE + M + row).to(tl.float32) - x1 *= rowscale - if HAS_DROPOUT: - # Compute dropout mask - # 7 rounds is good enough, and reduces register pressure - keep_mask = ( - tl.rand(tl.load(SEEDS + M + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - ) - x1 = tl.where(keep_mask, x1 / (1.0 - dropout_p), 0.0) - if STORE_DROPOUT_MASK: - tl.store(DROPOUT_MASK1 + row * N + cols, keep_mask, mask=cols < N) - x += x1 - if HAS_RESIDUAL: - residual = tl.load(RESIDUAL + cols, mask=cols < N, other=0.0).to(tl.float32) - x += residual - if STORE_RESIDUAL_OUT: - tl.store(RESIDUAL_OUT + cols, x, mask=cols < N) - if not IS_RMS_NORM: - mean = tl.sum(x, axis=0) / N - tl.store(Mean + row, mean) - xbar = tl.where(cols < N, x - mean, 0.0) - var = tl.sum(xbar * xbar, axis=0) / N - else: - xbar = tl.where(cols < N, x, 0.0) - var = tl.sum(xbar * xbar, axis=0) / N - rstd = 1 / tl.sqrt(var + eps) - tl.store(Rstd + row, rstd) - # Normalize and apply linear transformation - mask = cols < N - w = tl.load(W + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w += 1.0 - if HAS_BIAS: - b = tl.load(B + cols, mask=mask).to(tl.float32) - x_hat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd - y = x_hat * w + b if HAS_BIAS else x_hat * w - # Write output - tl.store(Y + cols, y, mask=mask) - if HAS_W1: - w1 = tl.load(W1 + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w1 += 1.0 - if HAS_B1: - b1 = tl.load(B1 + cols, mask=mask).to(tl.float32) - y1 = x_hat * w1 + b1 if HAS_B1 else x_hat * w1 - tl.store(Y1 + cols, y1, mask=mask) - - -def _layer_norm_fwd( - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - residual: Optional[Tensor] = None, - x1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - out_dtype: Optional[torch.dtype] = None, - residual_dtype: Optional[torch.dtype] = None, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - return_dropout_mask: bool = False, - out: Optional[Tensor] = None, - residual_out: Optional[Tensor] = None -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - # Need to wrap to handle the case where residual_out is a alias of x, which makes torch.library - # and torch.compile unhappy. Also allocate memory for out and residual_out if they are None - # so that _layer_norm_fwd_impl doesn't have to return them. - if out is None: - out = torch.empty_like(x, dtype=x.dtype if out_dtype is None else out_dtype) - if residual is not None: - residual_dtype = residual.dtype - if residual_out is None and ( - residual is not None - or (residual_dtype is not None and residual_dtype != x.dtype) - or dropout_p > 0.0 - or rowscale is not None - or x1 is not None - ): - residual_out = torch.empty_like( - x, dtype=residual_dtype if residual_dtype is not None else x.dtype - ) - else: - residual_out = None - y1, mean, rstd, seeds, dropout_mask, dropout_mask1 = _layer_norm_fwd_impl( - x, - weight, - bias, - eps, - out, - residual=residual, - x1=x1, - weight1=weight1, - bias1=bias1, - dropout_p=dropout_p, - rowscale=rowscale, - zero_centered_weight=zero_centered_weight, - is_rms_norm=is_rms_norm, - return_dropout_mask=return_dropout_mask, - residual_out=residual_out, - ) - # residual_out is None if residual is None and residual_dtype == input_dtype and dropout_p == 0.0 - if residual_out is None: - residual_out = x - return out, y1, mean, rstd, residual_out, seeds, dropout_mask, dropout_mask1 - - -# [2025-04-28] torch.library.triton_op ignores the schema argument, but here we need the schema -# since we're returning a tuple of tensors -@triton_op("flash_attn::layer_norm_fwd_impl", mutates_args={"out", "residual_out"}, - schema="(Tensor x, Tensor weight, Tensor bias, float eps, Tensor(a!) out, Tensor? residual, Tensor? x1, Tensor? weight1, Tensor? bias1, float dropout_p, Tensor? rowscale, bool zero_centered_weight, bool is_rms_norm, bool return_dropout_mask, Tensor(a!)? residual_out) -> (Tensor y1, Tensor mean, Tensor rstd, Tensor seeds, Tensor dropout_mask, Tensor dropout_mask1)") -def _layer_norm_fwd_impl( - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - out: Tensor, - residual: Optional[Tensor] = None, - x1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - return_dropout_mask: bool = False, - residual_out: Optional[Tensor] = None -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - M, N = x.shape - assert x.stride(-1) == 1 - if residual is not None: - assert residual.stride(-1) == 1 - assert residual.shape == (M, N) - assert weight.shape == (N,) - assert weight.stride(-1) == 1 - if bias is not None: - assert bias.stride(-1) == 1 - assert bias.shape == (N,) - if x1 is not None: - assert x1.shape == x.shape - assert rowscale is None - assert x1.stride(-1) == 1 - if weight1 is not None: - assert weight1.shape == (N,) - assert weight1.stride(-1) == 1 - if bias1 is not None: - assert bias1.shape == (N,) - assert bias1.stride(-1) == 1 - if rowscale is not None: - assert rowscale.is_contiguous() - assert rowscale.shape == (M,) - assert out.shape == x.shape - assert out.stride(-1) == 1 - if residual_out is not None: - assert residual_out.shape == x.shape - assert residual_out.stride(-1) == 1 - if weight1 is not None: - y1 = torch.empty_like(out) - assert y1.stride(-1) == 1 - else: - y1 = None - mean = torch.empty((M,), dtype=torch.float32, device=x.device) if not is_rms_norm else None - rstd = torch.empty((M,), dtype=torch.float32, device=x.device) - if dropout_p > 0.0: - seeds = torch.randint( - 2**32, (M if x1 is None else 2 * M,), device=x.device, dtype=torch.int64 - ) - else: - seeds = None - if return_dropout_mask and dropout_p > 0.0: - dropout_mask = torch.empty(M, N, device=x.device, dtype=torch.bool) - if x1 is not None: - dropout_mask1 = torch.empty(M, N, device=x.device, dtype=torch.bool) - else: - dropout_mask1 = None - else: - dropout_mask, dropout_mask1 = None, None - # Less than 64KB per feature: enqueue fused kernel - MAX_FUSED_SIZE = 65536 // x.element_size() - BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) - if N > BLOCK_N: - raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") - with torch.cuda.device(x.device.index): - torch.library.wrap_triton(_layer_norm_fwd_1pass_kernel)[(M,)]( - x, - out, - weight, - bias, - residual, - x1, - weight1, - bias1, - y1, - residual_out, - rowscale, - seeds, - dropout_mask, - dropout_mask1, - mean, - rstd, - x.stride(0), - out.stride(0), - residual.stride(0) if residual is not None else 0, - residual_out.stride(0) if residual_out is not None else 0, - x1.stride(0) if x1 is not None else 0, - y1.stride(0) if y1 is not None else 0, - M, - N, - eps, - dropout_p, - # Passing bool make torch inductor very unhappy since it then tries to compare to int_max - int(zero_centered_weight), - is_rms_norm, - BLOCK_N, - residual is not None, - residual_out is not None, - bias is not None, - dropout_p > 0.0, - dropout_mask is not None, - rowscale is not None, - HAS_X1=x1 is not None, - HAS_W1=weight1 is not None, - HAS_B1=bias1 is not None, - ) - return y1, mean, rstd, seeds, dropout_mask, dropout_mask1 - - -@triton.autotune( - configs=triton_autotune_configs(), - key=["N", "HAS_DRESIDUAL", "STORE_DRESIDUAL", "IS_RMS_NORM", "HAS_BIAS", "HAS_DROPOUT"], -) -# torch compile doesn't like triton.heuristics, so we set these manually when calling the kernel -# @triton.heuristics({"HAS_BIAS": lambda args: args["B"] is not None}) -# @triton.heuristics({"HAS_DRESIDUAL": lambda args: args["DRESIDUAL"] is not None}) -# @triton.heuristics({"STORE_DRESIDUAL": lambda args: args["DRESIDUAL_IN"] is not None}) -# @triton.heuristics({"HAS_ROWSCALE": lambda args: args["ROWSCALE"] is not None}) -# @triton.heuristics({"HAS_DY1": lambda args: args["DY1"] is not None}) -# @triton.heuristics({"HAS_DX1": lambda args: args["DX1"] is not None}) -# @triton.heuristics({"HAS_B1": lambda args: args["DB1"] is not None}) -# @triton.heuristics({"RECOMPUTE_OUTPUT": lambda args: args["Y"] is not None}) -@triton.jit -def _layer_norm_bwd_kernel( - X, # pointer to the input - W, # pointer to the weights - B, # pointer to the biases - Y, # pointer to the output to be recomputed - DY, # pointer to the output gradient - DX, # pointer to the input gradient - DW, # pointer to the partial sum of weights gradient - DB, # pointer to the partial sum of biases gradient - DRESIDUAL, - W1, - DY1, - DX1, - DW1, - DB1, - DRESIDUAL_IN, - ROWSCALE, - SEEDS, - Mean, # pointer to the mean - Rstd, # pointer to the 1/std - stride_x_row, # how much to increase the pointer when moving by 1 row - stride_y_row, - stride_dy_row, - stride_dx_row, - stride_dres_row, - stride_dy1_row, - stride_dx1_row, - stride_dres_in_row, - M, # number of rows in X - N, # number of columns in X - eps, # epsilon to avoid division by zero - dropout_p, - zero_centered_weight, - rows_per_program, - IS_RMS_NORM: tl.constexpr, - BLOCK_N: tl.constexpr, - HAS_DRESIDUAL: tl.constexpr, - STORE_DRESIDUAL: tl.constexpr, - HAS_BIAS: tl.constexpr, - HAS_DROPOUT: tl.constexpr, - HAS_ROWSCALE: tl.constexpr, - HAS_DY1: tl.constexpr, - HAS_DX1: tl.constexpr, - HAS_B1: tl.constexpr, - RECOMPUTE_OUTPUT: tl.constexpr, -): - # Map the program id to the elements of X, DX, and DY it should compute. - row_block_id = tl.program_id(0) - row_start = row_block_id * rows_per_program - # Do not early exit if row_start >= M, because we need to write DW and DB - cols = tl.arange(0, BLOCK_N) - mask = cols < N - X += row_start * stride_x_row - if HAS_DRESIDUAL: - DRESIDUAL += row_start * stride_dres_row - if STORE_DRESIDUAL: - DRESIDUAL_IN += row_start * stride_dres_in_row - DY += row_start * stride_dy_row - DX += row_start * stride_dx_row - if HAS_DY1: - DY1 += row_start * stride_dy1_row - if HAS_DX1: - DX1 += row_start * stride_dx1_row - if RECOMPUTE_OUTPUT: - Y += row_start * stride_y_row - w = tl.load(W + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w += 1.0 - if RECOMPUTE_OUTPUT and HAS_BIAS: - b = tl.load(B + cols, mask=mask, other=0.0).to(tl.float32) - if HAS_DY1: - w1 = tl.load(W1 + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w1 += 1.0 - dw = tl.zeros((BLOCK_N,), dtype=tl.float32) - if HAS_BIAS: - db = tl.zeros((BLOCK_N,), dtype=tl.float32) - if HAS_DY1: - dw1 = tl.zeros((BLOCK_N,), dtype=tl.float32) - if HAS_B1: - db1 = tl.zeros((BLOCK_N,), dtype=tl.float32) - row_end = min((row_block_id + 1) * rows_per_program, M) - for row in range(row_start, row_end): - # Load data to SRAM - x = tl.load(X + cols, mask=mask, other=0).to(tl.float32) - dy = tl.load(DY + cols, mask=mask, other=0).to(tl.float32) - if HAS_DY1: - dy1 = tl.load(DY1 + cols, mask=mask, other=0).to(tl.float32) - if not IS_RMS_NORM: - mean = tl.load(Mean + row) - rstd = tl.load(Rstd + row) - # Compute dx - xhat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd - xhat = tl.where(mask, xhat, 0.0) - if RECOMPUTE_OUTPUT: - y = xhat * w + b if HAS_BIAS else xhat * w - tl.store(Y + cols, y, mask=mask) - wdy = w * dy - dw += dy * xhat - if HAS_BIAS: - db += dy - if HAS_DY1: - wdy += w1 * dy1 - dw1 += dy1 * xhat - if HAS_B1: - db1 += dy1 - if not IS_RMS_NORM: - c1 = tl.sum(xhat * wdy, axis=0) / N - c2 = tl.sum(wdy, axis=0) / N - dx = (wdy - (xhat * c1 + c2)) * rstd - else: - c1 = tl.sum(xhat * wdy, axis=0) / N - dx = (wdy - xhat * c1) * rstd - if HAS_DRESIDUAL: - dres = tl.load(DRESIDUAL + cols, mask=mask, other=0).to(tl.float32) - dx += dres - # Write dx - if STORE_DRESIDUAL: - tl.store(DRESIDUAL_IN + cols, dx, mask=mask) - if HAS_DX1: - if HAS_DROPOUT: - keep_mask = ( - tl.rand(tl.load(SEEDS + M + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - ) - dx1 = tl.where(keep_mask, dx / (1.0 - dropout_p), 0.0) - else: - dx1 = dx - tl.store(DX1 + cols, dx1, mask=mask) - if HAS_DROPOUT: - keep_mask = tl.rand(tl.load(SEEDS + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - dx = tl.where(keep_mask, dx / (1.0 - dropout_p), 0.0) - if HAS_ROWSCALE: - rowscale = tl.load(ROWSCALE + row).to(tl.float32) - dx *= rowscale - tl.store(DX + cols, dx, mask=mask) - - X += stride_x_row - if HAS_DRESIDUAL: - DRESIDUAL += stride_dres_row - if STORE_DRESIDUAL: - DRESIDUAL_IN += stride_dres_in_row - if RECOMPUTE_OUTPUT: - Y += stride_y_row - DY += stride_dy_row - DX += stride_dx_row - if HAS_DY1: - DY1 += stride_dy1_row - if HAS_DX1: - DX1 += stride_dx1_row - tl.store(DW + row_block_id * N + cols, dw, mask=mask) - if HAS_BIAS: - tl.store(DB + row_block_id * N + cols, db, mask=mask) - if HAS_DY1: - tl.store(DW1 + row_block_id * N + cols, dw1, mask=mask) - if HAS_B1: - tl.store(DB1 + row_block_id * N + cols, db1, mask=mask) - - -def _layer_norm_bwd( - dy: Tensor, - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - mean: Tensor, - rstd: Tensor, - dresidual: Optional[Tensor] = None, - dy1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - seeds: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - has_residual: bool = False, - has_x1: bool = False, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - x_dtype: Optional[torch.dtype] = None, - recompute_output: bool = False, -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - # Need to wrap to handle the case where dresidual_in or dx1 are aliases of x, - # which makes torch.library unhappy - dx, dw, db, dresidual_in, dx1, dw1, db1, y = _layer_norm_bwd_impl( - dy, - x, - weight, - bias, - eps, - mean, - rstd, - dresidual, - dy1, - weight1, - bias1, - seeds, - dropout_p, - rowscale, - has_residual, - has_x1, - zero_centered_weight, - is_rms_norm, - x_dtype=x_dtype, - recompute_output=recompute_output, - ) - # Don't need to compute dresidual_in separately in this case - if has_residual and dx.dtype == x.dtype and dropout_p == 0.0 and rowscale is None: - dresidual_in = dx - if has_x1 and dropout_p == 0.0: - dx1 = dx - return dx, dw, db, dresidual_in, dx1, dw1, db1, y - - - -@triton_op("flash_attn::layer_norm_bwd_impl", mutates_args={}, - schema="(Tensor dy, Tensor x, Tensor weight, Tensor bias, float eps, Tensor mean, Tensor rstd, Tensor? dresidual, Tensor? dy1, Tensor? weight1, Tensor? bias1, Tensor? seeds, float dropout_p, Tensor? rowscale, bool has_residual, bool has_x1, bool zero_centered_weight, bool is_rms_norm, ScalarType? x_dtype, bool recompute_output) -> (Tensor dx, Tensor dw, Tensor db, Tensor dresidual_in, Tensor dx1, Tensor dw1, Tensor db1, Tensor y)", - allow_decomposition=False, # Don't let torch.compile trace inside - ) -def _layer_norm_bwd_impl( - dy: Tensor, - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - mean: Tensor, - rstd: Tensor, - dresidual: Optional[Tensor] = None, - dy1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - seeds: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - has_residual: bool = False, - has_x1: bool = False, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - x_dtype: Optional[torch.dtype] = None, - recompute_output: bool = False, -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - M, N = x.shape - assert x.stride(-1) == 1 - dy = maybe_contiguous_lastdim(dy) - assert dy.stride(-1) == 1 - assert dy.shape == (M, N) - if dresidual is not None: - dresidual = maybe_contiguous_lastdim(dresidual) - assert dresidual.stride(-1) == 1 - assert dresidual.shape == (M, N) - assert weight.shape == (N,) - assert weight.stride(-1) == 1 - if bias is not None: - assert bias.stride(-1) == 1 - assert bias.shape == (N,) - if dy1 is not None: - dy1 = maybe_contiguous_lastdim(dy1) - assert weight1 is not None - assert dy1.shape == dy.shape - assert dy1.stride(-1) == 1 - if weight1 is not None: - assert weight1.shape == (N,) - assert weight1.stride(-1) == 1 - if bias1 is not None: - assert bias1.shape == (N,) - assert bias1.stride(-1) == 1 - if seeds is not None: - assert seeds.is_contiguous() - assert seeds.shape == (M if not has_x1 else M * 2,) - if rowscale is not None: - assert rowscale.is_contiguous() - assert rowscale.shape == (M,) - # allocate output - dx = ( - torch.empty_like(x) - if x_dtype is None - else torch.empty(M, N, dtype=x_dtype, device=x.device) - ) - dresidual_in = ( - torch.empty_like(x) - if has_residual - and (dx.dtype != x.dtype or dropout_p > 0.0 or rowscale is not None or has_x1) - else None - ) - dx1 = torch.empty_like(dx) if (has_x1 and dropout_p > 0.0) else None - y = torch.empty(M, N, dtype=dy.dtype, device=dy.device) if recompute_output else None - if recompute_output: - assert weight1 is None, "recompute_output is not supported with parallel LayerNorm" - - # Less than 64KB per feature: enqueue fused kernel - MAX_FUSED_SIZE = 65536 // x.element_size() - BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) - if N > BLOCK_N: - raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") - # Increasing the multiple (e.g. 8) will allow more thread blocks to be launched and hide the - # latency of the gmem reads/writes, but will increase the time of summing up dw / db. - sm_count = torch.cuda.get_device_properties(x.device).multi_processor_count * 8 - _dw = torch.empty((sm_count, N), dtype=torch.float32, device=weight.device) - _db = ( - torch.empty((sm_count, N), dtype=torch.float32, device=bias.device) - if bias is not None - else None - ) - _dw1 = torch.empty_like(_dw) if weight1 is not None else None - _db1 = torch.empty_like(_db) if bias1 is not None else None - rows_per_program = math.ceil(M / sm_count) - grid = (sm_count,) - with torch.cuda.device(x.device.index): - torch.library.wrap_triton(_layer_norm_bwd_kernel)[grid]( - x, - weight, - bias, - y, - dy, - dx, - _dw, - _db, - dresidual, - weight1, - dy1, - dx1, - _dw1, - _db1, - dresidual_in, - rowscale, - seeds, - mean, - rstd, - x.stride(0), - 0 if not recompute_output else y.stride(0), - dy.stride(0), - dx.stride(0), - dresidual.stride(0) if dresidual is not None else 0, - dy1.stride(0) if dy1 is not None else 0, - dx1.stride(0) if dx1 is not None else 0, - dresidual_in.stride(0) if dresidual_in is not None else 0, - M, - N, - eps, - dropout_p, - # Passing bool make torch inductor very unhappy since it then tries to compare to int_max - int(zero_centered_weight), - rows_per_program, - is_rms_norm, - BLOCK_N, - dresidual is not None, - dresidual_in is not None, - bias is not None, - dropout_p > 0.0, - HAS_ROWSCALE=rowscale is not None, - HAS_DY1=dy1 is not None, - HAS_DX1=dx1 is not None, - HAS_B1=bias1 is not None, - RECOMPUTE_OUTPUT=y is not None, - ) - dw = _dw.sum(0).to(weight.dtype) - db = _db.sum(0).to(bias.dtype) if bias is not None else None - dw1 = _dw1.sum(0).to(weight1.dtype) if weight1 is not None else None - db1 = _db1.sum(0).to(bias1.dtype) if bias1 is not None else None - # dresidual_in and dx1 could be None, the wrapper will handle assigning them from dx - return dx, dw, db, dresidual_in, dx1, dw1, db1, y - - -class LayerNormFn(torch.autograd.Function): - - @staticmethod - def forward( - ctx, - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - residual_in_fp32=False, - zero_centered_weight=False, - is_rms_norm=False, - return_dropout_mask=False, - out_dtype=None, - out=None, - residual_out=None - ): - x_shape_og = x.shape - # reshape input data into 2D tensor - x = maybe_contiguous_lastdim(x.reshape(-1, x.shape[-1])) - if residual is not None: - assert residual.shape == x_shape_og - residual = maybe_contiguous_lastdim(residual.reshape(-1, residual.shape[-1])) - if x1 is not None: - assert x1.shape == x_shape_og - assert rowscale is None, "rowscale is not supported with parallel LayerNorm" - x1 = maybe_contiguous_lastdim(x1.reshape(-1, x1.shape[-1])) - weight = weight.contiguous() - bias = maybe_contiguous(bias) - weight1 = maybe_contiguous(weight1) - bias1 = maybe_contiguous(bias1) - if rowscale is not None: - rowscale = rowscale.reshape(-1).contiguous() - residual_dtype = ( - residual.dtype - if residual is not None - else (torch.float32 if residual_in_fp32 else None) - ) - if out is not None: - out = out.reshape(-1, out.shape[-1]) - if residual_out is not None: - residual_out = residual_out.reshape(-1, residual_out.shape[-1]) - y, y1, mean, rstd, residual_out, seeds, dropout_mask, dropout_mask1 = _layer_norm_fwd( - x, - weight, - bias, - eps, - residual, - x1, - weight1, - bias1, - dropout_p=dropout_p, - rowscale=rowscale, - out_dtype=out_dtype, - residual_dtype=residual_dtype, - zero_centered_weight=zero_centered_weight, - is_rms_norm=is_rms_norm, - return_dropout_mask=return_dropout_mask, - out=out, - residual_out=residual_out, - ) - ctx.save_for_backward( - residual_out, weight, bias, weight1, bias1, rowscale, seeds, mean, rstd - ) - ctx.x_shape_og = x_shape_og - ctx.eps = eps - ctx.dropout_p = dropout_p - ctx.is_rms_norm = is_rms_norm - ctx.has_residual = residual is not None - ctx.has_x1 = x1 is not None - ctx.prenorm = prenorm - ctx.x_dtype = x.dtype - ctx.zero_centered_weight = zero_centered_weight - y = y.reshape(x_shape_og) - y1 = y1.reshape(x_shape_og) if y1 is not None else None - residual_out = residual_out.reshape(x_shape_og) if residual_out is not None else None - dropout_mask = dropout_mask.reshape(x_shape_og) if dropout_mask is not None else None - dropout_mask1 = dropout_mask1.reshape(x_shape_og) if dropout_mask1 is not None else None - if not return_dropout_mask: - if weight1 is None: - return y if not prenorm else (y, residual_out) - else: - return (y, y1) if not prenorm else (y, y1, residual_out) - else: - if weight1 is None: - return ( - (y, dropout_mask, dropout_mask1) - if not prenorm - else (y, residual_out, dropout_mask, dropout_mask1) - ) - else: - return ( - (y, y1, dropout_mask, dropout_mask1) - if not prenorm - else (y, y1, residual_out, dropout_mask, dropout_mask1) - ) - - @staticmethod - def backward(ctx, dy, *args): - x, weight, bias, weight1, bias1, rowscale, seeds, mean, rstd = ctx.saved_tensors - dy = dy.reshape(-1, dy.shape[-1]) - if weight1 is not None: - dy1, args = args[0], args[1:] - dy1 = dy1.reshape(-1, dy1.shape[-1]) - assert dy1.shape == x.shape - else: - dy1 = None - if ctx.prenorm: - dresidual = args[0] - dresidual = dresidual.reshape(-1, dresidual.shape[-1]) - assert dresidual.shape == x.shape - else: - dresidual = None - dx, dw, db, dresidual_in, dx1, dw1, db1, _ = _layer_norm_bwd( - dy, - x, - weight, - bias, - ctx.eps, - mean, - rstd, - dresidual, - dy1, - weight1, - bias1, - seeds, - ctx.dropout_p, - rowscale, - ctx.has_residual, - ctx.has_x1, - ctx.zero_centered_weight, - ctx.is_rms_norm, - x_dtype=ctx.x_dtype, - recompute_output=False, - ) - return ( - dx.reshape(ctx.x_shape_og), - dw, - db, - dresidual_in.reshape(ctx.x_shape_og) if ctx.has_residual else None, - dx1.reshape(ctx.x_shape_og) if dx1 is not None else None, - dw1, - db1, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) - - -def layer_norm_fn( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - residual_in_fp32=False, - zero_centered_weight=False, - is_rms_norm=False, - return_dropout_mask=False, - out_dtype=None, - out=None, - residual_out=None -): - return LayerNormFn.apply( - x, - weight, - bias, - residual, - x1, - weight1, - bias1, - eps, - dropout_p, - rowscale, - prenorm, - residual_in_fp32, - zero_centered_weight, - is_rms_norm, - return_dropout_mask, - out_dtype, - out, - residual_out - ) - - -def rms_norm_fn( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - residual_in_fp32=False, - zero_centered_weight=False, - return_dropout_mask=False, - out_dtype=None, - out=None, - residual_out=None -): - return LayerNormFn.apply( - x, - weight, - bias, - residual, - x1, - weight1, - bias1, - eps, - dropout_p, - rowscale, - prenorm, - residual_in_fp32, - zero_centered_weight, - True, - return_dropout_mask, - out_dtype, - out, - residual_out - ) - - -class RMSNorm(torch.nn.Module): - - def __init__(self, hidden_size, eps=1e-5, dropout_p=0.0, zero_centered_weight=False, - device=None, dtype=None): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.eps = eps - if dropout_p > 0.0: - self.drop = torch.nn.Dropout(dropout_p) - else: - self.drop = None - self.zero_centered_weight = zero_centered_weight - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.register_parameter("bias", None) - self.reset_parameters() - - def reset_parameters(self): - if not self.zero_centered_weight: - torch.nn.init.ones_(self.weight) - else: - torch.nn.init.zeros_(self.weight) - - def forward(self, x, residual=None, prenorm=False, residual_in_fp32=False): - return rms_norm_fn( - x, - self.weight, - self.bias, - residual=residual, - eps=self.eps, - dropout_p=self.drop.p if self.drop is not None and self.training else 0.0, - prenorm=prenorm, - residual_in_fp32=residual_in_fp32, - zero_centered_weight=self.zero_centered_weight, - ) - - -class LayerNormLinearFn(torch.autograd.Function): - - @staticmethod - @custom_fwd - def forward( - ctx, - x, - norm_weight, - norm_bias, - linear_weight, - linear_bias, - residual=None, - eps=1e-6, - prenorm=False, - residual_in_fp32=False, - is_rms_norm=False, - ): - x_shape_og = x.shape - # reshape input data into 2D tensor - x = maybe_contiguous_lastdim(x.reshape(-1, x.shape[-1])) - if residual is not None: - assert residual.shape == x_shape_og - residual = maybe_contiguous_lastdim(residual.reshape(-1, residual.shape[-1])) - norm_weight = norm_weight.contiguous() - norm_bias = maybe_contiguous(norm_bias) - residual_dtype = ( - residual.dtype - if residual is not None - else (torch.float32 if residual_in_fp32 else None) - ) - y, _, mean, rstd, residual_out, *rest = _layer_norm_fwd( - x, - norm_weight, - norm_bias, - eps, - residual, - out_dtype=None if not torch.is_autocast_enabled() else torch.get_autocast_dtype("cuda"), - residual_dtype=residual_dtype, - is_rms_norm=is_rms_norm, - ) - y = y.reshape(x_shape_og) - dtype = torch.get_autocast_dtype("cuda") if torch.is_autocast_enabled() else y.dtype - linear_weight = linear_weight.to(dtype) - linear_bias = linear_bias.to(dtype) if linear_bias is not None else None - out = F.linear(y.to(linear_weight.dtype), linear_weight, linear_bias) - # We don't store y, will be recomputed in the backward pass to save memory - ctx.save_for_backward(residual_out, norm_weight, norm_bias, linear_weight, mean, rstd) - ctx.x_shape_og = x_shape_og - ctx.eps = eps - ctx.is_rms_norm = is_rms_norm - ctx.has_residual = residual is not None - ctx.prenorm = prenorm - ctx.x_dtype = x.dtype - ctx.linear_bias_is_none = linear_bias is None - return out if not prenorm else (out, residual_out.reshape(x_shape_og)) - - @staticmethod - @custom_bwd - def backward(ctx, dout, *args): - x, norm_weight, norm_bias, linear_weight, mean, rstd = ctx.saved_tensors - dout = dout.reshape(-1, dout.shape[-1]) - dy = F.linear(dout, linear_weight.t()) - dlinear_bias = None if ctx.linear_bias_is_none else dout.sum(0) - dy = maybe_contiguous_lastdim(dy) - assert dy.shape == x.shape - if ctx.prenorm: - dresidual = args[0] - dresidual = maybe_contiguous_lastdim(dresidual.reshape(-1, dresidual.shape[-1])) - assert dresidual.shape == x.shape - else: - dresidual = None - dx, dnorm_weight, dnorm_bias, dresidual_in, _, _, _, y = _layer_norm_bwd( - dy, - x, - norm_weight, - norm_bias, - ctx.eps, - mean, - rstd, - dresidual=dresidual, - has_residual=ctx.has_residual, - is_rms_norm=ctx.is_rms_norm, - x_dtype=ctx.x_dtype, - recompute_output=True, - ) - dlinear_weight = torch.einsum("bo,bi->oi", dout, y) - return ( - dx.reshape(ctx.x_shape_og), - dnorm_weight, - dnorm_bias, - dlinear_weight, - dlinear_bias, - dresidual_in.reshape(ctx.x_shape_og) if ctx.has_residual else None, - None, - None, - None, - None, - ) - - -def layer_norm_linear_fn( - x, - norm_weight, - norm_bias, - linear_weight, - linear_bias, - residual=None, - eps=1e-6, - prenorm=False, - residual_in_fp32=False, - is_rms_norm=False, -): - return LayerNormLinearFn.apply( - x, - norm_weight, - norm_bias, - linear_weight, - linear_bias, - residual, - eps, - prenorm, - residual_in_fp32, - is_rms_norm, - ) diff --git a/build/torch210-cu130-x86_64-windows/ops/triton/linear.py b/build/torch210-cu130-x86_64-windows/ops/triton/linear.py deleted file mode 100644 index 0d800851ebdd5a589bb03e5374b807d2c4bd148c..0000000000000000000000000000000000000000 --- a/build/torch210-cu130-x86_64-windows/ops/triton/linear.py +++ /dev/null @@ -1,594 +0,0 @@ -# Adapted from https://github.com/ELS-RD/kernl/blob/main/src/kernl/implementations/linear_layer.py -# and https://github.com/openai/triton/blob/master/python/triton/ops/matmul.py -from typing import Optional - -import torch -import triton -import triton.language as tl -from triton.ops.matmul_perf_model import early_config_prune, estimate_matmul_time - -from flash_attn.ops.triton.k_activations import ( - gelu, - gelu_approx, - gelu_approx_grad, - gelu_grad, - squared_relu, - squared_relu_grad, -) - -# CREDITS: Initially inspired by the Triton tutorial on matrix multiplications - - -def init_to_zero(name): - return lambda nargs: nargs[name].zero_() - - -def get_configs_io_bound(): - configs = [] - for num_stages in [2, 3, 4, 5, 6]: - for block_m in [16, 32]: - for block_k in [32, 64]: - for block_n in [32, 64, 128, 256]: - num_warps = 2 if block_n <= 64 else 4 - configs.append( - triton.Config( - { - "BLOCK_M": block_m, - "BLOCK_N": block_n, - "BLOCK_K": block_k, - "SPLIT_K": 1, - }, - num_stages=num_stages, - num_warps=num_warps, - ) - ) - # split_k not used - # for split_k in [2, 4, 8, 16]: - # configs.append(triton.Config( - # {'BLOCK_M': block_m, 'BLOCK_N': block_n, 'BLOCK_K': block_k, 'SPLIT_K': split_k}, - # num_stages=num_stages, num_warps=num_warps, pre_hook=init_to_zero('C'))) - return configs - - -@triton.autotune( - configs=[ - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - # good for int8 - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=4, - num_warps=4, - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - ] - + get_configs_io_bound(), - key=["CACHE_KEY_M", "CACHE_KEY_N", "CACHE_KEY_K"], - prune_configs_by={ - "early_config_prune": early_config_prune, - "perf_model": estimate_matmul_time, - "top_k": 10, - }, -) -@triton.heuristics( - { - "EVEN_K": lambda args: args["K"] % (args["BLOCK_K"] * args["SPLIT_K"]) == 0, - } -) -@triton.jit -def kernel_fwd( - C, # Pointers to matrices - ACT_INPUT, - A, - B, - bias, - # Matrix dimensions - M, - N, - K, - CACHE_KEY_M, - CACHE_KEY_N, - CACHE_KEY_K, - # The stride variables represent how much to increase the ptr by when moving by 1 - # element in a particular dimension. E.g. stride_am is how much to increase a_ptr - # by to get the element one row down (A has M rows) - stride_cm, - # stride_cn, # Assume that stride_cn == 1 - stride_am, - stride_ak, - stride_bn, - stride_bk, - # Meta-parameters - BLOCK_M: tl.constexpr, - GROUP_M: tl.constexpr, - BLOCK_N: tl.constexpr, - BLOCK_K: tl.constexpr, - # split k not used, not performant with activation, kept because early_config_prune is expecting it - SPLIT_K: tl.constexpr, - EVEN_K: tl.constexpr, - A_ROWMAJOR: tl.constexpr, - B_COLMAJOR: tl.constexpr, - BIAS: tl.constexpr, - SAVE_ACT_INPUT: tl.constexpr, - ACTIVATION: tl.constexpr, -): - - """ - Kernel for computing Out = activation(A x W + C) - - Input has shape (M, K) - - Weight has shape (K, N) - - Bias has shape (N,) - - Output has shape (M, N) - - ActInputs (optional) has shape (M, N) - 'ActInputs' optionally saves the A x W + C intermediate for backward computations - This kernel will consolidate over K - """ - - pid = tl.program_id(axis=0) - - grid_m = (M + BLOCK_M - 1) // BLOCK_M - grid_n = (N + BLOCK_N - 1) // BLOCK_N - # re-order program ID for better L2 performance - width = GROUP_M * grid_n - group_id = pid // width - group_size = min(grid_m - group_id * GROUP_M, GROUP_M) - pid_m = group_id * GROUP_M + (pid % group_size) - pid_n = (pid % width) // (group_size) - - # now compute the block that each program will go through - # rm (resp. rn) denotes a range of indices - # for rows (resp. col) of C - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - # trick to avoid masking on M and N axis - ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M) - rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N) - rk = tl.arange(0, BLOCK_K) - - if A_ROWMAJOR: - A = A + (ram[:, None] * stride_am + rk[None, :]) - else: - A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak) - if B_COLMAJOR: - B = B + (rk[:, None] + rbn[None, :] * stride_bn) - else: - B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn) - - acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) - - for k in range(K, 0, -BLOCK_K): - if EVEN_K: - a = tl.load(A) - b = tl.load(B) - else: - a = tl.load(A, mask=rk[None, :] < k, other=0.0) - b = tl.load(B, mask=rk[:, None] < k, other=0.0) - acc += tl.dot(a, b) - - if A_ROWMAJOR: - A += BLOCK_K - else: - A += BLOCK_K * stride_ak - if B_COLMAJOR: - B += BLOCK_K - else: - B += BLOCK_K * stride_bk - - # Putting bias after the matmul (instead of before) is faster, idk why - if BIAS: - bias = tl.load(bias + rn, mask=rn < N, other=0.0).to(tl.float32) - acc += bias[None, :] - - # optional: save the activation inputs - if SAVE_ACT_INPUT: - # act_in_ptrs = ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :] * stride_cn - act_in_ptrs = ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :] - tl.store(act_in_ptrs, acc) - - # optional: fused activation (while the data is in shared memory) - if ACTIVATION == "gelu": - acc = gelu(acc) - elif ACTIVATION == "gelu_approx": - acc = gelu_approx(acc) - elif ACTIVATION == "squared_relu": - acc = squared_relu(acc) - # rematerialize rm and rn to save registers - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - - # write back result - # C = C + rm[:, None] * stride_cm + rn[None, :] * stride_cn - C = C + rm[:, None] * stride_cm + rn[None, :] - mask = (rm < M)[:, None] & (rn < N)[None, :] - tl.store(C, acc) - - -def triton_linear_act( - x: torch.Tensor, - weight: torch.Tensor, - bias: Optional[torch.Tensor] = None, - activation: str = "id", - save_act_input: bool = False, -) -> torch.Tensor: - """ - Compute e = activation(x @ weight.T + bias). - This wrapper kicks the `kernel_fwd` Triton kernel - :param x: input tensor - :param weight: weight matrix - :param bias: an optional bias tensor - :param activation: Activation name. Needs to be a Triton kernel. - :param act_input: an optional tensor to save the activation inputs (for backward) - :return: result tensor - """ - # if torch.is_autocast_enabled(): - # dtype = torch.get_autocast_gpu_dtype() - # x, weight, bias = [a.to(dtype=dtype) for a in [x, weight, bias]] - - assert activation in ["id", "gelu", "gelu_approx", "squared_relu"] - - batch_shape, n = x.shape[:-1], x.shape[-1] - batch_dim = batch_shape.numel() - x_reshaped = x.reshape(batch_dim, n) - - if x_reshaped.stride(0) > 1 and x_reshaped.stride(1) > 1: - x_reshaped = x_reshaped.contiguous() - if weight.stride(0) > 1 and weight.stride(1) > 1: - weight = weight.contiguous() - bias = bias.contiguous() if bias is not None else None - - assert ( - x.dtype == weight.dtype - ), f"Input and weight must have the same dtype, got {x.dtype} and {weight.dtype}" - if bias is not None: - assert ( - x.dtype == bias.dtype - ), f"Input and bias must have the same dtype, got {x.dtype} and {bias.dtype}" - assert ( - x_reshaped.shape[1] == weight.shape[1] - ), f"Incompatible dimensions: {x_reshaped.shape} - {weight.shape}" - - assert ( - bias is None or bias.shape[0] == weight.shape[0] - ), "Incompatible dimensions in between weight and bias" - - M, K = x_reshaped.shape - N, K = weight.shape - - output = torch.empty((M, N), device=x.device, dtype=x.dtype) - act_input = torch.empty_like(output) if save_act_input else None - - # 1D launch kernel where each block gets its own program. - grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),) # noqa - - kernel_fwd[grid]( - output, - act_input, - x_reshaped, - weight, # data ptrs - bias if bias is not None else x, # auto skip bias if not present - M, # shapes - N, - K, - M // 32, # key for triton cache (limit number of compilations) - N // 32, - K // 32, - stride_cm=output.stride(0), # strides - # stride_cn=output.stride(1), - stride_am=x_reshaped.stride(0), - stride_ak=x_reshaped.stride(1), - stride_bk=weight.stride(1), - stride_bn=weight.stride(0), - BIAS=bias is not None, # optional fused bias - SAVE_ACT_INPUT=save_act_input, # optional save activation inputs - ACTIVATION=activation, # optional fused activation - A_ROWMAJOR=x_reshaped.stride(1) == 1, - B_COLMAJOR=weight.stride(1) == 1, - GROUP_M=8, # speed optimization: group the programs - ) - - if not save_act_input: - return output.reshape(*batch_shape, output.shape[-1]) - else: - return ( - output.reshape(*batch_shape, output.shape[-1]), - act_input.reshape(*batch_shape, act_input.shape[-1]), - ) - - -@triton.autotune( - configs=[ - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - # good for int8 - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=4, - num_warps=4, - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - ] - + get_configs_io_bound(), - key=["CACHE_KEY_M", "CACHE_KEY_N", "CACHE_KEY_K"], - prune_configs_by={ - "early_config_prune": early_config_prune, - "perf_model": estimate_matmul_time, - "top_k": 10, - }, -) -@triton.heuristics( - { - "EVEN_K": lambda args: args["K"] % (args["BLOCK_K"] * args["SPLIT_K"]) == 0, - } -) -@triton.jit -def kernel_bwd( - C, # Pointers to matrices - ACT_INPUT, - A, - B, - # Matrix dimensions - M, - N, - K, - CACHE_KEY_M, - CACHE_KEY_N, - CACHE_KEY_K, - # The stride variables represent how much to increase the ptr by when moving by 1 - # element in a particular dimension. E.g. stride_am is how much to increase a_ptr - # by to get the element one row down (A has M rows) - stride_cm, - # stride_cn, # Assume that stride_cn == 1 - stride_am, - stride_ak, - stride_bk, - stride_bn, - # Meta-parameters - BLOCK_M: tl.constexpr, - GROUP_M: tl.constexpr, - BLOCK_N: tl.constexpr, - BLOCK_K: tl.constexpr, - # split k not used, not performant with activation, kept because early_config_prune is expecting it - SPLIT_K: tl.constexpr, - EVEN_K: tl.constexpr, - ACTIVATION: tl.constexpr, -): - - """ - Kernel for computing Out = activation(A x W + C) - - Input has shape (M, K) - - Weight has shape (K, N) - - Output has shape (M, N) - - ActInputs (optional) has shape (M, N) - 'ActInputs' optionally saves the A x W + C intermediate for backward computations - This kernel will consolidate over K - """ - - pid = tl.program_id(axis=0) - - grid_m = (M + BLOCK_M - 1) // BLOCK_M - grid_n = (N + BLOCK_N - 1) // BLOCK_N - # re-order program ID for better L2 performance - width = GROUP_M * grid_n - group_id = pid // width - group_size = min(grid_m - group_id * GROUP_M, GROUP_M) - pid_m = group_id * GROUP_M + (pid % group_size) - pid_n = (pid % width) // (group_size) - - # now compute the block that each program will go through - # rm (resp. rn) denotes a range of indices - # for rows (resp. col) of C - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - # trick to avoid masking on M and N axis - ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M) - rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N) - rk = tl.arange(0, BLOCK_K) - - A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak) - B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn) - - acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) - - for k in range(K, 0, -BLOCK_K): - if EVEN_K: - a = tl.load(A) - b = tl.load(B) - else: - a = tl.load(A, mask=rk[None, :] < k, other=0.0) - b = tl.load(B, mask=rk[:, None] < k, other=0.0) - acc += tl.dot(a, b) - - A += BLOCK_K * stride_ak - B += BLOCK_K * stride_bk - - # optional: fused activation (while the data is in shared memory) - if ACTIVATION != "id": - act_in_ptrs = ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :] - act_input = tl.load(act_in_ptrs).to(acc.dtype) - if ACTIVATION == "gelu": - acc *= gelu_grad(act_input) - elif ACTIVATION == "gelu_approx": - acc *= gelu_approx_grad(act_input) - elif ACTIVATION == "squared_relu": - acc *= squared_relu_grad(act_input) - - # rematerialize rm and rn to save registers - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - - # write back result - C = C + rm[:, None] * stride_cm + rn[None, :] - mask = (rm < M)[:, None] & (rn < N)[None, :] - tl.store(C, acc, mask=mask) - - -def triton_dgrad_act( - grad_output: torch.Tensor, - weight: torch.Tensor, - activation: str = "id", - act_input: Optional[torch.Tensor] = None, -) -> torch.Tensor: - """ - Compute e = activation(grad_output @ weight + bias). - This wrapper kicks the `kernel_fwd` Triton kernel - :param grad_output: input tensor - :param weight: weight matrix - :param activation: Activation name. Needs to be a Triton kernel. - :param act_input: an optional tensor to save the activation inputs (for backward) - :return: result tensor - """ - assert activation in ["id", "gelu", "gelu_approx", "squared_relu"] - - batch_shape, n = grad_output.shape[:-1], grad_output.shape[-1] - batch_dim = batch_shape.numel() - grad_output_reshaped = grad_output.reshape(batch_dim, n) - - if grad_output_reshaped.stride(0) > 1 and grad_output_reshaped.stride(1) > 1: - grad_output_reshaped = grad_output_reshaped.contiguous() - if weight.stride(0) > 1 and weight.stride(1) > 1: - weight = weight.contiguous() - - assert ( - grad_output.dtype == weight.dtype - ), f"grad_output and weight must have the same dtype, got {grad_output.dtype} and {weight.dtype}" - assert ( - grad_output_reshaped.shape[1] == weight.shape[0] - ), f"Incompatible dimensions: {grad_output_reshaped.shape} - {weight.shape}" - if activation != "id": - assert act_input is not None, f"act_input is required for activation {activation}" - - # M, N, K in bwd are different from M, N, K in fwd - M, K = grad_output_reshaped.shape - K, N = weight.shape - - grad_input = torch.empty((M, N), device=grad_output.device, dtype=grad_output.dtype) - - # 1D launch kernel where each block gets its own program. - grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),) # noqa - - kernel_bwd[grid]( - grad_input, - act_input, - grad_output_reshaped, - weight, # data ptrs - M, # shapes - N, - K, - M // 32, # key for triton cache (limit number of compilations) - N // 32, - K // 32, - stride_cm=grad_input.stride(0), # strides - # stride_cn=grad_input.stride(1), - stride_am=grad_output_reshaped.stride(0), - stride_ak=grad_output_reshaped.stride(1), - stride_bk=weight.stride(0), - stride_bn=weight.stride(1), - ACTIVATION=activation, # optional fused activation - GROUP_M=8, # speed optimization: group the programs - ) - - return grad_input.reshape(*batch_shape, grad_input.shape[-1]) diff --git a/build/torch210-cu130-x86_64-windows/ops/triton/mlp.py b/build/torch210-cu130-x86_64-windows/ops/triton/mlp.py deleted file mode 100644 index ae7db36e3f1e79daca3dc8588ef891d904d79b35..0000000000000000000000000000000000000000 --- a/build/torch210-cu130-x86_64-windows/ops/triton/mlp.py +++ /dev/null @@ -1,149 +0,0 @@ -# The triton fused matmul + sqrelu is faster for fp16 but slower for bf16, compared -# to naive implementation. -import fused_dense_lib as fused_dense_cuda -import torch -import torch.nn as nn -import torch.nn.functional as F - -from flash_attn.utils.torch import custom_fwd, custom_bwd -from flash_attn.ops.activations import sqrelu_bwd, sqrelu_fwd -from flash_attn.ops.triton.linear import triton_dgrad_act, triton_linear_act - - -class FusedDenseSqreluDenseFunc(torch.autograd.Function): - @staticmethod - @custom_fwd - def forward(ctx, x, weight1, bias1, weight2, bias2, checkpoint_lvl=0): - """checkpoint_lvl: - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute act_input and gelu_out in the bwd - """ - if torch.is_autocast_enabled(): - dtype = torch.get_autocast_gpu_dtype() - x, weight1, bias1, weight2, bias2 = [ - a.to(dtype=dtype) for a in [x, weight1, bias1, weight2, bias2] - ] - is_bf16 = x.dtype == torch.bfloat16 - assert checkpoint_lvl in [0, 1, 2] - x = x.contiguous() - weight1 = weight1.contiguous() - bias1 = bias1.contiguous() - weight2 = weight2.contiguous() - bias2 = bias2.contiguous() - batch_shape, n = x.shape[:-1], x.shape[-1] - batch_dim = batch_shape.numel() - if is_bf16: - act_input = fused_dense_cuda.linear_bias_forward( - x.reshape(batch_dim, n), weight1, bias1 - ) - output1 = sqrelu_fwd(act_input) - else: - save_act_input = checkpoint_lvl != 2 - result = triton_linear_act( - x.reshape(batch_dim, n), - weight1, - bias1, - activation="squared_relu", - save_act_input=save_act_input, - ) - if save_act_input: - output1, act_input = result - else: - output1 = result - output2 = fused_dense_cuda.linear_bias_forward(output1, weight2, bias2) - ctx.checkpoint_lvl = checkpoint_lvl - if checkpoint_lvl == 0: - ctx.save_for_backward(x, weight1, bias1, weight2, act_input, output1) - elif checkpoint_lvl == 1: - ctx.save_for_backward(x, weight1, bias1, weight2, act_input) - elif checkpoint_lvl == 2: - ctx.save_for_backward(x, weight1, bias1, weight2) - return output2.reshape(*batch_shape, output2.shape[-1]) - - @staticmethod - @custom_bwd - def backward(ctx, grad_output): - grad_output = grad_output.contiguous() - checkpoint_lvl = ctx.checkpoint_lvl - x, weight1, bias1, weight2, *rest = ctx.saved_tensors - batch_shape, n = x.shape[:-1], x.shape[-1] - batch_dim = batch_shape.numel() - is_bf16 = x.dtype == torch.bfloat16 - if checkpoint_lvl == 0: - act_input, output1 = rest - elif checkpoint_lvl == 1: - (act_input,) = rest - output1 = sqrelu_fwd(act_input) - elif checkpoint_lvl == 2: - if is_bf16: - act_input = fused_dense_cuda.linear_bias_forward( - x.reshape(batch_dim, n), weight1, bias1 - ) - output1 = sqrelu_fwd(act_input) - else: - output1, act_input = triton_linear_act( - x.reshape(batch_dim, n), - weight1, - bias1, - activation="squared_relu", - save_act_input=True, - ) - - if is_bf16: - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad(output1, grad_output) - grad_output1 = grad_output @ weight2 - grad_act_input = sqrelu_bwd(grad_output1, act_input) - grad_input, grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_backward( - x.reshape(batch_dim, n), weight1, grad_act_input - ) - else: - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad(output1, grad_output) - grad_act_input = triton_dgrad_act( - grad_output, weight2, activation="squared_relu", act_input=act_input - ) - grad_input, grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_backward( - x.reshape(batch_dim, n), weight1, grad_act_input - ) - return grad_input.reshape_as(x), grad_weight1, grad_bias1, grad_weight2, grad_bias2, None - - -fused_dense_sqrelu_dense_function = FusedDenseSqreluDenseFunc.apply - - -class FusedDenseSqreluDense(nn.Module): - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - bias1=True, - bias2=True, - checkpoint_lvl=0, - device=None, - dtype=None, - ): - """ - checkpoint_lvl (increasing lvl means slower but more memory saving): - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute gelu_in and gelu_out in the bwd - """ - assert checkpoint_lvl in [0, 1, 2] - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features * 4 - assert bias1 == True, "DenseSqreluDense module without bias is currently not supported" - assert bias2 == True, "DenseSqreluDense module without bias is currently not supported" - self.checkpoint_lvl = checkpoint_lvl - self.fc1 = nn.Linear(in_features, hidden_features, bias=bias1, **factory_kwargs) - self.fc2 = nn.Linear(hidden_features, out_features, bias=bias2, **factory_kwargs) - - def forward(self, x): - assert x.is_cuda - return fused_dense_sqrelu_dense_function( - x, self.fc1.weight, self.fc1.bias, self.fc2.weight, self.fc2.bias, self.checkpoint_lvl - ) diff --git a/build/torch210-cu130-x86_64-windows/ops/triton/rotary.py b/build/torch210-cu130-x86_64-windows/ops/triton/rotary.py deleted file mode 100644 index ba1d164f4bf38d56e74169c8f1394e686d55271f..0000000000000000000000000000000000000000 --- a/build/torch210-cu130-x86_64-windows/ops/triton/rotary.py +++ /dev/null @@ -1,186 +0,0 @@ -# Copyright (c) 2025, Tri Dao. -# As of 2025-04-23, we require triton >= 3.0 - -from typing import Optional, Union - -import torch - -import triton -import triton.language as tl - - -@triton.jit -def rotary_kernel( - OUT, # Pointers to matrices - X, - COS, - SIN, - CU_SEQLENS, - SEQLEN_OFFSETS, # this could be int or a pointer - # Matrix dimensions - seqlen, - nheads, - seqlen_ro, - # strides - stride_out_batch, - stride_out_seqlen, - stride_out_nheads, - stride_out_headdim, - stride_x_batch, - stride_x_seqlen, - stride_x_nheads, - stride_x_headdim, - # Meta-parameters - # We want ROTARY_DIM to be constexpr, otherwise the triton compiler doesn't know that - # the mask is constant every 8 elements, and it will generate LDG.16 instead of LDG.128 - ROTARY_DIM: tl.constexpr, - IS_SEQLEN_OFFSETS_TENSOR: tl.constexpr, - IS_VARLEN: tl.constexpr, - INTERLEAVED: tl.constexpr, - CONJUGATE: tl.constexpr, - BLOCK_H: tl.constexpr, - BLOCK_M: tl.constexpr, -): - BLOCK_K: tl.constexpr = triton.next_power_of_2(ROTARY_DIM) - ROTARY_DIM_HALF = ROTARY_DIM // 2 - pid_head = tl.program_id(axis=0) - pid_m = tl.program_id(axis=1) - pid_batch = tl.program_id(axis=2) - - if not IS_VARLEN: - X = X + pid_batch * stride_x_batch - OUT = OUT + pid_batch * stride_out_batch - else: - start_idx = tl.load(CU_SEQLENS + pid_batch) - seqlen = tl.load(CU_SEQLENS + pid_batch + 1) - start_idx - X = X + start_idx * stride_x_seqlen - OUT = OUT + start_idx * stride_out_seqlen - - if pid_m * BLOCK_M >= seqlen: - return - - rh = pid_head * BLOCK_H + tl.arange(0, BLOCK_H) - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - if not IS_SEQLEN_OFFSETS_TENSOR: - rm_cs = rm + SEQLEN_OFFSETS - else: - rm_cs = rm + tl.load(SEQLEN_OFFSETS + pid_batch) - - rk_half = tl.arange(0, BLOCK_K // 2) - COS = COS + (rm_cs[:, None] * ROTARY_DIM_HALF + rk_half[None, :]) - SIN = SIN + (rm_cs[:, None] * ROTARY_DIM_HALF + rk_half[None, :]) - mask_cs = (rm_cs[:, None] < seqlen_ro) & (rk_half[None, :] < ROTARY_DIM_HALF) - cos = tl.load(COS, mask=mask_cs, other=1.0).to(tl.float32) - sin = tl.load(SIN, mask=mask_cs, other=0.0).to(tl.float32) - if CONJUGATE: - sin = -sin - - if not INTERLEAVED: - # Load the 1st and 2nd halves of X, do calculation, then store to 1st and 2nd halves of OUT - X = X + (rh[:, None, None] * stride_x_nheads + rm[None, :, None] * stride_x_seqlen + rk_half[None, None, :] * stride_x_headdim) - OUT = OUT + (rh[:, None, None] * stride_out_nheads + rm[None, :, None] * stride_out_seqlen + rk_half[None, None, :] * stride_out_headdim) - mask = (rh[:, None, None] < nheads) & (rm[None, :, None] < seqlen) & (rk_half[None, None, :] < ROTARY_DIM_HALF) - x0 = tl.load(X, mask=mask, other=0.0).to(tl.float32) - x1 = tl.load(X + ROTARY_DIM_HALF * stride_x_headdim, mask=mask, other=0.0,).to(tl.float32) - o0 = x0 * cos - x1 * sin - o1 = x0 * sin + x1 * cos - tl.store(OUT, o0, mask=mask) - tl.store(OUT + ROTARY_DIM_HALF * stride_out_headdim, o1, mask=mask) - else: - rk = tl.arange(0, BLOCK_K) - X = X + (rh[:, None, None] * stride_x_nheads + rm[None, :, None] * stride_x_seqlen + rk[None, None, :] * stride_x_headdim) - OUT = OUT + (rh[:, None, None] * stride_out_nheads + rm[None, :, None] * stride_out_seqlen + rk[None, None, :] * stride_out_headdim) - mask = (rh[:, None, None] < nheads) & (rm[None, :, None] < seqlen) & (rk[None, None, :] < ROTARY_DIM) - x = tl.load(X, mask=mask, other=0.0).to(tl.float32) - x0, x1 = tl.split(tl.reshape(x, [BLOCK_H, BLOCK_M, BLOCK_K // 2, 2])) - o0 = x0 * cos - x1 * sin - o1 = x0 * sin + x1 * cos - o = tl.reshape(tl.join(o0, o1), [BLOCK_H, BLOCK_M, BLOCK_K]) - tl.store(OUT, o, mask=mask) - - -def apply_rotary( - x: torch.Tensor, - cos: torch.Tensor, - sin: torch.Tensor, - seqlen_offsets: Union[int, torch.Tensor] = 0, - cu_seqlens: Optional[torch.Tensor] = None, - max_seqlen: Optional[int] = None, - interleaved=False, - inplace=False, - conjugate=False, -) -> torch.Tensor: - """ - Arguments: - x: (batch, seqlen, nheads, headdim) if cu_seqlens is None - else (total_seqlen, nheads, headdim). - cos: (seqlen_ro, rotary_dim / 2) - sin: (seqlen_ro, rotary_dim / 2) - seqlen_offsets: integer or integer tensor of size (batch,) - cu_seqlens: (batch + 1,) or None - max_seqlen: int - Returns: - y: (batch, seqlen, nheads, headdim) - """ - is_varlen = cu_seqlens is not None - if not is_varlen: - batch, seqlen, nheads, headdim = x.shape - else: - assert max_seqlen is not None, "If cu_seqlens is passed in, then max_seqlen must be passed" - total_seqlen, nheads, headdim = x.shape - batch_p_1 = cu_seqlens.shape[0] - batch = batch_p_1 - 1 - seqlen = max_seqlen - seqlen_ro, rotary_dim = cos.shape - assert sin.shape == cos.shape - rotary_dim *= 2 - assert rotary_dim <= headdim, "rotary_dim must be <= headdim" - assert headdim <= 256, "Only support headdim <= 256" - assert seqlen_ro >= seqlen, "seqlen_ro must be >= seqlen" - - cos, sin = cos.contiguous(), sin.contiguous() - if isinstance(seqlen_offsets, torch.Tensor): - assert seqlen_offsets.shape == (batch,) - assert seqlen_offsets.dtype in [torch.int32, torch.int64] - seqlen_offsets = seqlen_offsets.contiguous() - else: - assert seqlen_offsets + seqlen <= seqlen_ro - - output = torch.empty_like(x) if not inplace else x - if rotary_dim < headdim and not inplace: - output[..., rotary_dim:].copy_(x[..., rotary_dim:]) - - grid = lambda META: (triton.cdiv(nheads, META["BLOCK_H"]), triton.cdiv(seqlen, META["BLOCK_M"]), batch) # noqa - BLOCK_M = 8 if rotary_dim <= 128 else 4 - - # Need this, otherwise Triton tries to launch from cuda:0 and we get - # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?) - device_ctx = torch.cuda.device(x.device.index) if x.device.type == 'cuda' else torch.xpu.device(x.device.index) - with device_ctx: - torch.library.wrap_triton(rotary_kernel)[grid]( - output, # data ptrs - x, - cos, - sin, - cu_seqlens, - seqlen_offsets, - seqlen, # shapes - nheads, - seqlen_ro, - output.stride(0) if not is_varlen else 0, # batch_strides if not varlen else 0 - output.stride(-3), # seqlen_stride or total_seqlen_stride - output.stride(-2), # nheads_stride - output.stride(-1), # headdim_stride - x.stride(0) if not is_varlen else 0, # batch_strides if not varlen else 0 - x.stride(-3), # seqlen stride or total_seqlen_stride - x.stride(-2), # nheads stride - x.stride(-1), # headdim stride - rotary_dim, - isinstance(seqlen_offsets, torch.Tensor), - is_varlen, - interleaved, - conjugate, - BLOCK_M=BLOCK_M, - BLOCK_H=2, - ) - return output diff --git a/build/torch210-cxx11-cpu-x86_64-linux/__init__.py b/build/torch210-cxx11-cpu-x86_64-linux/__init__.py index 20ff0ff747abd7500913863b509925f414b1c973..ecc2f9d896b6c93f90b0a1499856dc0612177422 100644 --- a/build/torch210-cxx11-cpu-x86_64-linux/__init__.py +++ b/build/torch210-cxx11-cpu-x86_64-linux/__init__.py @@ -1,7 +1,5 @@ -from typing import List, Optional - +from typing import Optional, List import torch - from ._ops import ops as flash_attn_ops from .flash_attn_interface import ( flash_attn_func, @@ -13,23 +11,6 @@ from .flash_attn_interface import ( flash_attn_with_kvcache, ) -__all__ = [ - # High-level API (with autograd support) - "flash_attn_func", - "flash_attn_kvpacked_func", - "flash_attn_qkvpacked_func", - "flash_attn_varlen_func", - "flash_attn_varlen_kvpacked_func", - "flash_attn_varlen_qkvpacked_func", - "flash_attn_with_kvcache", - # Low-level ops (no autograd) - "fwd", - "varlen_fwd", - "bwd", - "varlen_bwd", - "fwd_kvcache", -] - def fwd( q: torch.Tensor, diff --git a/build/torch210-cxx11-cpu-x86_64-linux/_flash_attn2_588b404.abi3.so b/build/torch210-cxx11-cpu-x86_64-linux/_flash_attn2_588b404.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..d5fe2827907a352eb98f8e962de84a3c3696e2a3 --- /dev/null +++ b/build/torch210-cxx11-cpu-x86_64-linux/_flash_attn2_588b404.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:1d90d30dbcf574c7a50f2c9774884370e71e1e177062c6a233fcc7e1940fffcb +size 249504 diff --git a/build/torch210-cxx11-cpu-x86_64-linux/_flash_attn2_cpu_a511a4c.abi3.so b/build/torch210-cxx11-cpu-x86_64-linux/_flash_attn2_cpu_a511a4c.abi3.so deleted file mode 100644 index 5c25e37d8977b2c26a535ffe5888f4b23991e8d6..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cpu-x86_64-linux/_flash_attn2_cpu_a511a4c.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:145c40225f22ea012bbb986a80a1cbe08527ada0e26b5bbabff3952ed2bb58d0 -size 1942240 diff --git a/build/torch210-cxx11-cpu-x86_64-linux/_ops.py b/build/torch210-cxx11-cpu-x86_64-linux/_ops.py index 1d8ba74be8790a2dd7144aadaf20a5203b85810d..03506d8a18d17ca7296a5f1b7a18ec7dc4ad73bb 100644 --- a/build/torch210-cxx11-cpu-x86_64-linux/_ops.py +++ b/build/torch210-cxx11-cpu-x86_64-linux/_ops.py @@ -1,9 +1,9 @@ import torch -from . import _flash_attn2_cpu_a511a4c -ops = torch.ops._flash_attn2_cpu_a511a4c +from . import _flash_attn2_588b404 +ops = torch.ops._flash_attn2_588b404 def add_op_namespace_prefix(op_name: str): """ Prefix op by namespace. """ - return f"_flash_attn2_cpu_a511a4c::{op_name}" + return f"_flash_attn2_588b404::{op_name}" \ No newline at end of file diff --git a/build/torch210-cxx11-cpu-x86_64-linux/flash_attn2/__init__.py b/build/torch210-cxx11-cpu-x86_64-linux/flash_attn2/__init__.py index a9b2672c1cd85b74c1b3ded0fc0b2100e1aeac23..03dbc1afe1cf156661a2b1b22003cd5f599a0309 100644 --- a/build/torch210-cxx11-cpu-x86_64-linux/flash_attn2/__init__.py +++ b/build/torch210-cxx11-cpu-x86_64-linux/flash_attn2/__init__.py @@ -1,10 +1,10 @@ import ctypes -import importlib.util 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 diff --git a/build/torch210-cxx11-cpu-x86_64-linux/flash_attn_interface.py b/build/torch210-cxx11-cpu-x86_64-linux/flash_attn_interface.py index 0868572aa3324d9c96132859d123e3d48d296e34..66e8d3665834b141d703998496b3da6fe1f6fe06 100644 --- a/build/torch210-cxx11-cpu-x86_64-linux/flash_attn_interface.py +++ b/build/torch210-cxx11-cpu-x86_64-linux/flash_attn_interface.py @@ -16,11 +16,9 @@ import os from ._ops import ops as flash_attn -from ._ops import add_op_namespace_prefix # # isort: on - def maybe_contiguous(x): return x.contiguous() if x is not None and x.stride(-1) != 1 else x @@ -33,25 +31,14 @@ def _get_device(): else: return "cpu" +_XPU_AVAILABLE = torch.xpu.is_available() if hasattr(torch, "xpu") else False # TODO remove hasattr check when bwd is supported on XPU -def _get_block_size_n(device, head_dim, is_dropout, is_causal): - assert head_dim <= 256 - - if device.type == "xpu": - if head_dim <= 96: - return 64 - elif head_dim <= 128: - return 32 - elif head_dim <= 256: - return 64 - else: - return 32 +def _get_block_size_n(device, head_dim, is_dropout, is_causal): # This should match the block sizes in the CUDA kernel + assert head_dim <= 256 major, minor = torch.cuda.get_device_capability(device) - is_sm8x = ( - major == 8 and minor > 0 - ) # Only include sm86 and sm89, exclude sm80 (A100) + is_sm8x = major == 8 and minor > 0 # Only include sm86 and sm89, exclude sm80 (A100) is_sm80 = major == 8 and minor == 0 is_sm90 = major == 9 and minor == 0 if head_dim <= 32: @@ -77,11 +64,30 @@ def round_multiple(x, m): return (x + m - 1) // m * m -@torch.library.custom_op( - add_op_namespace_prefix("_flash_attn_forward"), - mutates_args=(), - device_types=_get_device(), -) +# torch.compile() support is only enabled for pytorch >= 2.4 +# The reason for this is that we are using the new custom_op and register_fake +# APIs, which support inplace modification of inputs in the function itself +if torch.__version__ >= "2.4.0": + _torch_custom_op_wrapper = torch.library.custom_op + _torch_register_fake_wrapper = torch.library.register_fake +else: + def noop_custom_op_wrapper(name, fn=None, /, *, mutates_args, device_types=None, schema=None): + def wrap(func): + return func + if fn is None: + return wrap + return fn + def noop_register_fake_wrapper(op, fn=None, /, *, lib=None, _stacklevel=1): + def wrap(func): + return func + if fn is None: + return wrap + return fn + _torch_custom_op_wrapper = noop_custom_op_wrapper + _torch_register_fake_wrapper = noop_register_fake_wrapper + + +@_torch_custom_op_wrapper("flash_attn::_flash_attn_forward", mutates_args=(), device_types=_get_device()) def _flash_attn_forward( q: torch.Tensor, k: torch.Tensor, @@ -93,7 +99,7 @@ def _flash_attn_forward( window_size_right: int, softcap: float, alibi_slopes: Optional[torch.Tensor], - return_softmax: bool, + return_softmax: bool ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: q, k, v = [maybe_contiguous(x) for x in (q, k, v)] out, softmax_lse, S_dmask, rng_state = flash_attn.fwd( @@ -114,7 +120,7 @@ def _flash_attn_forward( return out, softmax_lse, S_dmask, rng_state -@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_forward")) +@_torch_register_fake_wrapper("flash_attn::_flash_attn_forward") def _flash_attn_forward_fake( q: torch.Tensor, k: torch.Tensor, @@ -126,41 +132,28 @@ def _flash_attn_forward_fake( window_size_right: int, softcap: float, alibi_slopes: Optional[torch.Tensor], - return_softmax: bool, + return_softmax: bool ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: q, k, v = [maybe_contiguous(x) for x in (q, k, v)] batch_size, seqlen_q, num_heads, head_size = q.shape seqlen_k = k.shape[1] out = torch.empty_like(q) - softmax_lse = torch.empty( - (batch_size, num_heads, seqlen_q), - dtype=torch.float32, - device=q.device, - layout=q.layout, - ) + softmax_lse = torch.empty((batch_size, num_heads, seqlen_q), dtype=torch.float32, device=q.device, layout=q.layout) p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) if return_softmax: - p = torch.empty( - ( - batch_size, - num_heads, - round_multiple(seqlen_q, 128), - round_multiple(seqlen_k, 128), - ), - dtype=q.dtype, - device=q.device, - layout=q.layout, - ) + p = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128), round_multiple(seqlen_k, 128)), dtype=q.dtype, device=q.device, layout=q.layout) rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) return out, softmax_lse, p, rng_state -@torch.library.custom_op( - add_op_namespace_prefix("_flash_attn_varlen_forward"), - mutates_args=(), - device_types=_get_device(), -) +if torch.__version__ >= "2.4.0": + _wrapped_flash_attn_forward = torch.ops.flash_attn._flash_attn_forward +else: + _wrapped_flash_attn_forward = _flash_attn_forward + + +@_torch_custom_op_wrapper("flash_attn::_flash_attn_varlen_forward", mutates_args=(), device_types=_get_device()) def _flash_attn_varlen_forward( q: torch.Tensor, k: torch.Tensor, @@ -211,7 +204,7 @@ def _flash_attn_varlen_forward( return out, softmax_lse, S_dmask, rng_state -@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_varlen_forward")) +@_torch_register_fake_wrapper("flash_attn::_flash_attn_varlen_forward") def _flash_attn_varlen_forward_fake( q: torch.Tensor, k: torch.Tensor, @@ -237,30 +230,25 @@ def _flash_attn_varlen_forward_fake( paged_kv = block_table is not None batch_size = cu_seqlens_q.numel() - 1 total_q, num_heads, _ = q.shape - + out = torch.empty_like(q) - softmax_lse = torch.empty( - (num_heads, total_q), dtype=torch.float32, device=q.device, layout=q.layout - ) + softmax_lse = torch.empty((num_heads, total_q), dtype=torch.float32, device=q.device, layout=q.layout) p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) seqlen_q_rounded = round_multiple(max_seqlen_q, 128) seqlen_k_rounded = round_multiple(max_seqlen_k, 128) if return_softmax: - p = torch.empty( - (batch_size, num_heads, seqlen_q_rounded, seqlen_k_rounded), - dtype=q.dtype, - device=q.device, - layout=q.layout, - ) + p = torch.empty((batch_size, num_heads, seqlen_q_rounded, seqlen_k_rounded), dtype=q.dtype, device=q.device, layout=q.layout) rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) return out, softmax_lse, p, rng_state -@torch.library.custom_op( - add_op_namespace_prefix("_flash_attn_backward"), - mutates_args=("dq", "dk", "dv"), - device_types=_get_device(), -) +if torch.__version__ >= "2.4.0": + _wrapped_flash_attn_varlen_forward = torch.ops.flash_attn._flash_attn_varlen_forward +else: + _wrapped_flash_attn_varlen_forward = _flash_attn_varlen_forward + + +@_torch_custom_op_wrapper("flash_attn::_flash_attn_backward", mutates_args=("dq", "dk", "dv"), device_types=_get_device()) def _flash_attn_backward( dout: torch.Tensor, q: torch.Tensor, @@ -312,7 +300,7 @@ def _flash_attn_backward( return softmax_d -@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_backward")) +@_torch_register_fake_wrapper("flash_attn::_flash_attn_backward") def _flash_attn_backward_fake( dout: torch.Tensor, q: torch.Tensor, @@ -341,20 +329,18 @@ def _flash_attn_backward_fake( if dv is None: dv = torch.empty_like(v) batch_size, seqlen_q, num_heads, _ = q.shape - softmax_d = torch.empty( - (batch_size, num_heads, round_multiple(seqlen_q, 128)), - device=q.device, - dtype=torch.float32, - ) - + softmax_d = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128)), device=q.device, dtype=torch.float32) + return softmax_d -@torch.library.custom_op( - add_op_namespace_prefix("_flash_attn_varlen_backward"), - mutates_args=("dq", "dk", "dv"), - device_types=_get_device(), -) +if torch.__version__ >= "2.4.0": + _wrapped_flash_attn_backward = torch.ops.flash_attn._flash_attn_backward +else: + _wrapped_flash_attn_backward = _flash_attn_backward + + +@_torch_custom_op_wrapper("flash_attn::_flash_attn_varlen_backward", mutates_args=("dq", "dk", "dv"), device_types=_get_device()) def _flash_attn_varlen_backward( dout: torch.Tensor, q: torch.Tensor, @@ -418,7 +404,7 @@ def _flash_attn_varlen_backward( return softmax_d -@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_varlen_backward")) +@_torch_register_fake_wrapper("flash_attn::_flash_attn_varlen_backward") def _flash_attn_varlen_backward_fake( dout: torch.Tensor, q: torch.Tensor, @@ -454,13 +440,17 @@ def _flash_attn_varlen_backward_fake( dk = torch.empty_like(k) if dv is None: dv = torch.empty_like(v) - softmax_d = torch.empty( - (num_heads, total_q + 128 * batch_size), device=q.device, dtype=torch.float32 - ) - + softmax_d = torch.empty((num_heads, total_q + 128 * batch_size), device=q.device, dtype=torch.float32) + return softmax_d +if torch.__version__ >= "2.4.0": + _wrapped_flash_attn_varlen_backward = torch.ops.flash_attn._flash_attn_varlen_backward +else: + _wrapped_flash_attn_varlen_backward = _flash_attn_varlen_backward + + class FlashAttnQKVPackedFunc(torch.autograd.Function): @staticmethod def forward( @@ -485,7 +475,7 @@ class FlashAttnQKVPackedFunc(torch.autograd.Function): q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_forward( + out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( q, k, v, @@ -519,7 +509,7 @@ class FlashAttnQKVPackedFunc(torch.autograd.Function): dout_padded = dout if head_size_og % 8 != 0: dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _flash_attn_backward( + _wrapped_flash_attn_backward( dout_padded, q, k, @@ -569,7 +559,7 @@ class FlashAttnVarlenQKVPackedFunc(torch.autograd.Function): q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_varlen_forward( + out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( q, k, v, @@ -588,9 +578,7 @@ class FlashAttnVarlenQKVPackedFunc(torch.autograd.Function): block_table=None, ) if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens, rng_state - ) + ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens, rng_state) ctx.dropout_p = dropout_p ctx.max_seqlen = max_seqlen ctx.softmax_scale = softmax_scale @@ -611,7 +599,7 @@ class FlashAttnVarlenQKVPackedFunc(torch.autograd.Function): dout_padded = dout if head_size_og % 8 != 0: dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _flash_attn_varlen_backward( + _wrapped_flash_attn_varlen_backward( dout_padded, q, k, @@ -655,7 +643,9 @@ class FlashAttnKVPackedFunc(torch.autograd.Function): return_softmax, is_grad_enabled, ): - is_grad = is_grad_enabled and any(x.requires_grad for x in [q, kv]) + is_grad = is_grad_enabled and any( + x.requires_grad for x in [q, kv] + ) if softmax_scale is None: softmax_scale = q.shape[-1] ** (-0.5) k, v = kv[:, :, 0].detach(), kv[:, :, 1].detach() @@ -664,7 +654,7 @@ class FlashAttnKVPackedFunc(torch.autograd.Function): q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_forward( + out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( q, k, v, @@ -699,7 +689,7 @@ class FlashAttnKVPackedFunc(torch.autograd.Function): dout_padded = dout if head_size_og % 8 != 0: dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _flash_attn_backward( + _wrapped_flash_attn_backward( dout_padded, q, k, @@ -744,7 +734,9 @@ class FlashAttnVarlenKVPackedFunc(torch.autograd.Function): return_softmax, is_grad_enabled, ): - is_grad = is_grad_enabled and any(x.requires_grad for x in [q, kv]) + is_grad = is_grad_enabled and any( + x.requires_grad for x in [q, kv] + ) if softmax_scale is None: softmax_scale = q.shape[-1] ** (-0.5) k, v = kv[:, 0].detach(), kv[:, 1].detach() @@ -753,7 +745,7 @@ class FlashAttnVarlenKVPackedFunc(torch.autograd.Function): q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_varlen_forward( + out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( q, k, v, @@ -789,9 +781,7 @@ class FlashAttnVarlenKVPackedFunc(torch.autograd.Function): @staticmethod def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ( - ctx.saved_tensors - ) + q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors dq = torch.empty_like(q) kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) @@ -799,7 +789,7 @@ class FlashAttnVarlenKVPackedFunc(torch.autograd.Function): dout_padded = dout if head_size_og % 8 != 0: dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _flash_attn_varlen_backward( + _wrapped_flash_attn_varlen_backward( dout_padded, q, k, @@ -825,23 +815,7 @@ class FlashAttnVarlenKVPackedFunc(torch.autograd.Function): ) dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension dkv = dkv[..., : dout.shape[-1]] - return ( - dq, - dkv, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) + return dq, dkv, None, None, None, None, None, None, None, None, None, None, None, None, None class FlashAttnFunc(torch.autograd.Function): @@ -861,7 +835,9 @@ class FlashAttnFunc(torch.autograd.Function): return_softmax, is_grad_enabled, ): - is_grad = is_grad_enabled and any(x.requires_grad for x in [q, k, v]) + is_grad = is_grad_enabled and any( + x.requires_grad for x in [q, k, v] + ) if softmax_scale is None: softmax_scale = q.shape[-1] ** (-0.5) head_size_og = q.size(3) @@ -869,7 +845,7 @@ class FlashAttnFunc(torch.autograd.Function): q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_forward( + out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( q, k, v, @@ -902,7 +878,7 @@ class FlashAttnFunc(torch.autograd.Function): dout_padded = dout if head_size_og % 8 != 0: dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _flash_attn_backward( + _wrapped_flash_attn_backward( dout_padded, q, k, @@ -950,7 +926,9 @@ class FlashAttnVarlenFunc(torch.autograd.Function): block_table, is_grad_enabled, ): - is_grad = is_grad_enabled and any(x.requires_grad for x in [q, k, v]) + is_grad = is_grad_enabled and any( + x.requires_grad for x in [q, k, v] + ) if softmax_scale is None: softmax_scale = q.shape[-1] ** (-0.5) head_size_og = q.size(2) @@ -958,7 +936,7 @@ class FlashAttnVarlenFunc(torch.autograd.Function): q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_varlen_forward( + out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( q, k, v, @@ -995,15 +973,13 @@ class FlashAttnVarlenFunc(torch.autograd.Function): @staticmethod def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ( - ctx.saved_tensors - ) + q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) head_size_og = dout.size(2) dout_padded = dout if head_size_og % 8 != 0: dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _flash_attn_varlen_backward( + _wrapped_flash_attn_varlen_backward( dout_padded, q, k, @@ -1030,25 +1006,7 @@ class FlashAttnVarlenFunc(torch.autograd.Function): dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension dk = dk[..., : dout.shape[-1]] dv = dv[..., : dout.shape[-1]] - return ( - dq, - dk, - dv, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None def flash_attn_qkvpacked_func( @@ -1106,7 +1064,7 @@ def flash_attn_qkvpacked_func( alibi_slopes, deterministic, return_attn_probs, - torch.is_grad_enabled(), + False if _XPU_AVAILABLE else torch.is_grad_enabled(), ) @@ -1184,7 +1142,7 @@ def flash_attn_kvpacked_func( alibi_slopes, deterministic, return_attn_probs, - torch.is_grad_enabled(), + False if _XPU_AVAILABLE else torch.is_grad_enabled(), ) @@ -1196,7 +1154,7 @@ def flash_attn_func( softmax_scale=None, causal=False, window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated + softcap=0.0, # 0.0 means deactivated alibi_slopes=None, deterministic=False, return_attn_probs=False, @@ -1261,7 +1219,7 @@ def flash_attn_func( alibi_slopes, deterministic, return_attn_probs, - torch.is_grad_enabled(), + False if _XPU_AVAILABLE else torch.is_grad_enabled(), ) @@ -1273,7 +1231,7 @@ def flash_attn_varlen_qkvpacked_func( softmax_scale=None, causal=False, window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated + softcap=0.0, # 0.0 means deactivated alibi_slopes=None, deterministic=False, return_attn_probs=False, @@ -1327,7 +1285,7 @@ def flash_attn_varlen_qkvpacked_func( alibi_slopes, deterministic, return_attn_probs, - torch.is_grad_enabled(), + False if _XPU_AVAILABLE else torch.is_grad_enabled(), ) @@ -1342,7 +1300,7 @@ def flash_attn_varlen_kvpacked_func( softmax_scale=None, causal=False, window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated + softcap=0.0, # 0.0 means deactivated alibi_slopes=None, deterministic=False, return_attn_probs=False, @@ -1419,7 +1377,7 @@ def flash_attn_varlen_kvpacked_func( alibi_slopes, deterministic, return_attn_probs, - torch.is_grad_enabled(), + False if _XPU_AVAILABLE else torch.is_grad_enabled(), ) @@ -1435,7 +1393,7 @@ def flash_attn_varlen_func( softmax_scale=None, causal=False, window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated + softcap=0.0, # 0.0 means deactivated alibi_slopes=None, deterministic=False, return_attn_probs=False, @@ -1513,7 +1471,7 @@ def flash_attn_varlen_func( deterministic, return_attn_probs, block_table, - False if q.device.type == "cpu" else torch.is_grad_enabled(), + False if _XPU_AVAILABLE or q.device.type == "cpu" else torch.is_grad_enabled(), ) @@ -1532,7 +1490,7 @@ def flash_attn_with_kvcache( softmax_scale=None, causal=False, window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated + softcap=0.0, # 0.0 means deactivated rotary_interleaved=True, alibi_slopes=None, num_splits=0, diff --git a/build/torch210-cu128-x86_64-windows/layers/__init__.py b/build/torch210-cxx11-cpu-x86_64-linux/layers/__init__.py similarity index 100% rename from build/torch210-cu128-x86_64-windows/layers/__init__.py rename to build/torch210-cxx11-cpu-x86_64-linux/layers/__init__.py diff --git a/build/torch211-cxx11-cpu-x86_64-linux/layers/patch_embed.py b/build/torch210-cxx11-cpu-x86_64-linux/layers/patch_embed.py similarity index 100% rename from build/torch211-cxx11-cpu-x86_64-linux/layers/patch_embed.py rename to build/torch210-cxx11-cpu-x86_64-linux/layers/patch_embed.py diff --git a/build/torch211-cxx11-cpu-x86_64-linux/layers/rotary.py b/build/torch210-cxx11-cpu-x86_64-linux/layers/rotary.py similarity index 100% rename from build/torch211-cxx11-cpu-x86_64-linux/layers/rotary.py rename to build/torch210-cxx11-cpu-x86_64-linux/layers/rotary.py diff --git a/build/torch210-cxx11-cpu-x86_64-linux/metadata.json b/build/torch210-cxx11-cpu-x86_64-linux/metadata.json index 36fe50c2e205a669a17784861ef72c99418f9759..9cf5deed9898dce769f4cc73913d3530b92a0bd8 100644 --- a/build/torch210-cxx11-cpu-x86_64-linux/metadata.json +++ b/build/torch210-cxx11-cpu-x86_64-linux/metadata.json @@ -1,10 +1,4 @@ { - "name": "flash-attn2", - "id": "_flash_attn2_cpu_a511a4c", "version": 1, - "license": "BSD-3-Clause", - "python-depends": [], - "backend": { - "type": "cpu" - } -} + "python-depends": [] +} \ No newline at end of file diff --git a/build/torch210-cu128-x86_64-windows/ops/__init__.py b/build/torch210-cxx11-cpu-x86_64-linux/ops/__init__.py similarity index 100% rename from build/torch210-cu128-x86_64-windows/ops/__init__.py rename to build/torch210-cxx11-cpu-x86_64-linux/ops/__init__.py diff --git a/build/torch211-cxx11-cpu-x86_64-linux/ops/activations.py b/build/torch210-cxx11-cpu-x86_64-linux/ops/activations.py similarity index 100% rename from build/torch211-cxx11-cpu-x86_64-linux/ops/activations.py rename to build/torch210-cxx11-cpu-x86_64-linux/ops/activations.py diff --git a/build/torch211-cxx11-cpu-x86_64-linux/ops/fused_dense.py b/build/torch210-cxx11-cpu-x86_64-linux/ops/fused_dense.py similarity index 100% rename from build/torch211-cxx11-cpu-x86_64-linux/ops/fused_dense.py rename to build/torch210-cxx11-cpu-x86_64-linux/ops/fused_dense.py diff --git a/build/torch211-cxx11-cpu-x86_64-linux/ops/layer_norm.py b/build/torch210-cxx11-cpu-x86_64-linux/ops/layer_norm.py similarity index 100% rename from build/torch211-cxx11-cpu-x86_64-linux/ops/layer_norm.py rename to build/torch210-cxx11-cpu-x86_64-linux/ops/layer_norm.py diff --git a/build/torch211-cxx11-cpu-x86_64-linux/ops/rms_norm.py b/build/torch210-cxx11-cpu-x86_64-linux/ops/rms_norm.py similarity index 100% rename from build/torch211-cxx11-cpu-x86_64-linux/ops/rms_norm.py rename to build/torch210-cxx11-cpu-x86_64-linux/ops/rms_norm.py diff --git a/build/torch211-cxx11-cpu-x86_64-linux/ops/triton/__init__.py b/build/torch210-cxx11-cpu-x86_64-linux/ops/triton/__init__.py similarity index 100% rename from build/torch211-cxx11-cpu-x86_64-linux/ops/triton/__init__.py rename to build/torch210-cxx11-cpu-x86_64-linux/ops/triton/__init__.py diff --git a/build/torch211-cxx11-cpu-x86_64-linux/ops/triton/cross_entropy.py b/build/torch210-cxx11-cpu-x86_64-linux/ops/triton/cross_entropy.py similarity index 100% rename from build/torch211-cxx11-cpu-x86_64-linux/ops/triton/cross_entropy.py rename to build/torch210-cxx11-cpu-x86_64-linux/ops/triton/cross_entropy.py diff --git a/build/torch211-cxx11-cpu-x86_64-linux/ops/triton/k_activations.py b/build/torch210-cxx11-cpu-x86_64-linux/ops/triton/k_activations.py similarity index 100% rename from build/torch211-cxx11-cpu-x86_64-linux/ops/triton/k_activations.py rename to build/torch210-cxx11-cpu-x86_64-linux/ops/triton/k_activations.py diff --git a/build/torch29-cxx11-cu126-aarch64-linux/ops/triton/layer_norm.py b/build/torch210-cxx11-cpu-x86_64-linux/ops/triton/layer_norm.py similarity index 100% rename from build/torch29-cxx11-cu126-aarch64-linux/ops/triton/layer_norm.py rename to build/torch210-cxx11-cpu-x86_64-linux/ops/triton/layer_norm.py diff --git a/build/torch211-cxx11-cpu-x86_64-linux/ops/triton/linear.py b/build/torch210-cxx11-cpu-x86_64-linux/ops/triton/linear.py similarity index 100% rename from build/torch211-cxx11-cpu-x86_64-linux/ops/triton/linear.py rename to build/torch210-cxx11-cpu-x86_64-linux/ops/triton/linear.py diff --git a/build/torch211-cxx11-cpu-x86_64-linux/ops/triton/mlp.py b/build/torch210-cxx11-cpu-x86_64-linux/ops/triton/mlp.py similarity index 100% rename from build/torch211-cxx11-cpu-x86_64-linux/ops/triton/mlp.py rename to build/torch210-cxx11-cpu-x86_64-linux/ops/triton/mlp.py diff --git a/build/torch211-cxx11-cpu-x86_64-linux/ops/triton/rotary.py b/build/torch210-cxx11-cpu-x86_64-linux/ops/triton/rotary.py similarity index 100% rename from build/torch211-cxx11-cpu-x86_64-linux/ops/triton/rotary.py rename to build/torch210-cxx11-cpu-x86_64-linux/ops/triton/rotary.py diff --git a/build/torch210-cxx11-cu126-aarch64-linux/__init__.py b/build/torch210-cxx11-cu126-aarch64-linux/__init__.py deleted file mode 100644 index 20ff0ff747abd7500913863b509925f414b1c973..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu126-aarch64-linux/__init__.py +++ /dev/null @@ -1,412 +0,0 @@ -from typing import List, Optional - -import torch - -from ._ops import ops as flash_attn_ops -from .flash_attn_interface import ( - flash_attn_func, - flash_attn_kvpacked_func, - flash_attn_qkvpacked_func, - flash_attn_varlen_func, - flash_attn_varlen_kvpacked_func, - flash_attn_varlen_qkvpacked_func, - flash_attn_with_kvcache, -) - -__all__ = [ - # High-level API (with autograd support) - "flash_attn_func", - "flash_attn_kvpacked_func", - "flash_attn_qkvpacked_func", - "flash_attn_varlen_func", - "flash_attn_varlen_kvpacked_func", - "flash_attn_varlen_qkvpacked_func", - "flash_attn_with_kvcache", - # Low-level ops (no autograd) - "fwd", - "varlen_fwd", - "bwd", - "varlen_bwd", - "fwd_kvcache", -] - - -def fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - return_softmax: bool = False, - gen: Optional[torch.Generator] = None, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention. - - Args: - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Optional output tensor, same shape as q - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - return_softmax: Whether to return softmax weights - gen: Optional random number generator - - Returns: - List of tensors: [output, softmax_lse, (softmax if return_softmax)] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.fwd( - q, - k, - v, - out, - alibi_slopes, - p_dropout, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - gen, - ) - - -def varlen_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - out: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - block_table: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - max_seqlen_q: int = 0, - max_seqlen_k: int = 0, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - zero_tensors: bool = False, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - return_softmax: bool = False, - gen: Optional[torch.Generator] = None, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention with variable sequence lengths. - - Args: - q: Query tensor of shape [total_q, num_heads, head_size] - k: Key tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - v: Value tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1] - cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1] - out: Optional output tensor of shape [total_q, num_heads, head_size] - seqused_k: Optional tensor specifying how many keys to use per batch element [batch_size] - leftpad_k: Optional left padding for keys of shape [batch_size] - block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq] - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - max_seqlen_q: Maximum sequence length for queries - max_seqlen_k: Maximum sequence length for keys - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - zero_tensors: Whether to zero tensors before computation - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - return_softmax: Whether to return softmax weights - gen: Optional random number generator - - Returns: - List of tensors: [output, softmax_lse, (softmax if return_softmax)] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.varlen_fwd( - q, - k, - v, - out, - cu_seqlens_q, - cu_seqlens_k, - seqused_k, - leftpad_k, - block_table, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - p_dropout, - softmax_scale, - zero_tensors, - is_causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - gen, - ) - - -def bwd( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - p_dropout: float = 0.0, - 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, - gen: Optional[torch.Generator] = None, - rng_state: Optional[torch.Tensor] = None, -) -> List[torch.Tensor]: - """ - Backward pass for multi-head attention. - - Args: - dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size] - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size] - softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q] - dq: Optional gradient tensor for queries, same shape as q - dk: Optional gradient tensor for keys, same shape as k - dv: Optional gradient tensor for values, same shape as v - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - deterministic: Whether to use deterministic algorithms - gen: Optional random number generator - rng_state: Optional RNG state from forward pass - - Returns: - List of tensors: [dq, dk, dv] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - alibi_slopes, - p_dropout, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - deterministic, - gen, - rng_state, - ) - - -def varlen_bwd( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - max_seqlen_q: int = 0, - max_seqlen_k: int = 0, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - zero_tensors: bool = False, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - deterministic: bool = False, - gen: Optional[torch.Generator] = None, - rng_state: Optional[torch.Tensor] = None, -) -> List[torch.Tensor]: - """ - Backward pass for multi-head attention with variable sequence lengths. - - Args: - dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size] - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size] - softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q] - cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1] - cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1] - dq: Optional gradient tensor for queries, same shape as q - dk: Optional gradient tensor for keys, same shape as k - dv: Optional gradient tensor for values, same shape as v - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - max_seqlen_q: Maximum sequence length for queries - max_seqlen_k: Maximum sequence length for keys - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - zero_tensors: Whether to zero tensors before computation - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - deterministic: Whether to use deterministic algorithms - gen: Optional random number generator - rng_state: Optional RNG state from forward pass - - Returns: - List of tensors: [dq, dk, dv] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.varlen_bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - p_dropout, - softmax_scale, - zero_tensors, - is_causal, - window_size_left, - window_size_right, - softcap, - deterministic, - gen, - rng_state, - ) - - -def fwd_kvcache( - q: torch.Tensor, - kcache: torch.Tensor, - vcache: torch.Tensor, - k: Optional[torch.Tensor] = None, - v: Optional[torch.Tensor] = None, - seqlens_k: Optional[torch.Tensor] = None, - rotary_cos: Optional[torch.Tensor] = None, - rotary_sin: Optional[torch.Tensor] = None, - cache_batch_idx: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - block_table: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - out: 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, - is_rotary_interleaved: bool = False, - num_splits: int = 1, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention with KV cache. - - Args: - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - kcache: Key cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - vcache: Value cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - k: Optional new keys tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size] - v: Optional new values tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size] - seqlens_k: Optional sequence lengths for keys of shape [batch_size] - rotary_cos: Optional rotary cosine tensor of shape [seqlen_ro, rotary_dim/2] - rotary_sin: Optional rotary sine tensor of shape [seqlen_ro, rotary_dim/2] - cache_batch_idx: Optional indices to index into the KV cache - leftpad_k: Optional left padding for keys of shape [batch_size] - block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq] - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - out: Optional output tensor, same shape as q - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - is_rotary_interleaved: Whether rotary embeddings are interleaved - num_splits: Number of splits for computation - - Returns: - List of tensors: [output, softmax_lse] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.fwd_kvcache( - q, - kcache, - vcache, - k, - v, - seqlens_k, - rotary_cos, - rotary_sin, - cache_batch_idx, - leftpad_k, - block_table, - alibi_slopes, - out, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - is_rotary_interleaved, - num_splits, - ) diff --git a/build/torch210-cxx11-cu126-aarch64-linux/_flash_attn2_cuda_a511a4c.abi3.so b/build/torch210-cxx11-cu126-aarch64-linux/_flash_attn2_cuda_a511a4c.abi3.so deleted file mode 100644 index 9775c8e4bf5c4527fe72038a5e82e668e62dcf6d..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu126-aarch64-linux/_flash_attn2_cuda_a511a4c.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:27af83157113a13459d2666b0ebfb6da491af2029b33b5343d5684cc886bee8d -size 448608936 diff --git a/build/torch210-cxx11-cu126-aarch64-linux/_ops.py b/build/torch210-cxx11-cu126-aarch64-linux/_ops.py deleted file mode 100644 index 62b70a9b02c8f7663c7022564c1a7f58f46b6c65..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu126-aarch64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _flash_attn2_cuda_a511a4c -ops = torch.ops._flash_attn2_cuda_a511a4c - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_flash_attn2_cuda_a511a4c::{op_name}" diff --git a/build/torch210-cxx11-cu126-aarch64-linux/bert_padding.py b/build/torch210-cxx11-cu126-aarch64-linux/bert_padding.py deleted file mode 100644 index 3c2d35159a014a9d03aabead9e52e009168696ea..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu126-aarch64-linux/bert_padding.py +++ /dev/null @@ -1,218 +0,0 @@ -# Adapted from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/padding.py - -import torch -import torch.nn.functional as F -from einops import rearrange, repeat - - -class IndexFirstAxis(torch.autograd.Function): - @staticmethod - def forward(ctx, input, indices): - ctx.save_for_backward(indices) - assert input.ndim >= 2 - ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] - second_dim = other_shape.numel() - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - # return input[indices] - return torch.gather( - rearrange(input, "b ... -> b (...)"), 0, repeat(indices, "z -> z d", d=second_dim) - ).reshape(-1, *other_shape) - - @staticmethod - def backward(ctx, grad_output): - (indices,) = ctx.saved_tensors - assert grad_output.ndim >= 2 - other_shape = grad_output.shape[1:] - grad_output = rearrange(grad_output, "b ... -> b (...)") - grad_input = torch.zeros( - [ctx.first_axis_dim, grad_output.shape[1]], - device=grad_output.device, - dtype=grad_output.dtype, - ) - # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. - # grad_input[indices] = grad_output - grad_input.scatter_(0, repeat(indices, "z -> z d", d=grad_output.shape[1]), grad_output) - return grad_input.reshape(ctx.first_axis_dim, *other_shape), None - - -index_first_axis = IndexFirstAxis.apply - - -class IndexPutFirstAxis(torch.autograd.Function): - @staticmethod - def forward(ctx, values, indices, first_axis_dim): - ctx.save_for_backward(indices) - assert indices.ndim == 1 - assert values.ndim >= 2 - output = torch.zeros( - first_axis_dim, *values.shape[1:], device=values.device, dtype=values.dtype - ) - # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. - output[indices] = values - # output.scatter_(0, repeat(indices, 'z -> z d', d=values.shape[1]), values) - return output - - @staticmethod - def backward(ctx, grad_output): - (indices,) = ctx.saved_tensors - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - grad_values = grad_output[indices] - # grad_values = torch.gather(grad_output, 0, repeat(indices, 'z -> z d', d=grad_output.shape[1])) - return grad_values, None, None - - -index_put_first_axis = IndexPutFirstAxis.apply - - -class IndexFirstAxisResidual(torch.autograd.Function): - @staticmethod - def forward(ctx, input, indices): - ctx.save_for_backward(indices) - assert input.ndim >= 2 - ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] - second_dim = other_shape.numel() - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - output = input[indices] - # We don't want to reshape input (b ... -> b (...)) since it could change the channel_last - # memory format to channel_first. In other words, input might not be contiguous. - # If we don't detach, Pytorch complains about output being a view and is being modified inplace - return output, input.detach() - - @staticmethod - def backward(ctx, grad_output, grad_residual): - (indices,) = ctx.saved_tensors - assert grad_output.ndim >= 2 - other_shape = grad_output.shape[1:] - assert grad_residual.shape[1:] == other_shape - grad_input = grad_residual - # grad_input[indices] += grad_output - indices = indices.reshape(indices.shape[0], *((1,) * (grad_output.ndim - 1))) - indices = indices.expand_as(grad_output) - grad_input.scatter_add_(0, indices, grad_output) - return grad_input.reshape(ctx.first_axis_dim, *other_shape), None - - -index_first_axis_residual = IndexFirstAxisResidual.apply - - -def unpad_input(hidden_states, attention_mask, unused_mask=None): - """ - Arguments: - hidden_states: (batch, seqlen, ...) - attention_mask: (batch, seqlen), bool / int, 1 means valid and 0 means not valid. - unused_mask: (batch, seqlen), bool / int, 1 means the element is allocated but unused. - Return: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens selected in attention_mask + unused_mask. - indices: (total_nnz), the indices of masked tokens from the flattened input sequence. - cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. - max_seqlen_in_batch: int - seqused: (batch), returns the number of tokens selected in attention_mask + unused_mask. - """ - all_masks = (attention_mask + unused_mask) if unused_mask is not None else attention_mask - seqlens_in_batch = all_masks.sum(dim=-1, dtype=torch.int32) - used_seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) - indices = torch.nonzero(all_masks.flatten(), as_tuple=False).flatten() - max_seqlen_in_batch = seqlens_in_batch.max().item() - cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) - # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the - # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim - # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to - # index with integer indices. Moreover, torch's index is a bit slower than it needs to be, - # so we write custom forward and backward to make it a bit faster. - return ( - index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices), - indices, - cu_seqlens, - max_seqlen_in_batch, - used_seqlens_in_batch, - ) - - -def unpad_input_for_concatenated_sequences(hidden_states, attention_mask_in_length): - """ - Supports concatenating short samples in one sequence. The attention_mask_in_length is utilized to mask other short samples. It helps efficient training of variant lengths-based samples (e.g., the supervised fine-tuning task in large language model). - The motivation for this function is explained [here](https://github.com/Dao-AILab/flash-attention/issues/432#issuecomment-1668822286). - - For example, if batch = 3 and seqlen = 6, the attention_mask_in_length is: - ``` - [ - [2, 3, 0, 0, 0, 0], - [3, 2, 0, 0, 0, 0], - [6, 0, 0, 0, 0, 0] - ] - ``` - , which refers to the 3D-attention mask: - ``` - [ - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [0, 0, 1, 0, 0, 0], - [0, 0, 1, 1, 0, 0], - [0, 0, 1, 1, 1, 0], - [0, 0, 0, 0, 0, 1] - ], - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [0, 0, 0, 1, 0, 0], - [0, 0, 0, 1, 1, 0], - [0, 0, 0, 0, 0, 1] - ], - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [1, 1, 1, 1, 0, 0], - [1, 1, 1, 1, 1, 0], - [1, 1, 1, 1, 1, 1] - ] - ] - ```. - - Arguments: - hidden_states: (batch, seqlen, ...) - attention_mask_in_length: (batch, seqlen), int, a nonzero number (e.g., 1, 2, 3, etc.) means length of concatenated sequence in b-th batch, and 0 means none. - Return: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. - indices: (total_nnz), the indices of non-masked tokens from the flattened input sequence. - cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. - max_seqlen_in_batch: int - """ - length = attention_mask_in_length.sum(dim=-1) - seqlen = attention_mask_in_length.size(-1) - attention_mask_2d = torch.arange(seqlen, device=length.device, dtype=length.dtype).expand(len(length), seqlen) < length.unsqueeze(1) - real_indices_idx = torch.nonzero(attention_mask_in_length.flatten(), as_tuple=False).flatten() - seqlens_in_batch = attention_mask_in_length.flatten()[real_indices_idx] - indices = torch.nonzero(attention_mask_2d.flatten(), as_tuple=False).flatten() - max_seqlen_in_batch = seqlens_in_batch.max().item() - cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) - # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the - # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim - # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to - # index with integer indices. Moreover, torch's index is a bit slower than it needs to be, - # so we write custom forward and backward to make it a bit faster. - return ( - index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices), - indices, - cu_seqlens, - max_seqlen_in_batch, - ) - - -def pad_input(hidden_states, indices, batch, seqlen): - """ - Arguments: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. - indices: (total_nnz), the indices that represent the non-masked tokens of the original padded input sequence. - batch: int, batch size for the padded sequence. - seqlen: int, maximum sequence length for the padded sequence. - Return: - hidden_states: (batch, seqlen, ...) - """ - dim = hidden_states.shape[-1] - # output = torch.zeros((batch * seqlen), dim, device=hidden_states.device, dtype=hidden_states.dtype) - # output[indices] = hidden_states - output = index_put_first_axis(hidden_states, indices, batch * seqlen) - return rearrange(output, "(b s) ... -> b s ...", b=batch) diff --git a/build/torch210-cxx11-cu126-aarch64-linux/flash_attn2/__init__.py b/build/torch210-cxx11-cu126-aarch64-linux/flash_attn2/__init__.py deleted file mode 100644 index a9b2672c1cd85b74c1b3ded0fc0b2100e1aeac23..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu126-aarch64-linux/flash_attn2/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -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-cu126-aarch64-linux/flash_attn_interface.py b/build/torch210-cxx11-cu126-aarch64-linux/flash_attn_interface.py deleted file mode 100644 index 0868572aa3324d9c96132859d123e3d48d296e34..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu126-aarch64-linux/flash_attn_interface.py +++ /dev/null @@ -1,1662 +0,0 @@ -# Copyright (c) 2023, Tri Dao. - -from typing import Optional, Sequence, Tuple, Union - -import torch -import torch.nn as nn -import os - -# # isort: off -# # We need to import the CUDA kernels after importing torch -# USE_TRITON_ROCM = os.getenv("FLASH_ATTENTION_TRITON_AMD_ENABLE", "FALSE") == "TRUE" -# if USE_TRITON_ROCM: -# from .flash_attn_triton_amd import interface_fa as flash_attn -# else: -# import flash_attn_2_cuda as flash_attn - - -from ._ops import ops as flash_attn -from ._ops import add_op_namespace_prefix - -# # isort: on - - -def maybe_contiguous(x): - return x.contiguous() if x is not None and x.stride(-1) != 1 else x - - -def _get_device(): - if torch.xpu.is_available(): - return "xpu" - elif torch.cuda.is_available(): - return "cuda" - else: - return "cpu" - - -def _get_block_size_n(device, head_dim, is_dropout, is_causal): - assert head_dim <= 256 - - if device.type == "xpu": - if head_dim <= 96: - return 64 - elif head_dim <= 128: - return 32 - elif head_dim <= 256: - return 64 - else: - return 32 - - # This should match the block sizes in the CUDA kernel - major, minor = torch.cuda.get_device_capability(device) - is_sm8x = ( - major == 8 and minor > 0 - ) # Only include sm86 and sm89, exclude sm80 (A100) - is_sm80 = major == 8 and minor == 0 - is_sm90 = major == 9 and minor == 0 - if head_dim <= 32: - return 128 - if head_dim <= 64: - return 128 if not is_dropout else 64 - elif head_dim <= 96: - return 64 - elif head_dim <= 128: - if is_sm8x: - return 64 if (not is_dropout and is_causal) else 32 - else: - return 64 if not is_dropout else 32 - elif head_dim <= 192: - return 64 - elif head_dim <= 224: - return 64 - elif head_dim <= 256: - return 64 - - -def round_multiple(x, m): - return (x + m - 1) // m * m - - -@torch.library.custom_op( - add_op_namespace_prefix("_flash_attn_forward"), - mutates_args=(), - device_types=_get_device(), -) -def _flash_attn_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - return_softmax: bool, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - out, softmax_lse, S_dmask, rng_state = flash_attn.fwd( - q, - k, - v, - None, - alibi_slopes, - dropout_p, - softmax_scale, - causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - None, - ) - return out, softmax_lse, S_dmask, rng_state - - -@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, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - return_softmax: bool, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - batch_size, seqlen_q, num_heads, head_size = q.shape - seqlen_k = k.shape[1] - out = torch.empty_like(q) - softmax_lse = torch.empty( - (batch_size, num_heads, seqlen_q), - dtype=torch.float32, - device=q.device, - layout=q.layout, - ) - p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) - if return_softmax: - p = torch.empty( - ( - batch_size, - num_heads, - round_multiple(seqlen_q, 128), - round_multiple(seqlen_k, 128), - ), - dtype=q.dtype, - device=q.device, - layout=q.layout, - ) - rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) - - return out, softmax_lse, p, rng_state - - -@torch.library.custom_op( - add_op_namespace_prefix("_flash_attn_varlen_forward"), - mutates_args=(), - device_types=_get_device(), -) -def _flash_attn_varlen_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - alibi_slopes: Optional[torch.Tensor] = None, - return_softmax: bool = False, - block_table: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - out, softmax_lse, S_dmask, rng_state = flash_attn.varlen_fwd( - q, - k, - v, - None, - cu_seqlens_q, - cu_seqlens_k, - seqused_k, - leftpad_k, - block_table, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - zero_tensors, - causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - None, - ) - # if out.isnan().any() or softmax_lse.isnan().any(): - # breakpoint() - return out, softmax_lse, S_dmask, rng_state - - -@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_varlen_forward")) -def _flash_attn_varlen_forward_fake( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - alibi_slopes: Optional[torch.Tensor] = None, - return_softmax: bool = False, - block_table: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - paged_kv = block_table is not None - batch_size = cu_seqlens_q.numel() - 1 - total_q, num_heads, _ = q.shape - - out = torch.empty_like(q) - softmax_lse = torch.empty( - (num_heads, total_q), dtype=torch.float32, device=q.device, layout=q.layout - ) - p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) - seqlen_q_rounded = round_multiple(max_seqlen_q, 128) - seqlen_k_rounded = round_multiple(max_seqlen_k, 128) - if return_softmax: - p = torch.empty( - (batch_size, num_heads, seqlen_q_rounded, seqlen_k_rounded), - dtype=q.dtype, - device=q.device, - layout=q.layout, - ) - rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) - return out, softmax_lse, p, rng_state - - -@torch.library.custom_op( - add_op_namespace_prefix("_flash_attn_backward"), - mutates_args=("dq", "dk", "dv"), - device_types=_get_device(), -) -def _flash_attn_backward( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, -) -> 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)] - ( - dq, - dk, - dv, - softmax_d, - ) = flash_attn.bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - alibi_slopes, - dropout_p, - softmax_scale, - causal, - window_size_left, - window_size_right, - softcap, - deterministic, - None, - rng_state, - ) - 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, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, -) -> torch.Tensor: - dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] - if dq is None: - dq = torch.empty_like(q) - if dk is None: - dk = torch.empty_like(k) - if dv is None: - dv = torch.empty_like(v) - batch_size, seqlen_q, num_heads, _ = q.shape - softmax_d = torch.empty( - (batch_size, num_heads, round_multiple(seqlen_q, 128)), - device=q.device, - dtype=torch.float32, - ) - - return softmax_d - - -@torch.library.custom_op( - add_op_namespace_prefix("_flash_attn_varlen_backward"), - mutates_args=("dq", "dk", "dv"), - device_types=_get_device(), -) -def _flash_attn_varlen_backward( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> 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)] - ( - dq, - dk, - dv, - softmax_d, - ) = flash_attn.varlen_bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - zero_tensors, - causal, - window_size_left, - window_size_right, - softcap, - deterministic, - None, - rng_state, - ) - # if dk.isnan().any() or dk.isnan().any() or dv.isnan().any() or softmax_d.isnan().any(): - # breakpoint() - return softmax_d - - -@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_varlen_backward")) -def _flash_attn_varlen_backward_fake( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> torch.Tensor: - dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] - batch_size = cu_seqlens_q.numel() - 1 - total_q, num_heads, _ = q.shape - - if dq is None: - dq = torch.empty_like(q) - if dk is None: - dk = torch.empty_like(k) - if dv is None: - dv = torch.empty_like(v) - softmax_d = torch.empty( - (num_heads, total_q + 128 * batch_size), device=q.device, dtype=torch.float32 - ) - - return softmax_d - - -class FlashAttnQKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and qkv.requires_grad - if softmax_scale is None: - softmax_scale = qkv.shape[-1] ** (-0.5) - q, k, v = qkv[:, :, 0].detach(), qkv[:, :, 1].detach(), qkv[:, :, 2].detach() - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) - dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dqkv[:, :, 0], - dqkv[:, :, 1], - dqkv[:, :, 2], - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dqkv = dqkv[..., : dout.shape[-1]] # We could have padded the head dimension - return dqkv, None, None, None, None, None, None, None, None, None - - -class FlashAttnVarlenQKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - cu_seqlens, - max_seqlen, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and qkv.requires_grad - if softmax_scale is None: - softmax_scale = qkv.shape[-1] ** (-0.5) - q, k, v = qkv[:, 0].detach(), qkv[:, 1].detach(), qkv[:, 2].detach() - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens, - cu_seqlens, - max_seqlen, - max_seqlen, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=None, - ) - if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens, rng_state - ) - ctx.dropout_p = dropout_p - ctx.max_seqlen = max_seqlen - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens, rng_state = ctx.saved_tensors - qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) - dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dqkv[:, 0], - dqkv[:, 1], - dqkv[:, 2], - cu_seqlens, - cu_seqlens, - ctx.max_seqlen, - ctx.max_seqlen, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - 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 FlashAttnKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - kv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any(x.requires_grad for x in [q, kv]) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - k, v = kv[:, :, 0].detach(), kv[:, :, 1].detach() - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - dq = torch.empty_like(q) - kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) - dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dkv[:, :, 0], - dkv[:, :, 1], - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dkv = dkv[..., : dout.shape[-1]] - return dq, dkv, None, None, None, None, None, None, None, None, None - - -class FlashAttnVarlenKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any(x.requires_grad for x in [q, kv]) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - k, v = kv[:, 0].detach(), kv[:, 1].detach() - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=None, - ) - if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state - ) - ctx.dropout_p = dropout_p - 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.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ( - ctx.saved_tensors - ) - dq = torch.empty_like(q) - kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) - dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dkv[:, 0], - dkv[:, 1], - cu_seqlens_q, - cu_seqlens_k, - ctx.max_seqlen_q, - ctx.max_seqlen_k, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dkv = dkv[..., : dout.shape[-1]] - return ( - dq, - dkv, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) - - -class FlashAttnFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - k, - v, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any(x.requires_grad for x in [q, k, v]) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dk = dk[..., : dout.shape[-1]] - dv = dv[..., : dout.shape[-1]] - return dq, dk, dv, 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, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - block_table, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any(x.requires_grad for x in [q, k, v]) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=block_table, - ) - if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state - ) - ctx.dropout_p = dropout_p - 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.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ( - ctx.saved_tensors - ) - dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - ctx.max_seqlen_q, - ctx.max_seqlen_k, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dk = dk[..., : dout.shape[-1]] - dv = dv[..., : dout.shape[-1]] - return ( - dq, - dk, - dv, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) - - -def flash_attn_qkvpacked_func( - qkv, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # <=0.0 means deactivate - alibi_slopes=None, - deterministic=False, - 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, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_kvpacked_func( - q, - kv, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, -): - """dropout_p should be set to 0.0 during evaluation - If 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 K, V. - 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) - kv: (batch_size, seqlen, 2, 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. - softcap: float. Anything > 0 activates softcapping 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). - 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 FlashAttnKVPackedFunc.apply( - q, - kv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_func( - q, - k, - v, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - 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). - 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 FlashAttnFunc.apply( - q, - k, - v, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_qkvpacked_func( - qkv, - cu_seqlens, - max_seqlen, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - 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_varlen_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_varlen_kvpacked_func and flash_attn_varlen_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: (total, 3, nheads, headdim), where total = total number of tokens in the batch. - cu_seqlens: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into qkv. - max_seqlen: int. Maximum sequence length in the batch. - 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenQKVPackedFunc.apply( - qkv, - cu_seqlens, - max_seqlen, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_kvpacked_func( - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, -): - """dropout_p should be set to 0.0 during evaluation - If 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 K, V. - 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: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch. - kv: (total_k, 2, nheads_k, headdim), where total_k = total number of key tokens in the batch. - cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into q. - cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into kv. - max_seqlen_q: int. Maximum query sequence length in the batch. - max_seqlen_k: int. Maximum key sequence length in the batch. - 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 + 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenKVPackedFunc.apply( - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_func( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, - block_table=None, -): - """dropout_p should be set to 0.0 during evaluation - Supports multi-query and grouped-query attention (MQA/GQA) by passing in K, V 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: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch. - k: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. - v: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. - cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into q. - cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into kv. - max_seqlen_q: int. Maximum query sequence length in the batch. - max_seqlen_k: int. Maximum key sequence length in the batch. - 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 + 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenFunc.apply( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - block_table, - False if q.device.type == "cpu" else torch.is_grad_enabled(), - ) - - -def flash_attn_with_kvcache( - q, - k_cache, - v_cache, - k=None, - v=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, - block_table: Optional[torch.Tensor] = None, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - rotary_interleaved=True, - alibi_slopes=None, - num_splits=0, - 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 block_table, - or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_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) if there's no block_table, - or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_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). Similar to k. - 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. - block_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). - 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. - 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" - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-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) - cache_batch_idx = maybe_contiguous(cache_batch_idx) - block_table = maybe_contiguous(block_table) - out, softmax_lse = flash_attn.fwd_kvcache( - q, - k_cache, - v_cache, - k, - v, - cache_seqlens, - rotary_cos, - rotary_sin, - cache_batch_idx, - cache_leftpad, - block_table, - alibi_slopes, - None, - softmax_scale, - causal, - window_size[0], - window_size[1], - softcap, - rotary_interleaved, - num_splits, - ) - return (out, softmax_lse) if return_softmax_lse else out diff --git a/build/torch210-cxx11-cu126-aarch64-linux/metadata.json b/build/torch210-cxx11-cu126-aarch64-linux/metadata.json deleted file mode 100644 index 542c83a02afbfd411593a4c773f833fc674d037b..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu126-aarch64-linux/metadata.json +++ /dev/null @@ -1,14 +0,0 @@ -{ - "name": "flash-attn2", - "id": "_flash_attn2_cuda_a511a4c", - "version": 1, - "license": "BSD-3-Clause", - "python-depends": [], - "backend": { - "type": "cuda", - "archs": [ - "8.0", - "9.0" - ] - } -} diff --git a/build/torch210-cxx11-cu126-x86_64-linux/__init__.py b/build/torch210-cxx11-cu126-x86_64-linux/__init__.py index 20ff0ff747abd7500913863b509925f414b1c973..ecc2f9d896b6c93f90b0a1499856dc0612177422 100644 --- a/build/torch210-cxx11-cu126-x86_64-linux/__init__.py +++ b/build/torch210-cxx11-cu126-x86_64-linux/__init__.py @@ -1,7 +1,5 @@ -from typing import List, Optional - +from typing import Optional, List import torch - from ._ops import ops as flash_attn_ops from .flash_attn_interface import ( flash_attn_func, @@ -13,23 +11,6 @@ from .flash_attn_interface import ( flash_attn_with_kvcache, ) -__all__ = [ - # High-level API (with autograd support) - "flash_attn_func", - "flash_attn_kvpacked_func", - "flash_attn_qkvpacked_func", - "flash_attn_varlen_func", - "flash_attn_varlen_kvpacked_func", - "flash_attn_varlen_qkvpacked_func", - "flash_attn_with_kvcache", - # Low-level ops (no autograd) - "fwd", - "varlen_fwd", - "bwd", - "varlen_bwd", - "fwd_kvcache", -] - def fwd( q: torch.Tensor, diff --git a/build/torch210-cxx11-cu126-x86_64-linux/_flash_attn2_588b404.abi3.so b/build/torch210-cxx11-cu126-x86_64-linux/_flash_attn2_588b404.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..90afba4f7dcc6f7e5803b4567c938998d5a7e777 --- /dev/null +++ b/build/torch210-cxx11-cu126-x86_64-linux/_flash_attn2_588b404.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:247ade2063814573447dcb697fd39e738bcf5f0f5d40ac87eaf6cf6dba29298f +size 448708992 diff --git a/build/torch210-cxx11-cu126-x86_64-linux/_flash_attn2_cuda_a511a4c.abi3.so b/build/torch210-cxx11-cu126-x86_64-linux/_flash_attn2_cuda_a511a4c.abi3.so deleted file mode 100644 index 2a6895f40baae6c9bc48ef821f81d2f764bab66d..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu126-x86_64-linux/_flash_attn2_cuda_a511a4c.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:67c0e9fa9b8ff7d99ffd10f0b4bac9f9c81112d25de5798750a63996ca777e6d -size 448709080 diff --git a/build/torch210-cxx11-cu126-x86_64-linux/_ops.py b/build/torch210-cxx11-cu126-x86_64-linux/_ops.py index 62b70a9b02c8f7663c7022564c1a7f58f46b6c65..03506d8a18d17ca7296a5f1b7a18ec7dc4ad73bb 100644 --- a/build/torch210-cxx11-cu126-x86_64-linux/_ops.py +++ b/build/torch210-cxx11-cu126-x86_64-linux/_ops.py @@ -1,9 +1,9 @@ import torch -from . import _flash_attn2_cuda_a511a4c -ops = torch.ops._flash_attn2_cuda_a511a4c +from . import _flash_attn2_588b404 +ops = torch.ops._flash_attn2_588b404 def add_op_namespace_prefix(op_name: str): """ Prefix op by namespace. """ - return f"_flash_attn2_cuda_a511a4c::{op_name}" + return f"_flash_attn2_588b404::{op_name}" \ No newline at end of file diff --git a/build/torch210-cxx11-cu126-x86_64-linux/flash_attn2/__init__.py b/build/torch210-cxx11-cu126-x86_64-linux/flash_attn2/__init__.py index a9b2672c1cd85b74c1b3ded0fc0b2100e1aeac23..03dbc1afe1cf156661a2b1b22003cd5f599a0309 100644 --- a/build/torch210-cxx11-cu126-x86_64-linux/flash_attn2/__init__.py +++ b/build/torch210-cxx11-cu126-x86_64-linux/flash_attn2/__init__.py @@ -1,10 +1,10 @@ import ctypes -import importlib.util 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 diff --git a/build/torch210-cxx11-cu126-x86_64-linux/flash_attn_interface.py b/build/torch210-cxx11-cu126-x86_64-linux/flash_attn_interface.py index 0868572aa3324d9c96132859d123e3d48d296e34..66e8d3665834b141d703998496b3da6fe1f6fe06 100644 --- a/build/torch210-cxx11-cu126-x86_64-linux/flash_attn_interface.py +++ b/build/torch210-cxx11-cu126-x86_64-linux/flash_attn_interface.py @@ -16,11 +16,9 @@ import os from ._ops import ops as flash_attn -from ._ops import add_op_namespace_prefix # # isort: on - def maybe_contiguous(x): return x.contiguous() if x is not None and x.stride(-1) != 1 else x @@ -33,25 +31,14 @@ def _get_device(): else: return "cpu" +_XPU_AVAILABLE = torch.xpu.is_available() if hasattr(torch, "xpu") else False # TODO remove hasattr check when bwd is supported on XPU -def _get_block_size_n(device, head_dim, is_dropout, is_causal): - assert head_dim <= 256 - - if device.type == "xpu": - if head_dim <= 96: - return 64 - elif head_dim <= 128: - return 32 - elif head_dim <= 256: - return 64 - else: - return 32 +def _get_block_size_n(device, head_dim, is_dropout, is_causal): # This should match the block sizes in the CUDA kernel + assert head_dim <= 256 major, minor = torch.cuda.get_device_capability(device) - is_sm8x = ( - major == 8 and minor > 0 - ) # Only include sm86 and sm89, exclude sm80 (A100) + is_sm8x = major == 8 and minor > 0 # Only include sm86 and sm89, exclude sm80 (A100) is_sm80 = major == 8 and minor == 0 is_sm90 = major == 9 and minor == 0 if head_dim <= 32: @@ -77,11 +64,30 @@ def round_multiple(x, m): return (x + m - 1) // m * m -@torch.library.custom_op( - add_op_namespace_prefix("_flash_attn_forward"), - mutates_args=(), - device_types=_get_device(), -) +# torch.compile() support is only enabled for pytorch >= 2.4 +# The reason for this is that we are using the new custom_op and register_fake +# APIs, which support inplace modification of inputs in the function itself +if torch.__version__ >= "2.4.0": + _torch_custom_op_wrapper = torch.library.custom_op + _torch_register_fake_wrapper = torch.library.register_fake +else: + def noop_custom_op_wrapper(name, fn=None, /, *, mutates_args, device_types=None, schema=None): + def wrap(func): + return func + if fn is None: + return wrap + return fn + def noop_register_fake_wrapper(op, fn=None, /, *, lib=None, _stacklevel=1): + def wrap(func): + return func + if fn is None: + return wrap + return fn + _torch_custom_op_wrapper = noop_custom_op_wrapper + _torch_register_fake_wrapper = noop_register_fake_wrapper + + +@_torch_custom_op_wrapper("flash_attn::_flash_attn_forward", mutates_args=(), device_types=_get_device()) def _flash_attn_forward( q: torch.Tensor, k: torch.Tensor, @@ -93,7 +99,7 @@ def _flash_attn_forward( window_size_right: int, softcap: float, alibi_slopes: Optional[torch.Tensor], - return_softmax: bool, + return_softmax: bool ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: q, k, v = [maybe_contiguous(x) for x in (q, k, v)] out, softmax_lse, S_dmask, rng_state = flash_attn.fwd( @@ -114,7 +120,7 @@ def _flash_attn_forward( return out, softmax_lse, S_dmask, rng_state -@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_forward")) +@_torch_register_fake_wrapper("flash_attn::_flash_attn_forward") def _flash_attn_forward_fake( q: torch.Tensor, k: torch.Tensor, @@ -126,41 +132,28 @@ def _flash_attn_forward_fake( window_size_right: int, softcap: float, alibi_slopes: Optional[torch.Tensor], - return_softmax: bool, + return_softmax: bool ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: q, k, v = [maybe_contiguous(x) for x in (q, k, v)] batch_size, seqlen_q, num_heads, head_size = q.shape seqlen_k = k.shape[1] out = torch.empty_like(q) - softmax_lse = torch.empty( - (batch_size, num_heads, seqlen_q), - dtype=torch.float32, - device=q.device, - layout=q.layout, - ) + softmax_lse = torch.empty((batch_size, num_heads, seqlen_q), dtype=torch.float32, device=q.device, layout=q.layout) p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) if return_softmax: - p = torch.empty( - ( - batch_size, - num_heads, - round_multiple(seqlen_q, 128), - round_multiple(seqlen_k, 128), - ), - dtype=q.dtype, - device=q.device, - layout=q.layout, - ) + p = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128), round_multiple(seqlen_k, 128)), dtype=q.dtype, device=q.device, layout=q.layout) rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) return out, softmax_lse, p, rng_state -@torch.library.custom_op( - add_op_namespace_prefix("_flash_attn_varlen_forward"), - mutates_args=(), - device_types=_get_device(), -) +if torch.__version__ >= "2.4.0": + _wrapped_flash_attn_forward = torch.ops.flash_attn._flash_attn_forward +else: + _wrapped_flash_attn_forward = _flash_attn_forward + + +@_torch_custom_op_wrapper("flash_attn::_flash_attn_varlen_forward", mutates_args=(), device_types=_get_device()) def _flash_attn_varlen_forward( q: torch.Tensor, k: torch.Tensor, @@ -211,7 +204,7 @@ def _flash_attn_varlen_forward( return out, softmax_lse, S_dmask, rng_state -@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_varlen_forward")) +@_torch_register_fake_wrapper("flash_attn::_flash_attn_varlen_forward") def _flash_attn_varlen_forward_fake( q: torch.Tensor, k: torch.Tensor, @@ -237,30 +230,25 @@ def _flash_attn_varlen_forward_fake( paged_kv = block_table is not None batch_size = cu_seqlens_q.numel() - 1 total_q, num_heads, _ = q.shape - + out = torch.empty_like(q) - softmax_lse = torch.empty( - (num_heads, total_q), dtype=torch.float32, device=q.device, layout=q.layout - ) + softmax_lse = torch.empty((num_heads, total_q), dtype=torch.float32, device=q.device, layout=q.layout) p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) seqlen_q_rounded = round_multiple(max_seqlen_q, 128) seqlen_k_rounded = round_multiple(max_seqlen_k, 128) if return_softmax: - p = torch.empty( - (batch_size, num_heads, seqlen_q_rounded, seqlen_k_rounded), - dtype=q.dtype, - device=q.device, - layout=q.layout, - ) + p = torch.empty((batch_size, num_heads, seqlen_q_rounded, seqlen_k_rounded), dtype=q.dtype, device=q.device, layout=q.layout) rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) return out, softmax_lse, p, rng_state -@torch.library.custom_op( - add_op_namespace_prefix("_flash_attn_backward"), - mutates_args=("dq", "dk", "dv"), - device_types=_get_device(), -) +if torch.__version__ >= "2.4.0": + _wrapped_flash_attn_varlen_forward = torch.ops.flash_attn._flash_attn_varlen_forward +else: + _wrapped_flash_attn_varlen_forward = _flash_attn_varlen_forward + + +@_torch_custom_op_wrapper("flash_attn::_flash_attn_backward", mutates_args=("dq", "dk", "dv"), device_types=_get_device()) def _flash_attn_backward( dout: torch.Tensor, q: torch.Tensor, @@ -312,7 +300,7 @@ def _flash_attn_backward( return softmax_d -@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_backward")) +@_torch_register_fake_wrapper("flash_attn::_flash_attn_backward") def _flash_attn_backward_fake( dout: torch.Tensor, q: torch.Tensor, @@ -341,20 +329,18 @@ def _flash_attn_backward_fake( if dv is None: dv = torch.empty_like(v) batch_size, seqlen_q, num_heads, _ = q.shape - softmax_d = torch.empty( - (batch_size, num_heads, round_multiple(seqlen_q, 128)), - device=q.device, - dtype=torch.float32, - ) - + softmax_d = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128)), device=q.device, dtype=torch.float32) + return softmax_d -@torch.library.custom_op( - add_op_namespace_prefix("_flash_attn_varlen_backward"), - mutates_args=("dq", "dk", "dv"), - device_types=_get_device(), -) +if torch.__version__ >= "2.4.0": + _wrapped_flash_attn_backward = torch.ops.flash_attn._flash_attn_backward +else: + _wrapped_flash_attn_backward = _flash_attn_backward + + +@_torch_custom_op_wrapper("flash_attn::_flash_attn_varlen_backward", mutates_args=("dq", "dk", "dv"), device_types=_get_device()) def _flash_attn_varlen_backward( dout: torch.Tensor, q: torch.Tensor, @@ -418,7 +404,7 @@ def _flash_attn_varlen_backward( return softmax_d -@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_varlen_backward")) +@_torch_register_fake_wrapper("flash_attn::_flash_attn_varlen_backward") def _flash_attn_varlen_backward_fake( dout: torch.Tensor, q: torch.Tensor, @@ -454,13 +440,17 @@ def _flash_attn_varlen_backward_fake( dk = torch.empty_like(k) if dv is None: dv = torch.empty_like(v) - softmax_d = torch.empty( - (num_heads, total_q + 128 * batch_size), device=q.device, dtype=torch.float32 - ) - + softmax_d = torch.empty((num_heads, total_q + 128 * batch_size), device=q.device, dtype=torch.float32) + return softmax_d +if torch.__version__ >= "2.4.0": + _wrapped_flash_attn_varlen_backward = torch.ops.flash_attn._flash_attn_varlen_backward +else: + _wrapped_flash_attn_varlen_backward = _flash_attn_varlen_backward + + class FlashAttnQKVPackedFunc(torch.autograd.Function): @staticmethod def forward( @@ -485,7 +475,7 @@ class FlashAttnQKVPackedFunc(torch.autograd.Function): q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_forward( + out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( q, k, v, @@ -519,7 +509,7 @@ class FlashAttnQKVPackedFunc(torch.autograd.Function): dout_padded = dout if head_size_og % 8 != 0: dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _flash_attn_backward( + _wrapped_flash_attn_backward( dout_padded, q, k, @@ -569,7 +559,7 @@ class FlashAttnVarlenQKVPackedFunc(torch.autograd.Function): q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_varlen_forward( + out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( q, k, v, @@ -588,9 +578,7 @@ class FlashAttnVarlenQKVPackedFunc(torch.autograd.Function): block_table=None, ) if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens, rng_state - ) + ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens, rng_state) ctx.dropout_p = dropout_p ctx.max_seqlen = max_seqlen ctx.softmax_scale = softmax_scale @@ -611,7 +599,7 @@ class FlashAttnVarlenQKVPackedFunc(torch.autograd.Function): dout_padded = dout if head_size_og % 8 != 0: dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _flash_attn_varlen_backward( + _wrapped_flash_attn_varlen_backward( dout_padded, q, k, @@ -655,7 +643,9 @@ class FlashAttnKVPackedFunc(torch.autograd.Function): return_softmax, is_grad_enabled, ): - is_grad = is_grad_enabled and any(x.requires_grad for x in [q, kv]) + is_grad = is_grad_enabled and any( + x.requires_grad for x in [q, kv] + ) if softmax_scale is None: softmax_scale = q.shape[-1] ** (-0.5) k, v = kv[:, :, 0].detach(), kv[:, :, 1].detach() @@ -664,7 +654,7 @@ class FlashAttnKVPackedFunc(torch.autograd.Function): q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_forward( + out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( q, k, v, @@ -699,7 +689,7 @@ class FlashAttnKVPackedFunc(torch.autograd.Function): dout_padded = dout if head_size_og % 8 != 0: dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _flash_attn_backward( + _wrapped_flash_attn_backward( dout_padded, q, k, @@ -744,7 +734,9 @@ class FlashAttnVarlenKVPackedFunc(torch.autograd.Function): return_softmax, is_grad_enabled, ): - is_grad = is_grad_enabled and any(x.requires_grad for x in [q, kv]) + is_grad = is_grad_enabled and any( + x.requires_grad for x in [q, kv] + ) if softmax_scale is None: softmax_scale = q.shape[-1] ** (-0.5) k, v = kv[:, 0].detach(), kv[:, 1].detach() @@ -753,7 +745,7 @@ class FlashAttnVarlenKVPackedFunc(torch.autograd.Function): q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_varlen_forward( + out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( q, k, v, @@ -789,9 +781,7 @@ class FlashAttnVarlenKVPackedFunc(torch.autograd.Function): @staticmethod def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ( - ctx.saved_tensors - ) + q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors dq = torch.empty_like(q) kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) @@ -799,7 +789,7 @@ class FlashAttnVarlenKVPackedFunc(torch.autograd.Function): dout_padded = dout if head_size_og % 8 != 0: dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _flash_attn_varlen_backward( + _wrapped_flash_attn_varlen_backward( dout_padded, q, k, @@ -825,23 +815,7 @@ class FlashAttnVarlenKVPackedFunc(torch.autograd.Function): ) dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension dkv = dkv[..., : dout.shape[-1]] - return ( - dq, - dkv, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) + return dq, dkv, None, None, None, None, None, None, None, None, None, None, None, None, None class FlashAttnFunc(torch.autograd.Function): @@ -861,7 +835,9 @@ class FlashAttnFunc(torch.autograd.Function): return_softmax, is_grad_enabled, ): - is_grad = is_grad_enabled and any(x.requires_grad for x in [q, k, v]) + is_grad = is_grad_enabled and any( + x.requires_grad for x in [q, k, v] + ) if softmax_scale is None: softmax_scale = q.shape[-1] ** (-0.5) head_size_og = q.size(3) @@ -869,7 +845,7 @@ class FlashAttnFunc(torch.autograd.Function): q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_forward( + out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( q, k, v, @@ -902,7 +878,7 @@ class FlashAttnFunc(torch.autograd.Function): dout_padded = dout if head_size_og % 8 != 0: dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _flash_attn_backward( + _wrapped_flash_attn_backward( dout_padded, q, k, @@ -950,7 +926,9 @@ class FlashAttnVarlenFunc(torch.autograd.Function): block_table, is_grad_enabled, ): - is_grad = is_grad_enabled and any(x.requires_grad for x in [q, k, v]) + is_grad = is_grad_enabled and any( + x.requires_grad for x in [q, k, v] + ) if softmax_scale is None: softmax_scale = q.shape[-1] ** (-0.5) head_size_og = q.size(2) @@ -958,7 +936,7 @@ class FlashAttnVarlenFunc(torch.autograd.Function): q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_varlen_forward( + out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( q, k, v, @@ -995,15 +973,13 @@ class FlashAttnVarlenFunc(torch.autograd.Function): @staticmethod def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ( - ctx.saved_tensors - ) + q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) head_size_og = dout.size(2) dout_padded = dout if head_size_og % 8 != 0: dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _flash_attn_varlen_backward( + _wrapped_flash_attn_varlen_backward( dout_padded, q, k, @@ -1030,25 +1006,7 @@ class FlashAttnVarlenFunc(torch.autograd.Function): dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension dk = dk[..., : dout.shape[-1]] dv = dv[..., : dout.shape[-1]] - return ( - dq, - dk, - dv, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None def flash_attn_qkvpacked_func( @@ -1106,7 +1064,7 @@ def flash_attn_qkvpacked_func( alibi_slopes, deterministic, return_attn_probs, - torch.is_grad_enabled(), + False if _XPU_AVAILABLE else torch.is_grad_enabled(), ) @@ -1184,7 +1142,7 @@ def flash_attn_kvpacked_func( alibi_slopes, deterministic, return_attn_probs, - torch.is_grad_enabled(), + False if _XPU_AVAILABLE else torch.is_grad_enabled(), ) @@ -1196,7 +1154,7 @@ def flash_attn_func( softmax_scale=None, causal=False, window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated + softcap=0.0, # 0.0 means deactivated alibi_slopes=None, deterministic=False, return_attn_probs=False, @@ -1261,7 +1219,7 @@ def flash_attn_func( alibi_slopes, deterministic, return_attn_probs, - torch.is_grad_enabled(), + False if _XPU_AVAILABLE else torch.is_grad_enabled(), ) @@ -1273,7 +1231,7 @@ def flash_attn_varlen_qkvpacked_func( softmax_scale=None, causal=False, window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated + softcap=0.0, # 0.0 means deactivated alibi_slopes=None, deterministic=False, return_attn_probs=False, @@ -1327,7 +1285,7 @@ def flash_attn_varlen_qkvpacked_func( alibi_slopes, deterministic, return_attn_probs, - torch.is_grad_enabled(), + False if _XPU_AVAILABLE else torch.is_grad_enabled(), ) @@ -1342,7 +1300,7 @@ def flash_attn_varlen_kvpacked_func( softmax_scale=None, causal=False, window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated + softcap=0.0, # 0.0 means deactivated alibi_slopes=None, deterministic=False, return_attn_probs=False, @@ -1419,7 +1377,7 @@ def flash_attn_varlen_kvpacked_func( alibi_slopes, deterministic, return_attn_probs, - torch.is_grad_enabled(), + False if _XPU_AVAILABLE else torch.is_grad_enabled(), ) @@ -1435,7 +1393,7 @@ def flash_attn_varlen_func( softmax_scale=None, causal=False, window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated + softcap=0.0, # 0.0 means deactivated alibi_slopes=None, deterministic=False, return_attn_probs=False, @@ -1513,7 +1471,7 @@ def flash_attn_varlen_func( deterministic, return_attn_probs, block_table, - False if q.device.type == "cpu" else torch.is_grad_enabled(), + False if _XPU_AVAILABLE or q.device.type == "cpu" else torch.is_grad_enabled(), ) @@ -1532,7 +1490,7 @@ def flash_attn_with_kvcache( softmax_scale=None, causal=False, window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated + softcap=0.0, # 0.0 means deactivated rotary_interleaved=True, alibi_slopes=None, num_splits=0, diff --git a/build/torch210-cu130-x86_64-windows/layers/__init__.py b/build/torch210-cxx11-cu126-x86_64-linux/layers/__init__.py similarity index 100% rename from build/torch210-cu130-x86_64-windows/layers/__init__.py rename to build/torch210-cxx11-cu126-x86_64-linux/layers/__init__.py diff --git a/build/torch211-cxx11-cu126-aarch64-linux/layers/patch_embed.py b/build/torch210-cxx11-cu126-x86_64-linux/layers/patch_embed.py similarity index 100% rename from build/torch211-cxx11-cu126-aarch64-linux/layers/patch_embed.py rename to build/torch210-cxx11-cu126-x86_64-linux/layers/patch_embed.py diff --git a/build/torch211-cxx11-cu126-aarch64-linux/layers/rotary.py b/build/torch210-cxx11-cu126-x86_64-linux/layers/rotary.py similarity index 100% rename from build/torch211-cxx11-cu126-aarch64-linux/layers/rotary.py rename to build/torch210-cxx11-cu126-x86_64-linux/layers/rotary.py diff --git a/build/torch210-cxx11-cu126-x86_64-linux/metadata.json b/build/torch210-cxx11-cu126-x86_64-linux/metadata.json index 542c83a02afbfd411593a4c773f833fc674d037b..9cf5deed9898dce769f4cc73913d3530b92a0bd8 100644 --- a/build/torch210-cxx11-cu126-x86_64-linux/metadata.json +++ b/build/torch210-cxx11-cu126-x86_64-linux/metadata.json @@ -1,14 +1,4 @@ { - "name": "flash-attn2", - "id": "_flash_attn2_cuda_a511a4c", "version": 1, - "license": "BSD-3-Clause", - "python-depends": [], - "backend": { - "type": "cuda", - "archs": [ - "8.0", - "9.0" - ] - } -} + "python-depends": [] +} \ No newline at end of file diff --git a/build/torch210-cu130-x86_64-windows/ops/__init__.py b/build/torch210-cxx11-cu126-x86_64-linux/ops/__init__.py similarity index 100% rename from build/torch210-cu130-x86_64-windows/ops/__init__.py rename to build/torch210-cxx11-cu126-x86_64-linux/ops/__init__.py diff --git a/build/torch211-cxx11-cu126-aarch64-linux/ops/activations.py b/build/torch210-cxx11-cu126-x86_64-linux/ops/activations.py similarity index 100% rename from build/torch211-cxx11-cu126-aarch64-linux/ops/activations.py rename to build/torch210-cxx11-cu126-x86_64-linux/ops/activations.py diff --git a/build/torch211-cxx11-cu126-aarch64-linux/ops/fused_dense.py b/build/torch210-cxx11-cu126-x86_64-linux/ops/fused_dense.py similarity index 100% rename from build/torch211-cxx11-cu126-aarch64-linux/ops/fused_dense.py rename to build/torch210-cxx11-cu126-x86_64-linux/ops/fused_dense.py diff --git a/build/torch211-cxx11-cu126-aarch64-linux/ops/layer_norm.py b/build/torch210-cxx11-cu126-x86_64-linux/ops/layer_norm.py similarity index 100% rename from build/torch211-cxx11-cu126-aarch64-linux/ops/layer_norm.py rename to build/torch210-cxx11-cu126-x86_64-linux/ops/layer_norm.py diff --git a/build/torch211-cxx11-cu126-aarch64-linux/ops/rms_norm.py b/build/torch210-cxx11-cu126-x86_64-linux/ops/rms_norm.py similarity index 100% rename from build/torch211-cxx11-cu126-aarch64-linux/ops/rms_norm.py rename to build/torch210-cxx11-cu126-x86_64-linux/ops/rms_norm.py diff --git a/build/torch211-cxx11-cu126-aarch64-linux/ops/triton/__init__.py b/build/torch210-cxx11-cu126-x86_64-linux/ops/triton/__init__.py similarity index 100% rename from build/torch211-cxx11-cu126-aarch64-linux/ops/triton/__init__.py rename to build/torch210-cxx11-cu126-x86_64-linux/ops/triton/__init__.py diff --git a/build/torch211-cxx11-cu126-aarch64-linux/ops/triton/cross_entropy.py b/build/torch210-cxx11-cu126-x86_64-linux/ops/triton/cross_entropy.py similarity index 100% rename from build/torch211-cxx11-cu126-aarch64-linux/ops/triton/cross_entropy.py rename to build/torch210-cxx11-cu126-x86_64-linux/ops/triton/cross_entropy.py diff --git a/build/torch211-cxx11-cu126-aarch64-linux/ops/triton/k_activations.py b/build/torch210-cxx11-cu126-x86_64-linux/ops/triton/k_activations.py similarity index 100% rename from build/torch211-cxx11-cu126-aarch64-linux/ops/triton/k_activations.py rename to build/torch210-cxx11-cu126-x86_64-linux/ops/triton/k_activations.py diff --git a/build/torch29-cxx11-cu128-aarch64-linux/ops/triton/layer_norm.py b/build/torch210-cxx11-cu126-x86_64-linux/ops/triton/layer_norm.py similarity index 100% rename from build/torch29-cxx11-cu128-aarch64-linux/ops/triton/layer_norm.py rename to build/torch210-cxx11-cu126-x86_64-linux/ops/triton/layer_norm.py diff --git a/build/torch211-cxx11-cu126-aarch64-linux/ops/triton/linear.py b/build/torch210-cxx11-cu126-x86_64-linux/ops/triton/linear.py similarity index 100% rename from build/torch211-cxx11-cu126-aarch64-linux/ops/triton/linear.py rename to build/torch210-cxx11-cu126-x86_64-linux/ops/triton/linear.py diff --git a/build/torch211-cxx11-cu126-aarch64-linux/ops/triton/mlp.py b/build/torch210-cxx11-cu126-x86_64-linux/ops/triton/mlp.py similarity index 100% rename from build/torch211-cxx11-cu126-aarch64-linux/ops/triton/mlp.py rename to build/torch210-cxx11-cu126-x86_64-linux/ops/triton/mlp.py diff --git a/build/torch211-cxx11-cu126-aarch64-linux/ops/triton/rotary.py b/build/torch210-cxx11-cu126-x86_64-linux/ops/triton/rotary.py similarity index 100% rename from build/torch211-cxx11-cu126-aarch64-linux/ops/triton/rotary.py rename to build/torch210-cxx11-cu126-x86_64-linux/ops/triton/rotary.py diff --git a/build/torch210-cxx11-cu128-aarch64-linux/__init__.py b/build/torch210-cxx11-cu128-aarch64-linux/__init__.py deleted file mode 100644 index 20ff0ff747abd7500913863b509925f414b1c973..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu128-aarch64-linux/__init__.py +++ /dev/null @@ -1,412 +0,0 @@ -from typing import List, Optional - -import torch - -from ._ops import ops as flash_attn_ops -from .flash_attn_interface import ( - flash_attn_func, - flash_attn_kvpacked_func, - flash_attn_qkvpacked_func, - flash_attn_varlen_func, - flash_attn_varlen_kvpacked_func, - flash_attn_varlen_qkvpacked_func, - flash_attn_with_kvcache, -) - -__all__ = [ - # High-level API (with autograd support) - "flash_attn_func", - "flash_attn_kvpacked_func", - "flash_attn_qkvpacked_func", - "flash_attn_varlen_func", - "flash_attn_varlen_kvpacked_func", - "flash_attn_varlen_qkvpacked_func", - "flash_attn_with_kvcache", - # Low-level ops (no autograd) - "fwd", - "varlen_fwd", - "bwd", - "varlen_bwd", - "fwd_kvcache", -] - - -def fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - return_softmax: bool = False, - gen: Optional[torch.Generator] = None, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention. - - Args: - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Optional output tensor, same shape as q - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - return_softmax: Whether to return softmax weights - gen: Optional random number generator - - Returns: - List of tensors: [output, softmax_lse, (softmax if return_softmax)] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.fwd( - q, - k, - v, - out, - alibi_slopes, - p_dropout, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - gen, - ) - - -def varlen_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - out: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - block_table: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - max_seqlen_q: int = 0, - max_seqlen_k: int = 0, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - zero_tensors: bool = False, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - return_softmax: bool = False, - gen: Optional[torch.Generator] = None, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention with variable sequence lengths. - - Args: - q: Query tensor of shape [total_q, num_heads, head_size] - k: Key tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - v: Value tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1] - cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1] - out: Optional output tensor of shape [total_q, num_heads, head_size] - seqused_k: Optional tensor specifying how many keys to use per batch element [batch_size] - leftpad_k: Optional left padding for keys of shape [batch_size] - block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq] - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - max_seqlen_q: Maximum sequence length for queries - max_seqlen_k: Maximum sequence length for keys - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - zero_tensors: Whether to zero tensors before computation - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - return_softmax: Whether to return softmax weights - gen: Optional random number generator - - Returns: - List of tensors: [output, softmax_lse, (softmax if return_softmax)] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.varlen_fwd( - q, - k, - v, - out, - cu_seqlens_q, - cu_seqlens_k, - seqused_k, - leftpad_k, - block_table, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - p_dropout, - softmax_scale, - zero_tensors, - is_causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - gen, - ) - - -def bwd( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - p_dropout: float = 0.0, - 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, - gen: Optional[torch.Generator] = None, - rng_state: Optional[torch.Tensor] = None, -) -> List[torch.Tensor]: - """ - Backward pass for multi-head attention. - - Args: - dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size] - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size] - softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q] - dq: Optional gradient tensor for queries, same shape as q - dk: Optional gradient tensor for keys, same shape as k - dv: Optional gradient tensor for values, same shape as v - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - deterministic: Whether to use deterministic algorithms - gen: Optional random number generator - rng_state: Optional RNG state from forward pass - - Returns: - List of tensors: [dq, dk, dv] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - alibi_slopes, - p_dropout, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - deterministic, - gen, - rng_state, - ) - - -def varlen_bwd( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - max_seqlen_q: int = 0, - max_seqlen_k: int = 0, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - zero_tensors: bool = False, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - deterministic: bool = False, - gen: Optional[torch.Generator] = None, - rng_state: Optional[torch.Tensor] = None, -) -> List[torch.Tensor]: - """ - Backward pass for multi-head attention with variable sequence lengths. - - Args: - dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size] - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size] - softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q] - cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1] - cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1] - dq: Optional gradient tensor for queries, same shape as q - dk: Optional gradient tensor for keys, same shape as k - dv: Optional gradient tensor for values, same shape as v - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - max_seqlen_q: Maximum sequence length for queries - max_seqlen_k: Maximum sequence length for keys - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - zero_tensors: Whether to zero tensors before computation - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - deterministic: Whether to use deterministic algorithms - gen: Optional random number generator - rng_state: Optional RNG state from forward pass - - Returns: - List of tensors: [dq, dk, dv] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.varlen_bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - p_dropout, - softmax_scale, - zero_tensors, - is_causal, - window_size_left, - window_size_right, - softcap, - deterministic, - gen, - rng_state, - ) - - -def fwd_kvcache( - q: torch.Tensor, - kcache: torch.Tensor, - vcache: torch.Tensor, - k: Optional[torch.Tensor] = None, - v: Optional[torch.Tensor] = None, - seqlens_k: Optional[torch.Tensor] = None, - rotary_cos: Optional[torch.Tensor] = None, - rotary_sin: Optional[torch.Tensor] = None, - cache_batch_idx: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - block_table: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - out: 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, - is_rotary_interleaved: bool = False, - num_splits: int = 1, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention with KV cache. - - Args: - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - kcache: Key cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - vcache: Value cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - k: Optional new keys tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size] - v: Optional new values tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size] - seqlens_k: Optional sequence lengths for keys of shape [batch_size] - rotary_cos: Optional rotary cosine tensor of shape [seqlen_ro, rotary_dim/2] - rotary_sin: Optional rotary sine tensor of shape [seqlen_ro, rotary_dim/2] - cache_batch_idx: Optional indices to index into the KV cache - leftpad_k: Optional left padding for keys of shape [batch_size] - block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq] - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - out: Optional output tensor, same shape as q - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - is_rotary_interleaved: Whether rotary embeddings are interleaved - num_splits: Number of splits for computation - - Returns: - List of tensors: [output, softmax_lse] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.fwd_kvcache( - q, - kcache, - vcache, - k, - v, - seqlens_k, - rotary_cos, - rotary_sin, - cache_batch_idx, - leftpad_k, - block_table, - alibi_slopes, - out, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - is_rotary_interleaved, - num_splits, - ) diff --git a/build/torch210-cxx11-cu128-aarch64-linux/_flash_attn2_cuda_a511a4c.abi3.so b/build/torch210-cxx11-cu128-aarch64-linux/_flash_attn2_cuda_a511a4c.abi3.so deleted file mode 100644 index 1a127a80969102dc9c6df072143f5e7f3c661a30..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu128-aarch64-linux/_flash_attn2_cuda_a511a4c.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:54a3e0d7669eb5e39804fab3c6c3206ceff08f56a4502c280a428f9d9f166201 -size 1038067096 diff --git a/build/torch210-cxx11-cu128-aarch64-linux/_ops.py b/build/torch210-cxx11-cu128-aarch64-linux/_ops.py deleted file mode 100644 index 62b70a9b02c8f7663c7022564c1a7f58f46b6c65..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu128-aarch64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _flash_attn2_cuda_a511a4c -ops = torch.ops._flash_attn2_cuda_a511a4c - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_flash_attn2_cuda_a511a4c::{op_name}" diff --git a/build/torch210-cxx11-cu128-aarch64-linux/bert_padding.py b/build/torch210-cxx11-cu128-aarch64-linux/bert_padding.py deleted file mode 100644 index 3c2d35159a014a9d03aabead9e52e009168696ea..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu128-aarch64-linux/bert_padding.py +++ /dev/null @@ -1,218 +0,0 @@ -# Adapted from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/padding.py - -import torch -import torch.nn.functional as F -from einops import rearrange, repeat - - -class IndexFirstAxis(torch.autograd.Function): - @staticmethod - def forward(ctx, input, indices): - ctx.save_for_backward(indices) - assert input.ndim >= 2 - ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] - second_dim = other_shape.numel() - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - # return input[indices] - return torch.gather( - rearrange(input, "b ... -> b (...)"), 0, repeat(indices, "z -> z d", d=second_dim) - ).reshape(-1, *other_shape) - - @staticmethod - def backward(ctx, grad_output): - (indices,) = ctx.saved_tensors - assert grad_output.ndim >= 2 - other_shape = grad_output.shape[1:] - grad_output = rearrange(grad_output, "b ... -> b (...)") - grad_input = torch.zeros( - [ctx.first_axis_dim, grad_output.shape[1]], - device=grad_output.device, - dtype=grad_output.dtype, - ) - # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. - # grad_input[indices] = grad_output - grad_input.scatter_(0, repeat(indices, "z -> z d", d=grad_output.shape[1]), grad_output) - return grad_input.reshape(ctx.first_axis_dim, *other_shape), None - - -index_first_axis = IndexFirstAxis.apply - - -class IndexPutFirstAxis(torch.autograd.Function): - @staticmethod - def forward(ctx, values, indices, first_axis_dim): - ctx.save_for_backward(indices) - assert indices.ndim == 1 - assert values.ndim >= 2 - output = torch.zeros( - first_axis_dim, *values.shape[1:], device=values.device, dtype=values.dtype - ) - # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. - output[indices] = values - # output.scatter_(0, repeat(indices, 'z -> z d', d=values.shape[1]), values) - return output - - @staticmethod - def backward(ctx, grad_output): - (indices,) = ctx.saved_tensors - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - grad_values = grad_output[indices] - # grad_values = torch.gather(grad_output, 0, repeat(indices, 'z -> z d', d=grad_output.shape[1])) - return grad_values, None, None - - -index_put_first_axis = IndexPutFirstAxis.apply - - -class IndexFirstAxisResidual(torch.autograd.Function): - @staticmethod - def forward(ctx, input, indices): - ctx.save_for_backward(indices) - assert input.ndim >= 2 - ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] - second_dim = other_shape.numel() - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - output = input[indices] - # We don't want to reshape input (b ... -> b (...)) since it could change the channel_last - # memory format to channel_first. In other words, input might not be contiguous. - # If we don't detach, Pytorch complains about output being a view and is being modified inplace - return output, input.detach() - - @staticmethod - def backward(ctx, grad_output, grad_residual): - (indices,) = ctx.saved_tensors - assert grad_output.ndim >= 2 - other_shape = grad_output.shape[1:] - assert grad_residual.shape[1:] == other_shape - grad_input = grad_residual - # grad_input[indices] += grad_output - indices = indices.reshape(indices.shape[0], *((1,) * (grad_output.ndim - 1))) - indices = indices.expand_as(grad_output) - grad_input.scatter_add_(0, indices, grad_output) - return grad_input.reshape(ctx.first_axis_dim, *other_shape), None - - -index_first_axis_residual = IndexFirstAxisResidual.apply - - -def unpad_input(hidden_states, attention_mask, unused_mask=None): - """ - Arguments: - hidden_states: (batch, seqlen, ...) - attention_mask: (batch, seqlen), bool / int, 1 means valid and 0 means not valid. - unused_mask: (batch, seqlen), bool / int, 1 means the element is allocated but unused. - Return: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens selected in attention_mask + unused_mask. - indices: (total_nnz), the indices of masked tokens from the flattened input sequence. - cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. - max_seqlen_in_batch: int - seqused: (batch), returns the number of tokens selected in attention_mask + unused_mask. - """ - all_masks = (attention_mask + unused_mask) if unused_mask is not None else attention_mask - seqlens_in_batch = all_masks.sum(dim=-1, dtype=torch.int32) - used_seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) - indices = torch.nonzero(all_masks.flatten(), as_tuple=False).flatten() - max_seqlen_in_batch = seqlens_in_batch.max().item() - cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) - # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the - # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim - # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to - # index with integer indices. Moreover, torch's index is a bit slower than it needs to be, - # so we write custom forward and backward to make it a bit faster. - return ( - index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices), - indices, - cu_seqlens, - max_seqlen_in_batch, - used_seqlens_in_batch, - ) - - -def unpad_input_for_concatenated_sequences(hidden_states, attention_mask_in_length): - """ - Supports concatenating short samples in one sequence. The attention_mask_in_length is utilized to mask other short samples. It helps efficient training of variant lengths-based samples (e.g., the supervised fine-tuning task in large language model). - The motivation for this function is explained [here](https://github.com/Dao-AILab/flash-attention/issues/432#issuecomment-1668822286). - - For example, if batch = 3 and seqlen = 6, the attention_mask_in_length is: - ``` - [ - [2, 3, 0, 0, 0, 0], - [3, 2, 0, 0, 0, 0], - [6, 0, 0, 0, 0, 0] - ] - ``` - , which refers to the 3D-attention mask: - ``` - [ - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [0, 0, 1, 0, 0, 0], - [0, 0, 1, 1, 0, 0], - [0, 0, 1, 1, 1, 0], - [0, 0, 0, 0, 0, 1] - ], - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [0, 0, 0, 1, 0, 0], - [0, 0, 0, 1, 1, 0], - [0, 0, 0, 0, 0, 1] - ], - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [1, 1, 1, 1, 0, 0], - [1, 1, 1, 1, 1, 0], - [1, 1, 1, 1, 1, 1] - ] - ] - ```. - - Arguments: - hidden_states: (batch, seqlen, ...) - attention_mask_in_length: (batch, seqlen), int, a nonzero number (e.g., 1, 2, 3, etc.) means length of concatenated sequence in b-th batch, and 0 means none. - Return: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. - indices: (total_nnz), the indices of non-masked tokens from the flattened input sequence. - cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. - max_seqlen_in_batch: int - """ - length = attention_mask_in_length.sum(dim=-1) - seqlen = attention_mask_in_length.size(-1) - attention_mask_2d = torch.arange(seqlen, device=length.device, dtype=length.dtype).expand(len(length), seqlen) < length.unsqueeze(1) - real_indices_idx = torch.nonzero(attention_mask_in_length.flatten(), as_tuple=False).flatten() - seqlens_in_batch = attention_mask_in_length.flatten()[real_indices_idx] - indices = torch.nonzero(attention_mask_2d.flatten(), as_tuple=False).flatten() - max_seqlen_in_batch = seqlens_in_batch.max().item() - cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) - # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the - # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim - # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to - # index with integer indices. Moreover, torch's index is a bit slower than it needs to be, - # so we write custom forward and backward to make it a bit faster. - return ( - index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices), - indices, - cu_seqlens, - max_seqlen_in_batch, - ) - - -def pad_input(hidden_states, indices, batch, seqlen): - """ - Arguments: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. - indices: (total_nnz), the indices that represent the non-masked tokens of the original padded input sequence. - batch: int, batch size for the padded sequence. - seqlen: int, maximum sequence length for the padded sequence. - Return: - hidden_states: (batch, seqlen, ...) - """ - dim = hidden_states.shape[-1] - # output = torch.zeros((batch * seqlen), dim, device=hidden_states.device, dtype=hidden_states.dtype) - # output[indices] = hidden_states - output = index_put_first_axis(hidden_states, indices, batch * seqlen) - return rearrange(output, "(b s) ... -> b s ...", b=batch) diff --git a/build/torch210-cxx11-cu128-aarch64-linux/flash_attn2/__init__.py b/build/torch210-cxx11-cu128-aarch64-linux/flash_attn2/__init__.py deleted file mode 100644 index a9b2672c1cd85b74c1b3ded0fc0b2100e1aeac23..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu128-aarch64-linux/flash_attn2/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -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-aarch64-linux/flash_attn_interface.py b/build/torch210-cxx11-cu128-aarch64-linux/flash_attn_interface.py deleted file mode 100644 index 0868572aa3324d9c96132859d123e3d48d296e34..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu128-aarch64-linux/flash_attn_interface.py +++ /dev/null @@ -1,1662 +0,0 @@ -# Copyright (c) 2023, Tri Dao. - -from typing import Optional, Sequence, Tuple, Union - -import torch -import torch.nn as nn -import os - -# # isort: off -# # We need to import the CUDA kernels after importing torch -# USE_TRITON_ROCM = os.getenv("FLASH_ATTENTION_TRITON_AMD_ENABLE", "FALSE") == "TRUE" -# if USE_TRITON_ROCM: -# from .flash_attn_triton_amd import interface_fa as flash_attn -# else: -# import flash_attn_2_cuda as flash_attn - - -from ._ops import ops as flash_attn -from ._ops import add_op_namespace_prefix - -# # isort: on - - -def maybe_contiguous(x): - return x.contiguous() if x is not None and x.stride(-1) != 1 else x - - -def _get_device(): - if torch.xpu.is_available(): - return "xpu" - elif torch.cuda.is_available(): - return "cuda" - else: - return "cpu" - - -def _get_block_size_n(device, head_dim, is_dropout, is_causal): - assert head_dim <= 256 - - if device.type == "xpu": - if head_dim <= 96: - return 64 - elif head_dim <= 128: - return 32 - elif head_dim <= 256: - return 64 - else: - return 32 - - # This should match the block sizes in the CUDA kernel - major, minor = torch.cuda.get_device_capability(device) - is_sm8x = ( - major == 8 and minor > 0 - ) # Only include sm86 and sm89, exclude sm80 (A100) - is_sm80 = major == 8 and minor == 0 - is_sm90 = major == 9 and minor == 0 - if head_dim <= 32: - return 128 - if head_dim <= 64: - return 128 if not is_dropout else 64 - elif head_dim <= 96: - return 64 - elif head_dim <= 128: - if is_sm8x: - return 64 if (not is_dropout and is_causal) else 32 - else: - return 64 if not is_dropout else 32 - elif head_dim <= 192: - return 64 - elif head_dim <= 224: - return 64 - elif head_dim <= 256: - return 64 - - -def round_multiple(x, m): - return (x + m - 1) // m * m - - -@torch.library.custom_op( - add_op_namespace_prefix("_flash_attn_forward"), - mutates_args=(), - device_types=_get_device(), -) -def _flash_attn_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - return_softmax: bool, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - out, softmax_lse, S_dmask, rng_state = flash_attn.fwd( - q, - k, - v, - None, - alibi_slopes, - dropout_p, - softmax_scale, - causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - None, - ) - return out, softmax_lse, S_dmask, rng_state - - -@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, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - return_softmax: bool, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - batch_size, seqlen_q, num_heads, head_size = q.shape - seqlen_k = k.shape[1] - out = torch.empty_like(q) - softmax_lse = torch.empty( - (batch_size, num_heads, seqlen_q), - dtype=torch.float32, - device=q.device, - layout=q.layout, - ) - p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) - if return_softmax: - p = torch.empty( - ( - batch_size, - num_heads, - round_multiple(seqlen_q, 128), - round_multiple(seqlen_k, 128), - ), - dtype=q.dtype, - device=q.device, - layout=q.layout, - ) - rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) - - return out, softmax_lse, p, rng_state - - -@torch.library.custom_op( - add_op_namespace_prefix("_flash_attn_varlen_forward"), - mutates_args=(), - device_types=_get_device(), -) -def _flash_attn_varlen_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - alibi_slopes: Optional[torch.Tensor] = None, - return_softmax: bool = False, - block_table: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - out, softmax_lse, S_dmask, rng_state = flash_attn.varlen_fwd( - q, - k, - v, - None, - cu_seqlens_q, - cu_seqlens_k, - seqused_k, - leftpad_k, - block_table, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - zero_tensors, - causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - None, - ) - # if out.isnan().any() or softmax_lse.isnan().any(): - # breakpoint() - return out, softmax_lse, S_dmask, rng_state - - -@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_varlen_forward")) -def _flash_attn_varlen_forward_fake( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - alibi_slopes: Optional[torch.Tensor] = None, - return_softmax: bool = False, - block_table: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - paged_kv = block_table is not None - batch_size = cu_seqlens_q.numel() - 1 - total_q, num_heads, _ = q.shape - - out = torch.empty_like(q) - softmax_lse = torch.empty( - (num_heads, total_q), dtype=torch.float32, device=q.device, layout=q.layout - ) - p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) - seqlen_q_rounded = round_multiple(max_seqlen_q, 128) - seqlen_k_rounded = round_multiple(max_seqlen_k, 128) - if return_softmax: - p = torch.empty( - (batch_size, num_heads, seqlen_q_rounded, seqlen_k_rounded), - dtype=q.dtype, - device=q.device, - layout=q.layout, - ) - rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) - return out, softmax_lse, p, rng_state - - -@torch.library.custom_op( - add_op_namespace_prefix("_flash_attn_backward"), - mutates_args=("dq", "dk", "dv"), - device_types=_get_device(), -) -def _flash_attn_backward( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, -) -> 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)] - ( - dq, - dk, - dv, - softmax_d, - ) = flash_attn.bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - alibi_slopes, - dropout_p, - softmax_scale, - causal, - window_size_left, - window_size_right, - softcap, - deterministic, - None, - rng_state, - ) - 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, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, -) -> torch.Tensor: - dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] - if dq is None: - dq = torch.empty_like(q) - if dk is None: - dk = torch.empty_like(k) - if dv is None: - dv = torch.empty_like(v) - batch_size, seqlen_q, num_heads, _ = q.shape - softmax_d = torch.empty( - (batch_size, num_heads, round_multiple(seqlen_q, 128)), - device=q.device, - dtype=torch.float32, - ) - - return softmax_d - - -@torch.library.custom_op( - add_op_namespace_prefix("_flash_attn_varlen_backward"), - mutates_args=("dq", "dk", "dv"), - device_types=_get_device(), -) -def _flash_attn_varlen_backward( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> 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)] - ( - dq, - dk, - dv, - softmax_d, - ) = flash_attn.varlen_bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - zero_tensors, - causal, - window_size_left, - window_size_right, - softcap, - deterministic, - None, - rng_state, - ) - # if dk.isnan().any() or dk.isnan().any() or dv.isnan().any() or softmax_d.isnan().any(): - # breakpoint() - return softmax_d - - -@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_varlen_backward")) -def _flash_attn_varlen_backward_fake( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> torch.Tensor: - dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] - batch_size = cu_seqlens_q.numel() - 1 - total_q, num_heads, _ = q.shape - - if dq is None: - dq = torch.empty_like(q) - if dk is None: - dk = torch.empty_like(k) - if dv is None: - dv = torch.empty_like(v) - softmax_d = torch.empty( - (num_heads, total_q + 128 * batch_size), device=q.device, dtype=torch.float32 - ) - - return softmax_d - - -class FlashAttnQKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and qkv.requires_grad - if softmax_scale is None: - softmax_scale = qkv.shape[-1] ** (-0.5) - q, k, v = qkv[:, :, 0].detach(), qkv[:, :, 1].detach(), qkv[:, :, 2].detach() - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) - dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dqkv[:, :, 0], - dqkv[:, :, 1], - dqkv[:, :, 2], - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dqkv = dqkv[..., : dout.shape[-1]] # We could have padded the head dimension - return dqkv, None, None, None, None, None, None, None, None, None - - -class FlashAttnVarlenQKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - cu_seqlens, - max_seqlen, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and qkv.requires_grad - if softmax_scale is None: - softmax_scale = qkv.shape[-1] ** (-0.5) - q, k, v = qkv[:, 0].detach(), qkv[:, 1].detach(), qkv[:, 2].detach() - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens, - cu_seqlens, - max_seqlen, - max_seqlen, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=None, - ) - if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens, rng_state - ) - ctx.dropout_p = dropout_p - ctx.max_seqlen = max_seqlen - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens, rng_state = ctx.saved_tensors - qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) - dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dqkv[:, 0], - dqkv[:, 1], - dqkv[:, 2], - cu_seqlens, - cu_seqlens, - ctx.max_seqlen, - ctx.max_seqlen, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - 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 FlashAttnKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - kv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any(x.requires_grad for x in [q, kv]) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - k, v = kv[:, :, 0].detach(), kv[:, :, 1].detach() - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - dq = torch.empty_like(q) - kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) - dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dkv[:, :, 0], - dkv[:, :, 1], - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dkv = dkv[..., : dout.shape[-1]] - return dq, dkv, None, None, None, None, None, None, None, None, None - - -class FlashAttnVarlenKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any(x.requires_grad for x in [q, kv]) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - k, v = kv[:, 0].detach(), kv[:, 1].detach() - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=None, - ) - if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state - ) - ctx.dropout_p = dropout_p - 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.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ( - ctx.saved_tensors - ) - dq = torch.empty_like(q) - kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) - dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dkv[:, 0], - dkv[:, 1], - cu_seqlens_q, - cu_seqlens_k, - ctx.max_seqlen_q, - ctx.max_seqlen_k, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dkv = dkv[..., : dout.shape[-1]] - return ( - dq, - dkv, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) - - -class FlashAttnFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - k, - v, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any(x.requires_grad for x in [q, k, v]) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dk = dk[..., : dout.shape[-1]] - dv = dv[..., : dout.shape[-1]] - return dq, dk, dv, 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, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - block_table, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any(x.requires_grad for x in [q, k, v]) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=block_table, - ) - if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state - ) - ctx.dropout_p = dropout_p - 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.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ( - ctx.saved_tensors - ) - dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - ctx.max_seqlen_q, - ctx.max_seqlen_k, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dk = dk[..., : dout.shape[-1]] - dv = dv[..., : dout.shape[-1]] - return ( - dq, - dk, - dv, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) - - -def flash_attn_qkvpacked_func( - qkv, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # <=0.0 means deactivate - alibi_slopes=None, - deterministic=False, - 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, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_kvpacked_func( - q, - kv, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, -): - """dropout_p should be set to 0.0 during evaluation - If 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 K, V. - 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) - kv: (batch_size, seqlen, 2, 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. - softcap: float. Anything > 0 activates softcapping 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). - 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 FlashAttnKVPackedFunc.apply( - q, - kv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_func( - q, - k, - v, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - 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). - 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 FlashAttnFunc.apply( - q, - k, - v, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_qkvpacked_func( - qkv, - cu_seqlens, - max_seqlen, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - 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_varlen_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_varlen_kvpacked_func and flash_attn_varlen_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: (total, 3, nheads, headdim), where total = total number of tokens in the batch. - cu_seqlens: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into qkv. - max_seqlen: int. Maximum sequence length in the batch. - 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenQKVPackedFunc.apply( - qkv, - cu_seqlens, - max_seqlen, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_kvpacked_func( - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, -): - """dropout_p should be set to 0.0 during evaluation - If 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 K, V. - 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: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch. - kv: (total_k, 2, nheads_k, headdim), where total_k = total number of key tokens in the batch. - cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into q. - cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into kv. - max_seqlen_q: int. Maximum query sequence length in the batch. - max_seqlen_k: int. Maximum key sequence length in the batch. - 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 + 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenKVPackedFunc.apply( - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_func( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, - block_table=None, -): - """dropout_p should be set to 0.0 during evaluation - Supports multi-query and grouped-query attention (MQA/GQA) by passing in K, V 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: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch. - k: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. - v: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. - cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into q. - cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into kv. - max_seqlen_q: int. Maximum query sequence length in the batch. - max_seqlen_k: int. Maximum key sequence length in the batch. - 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 + 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenFunc.apply( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - block_table, - False if q.device.type == "cpu" else torch.is_grad_enabled(), - ) - - -def flash_attn_with_kvcache( - q, - k_cache, - v_cache, - k=None, - v=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, - block_table: Optional[torch.Tensor] = None, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - rotary_interleaved=True, - alibi_slopes=None, - num_splits=0, - 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 block_table, - or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_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) if there's no block_table, - or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_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). Similar to k. - 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. - block_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). - 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. - 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" - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-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) - cache_batch_idx = maybe_contiguous(cache_batch_idx) - block_table = maybe_contiguous(block_table) - out, softmax_lse = flash_attn.fwd_kvcache( - q, - k_cache, - v_cache, - k, - v, - cache_seqlens, - rotary_cos, - rotary_sin, - cache_batch_idx, - cache_leftpad, - block_table, - alibi_slopes, - None, - softmax_scale, - causal, - window_size[0], - window_size[1], - softcap, - rotary_interleaved, - num_splits, - ) - return (out, softmax_lse) if return_softmax_lse else out diff --git a/build/torch210-cxx11-cu128-aarch64-linux/metadata.json b/build/torch210-cxx11-cu128-aarch64-linux/metadata.json deleted file mode 100644 index 5912a9cf2f79a4ad05d354c2cf7559adff1b1ac2..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu128-aarch64-linux/metadata.json +++ /dev/null @@ -1,16 +0,0 @@ -{ - "name": "flash-attn2", - "id": "_flash_attn2_cuda_a511a4c", - "version": 1, - "license": "BSD-3-Clause", - "python-depends": [], - "backend": { - "type": "cuda", - "archs": [ - "10.0", - "12.0", - "8.0", - "9.0" - ] - } -} diff --git a/build/torch210-cxx11-cu128-x86_64-linux/__init__.py b/build/torch210-cxx11-cu128-x86_64-linux/__init__.py index 20ff0ff747abd7500913863b509925f414b1c973..ecc2f9d896b6c93f90b0a1499856dc0612177422 100644 --- a/build/torch210-cxx11-cu128-x86_64-linux/__init__.py +++ b/build/torch210-cxx11-cu128-x86_64-linux/__init__.py @@ -1,7 +1,5 @@ -from typing import List, Optional - +from typing import Optional, List import torch - from ._ops import ops as flash_attn_ops from .flash_attn_interface import ( flash_attn_func, @@ -13,23 +11,6 @@ from .flash_attn_interface import ( flash_attn_with_kvcache, ) -__all__ = [ - # High-level API (with autograd support) - "flash_attn_func", - "flash_attn_kvpacked_func", - "flash_attn_qkvpacked_func", - "flash_attn_varlen_func", - "flash_attn_varlen_kvpacked_func", - "flash_attn_varlen_qkvpacked_func", - "flash_attn_with_kvcache", - # Low-level ops (no autograd) - "fwd", - "varlen_fwd", - "bwd", - "varlen_bwd", - "fwd_kvcache", -] - def fwd( q: torch.Tensor, diff --git a/build/torch210-cxx11-cu128-x86_64-linux/_flash_attn2_588b404.abi3.so b/build/torch210-cxx11-cu128-x86_64-linux/_flash_attn2_588b404.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..dc5528ea731e761511650f7535b9e4c1a1b90e20 --- /dev/null +++ b/build/torch210-cxx11-cu128-x86_64-linux/_flash_attn2_588b404.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:09cfe096dc8f0010e99225d44263e4d9172d4b542d48d656b3b9fd718ca55b7d +size 1037803376 diff --git a/build/torch210-cxx11-cu128-x86_64-linux/_flash_attn2_cuda_a511a4c.abi3.so b/build/torch210-cxx11-cu128-x86_64-linux/_flash_attn2_cuda_a511a4c.abi3.so deleted file mode 100644 index 1bd4e3ea0326d594f0189f96120f4d27481bfc28..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu128-x86_64-linux/_flash_attn2_cuda_a511a4c.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:c2acd487c5d1d77b7f6721eabc0fee4f8c5c8f9e682aadf43bfa99b394eecba7 -size 1037795600 diff --git a/build/torch210-cxx11-cu128-x86_64-linux/_ops.py b/build/torch210-cxx11-cu128-x86_64-linux/_ops.py index 62b70a9b02c8f7663c7022564c1a7f58f46b6c65..03506d8a18d17ca7296a5f1b7a18ec7dc4ad73bb 100644 --- a/build/torch210-cxx11-cu128-x86_64-linux/_ops.py +++ b/build/torch210-cxx11-cu128-x86_64-linux/_ops.py @@ -1,9 +1,9 @@ import torch -from . import _flash_attn2_cuda_a511a4c -ops = torch.ops._flash_attn2_cuda_a511a4c +from . import _flash_attn2_588b404 +ops = torch.ops._flash_attn2_588b404 def add_op_namespace_prefix(op_name: str): """ Prefix op by namespace. """ - return f"_flash_attn2_cuda_a511a4c::{op_name}" + return f"_flash_attn2_588b404::{op_name}" \ No newline at end of file diff --git a/build/torch210-cxx11-cu128-x86_64-linux/flash_attn2/__init__.py b/build/torch210-cxx11-cu128-x86_64-linux/flash_attn2/__init__.py index a9b2672c1cd85b74c1b3ded0fc0b2100e1aeac23..03dbc1afe1cf156661a2b1b22003cd5f599a0309 100644 --- a/build/torch210-cxx11-cu128-x86_64-linux/flash_attn2/__init__.py +++ b/build/torch210-cxx11-cu128-x86_64-linux/flash_attn2/__init__.py @@ -1,10 +1,10 @@ import ctypes -import importlib.util 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 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 index 0868572aa3324d9c96132859d123e3d48d296e34..66e8d3665834b141d703998496b3da6fe1f6fe06 100644 --- a/build/torch210-cxx11-cu128-x86_64-linux/flash_attn_interface.py +++ b/build/torch210-cxx11-cu128-x86_64-linux/flash_attn_interface.py @@ -16,11 +16,9 @@ import os from ._ops import ops as flash_attn -from ._ops import add_op_namespace_prefix # # isort: on - def maybe_contiguous(x): return x.contiguous() if x is not None and x.stride(-1) != 1 else x @@ -33,25 +31,14 @@ def _get_device(): else: return "cpu" +_XPU_AVAILABLE = torch.xpu.is_available() if hasattr(torch, "xpu") else False # TODO remove hasattr check when bwd is supported on XPU -def _get_block_size_n(device, head_dim, is_dropout, is_causal): - assert head_dim <= 256 - - if device.type == "xpu": - if head_dim <= 96: - return 64 - elif head_dim <= 128: - return 32 - elif head_dim <= 256: - return 64 - else: - return 32 +def _get_block_size_n(device, head_dim, is_dropout, is_causal): # This should match the block sizes in the CUDA kernel + assert head_dim <= 256 major, minor = torch.cuda.get_device_capability(device) - is_sm8x = ( - major == 8 and minor > 0 - ) # Only include sm86 and sm89, exclude sm80 (A100) + is_sm8x = major == 8 and minor > 0 # Only include sm86 and sm89, exclude sm80 (A100) is_sm80 = major == 8 and minor == 0 is_sm90 = major == 9 and minor == 0 if head_dim <= 32: @@ -77,11 +64,30 @@ def round_multiple(x, m): return (x + m - 1) // m * m -@torch.library.custom_op( - add_op_namespace_prefix("_flash_attn_forward"), - mutates_args=(), - device_types=_get_device(), -) +# torch.compile() support is only enabled for pytorch >= 2.4 +# The reason for this is that we are using the new custom_op and register_fake +# APIs, which support inplace modification of inputs in the function itself +if torch.__version__ >= "2.4.0": + _torch_custom_op_wrapper = torch.library.custom_op + _torch_register_fake_wrapper = torch.library.register_fake +else: + def noop_custom_op_wrapper(name, fn=None, /, *, mutates_args, device_types=None, schema=None): + def wrap(func): + return func + if fn is None: + return wrap + return fn + def noop_register_fake_wrapper(op, fn=None, /, *, lib=None, _stacklevel=1): + def wrap(func): + return func + if fn is None: + return wrap + return fn + _torch_custom_op_wrapper = noop_custom_op_wrapper + _torch_register_fake_wrapper = noop_register_fake_wrapper + + +@_torch_custom_op_wrapper("flash_attn::_flash_attn_forward", mutates_args=(), device_types=_get_device()) def _flash_attn_forward( q: torch.Tensor, k: torch.Tensor, @@ -93,7 +99,7 @@ def _flash_attn_forward( window_size_right: int, softcap: float, alibi_slopes: Optional[torch.Tensor], - return_softmax: bool, + return_softmax: bool ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: q, k, v = [maybe_contiguous(x) for x in (q, k, v)] out, softmax_lse, S_dmask, rng_state = flash_attn.fwd( @@ -114,7 +120,7 @@ def _flash_attn_forward( return out, softmax_lse, S_dmask, rng_state -@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_forward")) +@_torch_register_fake_wrapper("flash_attn::_flash_attn_forward") def _flash_attn_forward_fake( q: torch.Tensor, k: torch.Tensor, @@ -126,41 +132,28 @@ def _flash_attn_forward_fake( window_size_right: int, softcap: float, alibi_slopes: Optional[torch.Tensor], - return_softmax: bool, + return_softmax: bool ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: q, k, v = [maybe_contiguous(x) for x in (q, k, v)] batch_size, seqlen_q, num_heads, head_size = q.shape seqlen_k = k.shape[1] out = torch.empty_like(q) - softmax_lse = torch.empty( - (batch_size, num_heads, seqlen_q), - dtype=torch.float32, - device=q.device, - layout=q.layout, - ) + softmax_lse = torch.empty((batch_size, num_heads, seqlen_q), dtype=torch.float32, device=q.device, layout=q.layout) p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) if return_softmax: - p = torch.empty( - ( - batch_size, - num_heads, - round_multiple(seqlen_q, 128), - round_multiple(seqlen_k, 128), - ), - dtype=q.dtype, - device=q.device, - layout=q.layout, - ) + p = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128), round_multiple(seqlen_k, 128)), dtype=q.dtype, device=q.device, layout=q.layout) rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) return out, softmax_lse, p, rng_state -@torch.library.custom_op( - add_op_namespace_prefix("_flash_attn_varlen_forward"), - mutates_args=(), - device_types=_get_device(), -) +if torch.__version__ >= "2.4.0": + _wrapped_flash_attn_forward = torch.ops.flash_attn._flash_attn_forward +else: + _wrapped_flash_attn_forward = _flash_attn_forward + + +@_torch_custom_op_wrapper("flash_attn::_flash_attn_varlen_forward", mutates_args=(), device_types=_get_device()) def _flash_attn_varlen_forward( q: torch.Tensor, k: torch.Tensor, @@ -211,7 +204,7 @@ def _flash_attn_varlen_forward( return out, softmax_lse, S_dmask, rng_state -@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_varlen_forward")) +@_torch_register_fake_wrapper("flash_attn::_flash_attn_varlen_forward") def _flash_attn_varlen_forward_fake( q: torch.Tensor, k: torch.Tensor, @@ -237,30 +230,25 @@ def _flash_attn_varlen_forward_fake( paged_kv = block_table is not None batch_size = cu_seqlens_q.numel() - 1 total_q, num_heads, _ = q.shape - + out = torch.empty_like(q) - softmax_lse = torch.empty( - (num_heads, total_q), dtype=torch.float32, device=q.device, layout=q.layout - ) + softmax_lse = torch.empty((num_heads, total_q), dtype=torch.float32, device=q.device, layout=q.layout) p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) seqlen_q_rounded = round_multiple(max_seqlen_q, 128) seqlen_k_rounded = round_multiple(max_seqlen_k, 128) if return_softmax: - p = torch.empty( - (batch_size, num_heads, seqlen_q_rounded, seqlen_k_rounded), - dtype=q.dtype, - device=q.device, - layout=q.layout, - ) + p = torch.empty((batch_size, num_heads, seqlen_q_rounded, seqlen_k_rounded), dtype=q.dtype, device=q.device, layout=q.layout) rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) return out, softmax_lse, p, rng_state -@torch.library.custom_op( - add_op_namespace_prefix("_flash_attn_backward"), - mutates_args=("dq", "dk", "dv"), - device_types=_get_device(), -) +if torch.__version__ >= "2.4.0": + _wrapped_flash_attn_varlen_forward = torch.ops.flash_attn._flash_attn_varlen_forward +else: + _wrapped_flash_attn_varlen_forward = _flash_attn_varlen_forward + + +@_torch_custom_op_wrapper("flash_attn::_flash_attn_backward", mutates_args=("dq", "dk", "dv"), device_types=_get_device()) def _flash_attn_backward( dout: torch.Tensor, q: torch.Tensor, @@ -312,7 +300,7 @@ def _flash_attn_backward( return softmax_d -@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_backward")) +@_torch_register_fake_wrapper("flash_attn::_flash_attn_backward") def _flash_attn_backward_fake( dout: torch.Tensor, q: torch.Tensor, @@ -341,20 +329,18 @@ def _flash_attn_backward_fake( if dv is None: dv = torch.empty_like(v) batch_size, seqlen_q, num_heads, _ = q.shape - softmax_d = torch.empty( - (batch_size, num_heads, round_multiple(seqlen_q, 128)), - device=q.device, - dtype=torch.float32, - ) - + softmax_d = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128)), device=q.device, dtype=torch.float32) + return softmax_d -@torch.library.custom_op( - add_op_namespace_prefix("_flash_attn_varlen_backward"), - mutates_args=("dq", "dk", "dv"), - device_types=_get_device(), -) +if torch.__version__ >= "2.4.0": + _wrapped_flash_attn_backward = torch.ops.flash_attn._flash_attn_backward +else: + _wrapped_flash_attn_backward = _flash_attn_backward + + +@_torch_custom_op_wrapper("flash_attn::_flash_attn_varlen_backward", mutates_args=("dq", "dk", "dv"), device_types=_get_device()) def _flash_attn_varlen_backward( dout: torch.Tensor, q: torch.Tensor, @@ -418,7 +404,7 @@ def _flash_attn_varlen_backward( return softmax_d -@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_varlen_backward")) +@_torch_register_fake_wrapper("flash_attn::_flash_attn_varlen_backward") def _flash_attn_varlen_backward_fake( dout: torch.Tensor, q: torch.Tensor, @@ -454,13 +440,17 @@ def _flash_attn_varlen_backward_fake( dk = torch.empty_like(k) if dv is None: dv = torch.empty_like(v) - softmax_d = torch.empty( - (num_heads, total_q + 128 * batch_size), device=q.device, dtype=torch.float32 - ) - + softmax_d = torch.empty((num_heads, total_q + 128 * batch_size), device=q.device, dtype=torch.float32) + return softmax_d +if torch.__version__ >= "2.4.0": + _wrapped_flash_attn_varlen_backward = torch.ops.flash_attn._flash_attn_varlen_backward +else: + _wrapped_flash_attn_varlen_backward = _flash_attn_varlen_backward + + class FlashAttnQKVPackedFunc(torch.autograd.Function): @staticmethod def forward( @@ -485,7 +475,7 @@ class FlashAttnQKVPackedFunc(torch.autograd.Function): q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_forward( + out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( q, k, v, @@ -519,7 +509,7 @@ class FlashAttnQKVPackedFunc(torch.autograd.Function): dout_padded = dout if head_size_og % 8 != 0: dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _flash_attn_backward( + _wrapped_flash_attn_backward( dout_padded, q, k, @@ -569,7 +559,7 @@ class FlashAttnVarlenQKVPackedFunc(torch.autograd.Function): q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_varlen_forward( + out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( q, k, v, @@ -588,9 +578,7 @@ class FlashAttnVarlenQKVPackedFunc(torch.autograd.Function): block_table=None, ) if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens, rng_state - ) + ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens, rng_state) ctx.dropout_p = dropout_p ctx.max_seqlen = max_seqlen ctx.softmax_scale = softmax_scale @@ -611,7 +599,7 @@ class FlashAttnVarlenQKVPackedFunc(torch.autograd.Function): dout_padded = dout if head_size_og % 8 != 0: dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _flash_attn_varlen_backward( + _wrapped_flash_attn_varlen_backward( dout_padded, q, k, @@ -655,7 +643,9 @@ class FlashAttnKVPackedFunc(torch.autograd.Function): return_softmax, is_grad_enabled, ): - is_grad = is_grad_enabled and any(x.requires_grad for x in [q, kv]) + is_grad = is_grad_enabled and any( + x.requires_grad for x in [q, kv] + ) if softmax_scale is None: softmax_scale = q.shape[-1] ** (-0.5) k, v = kv[:, :, 0].detach(), kv[:, :, 1].detach() @@ -664,7 +654,7 @@ class FlashAttnKVPackedFunc(torch.autograd.Function): q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_forward( + out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( q, k, v, @@ -699,7 +689,7 @@ class FlashAttnKVPackedFunc(torch.autograd.Function): dout_padded = dout if head_size_og % 8 != 0: dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _flash_attn_backward( + _wrapped_flash_attn_backward( dout_padded, q, k, @@ -744,7 +734,9 @@ class FlashAttnVarlenKVPackedFunc(torch.autograd.Function): return_softmax, is_grad_enabled, ): - is_grad = is_grad_enabled and any(x.requires_grad for x in [q, kv]) + is_grad = is_grad_enabled and any( + x.requires_grad for x in [q, kv] + ) if softmax_scale is None: softmax_scale = q.shape[-1] ** (-0.5) k, v = kv[:, 0].detach(), kv[:, 1].detach() @@ -753,7 +745,7 @@ class FlashAttnVarlenKVPackedFunc(torch.autograd.Function): q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_varlen_forward( + out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( q, k, v, @@ -789,9 +781,7 @@ class FlashAttnVarlenKVPackedFunc(torch.autograd.Function): @staticmethod def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ( - ctx.saved_tensors - ) + q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors dq = torch.empty_like(q) kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) @@ -799,7 +789,7 @@ class FlashAttnVarlenKVPackedFunc(torch.autograd.Function): dout_padded = dout if head_size_og % 8 != 0: dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _flash_attn_varlen_backward( + _wrapped_flash_attn_varlen_backward( dout_padded, q, k, @@ -825,23 +815,7 @@ class FlashAttnVarlenKVPackedFunc(torch.autograd.Function): ) dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension dkv = dkv[..., : dout.shape[-1]] - return ( - dq, - dkv, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) + return dq, dkv, None, None, None, None, None, None, None, None, None, None, None, None, None class FlashAttnFunc(torch.autograd.Function): @@ -861,7 +835,9 @@ class FlashAttnFunc(torch.autograd.Function): return_softmax, is_grad_enabled, ): - is_grad = is_grad_enabled and any(x.requires_grad for x in [q, k, v]) + is_grad = is_grad_enabled and any( + x.requires_grad for x in [q, k, v] + ) if softmax_scale is None: softmax_scale = q.shape[-1] ** (-0.5) head_size_og = q.size(3) @@ -869,7 +845,7 @@ class FlashAttnFunc(torch.autograd.Function): q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_forward( + out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( q, k, v, @@ -902,7 +878,7 @@ class FlashAttnFunc(torch.autograd.Function): dout_padded = dout if head_size_og % 8 != 0: dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _flash_attn_backward( + _wrapped_flash_attn_backward( dout_padded, q, k, @@ -950,7 +926,9 @@ class FlashAttnVarlenFunc(torch.autograd.Function): block_table, is_grad_enabled, ): - is_grad = is_grad_enabled and any(x.requires_grad for x in [q, k, v]) + is_grad = is_grad_enabled and any( + x.requires_grad for x in [q, k, v] + ) if softmax_scale is None: softmax_scale = q.shape[-1] ** (-0.5) head_size_og = q.size(2) @@ -958,7 +936,7 @@ class FlashAttnVarlenFunc(torch.autograd.Function): q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_varlen_forward( + out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( q, k, v, @@ -995,15 +973,13 @@ class FlashAttnVarlenFunc(torch.autograd.Function): @staticmethod def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ( - ctx.saved_tensors - ) + q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) head_size_og = dout.size(2) dout_padded = dout if head_size_og % 8 != 0: dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _flash_attn_varlen_backward( + _wrapped_flash_attn_varlen_backward( dout_padded, q, k, @@ -1030,25 +1006,7 @@ class FlashAttnVarlenFunc(torch.autograd.Function): dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension dk = dk[..., : dout.shape[-1]] dv = dv[..., : dout.shape[-1]] - return ( - dq, - dk, - dv, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None def flash_attn_qkvpacked_func( @@ -1106,7 +1064,7 @@ def flash_attn_qkvpacked_func( alibi_slopes, deterministic, return_attn_probs, - torch.is_grad_enabled(), + False if _XPU_AVAILABLE else torch.is_grad_enabled(), ) @@ -1184,7 +1142,7 @@ def flash_attn_kvpacked_func( alibi_slopes, deterministic, return_attn_probs, - torch.is_grad_enabled(), + False if _XPU_AVAILABLE else torch.is_grad_enabled(), ) @@ -1196,7 +1154,7 @@ def flash_attn_func( softmax_scale=None, causal=False, window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated + softcap=0.0, # 0.0 means deactivated alibi_slopes=None, deterministic=False, return_attn_probs=False, @@ -1261,7 +1219,7 @@ def flash_attn_func( alibi_slopes, deterministic, return_attn_probs, - torch.is_grad_enabled(), + False if _XPU_AVAILABLE else torch.is_grad_enabled(), ) @@ -1273,7 +1231,7 @@ def flash_attn_varlen_qkvpacked_func( softmax_scale=None, causal=False, window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated + softcap=0.0, # 0.0 means deactivated alibi_slopes=None, deterministic=False, return_attn_probs=False, @@ -1327,7 +1285,7 @@ def flash_attn_varlen_qkvpacked_func( alibi_slopes, deterministic, return_attn_probs, - torch.is_grad_enabled(), + False if _XPU_AVAILABLE else torch.is_grad_enabled(), ) @@ -1342,7 +1300,7 @@ def flash_attn_varlen_kvpacked_func( softmax_scale=None, causal=False, window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated + softcap=0.0, # 0.0 means deactivated alibi_slopes=None, deterministic=False, return_attn_probs=False, @@ -1419,7 +1377,7 @@ def flash_attn_varlen_kvpacked_func( alibi_slopes, deterministic, return_attn_probs, - torch.is_grad_enabled(), + False if _XPU_AVAILABLE else torch.is_grad_enabled(), ) @@ -1435,7 +1393,7 @@ def flash_attn_varlen_func( softmax_scale=None, causal=False, window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated + softcap=0.0, # 0.0 means deactivated alibi_slopes=None, deterministic=False, return_attn_probs=False, @@ -1513,7 +1471,7 @@ def flash_attn_varlen_func( deterministic, return_attn_probs, block_table, - False if q.device.type == "cpu" else torch.is_grad_enabled(), + False if _XPU_AVAILABLE or q.device.type == "cpu" else torch.is_grad_enabled(), ) @@ -1532,7 +1490,7 @@ def flash_attn_with_kvcache( softmax_scale=None, causal=False, window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated + softcap=0.0, # 0.0 means deactivated rotary_interleaved=True, alibi_slopes=None, num_splits=0, diff --git a/build/torch211-cxx11-cpu-x86_64-linux/layers/__init__.py b/build/torch210-cxx11-cu128-x86_64-linux/layers/__init__.py similarity index 100% rename from build/torch211-cxx11-cpu-x86_64-linux/layers/__init__.py rename to build/torch210-cxx11-cu128-x86_64-linux/layers/__init__.py diff --git a/build/torch211-cxx11-cu126-x86_64-linux/layers/patch_embed.py b/build/torch210-cxx11-cu128-x86_64-linux/layers/patch_embed.py similarity index 100% rename from build/torch211-cxx11-cu126-x86_64-linux/layers/patch_embed.py rename to build/torch210-cxx11-cu128-x86_64-linux/layers/patch_embed.py diff --git a/build/torch211-cxx11-cu126-x86_64-linux/layers/rotary.py b/build/torch210-cxx11-cu128-x86_64-linux/layers/rotary.py similarity index 100% rename from build/torch211-cxx11-cu126-x86_64-linux/layers/rotary.py rename to build/torch210-cxx11-cu128-x86_64-linux/layers/rotary.py diff --git a/build/torch210-cxx11-cu128-x86_64-linux/metadata.json b/build/torch210-cxx11-cu128-x86_64-linux/metadata.json index 5912a9cf2f79a4ad05d354c2cf7559adff1b1ac2..9cf5deed9898dce769f4cc73913d3530b92a0bd8 100644 --- a/build/torch210-cxx11-cu128-x86_64-linux/metadata.json +++ b/build/torch210-cxx11-cu128-x86_64-linux/metadata.json @@ -1,16 +1,4 @@ { - "name": "flash-attn2", - "id": "_flash_attn2_cuda_a511a4c", "version": 1, - "license": "BSD-3-Clause", - "python-depends": [], - "backend": { - "type": "cuda", - "archs": [ - "10.0", - "12.0", - "8.0", - "9.0" - ] - } -} + "python-depends": [] +} \ No newline at end of file diff --git a/build/torch211-cxx11-cpu-x86_64-linux/ops/__init__.py b/build/torch210-cxx11-cu128-x86_64-linux/ops/__init__.py similarity index 100% rename from build/torch211-cxx11-cpu-x86_64-linux/ops/__init__.py rename to build/torch210-cxx11-cu128-x86_64-linux/ops/__init__.py diff --git a/build/torch211-cxx11-cu126-x86_64-linux/ops/activations.py b/build/torch210-cxx11-cu128-x86_64-linux/ops/activations.py similarity index 100% rename from build/torch211-cxx11-cu126-x86_64-linux/ops/activations.py rename to build/torch210-cxx11-cu128-x86_64-linux/ops/activations.py diff --git a/build/torch211-cxx11-cu126-x86_64-linux/ops/fused_dense.py b/build/torch210-cxx11-cu128-x86_64-linux/ops/fused_dense.py similarity index 100% rename from build/torch211-cxx11-cu126-x86_64-linux/ops/fused_dense.py rename to build/torch210-cxx11-cu128-x86_64-linux/ops/fused_dense.py diff --git a/build/torch211-cxx11-cu126-x86_64-linux/ops/layer_norm.py b/build/torch210-cxx11-cu128-x86_64-linux/ops/layer_norm.py similarity index 100% rename from build/torch211-cxx11-cu126-x86_64-linux/ops/layer_norm.py rename to build/torch210-cxx11-cu128-x86_64-linux/ops/layer_norm.py diff --git a/build/torch211-cxx11-cu126-x86_64-linux/ops/rms_norm.py b/build/torch210-cxx11-cu128-x86_64-linux/ops/rms_norm.py similarity index 100% rename from build/torch211-cxx11-cu126-x86_64-linux/ops/rms_norm.py rename to build/torch210-cxx11-cu128-x86_64-linux/ops/rms_norm.py diff --git a/build/torch211-cxx11-cu126-x86_64-linux/ops/triton/__init__.py b/build/torch210-cxx11-cu128-x86_64-linux/ops/triton/__init__.py similarity index 100% rename from build/torch211-cxx11-cu126-x86_64-linux/ops/triton/__init__.py rename to build/torch210-cxx11-cu128-x86_64-linux/ops/triton/__init__.py diff --git a/build/torch211-cxx11-cu126-x86_64-linux/ops/triton/cross_entropy.py b/build/torch210-cxx11-cu128-x86_64-linux/ops/triton/cross_entropy.py similarity index 100% rename from build/torch211-cxx11-cu126-x86_64-linux/ops/triton/cross_entropy.py rename to build/torch210-cxx11-cu128-x86_64-linux/ops/triton/cross_entropy.py diff --git a/build/torch211-cxx11-cu126-x86_64-linux/ops/triton/k_activations.py b/build/torch210-cxx11-cu128-x86_64-linux/ops/triton/k_activations.py similarity index 100% rename from build/torch211-cxx11-cu126-x86_64-linux/ops/triton/k_activations.py rename to build/torch210-cxx11-cu128-x86_64-linux/ops/triton/k_activations.py diff --git a/build/torch29-cxx11-cu129-aarch64-linux/ops/triton/layer_norm.py b/build/torch210-cxx11-cu128-x86_64-linux/ops/triton/layer_norm.py similarity index 100% rename from build/torch29-cxx11-cu129-aarch64-linux/ops/triton/layer_norm.py rename to build/torch210-cxx11-cu128-x86_64-linux/ops/triton/layer_norm.py diff --git a/build/torch211-cxx11-cu126-x86_64-linux/ops/triton/linear.py b/build/torch210-cxx11-cu128-x86_64-linux/ops/triton/linear.py similarity index 100% rename from build/torch211-cxx11-cu126-x86_64-linux/ops/triton/linear.py rename to build/torch210-cxx11-cu128-x86_64-linux/ops/triton/linear.py diff --git a/build/torch211-cxx11-cu126-x86_64-linux/ops/triton/mlp.py b/build/torch210-cxx11-cu128-x86_64-linux/ops/triton/mlp.py similarity index 100% rename from build/torch211-cxx11-cu126-x86_64-linux/ops/triton/mlp.py rename to build/torch210-cxx11-cu128-x86_64-linux/ops/triton/mlp.py diff --git a/build/torch211-cxx11-cu126-x86_64-linux/ops/triton/rotary.py b/build/torch210-cxx11-cu128-x86_64-linux/ops/triton/rotary.py similarity index 100% rename from build/torch211-cxx11-cu126-x86_64-linux/ops/triton/rotary.py rename to build/torch210-cxx11-cu128-x86_64-linux/ops/triton/rotary.py diff --git a/build/torch210-cxx11-cu130-aarch64-linux/__init__.py b/build/torch210-cxx11-cu130-aarch64-linux/__init__.py deleted file mode 100644 index 20ff0ff747abd7500913863b509925f414b1c973..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu130-aarch64-linux/__init__.py +++ /dev/null @@ -1,412 +0,0 @@ -from typing import List, Optional - -import torch - -from ._ops import ops as flash_attn_ops -from .flash_attn_interface import ( - flash_attn_func, - flash_attn_kvpacked_func, - flash_attn_qkvpacked_func, - flash_attn_varlen_func, - flash_attn_varlen_kvpacked_func, - flash_attn_varlen_qkvpacked_func, - flash_attn_with_kvcache, -) - -__all__ = [ - # High-level API (with autograd support) - "flash_attn_func", - "flash_attn_kvpacked_func", - "flash_attn_qkvpacked_func", - "flash_attn_varlen_func", - "flash_attn_varlen_kvpacked_func", - "flash_attn_varlen_qkvpacked_func", - "flash_attn_with_kvcache", - # Low-level ops (no autograd) - "fwd", - "varlen_fwd", - "bwd", - "varlen_bwd", - "fwd_kvcache", -] - - -def fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - return_softmax: bool = False, - gen: Optional[torch.Generator] = None, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention. - - Args: - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Optional output tensor, same shape as q - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - return_softmax: Whether to return softmax weights - gen: Optional random number generator - - Returns: - List of tensors: [output, softmax_lse, (softmax if return_softmax)] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.fwd( - q, - k, - v, - out, - alibi_slopes, - p_dropout, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - gen, - ) - - -def varlen_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - out: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - block_table: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - max_seqlen_q: int = 0, - max_seqlen_k: int = 0, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - zero_tensors: bool = False, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - return_softmax: bool = False, - gen: Optional[torch.Generator] = None, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention with variable sequence lengths. - - Args: - q: Query tensor of shape [total_q, num_heads, head_size] - k: Key tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - v: Value tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1] - cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1] - out: Optional output tensor of shape [total_q, num_heads, head_size] - seqused_k: Optional tensor specifying how many keys to use per batch element [batch_size] - leftpad_k: Optional left padding for keys of shape [batch_size] - block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq] - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - max_seqlen_q: Maximum sequence length for queries - max_seqlen_k: Maximum sequence length for keys - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - zero_tensors: Whether to zero tensors before computation - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - return_softmax: Whether to return softmax weights - gen: Optional random number generator - - Returns: - List of tensors: [output, softmax_lse, (softmax if return_softmax)] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.varlen_fwd( - q, - k, - v, - out, - cu_seqlens_q, - cu_seqlens_k, - seqused_k, - leftpad_k, - block_table, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - p_dropout, - softmax_scale, - zero_tensors, - is_causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - gen, - ) - - -def bwd( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - p_dropout: float = 0.0, - 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, - gen: Optional[torch.Generator] = None, - rng_state: Optional[torch.Tensor] = None, -) -> List[torch.Tensor]: - """ - Backward pass for multi-head attention. - - Args: - dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size] - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size] - softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q] - dq: Optional gradient tensor for queries, same shape as q - dk: Optional gradient tensor for keys, same shape as k - dv: Optional gradient tensor for values, same shape as v - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - deterministic: Whether to use deterministic algorithms - gen: Optional random number generator - rng_state: Optional RNG state from forward pass - - Returns: - List of tensors: [dq, dk, dv] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - alibi_slopes, - p_dropout, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - deterministic, - gen, - rng_state, - ) - - -def varlen_bwd( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - max_seqlen_q: int = 0, - max_seqlen_k: int = 0, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - zero_tensors: bool = False, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - deterministic: bool = False, - gen: Optional[torch.Generator] = None, - rng_state: Optional[torch.Tensor] = None, -) -> List[torch.Tensor]: - """ - Backward pass for multi-head attention with variable sequence lengths. - - Args: - dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size] - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size] - softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q] - cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1] - cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1] - dq: Optional gradient tensor for queries, same shape as q - dk: Optional gradient tensor for keys, same shape as k - dv: Optional gradient tensor for values, same shape as v - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - max_seqlen_q: Maximum sequence length for queries - max_seqlen_k: Maximum sequence length for keys - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - zero_tensors: Whether to zero tensors before computation - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - deterministic: Whether to use deterministic algorithms - gen: Optional random number generator - rng_state: Optional RNG state from forward pass - - Returns: - List of tensors: [dq, dk, dv] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.varlen_bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - p_dropout, - softmax_scale, - zero_tensors, - is_causal, - window_size_left, - window_size_right, - softcap, - deterministic, - gen, - rng_state, - ) - - -def fwd_kvcache( - q: torch.Tensor, - kcache: torch.Tensor, - vcache: torch.Tensor, - k: Optional[torch.Tensor] = None, - v: Optional[torch.Tensor] = None, - seqlens_k: Optional[torch.Tensor] = None, - rotary_cos: Optional[torch.Tensor] = None, - rotary_sin: Optional[torch.Tensor] = None, - cache_batch_idx: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - block_table: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - out: 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, - is_rotary_interleaved: bool = False, - num_splits: int = 1, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention with KV cache. - - Args: - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - kcache: Key cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - vcache: Value cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - k: Optional new keys tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size] - v: Optional new values tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size] - seqlens_k: Optional sequence lengths for keys of shape [batch_size] - rotary_cos: Optional rotary cosine tensor of shape [seqlen_ro, rotary_dim/2] - rotary_sin: Optional rotary sine tensor of shape [seqlen_ro, rotary_dim/2] - cache_batch_idx: Optional indices to index into the KV cache - leftpad_k: Optional left padding for keys of shape [batch_size] - block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq] - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - out: Optional output tensor, same shape as q - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - is_rotary_interleaved: Whether rotary embeddings are interleaved - num_splits: Number of splits for computation - - Returns: - List of tensors: [output, softmax_lse] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.fwd_kvcache( - q, - kcache, - vcache, - k, - v, - seqlens_k, - rotary_cos, - rotary_sin, - cache_batch_idx, - leftpad_k, - block_table, - alibi_slopes, - out, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - is_rotary_interleaved, - num_splits, - ) diff --git a/build/torch210-cxx11-cu130-aarch64-linux/_flash_attn2_cuda_a511a4c.abi3.so b/build/torch210-cxx11-cu130-aarch64-linux/_flash_attn2_cuda_a511a4c.abi3.so deleted file mode 100644 index 807f2858c28ec0c784d3c653b0049cb4c472805a..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu130-aarch64-linux/_flash_attn2_cuda_a511a4c.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:3880a1cab17ae5b2854d7578bde6ff3b649347c26d8ccfe7571eb2c00e7a84d9 -size 1008655376 diff --git a/build/torch210-cxx11-cu130-aarch64-linux/_ops.py b/build/torch210-cxx11-cu130-aarch64-linux/_ops.py deleted file mode 100644 index 62b70a9b02c8f7663c7022564c1a7f58f46b6c65..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu130-aarch64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _flash_attn2_cuda_a511a4c -ops = torch.ops._flash_attn2_cuda_a511a4c - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_flash_attn2_cuda_a511a4c::{op_name}" diff --git a/build/torch210-cxx11-cu130-aarch64-linux/bert_padding.py b/build/torch210-cxx11-cu130-aarch64-linux/bert_padding.py deleted file mode 100644 index 3c2d35159a014a9d03aabead9e52e009168696ea..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu130-aarch64-linux/bert_padding.py +++ /dev/null @@ -1,218 +0,0 @@ -# Adapted from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/padding.py - -import torch -import torch.nn.functional as F -from einops import rearrange, repeat - - -class IndexFirstAxis(torch.autograd.Function): - @staticmethod - def forward(ctx, input, indices): - ctx.save_for_backward(indices) - assert input.ndim >= 2 - ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] - second_dim = other_shape.numel() - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - # return input[indices] - return torch.gather( - rearrange(input, "b ... -> b (...)"), 0, repeat(indices, "z -> z d", d=second_dim) - ).reshape(-1, *other_shape) - - @staticmethod - def backward(ctx, grad_output): - (indices,) = ctx.saved_tensors - assert grad_output.ndim >= 2 - other_shape = grad_output.shape[1:] - grad_output = rearrange(grad_output, "b ... -> b (...)") - grad_input = torch.zeros( - [ctx.first_axis_dim, grad_output.shape[1]], - device=grad_output.device, - dtype=grad_output.dtype, - ) - # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. - # grad_input[indices] = grad_output - grad_input.scatter_(0, repeat(indices, "z -> z d", d=grad_output.shape[1]), grad_output) - return grad_input.reshape(ctx.first_axis_dim, *other_shape), None - - -index_first_axis = IndexFirstAxis.apply - - -class IndexPutFirstAxis(torch.autograd.Function): - @staticmethod - def forward(ctx, values, indices, first_axis_dim): - ctx.save_for_backward(indices) - assert indices.ndim == 1 - assert values.ndim >= 2 - output = torch.zeros( - first_axis_dim, *values.shape[1:], device=values.device, dtype=values.dtype - ) - # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. - output[indices] = values - # output.scatter_(0, repeat(indices, 'z -> z d', d=values.shape[1]), values) - return output - - @staticmethod - def backward(ctx, grad_output): - (indices,) = ctx.saved_tensors - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - grad_values = grad_output[indices] - # grad_values = torch.gather(grad_output, 0, repeat(indices, 'z -> z d', d=grad_output.shape[1])) - return grad_values, None, None - - -index_put_first_axis = IndexPutFirstAxis.apply - - -class IndexFirstAxisResidual(torch.autograd.Function): - @staticmethod - def forward(ctx, input, indices): - ctx.save_for_backward(indices) - assert input.ndim >= 2 - ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] - second_dim = other_shape.numel() - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - output = input[indices] - # We don't want to reshape input (b ... -> b (...)) since it could change the channel_last - # memory format to channel_first. In other words, input might not be contiguous. - # If we don't detach, Pytorch complains about output being a view and is being modified inplace - return output, input.detach() - - @staticmethod - def backward(ctx, grad_output, grad_residual): - (indices,) = ctx.saved_tensors - assert grad_output.ndim >= 2 - other_shape = grad_output.shape[1:] - assert grad_residual.shape[1:] == other_shape - grad_input = grad_residual - # grad_input[indices] += grad_output - indices = indices.reshape(indices.shape[0], *((1,) * (grad_output.ndim - 1))) - indices = indices.expand_as(grad_output) - grad_input.scatter_add_(0, indices, grad_output) - return grad_input.reshape(ctx.first_axis_dim, *other_shape), None - - -index_first_axis_residual = IndexFirstAxisResidual.apply - - -def unpad_input(hidden_states, attention_mask, unused_mask=None): - """ - Arguments: - hidden_states: (batch, seqlen, ...) - attention_mask: (batch, seqlen), bool / int, 1 means valid and 0 means not valid. - unused_mask: (batch, seqlen), bool / int, 1 means the element is allocated but unused. - Return: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens selected in attention_mask + unused_mask. - indices: (total_nnz), the indices of masked tokens from the flattened input sequence. - cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. - max_seqlen_in_batch: int - seqused: (batch), returns the number of tokens selected in attention_mask + unused_mask. - """ - all_masks = (attention_mask + unused_mask) if unused_mask is not None else attention_mask - seqlens_in_batch = all_masks.sum(dim=-1, dtype=torch.int32) - used_seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) - indices = torch.nonzero(all_masks.flatten(), as_tuple=False).flatten() - max_seqlen_in_batch = seqlens_in_batch.max().item() - cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) - # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the - # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim - # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to - # index with integer indices. Moreover, torch's index is a bit slower than it needs to be, - # so we write custom forward and backward to make it a bit faster. - return ( - index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices), - indices, - cu_seqlens, - max_seqlen_in_batch, - used_seqlens_in_batch, - ) - - -def unpad_input_for_concatenated_sequences(hidden_states, attention_mask_in_length): - """ - Supports concatenating short samples in one sequence. The attention_mask_in_length is utilized to mask other short samples. It helps efficient training of variant lengths-based samples (e.g., the supervised fine-tuning task in large language model). - The motivation for this function is explained [here](https://github.com/Dao-AILab/flash-attention/issues/432#issuecomment-1668822286). - - For example, if batch = 3 and seqlen = 6, the attention_mask_in_length is: - ``` - [ - [2, 3, 0, 0, 0, 0], - [3, 2, 0, 0, 0, 0], - [6, 0, 0, 0, 0, 0] - ] - ``` - , which refers to the 3D-attention mask: - ``` - [ - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [0, 0, 1, 0, 0, 0], - [0, 0, 1, 1, 0, 0], - [0, 0, 1, 1, 1, 0], - [0, 0, 0, 0, 0, 1] - ], - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [0, 0, 0, 1, 0, 0], - [0, 0, 0, 1, 1, 0], - [0, 0, 0, 0, 0, 1] - ], - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [1, 1, 1, 1, 0, 0], - [1, 1, 1, 1, 1, 0], - [1, 1, 1, 1, 1, 1] - ] - ] - ```. - - Arguments: - hidden_states: (batch, seqlen, ...) - attention_mask_in_length: (batch, seqlen), int, a nonzero number (e.g., 1, 2, 3, etc.) means length of concatenated sequence in b-th batch, and 0 means none. - Return: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. - indices: (total_nnz), the indices of non-masked tokens from the flattened input sequence. - cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. - max_seqlen_in_batch: int - """ - length = attention_mask_in_length.sum(dim=-1) - seqlen = attention_mask_in_length.size(-1) - attention_mask_2d = torch.arange(seqlen, device=length.device, dtype=length.dtype).expand(len(length), seqlen) < length.unsqueeze(1) - real_indices_idx = torch.nonzero(attention_mask_in_length.flatten(), as_tuple=False).flatten() - seqlens_in_batch = attention_mask_in_length.flatten()[real_indices_idx] - indices = torch.nonzero(attention_mask_2d.flatten(), as_tuple=False).flatten() - max_seqlen_in_batch = seqlens_in_batch.max().item() - cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) - # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the - # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim - # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to - # index with integer indices. Moreover, torch's index is a bit slower than it needs to be, - # so we write custom forward and backward to make it a bit faster. - return ( - index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices), - indices, - cu_seqlens, - max_seqlen_in_batch, - ) - - -def pad_input(hidden_states, indices, batch, seqlen): - """ - Arguments: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. - indices: (total_nnz), the indices that represent the non-masked tokens of the original padded input sequence. - batch: int, batch size for the padded sequence. - seqlen: int, maximum sequence length for the padded sequence. - Return: - hidden_states: (batch, seqlen, ...) - """ - dim = hidden_states.shape[-1] - # output = torch.zeros((batch * seqlen), dim, device=hidden_states.device, dtype=hidden_states.dtype) - # output[indices] = hidden_states - output = index_put_first_axis(hidden_states, indices, batch * seqlen) - return rearrange(output, "(b s) ... -> b s ...", b=batch) diff --git a/build/torch210-cxx11-cu130-aarch64-linux/flash_attn2/__init__.py b/build/torch210-cxx11-cu130-aarch64-linux/flash_attn2/__init__.py deleted file mode 100644 index a9b2672c1cd85b74c1b3ded0fc0b2100e1aeac23..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu130-aarch64-linux/flash_attn2/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -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-aarch64-linux/flash_attn_interface.py b/build/torch210-cxx11-cu130-aarch64-linux/flash_attn_interface.py deleted file mode 100644 index 0868572aa3324d9c96132859d123e3d48d296e34..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu130-aarch64-linux/flash_attn_interface.py +++ /dev/null @@ -1,1662 +0,0 @@ -# Copyright (c) 2023, Tri Dao. - -from typing import Optional, Sequence, Tuple, Union - -import torch -import torch.nn as nn -import os - -# # isort: off -# # We need to import the CUDA kernels after importing torch -# USE_TRITON_ROCM = os.getenv("FLASH_ATTENTION_TRITON_AMD_ENABLE", "FALSE") == "TRUE" -# if USE_TRITON_ROCM: -# from .flash_attn_triton_amd import interface_fa as flash_attn -# else: -# import flash_attn_2_cuda as flash_attn - - -from ._ops import ops as flash_attn -from ._ops import add_op_namespace_prefix - -# # isort: on - - -def maybe_contiguous(x): - return x.contiguous() if x is not None and x.stride(-1) != 1 else x - - -def _get_device(): - if torch.xpu.is_available(): - return "xpu" - elif torch.cuda.is_available(): - return "cuda" - else: - return "cpu" - - -def _get_block_size_n(device, head_dim, is_dropout, is_causal): - assert head_dim <= 256 - - if device.type == "xpu": - if head_dim <= 96: - return 64 - elif head_dim <= 128: - return 32 - elif head_dim <= 256: - return 64 - else: - return 32 - - # This should match the block sizes in the CUDA kernel - major, minor = torch.cuda.get_device_capability(device) - is_sm8x = ( - major == 8 and minor > 0 - ) # Only include sm86 and sm89, exclude sm80 (A100) - is_sm80 = major == 8 and minor == 0 - is_sm90 = major == 9 and minor == 0 - if head_dim <= 32: - return 128 - if head_dim <= 64: - return 128 if not is_dropout else 64 - elif head_dim <= 96: - return 64 - elif head_dim <= 128: - if is_sm8x: - return 64 if (not is_dropout and is_causal) else 32 - else: - return 64 if not is_dropout else 32 - elif head_dim <= 192: - return 64 - elif head_dim <= 224: - return 64 - elif head_dim <= 256: - return 64 - - -def round_multiple(x, m): - return (x + m - 1) // m * m - - -@torch.library.custom_op( - add_op_namespace_prefix("_flash_attn_forward"), - mutates_args=(), - device_types=_get_device(), -) -def _flash_attn_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - return_softmax: bool, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - out, softmax_lse, S_dmask, rng_state = flash_attn.fwd( - q, - k, - v, - None, - alibi_slopes, - dropout_p, - softmax_scale, - causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - None, - ) - return out, softmax_lse, S_dmask, rng_state - - -@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, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - return_softmax: bool, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - batch_size, seqlen_q, num_heads, head_size = q.shape - seqlen_k = k.shape[1] - out = torch.empty_like(q) - softmax_lse = torch.empty( - (batch_size, num_heads, seqlen_q), - dtype=torch.float32, - device=q.device, - layout=q.layout, - ) - p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) - if return_softmax: - p = torch.empty( - ( - batch_size, - num_heads, - round_multiple(seqlen_q, 128), - round_multiple(seqlen_k, 128), - ), - dtype=q.dtype, - device=q.device, - layout=q.layout, - ) - rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) - - return out, softmax_lse, p, rng_state - - -@torch.library.custom_op( - add_op_namespace_prefix("_flash_attn_varlen_forward"), - mutates_args=(), - device_types=_get_device(), -) -def _flash_attn_varlen_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - alibi_slopes: Optional[torch.Tensor] = None, - return_softmax: bool = False, - block_table: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - out, softmax_lse, S_dmask, rng_state = flash_attn.varlen_fwd( - q, - k, - v, - None, - cu_seqlens_q, - cu_seqlens_k, - seqused_k, - leftpad_k, - block_table, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - zero_tensors, - causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - None, - ) - # if out.isnan().any() or softmax_lse.isnan().any(): - # breakpoint() - return out, softmax_lse, S_dmask, rng_state - - -@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_varlen_forward")) -def _flash_attn_varlen_forward_fake( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - alibi_slopes: Optional[torch.Tensor] = None, - return_softmax: bool = False, - block_table: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - paged_kv = block_table is not None - batch_size = cu_seqlens_q.numel() - 1 - total_q, num_heads, _ = q.shape - - out = torch.empty_like(q) - softmax_lse = torch.empty( - (num_heads, total_q), dtype=torch.float32, device=q.device, layout=q.layout - ) - p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) - seqlen_q_rounded = round_multiple(max_seqlen_q, 128) - seqlen_k_rounded = round_multiple(max_seqlen_k, 128) - if return_softmax: - p = torch.empty( - (batch_size, num_heads, seqlen_q_rounded, seqlen_k_rounded), - dtype=q.dtype, - device=q.device, - layout=q.layout, - ) - rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) - return out, softmax_lse, p, rng_state - - -@torch.library.custom_op( - add_op_namespace_prefix("_flash_attn_backward"), - mutates_args=("dq", "dk", "dv"), - device_types=_get_device(), -) -def _flash_attn_backward( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, -) -> 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)] - ( - dq, - dk, - dv, - softmax_d, - ) = flash_attn.bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - alibi_slopes, - dropout_p, - softmax_scale, - causal, - window_size_left, - window_size_right, - softcap, - deterministic, - None, - rng_state, - ) - 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, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, -) -> torch.Tensor: - dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] - if dq is None: - dq = torch.empty_like(q) - if dk is None: - dk = torch.empty_like(k) - if dv is None: - dv = torch.empty_like(v) - batch_size, seqlen_q, num_heads, _ = q.shape - softmax_d = torch.empty( - (batch_size, num_heads, round_multiple(seqlen_q, 128)), - device=q.device, - dtype=torch.float32, - ) - - return softmax_d - - -@torch.library.custom_op( - add_op_namespace_prefix("_flash_attn_varlen_backward"), - mutates_args=("dq", "dk", "dv"), - device_types=_get_device(), -) -def _flash_attn_varlen_backward( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> 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)] - ( - dq, - dk, - dv, - softmax_d, - ) = flash_attn.varlen_bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - zero_tensors, - causal, - window_size_left, - window_size_right, - softcap, - deterministic, - None, - rng_state, - ) - # if dk.isnan().any() or dk.isnan().any() or dv.isnan().any() or softmax_d.isnan().any(): - # breakpoint() - return softmax_d - - -@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_varlen_backward")) -def _flash_attn_varlen_backward_fake( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> torch.Tensor: - dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] - batch_size = cu_seqlens_q.numel() - 1 - total_q, num_heads, _ = q.shape - - if dq is None: - dq = torch.empty_like(q) - if dk is None: - dk = torch.empty_like(k) - if dv is None: - dv = torch.empty_like(v) - softmax_d = torch.empty( - (num_heads, total_q + 128 * batch_size), device=q.device, dtype=torch.float32 - ) - - return softmax_d - - -class FlashAttnQKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and qkv.requires_grad - if softmax_scale is None: - softmax_scale = qkv.shape[-1] ** (-0.5) - q, k, v = qkv[:, :, 0].detach(), qkv[:, :, 1].detach(), qkv[:, :, 2].detach() - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) - dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dqkv[:, :, 0], - dqkv[:, :, 1], - dqkv[:, :, 2], - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dqkv = dqkv[..., : dout.shape[-1]] # We could have padded the head dimension - return dqkv, None, None, None, None, None, None, None, None, None - - -class FlashAttnVarlenQKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - cu_seqlens, - max_seqlen, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and qkv.requires_grad - if softmax_scale is None: - softmax_scale = qkv.shape[-1] ** (-0.5) - q, k, v = qkv[:, 0].detach(), qkv[:, 1].detach(), qkv[:, 2].detach() - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens, - cu_seqlens, - max_seqlen, - max_seqlen, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=None, - ) - if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens, rng_state - ) - ctx.dropout_p = dropout_p - ctx.max_seqlen = max_seqlen - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens, rng_state = ctx.saved_tensors - qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) - dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dqkv[:, 0], - dqkv[:, 1], - dqkv[:, 2], - cu_seqlens, - cu_seqlens, - ctx.max_seqlen, - ctx.max_seqlen, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - 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 FlashAttnKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - kv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any(x.requires_grad for x in [q, kv]) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - k, v = kv[:, :, 0].detach(), kv[:, :, 1].detach() - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - dq = torch.empty_like(q) - kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) - dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dkv[:, :, 0], - dkv[:, :, 1], - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dkv = dkv[..., : dout.shape[-1]] - return dq, dkv, None, None, None, None, None, None, None, None, None - - -class FlashAttnVarlenKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any(x.requires_grad for x in [q, kv]) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - k, v = kv[:, 0].detach(), kv[:, 1].detach() - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=None, - ) - if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state - ) - ctx.dropout_p = dropout_p - 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.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ( - ctx.saved_tensors - ) - dq = torch.empty_like(q) - kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) - dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dkv[:, 0], - dkv[:, 1], - cu_seqlens_q, - cu_seqlens_k, - ctx.max_seqlen_q, - ctx.max_seqlen_k, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dkv = dkv[..., : dout.shape[-1]] - return ( - dq, - dkv, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) - - -class FlashAttnFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - k, - v, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any(x.requires_grad for x in [q, k, v]) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dk = dk[..., : dout.shape[-1]] - dv = dv[..., : dout.shape[-1]] - return dq, dk, dv, 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, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - block_table, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any(x.requires_grad for x in [q, k, v]) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=block_table, - ) - if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state - ) - ctx.dropout_p = dropout_p - 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.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ( - ctx.saved_tensors - ) - dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - ctx.max_seqlen_q, - ctx.max_seqlen_k, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dk = dk[..., : dout.shape[-1]] - dv = dv[..., : dout.shape[-1]] - return ( - dq, - dk, - dv, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) - - -def flash_attn_qkvpacked_func( - qkv, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # <=0.0 means deactivate - alibi_slopes=None, - deterministic=False, - 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, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_kvpacked_func( - q, - kv, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, -): - """dropout_p should be set to 0.0 during evaluation - If 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 K, V. - 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) - kv: (batch_size, seqlen, 2, 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. - softcap: float. Anything > 0 activates softcapping 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). - 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 FlashAttnKVPackedFunc.apply( - q, - kv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_func( - q, - k, - v, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - 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). - 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 FlashAttnFunc.apply( - q, - k, - v, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_qkvpacked_func( - qkv, - cu_seqlens, - max_seqlen, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - 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_varlen_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_varlen_kvpacked_func and flash_attn_varlen_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: (total, 3, nheads, headdim), where total = total number of tokens in the batch. - cu_seqlens: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into qkv. - max_seqlen: int. Maximum sequence length in the batch. - 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenQKVPackedFunc.apply( - qkv, - cu_seqlens, - max_seqlen, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_kvpacked_func( - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, -): - """dropout_p should be set to 0.0 during evaluation - If 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 K, V. - 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: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch. - kv: (total_k, 2, nheads_k, headdim), where total_k = total number of key tokens in the batch. - cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into q. - cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into kv. - max_seqlen_q: int. Maximum query sequence length in the batch. - max_seqlen_k: int. Maximum key sequence length in the batch. - 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 + 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenKVPackedFunc.apply( - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_func( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, - block_table=None, -): - """dropout_p should be set to 0.0 during evaluation - Supports multi-query and grouped-query attention (MQA/GQA) by passing in K, V 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: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch. - k: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. - v: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. - cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into q. - cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into kv. - max_seqlen_q: int. Maximum query sequence length in the batch. - max_seqlen_k: int. Maximum key sequence length in the batch. - 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 + 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenFunc.apply( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - block_table, - False if q.device.type == "cpu" else torch.is_grad_enabled(), - ) - - -def flash_attn_with_kvcache( - q, - k_cache, - v_cache, - k=None, - v=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, - block_table: Optional[torch.Tensor] = None, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - rotary_interleaved=True, - alibi_slopes=None, - num_splits=0, - 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 block_table, - or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_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) if there's no block_table, - or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_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). Similar to k. - 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. - block_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). - 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. - 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" - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-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) - cache_batch_idx = maybe_contiguous(cache_batch_idx) - block_table = maybe_contiguous(block_table) - out, softmax_lse = flash_attn.fwd_kvcache( - q, - k_cache, - v_cache, - k, - v, - cache_seqlens, - rotary_cos, - rotary_sin, - cache_batch_idx, - cache_leftpad, - block_table, - alibi_slopes, - None, - softmax_scale, - causal, - window_size[0], - window_size[1], - softcap, - rotary_interleaved, - num_splits, - ) - return (out, softmax_lse) if return_softmax_lse else out diff --git a/build/torch210-cxx11-cu130-aarch64-linux/metadata.json b/build/torch210-cxx11-cu130-aarch64-linux/metadata.json deleted file mode 100644 index 5912a9cf2f79a4ad05d354c2cf7559adff1b1ac2..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu130-aarch64-linux/metadata.json +++ /dev/null @@ -1,16 +0,0 @@ -{ - "name": "flash-attn2", - "id": "_flash_attn2_cuda_a511a4c", - "version": 1, - "license": "BSD-3-Clause", - "python-depends": [], - "backend": { - "type": "cuda", - "archs": [ - "10.0", - "12.0", - "8.0", - "9.0" - ] - } -} diff --git a/build/torch210-cxx11-cu130-x86_64-linux/__init__.py b/build/torch210-cxx11-cu130-x86_64-linux/__init__.py index 20ff0ff747abd7500913863b509925f414b1c973..ecc2f9d896b6c93f90b0a1499856dc0612177422 100644 --- a/build/torch210-cxx11-cu130-x86_64-linux/__init__.py +++ b/build/torch210-cxx11-cu130-x86_64-linux/__init__.py @@ -1,7 +1,5 @@ -from typing import List, Optional - +from typing import Optional, List import torch - from ._ops import ops as flash_attn_ops from .flash_attn_interface import ( flash_attn_func, @@ -13,23 +11,6 @@ from .flash_attn_interface import ( flash_attn_with_kvcache, ) -__all__ = [ - # High-level API (with autograd support) - "flash_attn_func", - "flash_attn_kvpacked_func", - "flash_attn_qkvpacked_func", - "flash_attn_varlen_func", - "flash_attn_varlen_kvpacked_func", - "flash_attn_varlen_qkvpacked_func", - "flash_attn_with_kvcache", - # Low-level ops (no autograd) - "fwd", - "varlen_fwd", - "bwd", - "varlen_bwd", - "fwd_kvcache", -] - def fwd( q: torch.Tensor, diff --git a/build/torch210-cxx11-cu130-x86_64-linux/_flash_attn2_588b404.abi3.so b/build/torch210-cxx11-cu130-x86_64-linux/_flash_attn2_588b404.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..ee07c96066734631aa83e146852c588a4b3eb822 --- /dev/null +++ b/build/torch210-cxx11-cu130-x86_64-linux/_flash_attn2_588b404.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:196d3756a7d099f5e23ddd53ebc47aadf558a96e1d7873f5a14faec09bb7b707 +size 1009055064 diff --git a/build/torch210-cxx11-cu130-x86_64-linux/_flash_attn2_cuda_a511a4c.abi3.so b/build/torch210-cxx11-cu130-x86_64-linux/_flash_attn2_cuda_a511a4c.abi3.so deleted file mode 100644 index c67d7047bc61d0ebde15a0bb2629adcc01d742a7..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu130-x86_64-linux/_flash_attn2_cuda_a511a4c.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:01aa71482e4e1c20abcb2fa402a0d6f3386a93e023aad1fb0ffd1a7c4b218677 -size 1008994200 diff --git a/build/torch210-cxx11-cu130-x86_64-linux/_ops.py b/build/torch210-cxx11-cu130-x86_64-linux/_ops.py index 62b70a9b02c8f7663c7022564c1a7f58f46b6c65..03506d8a18d17ca7296a5f1b7a18ec7dc4ad73bb 100644 --- a/build/torch210-cxx11-cu130-x86_64-linux/_ops.py +++ b/build/torch210-cxx11-cu130-x86_64-linux/_ops.py @@ -1,9 +1,9 @@ import torch -from . import _flash_attn2_cuda_a511a4c -ops = torch.ops._flash_attn2_cuda_a511a4c +from . import _flash_attn2_588b404 +ops = torch.ops._flash_attn2_588b404 def add_op_namespace_prefix(op_name: str): """ Prefix op by namespace. """ - return f"_flash_attn2_cuda_a511a4c::{op_name}" + return f"_flash_attn2_588b404::{op_name}" \ No newline at end of file diff --git a/build/torch210-cxx11-cu130-x86_64-linux/flash_attn2/__init__.py b/build/torch210-cxx11-cu130-x86_64-linux/flash_attn2/__init__.py index a9b2672c1cd85b74c1b3ded0fc0b2100e1aeac23..03dbc1afe1cf156661a2b1b22003cd5f599a0309 100644 --- a/build/torch210-cxx11-cu130-x86_64-linux/flash_attn2/__init__.py +++ b/build/torch210-cxx11-cu130-x86_64-linux/flash_attn2/__init__.py @@ -1,10 +1,10 @@ import ctypes -import importlib.util 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 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 index 0868572aa3324d9c96132859d123e3d48d296e34..66e8d3665834b141d703998496b3da6fe1f6fe06 100644 --- a/build/torch210-cxx11-cu130-x86_64-linux/flash_attn_interface.py +++ b/build/torch210-cxx11-cu130-x86_64-linux/flash_attn_interface.py @@ -16,11 +16,9 @@ import os from ._ops import ops as flash_attn -from ._ops import add_op_namespace_prefix # # isort: on - def maybe_contiguous(x): return x.contiguous() if x is not None and x.stride(-1) != 1 else x @@ -33,25 +31,14 @@ def _get_device(): else: return "cpu" +_XPU_AVAILABLE = torch.xpu.is_available() if hasattr(torch, "xpu") else False # TODO remove hasattr check when bwd is supported on XPU -def _get_block_size_n(device, head_dim, is_dropout, is_causal): - assert head_dim <= 256 - - if device.type == "xpu": - if head_dim <= 96: - return 64 - elif head_dim <= 128: - return 32 - elif head_dim <= 256: - return 64 - else: - return 32 +def _get_block_size_n(device, head_dim, is_dropout, is_causal): # This should match the block sizes in the CUDA kernel + assert head_dim <= 256 major, minor = torch.cuda.get_device_capability(device) - is_sm8x = ( - major == 8 and minor > 0 - ) # Only include sm86 and sm89, exclude sm80 (A100) + is_sm8x = major == 8 and minor > 0 # Only include sm86 and sm89, exclude sm80 (A100) is_sm80 = major == 8 and minor == 0 is_sm90 = major == 9 and minor == 0 if head_dim <= 32: @@ -77,11 +64,30 @@ def round_multiple(x, m): return (x + m - 1) // m * m -@torch.library.custom_op( - add_op_namespace_prefix("_flash_attn_forward"), - mutates_args=(), - device_types=_get_device(), -) +# torch.compile() support is only enabled for pytorch >= 2.4 +# The reason for this is that we are using the new custom_op and register_fake +# APIs, which support inplace modification of inputs in the function itself +if torch.__version__ >= "2.4.0": + _torch_custom_op_wrapper = torch.library.custom_op + _torch_register_fake_wrapper = torch.library.register_fake +else: + def noop_custom_op_wrapper(name, fn=None, /, *, mutates_args, device_types=None, schema=None): + def wrap(func): + return func + if fn is None: + return wrap + return fn + def noop_register_fake_wrapper(op, fn=None, /, *, lib=None, _stacklevel=1): + def wrap(func): + return func + if fn is None: + return wrap + return fn + _torch_custom_op_wrapper = noop_custom_op_wrapper + _torch_register_fake_wrapper = noop_register_fake_wrapper + + +@_torch_custom_op_wrapper("flash_attn::_flash_attn_forward", mutates_args=(), device_types=_get_device()) def _flash_attn_forward( q: torch.Tensor, k: torch.Tensor, @@ -93,7 +99,7 @@ def _flash_attn_forward( window_size_right: int, softcap: float, alibi_slopes: Optional[torch.Tensor], - return_softmax: bool, + return_softmax: bool ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: q, k, v = [maybe_contiguous(x) for x in (q, k, v)] out, softmax_lse, S_dmask, rng_state = flash_attn.fwd( @@ -114,7 +120,7 @@ def _flash_attn_forward( return out, softmax_lse, S_dmask, rng_state -@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_forward")) +@_torch_register_fake_wrapper("flash_attn::_flash_attn_forward") def _flash_attn_forward_fake( q: torch.Tensor, k: torch.Tensor, @@ -126,41 +132,28 @@ def _flash_attn_forward_fake( window_size_right: int, softcap: float, alibi_slopes: Optional[torch.Tensor], - return_softmax: bool, + return_softmax: bool ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: q, k, v = [maybe_contiguous(x) for x in (q, k, v)] batch_size, seqlen_q, num_heads, head_size = q.shape seqlen_k = k.shape[1] out = torch.empty_like(q) - softmax_lse = torch.empty( - (batch_size, num_heads, seqlen_q), - dtype=torch.float32, - device=q.device, - layout=q.layout, - ) + softmax_lse = torch.empty((batch_size, num_heads, seqlen_q), dtype=torch.float32, device=q.device, layout=q.layout) p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) if return_softmax: - p = torch.empty( - ( - batch_size, - num_heads, - round_multiple(seqlen_q, 128), - round_multiple(seqlen_k, 128), - ), - dtype=q.dtype, - device=q.device, - layout=q.layout, - ) + p = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128), round_multiple(seqlen_k, 128)), dtype=q.dtype, device=q.device, layout=q.layout) rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) return out, softmax_lse, p, rng_state -@torch.library.custom_op( - add_op_namespace_prefix("_flash_attn_varlen_forward"), - mutates_args=(), - device_types=_get_device(), -) +if torch.__version__ >= "2.4.0": + _wrapped_flash_attn_forward = torch.ops.flash_attn._flash_attn_forward +else: + _wrapped_flash_attn_forward = _flash_attn_forward + + +@_torch_custom_op_wrapper("flash_attn::_flash_attn_varlen_forward", mutates_args=(), device_types=_get_device()) def _flash_attn_varlen_forward( q: torch.Tensor, k: torch.Tensor, @@ -211,7 +204,7 @@ def _flash_attn_varlen_forward( return out, softmax_lse, S_dmask, rng_state -@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_varlen_forward")) +@_torch_register_fake_wrapper("flash_attn::_flash_attn_varlen_forward") def _flash_attn_varlen_forward_fake( q: torch.Tensor, k: torch.Tensor, @@ -237,30 +230,25 @@ def _flash_attn_varlen_forward_fake( paged_kv = block_table is not None batch_size = cu_seqlens_q.numel() - 1 total_q, num_heads, _ = q.shape - + out = torch.empty_like(q) - softmax_lse = torch.empty( - (num_heads, total_q), dtype=torch.float32, device=q.device, layout=q.layout - ) + softmax_lse = torch.empty((num_heads, total_q), dtype=torch.float32, device=q.device, layout=q.layout) p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) seqlen_q_rounded = round_multiple(max_seqlen_q, 128) seqlen_k_rounded = round_multiple(max_seqlen_k, 128) if return_softmax: - p = torch.empty( - (batch_size, num_heads, seqlen_q_rounded, seqlen_k_rounded), - dtype=q.dtype, - device=q.device, - layout=q.layout, - ) + p = torch.empty((batch_size, num_heads, seqlen_q_rounded, seqlen_k_rounded), dtype=q.dtype, device=q.device, layout=q.layout) rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) return out, softmax_lse, p, rng_state -@torch.library.custom_op( - add_op_namespace_prefix("_flash_attn_backward"), - mutates_args=("dq", "dk", "dv"), - device_types=_get_device(), -) +if torch.__version__ >= "2.4.0": + _wrapped_flash_attn_varlen_forward = torch.ops.flash_attn._flash_attn_varlen_forward +else: + _wrapped_flash_attn_varlen_forward = _flash_attn_varlen_forward + + +@_torch_custom_op_wrapper("flash_attn::_flash_attn_backward", mutates_args=("dq", "dk", "dv"), device_types=_get_device()) def _flash_attn_backward( dout: torch.Tensor, q: torch.Tensor, @@ -312,7 +300,7 @@ def _flash_attn_backward( return softmax_d -@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_backward")) +@_torch_register_fake_wrapper("flash_attn::_flash_attn_backward") def _flash_attn_backward_fake( dout: torch.Tensor, q: torch.Tensor, @@ -341,20 +329,18 @@ def _flash_attn_backward_fake( if dv is None: dv = torch.empty_like(v) batch_size, seqlen_q, num_heads, _ = q.shape - softmax_d = torch.empty( - (batch_size, num_heads, round_multiple(seqlen_q, 128)), - device=q.device, - dtype=torch.float32, - ) - + softmax_d = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128)), device=q.device, dtype=torch.float32) + return softmax_d -@torch.library.custom_op( - add_op_namespace_prefix("_flash_attn_varlen_backward"), - mutates_args=("dq", "dk", "dv"), - device_types=_get_device(), -) +if torch.__version__ >= "2.4.0": + _wrapped_flash_attn_backward = torch.ops.flash_attn._flash_attn_backward +else: + _wrapped_flash_attn_backward = _flash_attn_backward + + +@_torch_custom_op_wrapper("flash_attn::_flash_attn_varlen_backward", mutates_args=("dq", "dk", "dv"), device_types=_get_device()) def _flash_attn_varlen_backward( dout: torch.Tensor, q: torch.Tensor, @@ -418,7 +404,7 @@ def _flash_attn_varlen_backward( return softmax_d -@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_varlen_backward")) +@_torch_register_fake_wrapper("flash_attn::_flash_attn_varlen_backward") def _flash_attn_varlen_backward_fake( dout: torch.Tensor, q: torch.Tensor, @@ -454,13 +440,17 @@ def _flash_attn_varlen_backward_fake( dk = torch.empty_like(k) if dv is None: dv = torch.empty_like(v) - softmax_d = torch.empty( - (num_heads, total_q + 128 * batch_size), device=q.device, dtype=torch.float32 - ) - + softmax_d = torch.empty((num_heads, total_q + 128 * batch_size), device=q.device, dtype=torch.float32) + return softmax_d +if torch.__version__ >= "2.4.0": + _wrapped_flash_attn_varlen_backward = torch.ops.flash_attn._flash_attn_varlen_backward +else: + _wrapped_flash_attn_varlen_backward = _flash_attn_varlen_backward + + class FlashAttnQKVPackedFunc(torch.autograd.Function): @staticmethod def forward( @@ -485,7 +475,7 @@ class FlashAttnQKVPackedFunc(torch.autograd.Function): q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_forward( + out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( q, k, v, @@ -519,7 +509,7 @@ class FlashAttnQKVPackedFunc(torch.autograd.Function): dout_padded = dout if head_size_og % 8 != 0: dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _flash_attn_backward( + _wrapped_flash_attn_backward( dout_padded, q, k, @@ -569,7 +559,7 @@ class FlashAttnVarlenQKVPackedFunc(torch.autograd.Function): q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_varlen_forward( + out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( q, k, v, @@ -588,9 +578,7 @@ class FlashAttnVarlenQKVPackedFunc(torch.autograd.Function): block_table=None, ) if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens, rng_state - ) + ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens, rng_state) ctx.dropout_p = dropout_p ctx.max_seqlen = max_seqlen ctx.softmax_scale = softmax_scale @@ -611,7 +599,7 @@ class FlashAttnVarlenQKVPackedFunc(torch.autograd.Function): dout_padded = dout if head_size_og % 8 != 0: dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _flash_attn_varlen_backward( + _wrapped_flash_attn_varlen_backward( dout_padded, q, k, @@ -655,7 +643,9 @@ class FlashAttnKVPackedFunc(torch.autograd.Function): return_softmax, is_grad_enabled, ): - is_grad = is_grad_enabled and any(x.requires_grad for x in [q, kv]) + is_grad = is_grad_enabled and any( + x.requires_grad for x in [q, kv] + ) if softmax_scale is None: softmax_scale = q.shape[-1] ** (-0.5) k, v = kv[:, :, 0].detach(), kv[:, :, 1].detach() @@ -664,7 +654,7 @@ class FlashAttnKVPackedFunc(torch.autograd.Function): q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_forward( + out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( q, k, v, @@ -699,7 +689,7 @@ class FlashAttnKVPackedFunc(torch.autograd.Function): dout_padded = dout if head_size_og % 8 != 0: dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _flash_attn_backward( + _wrapped_flash_attn_backward( dout_padded, q, k, @@ -744,7 +734,9 @@ class FlashAttnVarlenKVPackedFunc(torch.autograd.Function): return_softmax, is_grad_enabled, ): - is_grad = is_grad_enabled and any(x.requires_grad for x in [q, kv]) + is_grad = is_grad_enabled and any( + x.requires_grad for x in [q, kv] + ) if softmax_scale is None: softmax_scale = q.shape[-1] ** (-0.5) k, v = kv[:, 0].detach(), kv[:, 1].detach() @@ -753,7 +745,7 @@ class FlashAttnVarlenKVPackedFunc(torch.autograd.Function): q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_varlen_forward( + out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( q, k, v, @@ -789,9 +781,7 @@ class FlashAttnVarlenKVPackedFunc(torch.autograd.Function): @staticmethod def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ( - ctx.saved_tensors - ) + q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors dq = torch.empty_like(q) kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) @@ -799,7 +789,7 @@ class FlashAttnVarlenKVPackedFunc(torch.autograd.Function): dout_padded = dout if head_size_og % 8 != 0: dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _flash_attn_varlen_backward( + _wrapped_flash_attn_varlen_backward( dout_padded, q, k, @@ -825,23 +815,7 @@ class FlashAttnVarlenKVPackedFunc(torch.autograd.Function): ) dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension dkv = dkv[..., : dout.shape[-1]] - return ( - dq, - dkv, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) + return dq, dkv, None, None, None, None, None, None, None, None, None, None, None, None, None class FlashAttnFunc(torch.autograd.Function): @@ -861,7 +835,9 @@ class FlashAttnFunc(torch.autograd.Function): return_softmax, is_grad_enabled, ): - is_grad = is_grad_enabled and any(x.requires_grad for x in [q, k, v]) + is_grad = is_grad_enabled and any( + x.requires_grad for x in [q, k, v] + ) if softmax_scale is None: softmax_scale = q.shape[-1] ** (-0.5) head_size_og = q.size(3) @@ -869,7 +845,7 @@ class FlashAttnFunc(torch.autograd.Function): q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_forward( + out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( q, k, v, @@ -902,7 +878,7 @@ class FlashAttnFunc(torch.autograd.Function): dout_padded = dout if head_size_og % 8 != 0: dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _flash_attn_backward( + _wrapped_flash_attn_backward( dout_padded, q, k, @@ -950,7 +926,9 @@ class FlashAttnVarlenFunc(torch.autograd.Function): block_table, is_grad_enabled, ): - is_grad = is_grad_enabled and any(x.requires_grad for x in [q, k, v]) + is_grad = is_grad_enabled and any( + x.requires_grad for x in [q, k, v] + ) if softmax_scale is None: softmax_scale = q.shape[-1] ** (-0.5) head_size_og = q.size(2) @@ -958,7 +936,7 @@ class FlashAttnVarlenFunc(torch.autograd.Function): q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_varlen_forward( + out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( q, k, v, @@ -995,15 +973,13 @@ class FlashAttnVarlenFunc(torch.autograd.Function): @staticmethod def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ( - ctx.saved_tensors - ) + q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) head_size_og = dout.size(2) dout_padded = dout if head_size_og % 8 != 0: dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _flash_attn_varlen_backward( + _wrapped_flash_attn_varlen_backward( dout_padded, q, k, @@ -1030,25 +1006,7 @@ class FlashAttnVarlenFunc(torch.autograd.Function): dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension dk = dk[..., : dout.shape[-1]] dv = dv[..., : dout.shape[-1]] - return ( - dq, - dk, - dv, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None def flash_attn_qkvpacked_func( @@ -1106,7 +1064,7 @@ def flash_attn_qkvpacked_func( alibi_slopes, deterministic, return_attn_probs, - torch.is_grad_enabled(), + False if _XPU_AVAILABLE else torch.is_grad_enabled(), ) @@ -1184,7 +1142,7 @@ def flash_attn_kvpacked_func( alibi_slopes, deterministic, return_attn_probs, - torch.is_grad_enabled(), + False if _XPU_AVAILABLE else torch.is_grad_enabled(), ) @@ -1196,7 +1154,7 @@ def flash_attn_func( softmax_scale=None, causal=False, window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated + softcap=0.0, # 0.0 means deactivated alibi_slopes=None, deterministic=False, return_attn_probs=False, @@ -1261,7 +1219,7 @@ def flash_attn_func( alibi_slopes, deterministic, return_attn_probs, - torch.is_grad_enabled(), + False if _XPU_AVAILABLE else torch.is_grad_enabled(), ) @@ -1273,7 +1231,7 @@ def flash_attn_varlen_qkvpacked_func( softmax_scale=None, causal=False, window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated + softcap=0.0, # 0.0 means deactivated alibi_slopes=None, deterministic=False, return_attn_probs=False, @@ -1327,7 +1285,7 @@ def flash_attn_varlen_qkvpacked_func( alibi_slopes, deterministic, return_attn_probs, - torch.is_grad_enabled(), + False if _XPU_AVAILABLE else torch.is_grad_enabled(), ) @@ -1342,7 +1300,7 @@ def flash_attn_varlen_kvpacked_func( softmax_scale=None, causal=False, window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated + softcap=0.0, # 0.0 means deactivated alibi_slopes=None, deterministic=False, return_attn_probs=False, @@ -1419,7 +1377,7 @@ def flash_attn_varlen_kvpacked_func( alibi_slopes, deterministic, return_attn_probs, - torch.is_grad_enabled(), + False if _XPU_AVAILABLE else torch.is_grad_enabled(), ) @@ -1435,7 +1393,7 @@ def flash_attn_varlen_func( softmax_scale=None, causal=False, window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated + softcap=0.0, # 0.0 means deactivated alibi_slopes=None, deterministic=False, return_attn_probs=False, @@ -1513,7 +1471,7 @@ def flash_attn_varlen_func( deterministic, return_attn_probs, block_table, - False if q.device.type == "cpu" else torch.is_grad_enabled(), + False if _XPU_AVAILABLE or q.device.type == "cpu" else torch.is_grad_enabled(), ) @@ -1532,7 +1490,7 @@ def flash_attn_with_kvcache( softmax_scale=None, causal=False, window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated + softcap=0.0, # 0.0 means deactivated rotary_interleaved=True, alibi_slopes=None, num_splits=0, diff --git a/build/torch211-cxx11-cu126-aarch64-linux/layers/__init__.py b/build/torch210-cxx11-cu130-x86_64-linux/layers/__init__.py similarity index 100% rename from build/torch211-cxx11-cu126-aarch64-linux/layers/__init__.py rename to build/torch210-cxx11-cu130-x86_64-linux/layers/__init__.py diff --git a/build/torch211-cxx11-cu128-aarch64-linux/layers/patch_embed.py b/build/torch210-cxx11-cu130-x86_64-linux/layers/patch_embed.py similarity index 100% rename from build/torch211-cxx11-cu128-aarch64-linux/layers/patch_embed.py rename to build/torch210-cxx11-cu130-x86_64-linux/layers/patch_embed.py diff --git a/build/torch211-cxx11-cu128-aarch64-linux/layers/rotary.py b/build/torch210-cxx11-cu130-x86_64-linux/layers/rotary.py similarity index 100% rename from build/torch211-cxx11-cu128-aarch64-linux/layers/rotary.py rename to build/torch210-cxx11-cu130-x86_64-linux/layers/rotary.py diff --git a/build/torch210-cxx11-cu130-x86_64-linux/metadata.json b/build/torch210-cxx11-cu130-x86_64-linux/metadata.json index 5912a9cf2f79a4ad05d354c2cf7559adff1b1ac2..9cf5deed9898dce769f4cc73913d3530b92a0bd8 100644 --- a/build/torch210-cxx11-cu130-x86_64-linux/metadata.json +++ b/build/torch210-cxx11-cu130-x86_64-linux/metadata.json @@ -1,16 +1,4 @@ { - "name": "flash-attn2", - "id": "_flash_attn2_cuda_a511a4c", "version": 1, - "license": "BSD-3-Clause", - "python-depends": [], - "backend": { - "type": "cuda", - "archs": [ - "10.0", - "12.0", - "8.0", - "9.0" - ] - } -} + "python-depends": [] +} \ No newline at end of file diff --git a/build/torch211-cxx11-cu126-aarch64-linux/ops/__init__.py b/build/torch210-cxx11-cu130-x86_64-linux/ops/__init__.py similarity index 100% rename from build/torch211-cxx11-cu126-aarch64-linux/ops/__init__.py rename to build/torch210-cxx11-cu130-x86_64-linux/ops/__init__.py diff --git a/build/torch211-cxx11-cu128-aarch64-linux/ops/activations.py b/build/torch210-cxx11-cu130-x86_64-linux/ops/activations.py similarity index 100% rename from build/torch211-cxx11-cu128-aarch64-linux/ops/activations.py rename to build/torch210-cxx11-cu130-x86_64-linux/ops/activations.py diff --git a/build/torch211-cxx11-cu128-aarch64-linux/ops/fused_dense.py b/build/torch210-cxx11-cu130-x86_64-linux/ops/fused_dense.py similarity index 100% rename from build/torch211-cxx11-cu128-aarch64-linux/ops/fused_dense.py rename to build/torch210-cxx11-cu130-x86_64-linux/ops/fused_dense.py diff --git a/build/torch211-cxx11-cu128-aarch64-linux/ops/layer_norm.py b/build/torch210-cxx11-cu130-x86_64-linux/ops/layer_norm.py similarity index 100% rename from build/torch211-cxx11-cu128-aarch64-linux/ops/layer_norm.py rename to build/torch210-cxx11-cu130-x86_64-linux/ops/layer_norm.py diff --git a/build/torch211-cxx11-cu128-aarch64-linux/ops/rms_norm.py b/build/torch210-cxx11-cu130-x86_64-linux/ops/rms_norm.py similarity index 100% rename from build/torch211-cxx11-cu128-aarch64-linux/ops/rms_norm.py rename to build/torch210-cxx11-cu130-x86_64-linux/ops/rms_norm.py diff --git a/build/torch211-cxx11-cu128-aarch64-linux/ops/triton/__init__.py b/build/torch210-cxx11-cu130-x86_64-linux/ops/triton/__init__.py similarity index 100% rename from build/torch211-cxx11-cu128-aarch64-linux/ops/triton/__init__.py rename to build/torch210-cxx11-cu130-x86_64-linux/ops/triton/__init__.py diff --git a/build/torch211-cxx11-cu128-aarch64-linux/ops/triton/cross_entropy.py b/build/torch210-cxx11-cu130-x86_64-linux/ops/triton/cross_entropy.py similarity index 100% rename from build/torch211-cxx11-cu128-aarch64-linux/ops/triton/cross_entropy.py rename to build/torch210-cxx11-cu130-x86_64-linux/ops/triton/cross_entropy.py diff --git a/build/torch211-cxx11-cu128-aarch64-linux/ops/triton/k_activations.py b/build/torch210-cxx11-cu130-x86_64-linux/ops/triton/k_activations.py similarity index 100% rename from build/torch211-cxx11-cu128-aarch64-linux/ops/triton/k_activations.py rename to build/torch210-cxx11-cu130-x86_64-linux/ops/triton/k_activations.py diff --git a/build/torch29-cxx11-cu129-x86_64-linux/ops/triton/layer_norm.py b/build/torch210-cxx11-cu130-x86_64-linux/ops/triton/layer_norm.py similarity index 100% rename from build/torch29-cxx11-cu129-x86_64-linux/ops/triton/layer_norm.py rename to build/torch210-cxx11-cu130-x86_64-linux/ops/triton/layer_norm.py diff --git a/build/torch211-cxx11-cu128-aarch64-linux/ops/triton/linear.py b/build/torch210-cxx11-cu130-x86_64-linux/ops/triton/linear.py similarity index 100% rename from build/torch211-cxx11-cu128-aarch64-linux/ops/triton/linear.py rename to build/torch210-cxx11-cu130-x86_64-linux/ops/triton/linear.py diff --git a/build/torch211-cxx11-cu128-aarch64-linux/ops/triton/mlp.py b/build/torch210-cxx11-cu130-x86_64-linux/ops/triton/mlp.py similarity index 100% rename from build/torch211-cxx11-cu128-aarch64-linux/ops/triton/mlp.py rename to build/torch210-cxx11-cu130-x86_64-linux/ops/triton/mlp.py diff --git a/build/torch211-cxx11-cu128-aarch64-linux/ops/triton/rotary.py b/build/torch210-cxx11-cu130-x86_64-linux/ops/triton/rotary.py similarity index 100% rename from build/torch211-cxx11-cu128-aarch64-linux/ops/triton/rotary.py rename to build/torch210-cxx11-cu130-x86_64-linux/ops/triton/rotary.py diff --git a/build/torch210-cxx11-xpu20253-x86_64-linux/__init__.py b/build/torch210-cxx11-xpu20253-x86_64-linux/__init__.py index 20ff0ff747abd7500913863b509925f414b1c973..ecc2f9d896b6c93f90b0a1499856dc0612177422 100644 --- a/build/torch210-cxx11-xpu20253-x86_64-linux/__init__.py +++ b/build/torch210-cxx11-xpu20253-x86_64-linux/__init__.py @@ -1,7 +1,5 @@ -from typing import List, Optional - +from typing import Optional, List import torch - from ._ops import ops as flash_attn_ops from .flash_attn_interface import ( flash_attn_func, @@ -13,23 +11,6 @@ from .flash_attn_interface import ( flash_attn_with_kvcache, ) -__all__ = [ - # High-level API (with autograd support) - "flash_attn_func", - "flash_attn_kvpacked_func", - "flash_attn_qkvpacked_func", - "flash_attn_varlen_func", - "flash_attn_varlen_kvpacked_func", - "flash_attn_varlen_qkvpacked_func", - "flash_attn_with_kvcache", - # Low-level ops (no autograd) - "fwd", - "varlen_fwd", - "bwd", - "varlen_bwd", - "fwd_kvcache", -] - def fwd( q: torch.Tensor, diff --git a/build/torch210-cxx11-xpu20253-x86_64-linux/_flash_attn2_588b404.abi3.so b/build/torch210-cxx11-xpu20253-x86_64-linux/_flash_attn2_588b404.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..2595f94249e401c179fb2d7dc02422e9d2fa0837 --- /dev/null +++ b/build/torch210-cxx11-xpu20253-x86_64-linux/_flash_attn2_588b404.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:1b431c2d9fcbb6f0923c3dfcaab8ee5f6df980fd39877cd6ff5f44373cd02271 +size 10797184 diff --git a/build/torch210-cxx11-xpu20253-x86_64-linux/_flash_attn2_xpu_a511a4c.abi3.so b/build/torch210-cxx11-xpu20253-x86_64-linux/_flash_attn2_xpu_a511a4c.abi3.so deleted file mode 100644 index 65579aae428daec600d729d1d87cdc1d8bea3835..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-xpu20253-x86_64-linux/_flash_attn2_xpu_a511a4c.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:410153480117b258b7b2227e5217870822a84b5d29c58c4120a66f473eebae44 -size 84309376 diff --git a/build/torch210-cxx11-xpu20253-x86_64-linux/_ops.py b/build/torch210-cxx11-xpu20253-x86_64-linux/_ops.py index 9386c5279bcd528824c7ee73a4ff7fed9c08e2d3..03506d8a18d17ca7296a5f1b7a18ec7dc4ad73bb 100644 --- a/build/torch210-cxx11-xpu20253-x86_64-linux/_ops.py +++ b/build/torch210-cxx11-xpu20253-x86_64-linux/_ops.py @@ -1,9 +1,9 @@ import torch -from . import _flash_attn2_xpu_a511a4c -ops = torch.ops._flash_attn2_xpu_a511a4c +from . import _flash_attn2_588b404 +ops = torch.ops._flash_attn2_588b404 def add_op_namespace_prefix(op_name: str): """ Prefix op by namespace. """ - return f"_flash_attn2_xpu_a511a4c::{op_name}" + return f"_flash_attn2_588b404::{op_name}" \ No newline at end of file diff --git a/build/torch210-cxx11-xpu20253-x86_64-linux/flash_attn2/__init__.py b/build/torch210-cxx11-xpu20253-x86_64-linux/flash_attn2/__init__.py index a9b2672c1cd85b74c1b3ded0fc0b2100e1aeac23..03dbc1afe1cf156661a2b1b22003cd5f599a0309 100644 --- a/build/torch210-cxx11-xpu20253-x86_64-linux/flash_attn2/__init__.py +++ b/build/torch210-cxx11-xpu20253-x86_64-linux/flash_attn2/__init__.py @@ -1,10 +1,10 @@ import ctypes -import importlib.util 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 diff --git a/build/torch210-cxx11-xpu20253-x86_64-linux/flash_attn_interface.py b/build/torch210-cxx11-xpu20253-x86_64-linux/flash_attn_interface.py index 0868572aa3324d9c96132859d123e3d48d296e34..66e8d3665834b141d703998496b3da6fe1f6fe06 100644 --- a/build/torch210-cxx11-xpu20253-x86_64-linux/flash_attn_interface.py +++ b/build/torch210-cxx11-xpu20253-x86_64-linux/flash_attn_interface.py @@ -16,11 +16,9 @@ import os from ._ops import ops as flash_attn -from ._ops import add_op_namespace_prefix # # isort: on - def maybe_contiguous(x): return x.contiguous() if x is not None and x.stride(-1) != 1 else x @@ -33,25 +31,14 @@ def _get_device(): else: return "cpu" +_XPU_AVAILABLE = torch.xpu.is_available() if hasattr(torch, "xpu") else False # TODO remove hasattr check when bwd is supported on XPU -def _get_block_size_n(device, head_dim, is_dropout, is_causal): - assert head_dim <= 256 - - if device.type == "xpu": - if head_dim <= 96: - return 64 - elif head_dim <= 128: - return 32 - elif head_dim <= 256: - return 64 - else: - return 32 +def _get_block_size_n(device, head_dim, is_dropout, is_causal): # This should match the block sizes in the CUDA kernel + assert head_dim <= 256 major, minor = torch.cuda.get_device_capability(device) - is_sm8x = ( - major == 8 and minor > 0 - ) # Only include sm86 and sm89, exclude sm80 (A100) + is_sm8x = major == 8 and minor > 0 # Only include sm86 and sm89, exclude sm80 (A100) is_sm80 = major == 8 and minor == 0 is_sm90 = major == 9 and minor == 0 if head_dim <= 32: @@ -77,11 +64,30 @@ def round_multiple(x, m): return (x + m - 1) // m * m -@torch.library.custom_op( - add_op_namespace_prefix("_flash_attn_forward"), - mutates_args=(), - device_types=_get_device(), -) +# torch.compile() support is only enabled for pytorch >= 2.4 +# The reason for this is that we are using the new custom_op and register_fake +# APIs, which support inplace modification of inputs in the function itself +if torch.__version__ >= "2.4.0": + _torch_custom_op_wrapper = torch.library.custom_op + _torch_register_fake_wrapper = torch.library.register_fake +else: + def noop_custom_op_wrapper(name, fn=None, /, *, mutates_args, device_types=None, schema=None): + def wrap(func): + return func + if fn is None: + return wrap + return fn + def noop_register_fake_wrapper(op, fn=None, /, *, lib=None, _stacklevel=1): + def wrap(func): + return func + if fn is None: + return wrap + return fn + _torch_custom_op_wrapper = noop_custom_op_wrapper + _torch_register_fake_wrapper = noop_register_fake_wrapper + + +@_torch_custom_op_wrapper("flash_attn::_flash_attn_forward", mutates_args=(), device_types=_get_device()) def _flash_attn_forward( q: torch.Tensor, k: torch.Tensor, @@ -93,7 +99,7 @@ def _flash_attn_forward( window_size_right: int, softcap: float, alibi_slopes: Optional[torch.Tensor], - return_softmax: bool, + return_softmax: bool ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: q, k, v = [maybe_contiguous(x) for x in (q, k, v)] out, softmax_lse, S_dmask, rng_state = flash_attn.fwd( @@ -114,7 +120,7 @@ def _flash_attn_forward( return out, softmax_lse, S_dmask, rng_state -@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_forward")) +@_torch_register_fake_wrapper("flash_attn::_flash_attn_forward") def _flash_attn_forward_fake( q: torch.Tensor, k: torch.Tensor, @@ -126,41 +132,28 @@ def _flash_attn_forward_fake( window_size_right: int, softcap: float, alibi_slopes: Optional[torch.Tensor], - return_softmax: bool, + return_softmax: bool ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: q, k, v = [maybe_contiguous(x) for x in (q, k, v)] batch_size, seqlen_q, num_heads, head_size = q.shape seqlen_k = k.shape[1] out = torch.empty_like(q) - softmax_lse = torch.empty( - (batch_size, num_heads, seqlen_q), - dtype=torch.float32, - device=q.device, - layout=q.layout, - ) + softmax_lse = torch.empty((batch_size, num_heads, seqlen_q), dtype=torch.float32, device=q.device, layout=q.layout) p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) if return_softmax: - p = torch.empty( - ( - batch_size, - num_heads, - round_multiple(seqlen_q, 128), - round_multiple(seqlen_k, 128), - ), - dtype=q.dtype, - device=q.device, - layout=q.layout, - ) + p = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128), round_multiple(seqlen_k, 128)), dtype=q.dtype, device=q.device, layout=q.layout) rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) return out, softmax_lse, p, rng_state -@torch.library.custom_op( - add_op_namespace_prefix("_flash_attn_varlen_forward"), - mutates_args=(), - device_types=_get_device(), -) +if torch.__version__ >= "2.4.0": + _wrapped_flash_attn_forward = torch.ops.flash_attn._flash_attn_forward +else: + _wrapped_flash_attn_forward = _flash_attn_forward + + +@_torch_custom_op_wrapper("flash_attn::_flash_attn_varlen_forward", mutates_args=(), device_types=_get_device()) def _flash_attn_varlen_forward( q: torch.Tensor, k: torch.Tensor, @@ -211,7 +204,7 @@ def _flash_attn_varlen_forward( return out, softmax_lse, S_dmask, rng_state -@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_varlen_forward")) +@_torch_register_fake_wrapper("flash_attn::_flash_attn_varlen_forward") def _flash_attn_varlen_forward_fake( q: torch.Tensor, k: torch.Tensor, @@ -237,30 +230,25 @@ def _flash_attn_varlen_forward_fake( paged_kv = block_table is not None batch_size = cu_seqlens_q.numel() - 1 total_q, num_heads, _ = q.shape - + out = torch.empty_like(q) - softmax_lse = torch.empty( - (num_heads, total_q), dtype=torch.float32, device=q.device, layout=q.layout - ) + softmax_lse = torch.empty((num_heads, total_q), dtype=torch.float32, device=q.device, layout=q.layout) p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) seqlen_q_rounded = round_multiple(max_seqlen_q, 128) seqlen_k_rounded = round_multiple(max_seqlen_k, 128) if return_softmax: - p = torch.empty( - (batch_size, num_heads, seqlen_q_rounded, seqlen_k_rounded), - dtype=q.dtype, - device=q.device, - layout=q.layout, - ) + p = torch.empty((batch_size, num_heads, seqlen_q_rounded, seqlen_k_rounded), dtype=q.dtype, device=q.device, layout=q.layout) rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) return out, softmax_lse, p, rng_state -@torch.library.custom_op( - add_op_namespace_prefix("_flash_attn_backward"), - mutates_args=("dq", "dk", "dv"), - device_types=_get_device(), -) +if torch.__version__ >= "2.4.0": + _wrapped_flash_attn_varlen_forward = torch.ops.flash_attn._flash_attn_varlen_forward +else: + _wrapped_flash_attn_varlen_forward = _flash_attn_varlen_forward + + +@_torch_custom_op_wrapper("flash_attn::_flash_attn_backward", mutates_args=("dq", "dk", "dv"), device_types=_get_device()) def _flash_attn_backward( dout: torch.Tensor, q: torch.Tensor, @@ -312,7 +300,7 @@ def _flash_attn_backward( return softmax_d -@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_backward")) +@_torch_register_fake_wrapper("flash_attn::_flash_attn_backward") def _flash_attn_backward_fake( dout: torch.Tensor, q: torch.Tensor, @@ -341,20 +329,18 @@ def _flash_attn_backward_fake( if dv is None: dv = torch.empty_like(v) batch_size, seqlen_q, num_heads, _ = q.shape - softmax_d = torch.empty( - (batch_size, num_heads, round_multiple(seqlen_q, 128)), - device=q.device, - dtype=torch.float32, - ) - + softmax_d = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128)), device=q.device, dtype=torch.float32) + return softmax_d -@torch.library.custom_op( - add_op_namespace_prefix("_flash_attn_varlen_backward"), - mutates_args=("dq", "dk", "dv"), - device_types=_get_device(), -) +if torch.__version__ >= "2.4.0": + _wrapped_flash_attn_backward = torch.ops.flash_attn._flash_attn_backward +else: + _wrapped_flash_attn_backward = _flash_attn_backward + + +@_torch_custom_op_wrapper("flash_attn::_flash_attn_varlen_backward", mutates_args=("dq", "dk", "dv"), device_types=_get_device()) def _flash_attn_varlen_backward( dout: torch.Tensor, q: torch.Tensor, @@ -418,7 +404,7 @@ def _flash_attn_varlen_backward( return softmax_d -@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_varlen_backward")) +@_torch_register_fake_wrapper("flash_attn::_flash_attn_varlen_backward") def _flash_attn_varlen_backward_fake( dout: torch.Tensor, q: torch.Tensor, @@ -454,13 +440,17 @@ def _flash_attn_varlen_backward_fake( dk = torch.empty_like(k) if dv is None: dv = torch.empty_like(v) - softmax_d = torch.empty( - (num_heads, total_q + 128 * batch_size), device=q.device, dtype=torch.float32 - ) - + softmax_d = torch.empty((num_heads, total_q + 128 * batch_size), device=q.device, dtype=torch.float32) + return softmax_d +if torch.__version__ >= "2.4.0": + _wrapped_flash_attn_varlen_backward = torch.ops.flash_attn._flash_attn_varlen_backward +else: + _wrapped_flash_attn_varlen_backward = _flash_attn_varlen_backward + + class FlashAttnQKVPackedFunc(torch.autograd.Function): @staticmethod def forward( @@ -485,7 +475,7 @@ class FlashAttnQKVPackedFunc(torch.autograd.Function): q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_forward( + out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( q, k, v, @@ -519,7 +509,7 @@ class FlashAttnQKVPackedFunc(torch.autograd.Function): dout_padded = dout if head_size_og % 8 != 0: dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _flash_attn_backward( + _wrapped_flash_attn_backward( dout_padded, q, k, @@ -569,7 +559,7 @@ class FlashAttnVarlenQKVPackedFunc(torch.autograd.Function): q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_varlen_forward( + out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( q, k, v, @@ -588,9 +578,7 @@ class FlashAttnVarlenQKVPackedFunc(torch.autograd.Function): block_table=None, ) if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens, rng_state - ) + ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens, rng_state) ctx.dropout_p = dropout_p ctx.max_seqlen = max_seqlen ctx.softmax_scale = softmax_scale @@ -611,7 +599,7 @@ class FlashAttnVarlenQKVPackedFunc(torch.autograd.Function): dout_padded = dout if head_size_og % 8 != 0: dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _flash_attn_varlen_backward( + _wrapped_flash_attn_varlen_backward( dout_padded, q, k, @@ -655,7 +643,9 @@ class FlashAttnKVPackedFunc(torch.autograd.Function): return_softmax, is_grad_enabled, ): - is_grad = is_grad_enabled and any(x.requires_grad for x in [q, kv]) + is_grad = is_grad_enabled and any( + x.requires_grad for x in [q, kv] + ) if softmax_scale is None: softmax_scale = q.shape[-1] ** (-0.5) k, v = kv[:, :, 0].detach(), kv[:, :, 1].detach() @@ -664,7 +654,7 @@ class FlashAttnKVPackedFunc(torch.autograd.Function): q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_forward( + out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( q, k, v, @@ -699,7 +689,7 @@ class FlashAttnKVPackedFunc(torch.autograd.Function): dout_padded = dout if head_size_og % 8 != 0: dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _flash_attn_backward( + _wrapped_flash_attn_backward( dout_padded, q, k, @@ -744,7 +734,9 @@ class FlashAttnVarlenKVPackedFunc(torch.autograd.Function): return_softmax, is_grad_enabled, ): - is_grad = is_grad_enabled and any(x.requires_grad for x in [q, kv]) + is_grad = is_grad_enabled and any( + x.requires_grad for x in [q, kv] + ) if softmax_scale is None: softmax_scale = q.shape[-1] ** (-0.5) k, v = kv[:, 0].detach(), kv[:, 1].detach() @@ -753,7 +745,7 @@ class FlashAttnVarlenKVPackedFunc(torch.autograd.Function): q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_varlen_forward( + out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( q, k, v, @@ -789,9 +781,7 @@ class FlashAttnVarlenKVPackedFunc(torch.autograd.Function): @staticmethod def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ( - ctx.saved_tensors - ) + q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors dq = torch.empty_like(q) kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) @@ -799,7 +789,7 @@ class FlashAttnVarlenKVPackedFunc(torch.autograd.Function): dout_padded = dout if head_size_og % 8 != 0: dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _flash_attn_varlen_backward( + _wrapped_flash_attn_varlen_backward( dout_padded, q, k, @@ -825,23 +815,7 @@ class FlashAttnVarlenKVPackedFunc(torch.autograd.Function): ) dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension dkv = dkv[..., : dout.shape[-1]] - return ( - dq, - dkv, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) + return dq, dkv, None, None, None, None, None, None, None, None, None, None, None, None, None class FlashAttnFunc(torch.autograd.Function): @@ -861,7 +835,9 @@ class FlashAttnFunc(torch.autograd.Function): return_softmax, is_grad_enabled, ): - is_grad = is_grad_enabled and any(x.requires_grad for x in [q, k, v]) + is_grad = is_grad_enabled and any( + x.requires_grad for x in [q, k, v] + ) if softmax_scale is None: softmax_scale = q.shape[-1] ** (-0.5) head_size_og = q.size(3) @@ -869,7 +845,7 @@ class FlashAttnFunc(torch.autograd.Function): q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_forward( + out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( q, k, v, @@ -902,7 +878,7 @@ class FlashAttnFunc(torch.autograd.Function): dout_padded = dout if head_size_og % 8 != 0: dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _flash_attn_backward( + _wrapped_flash_attn_backward( dout_padded, q, k, @@ -950,7 +926,9 @@ class FlashAttnVarlenFunc(torch.autograd.Function): block_table, is_grad_enabled, ): - is_grad = is_grad_enabled and any(x.requires_grad for x in [q, k, v]) + is_grad = is_grad_enabled and any( + x.requires_grad for x in [q, k, v] + ) if softmax_scale is None: softmax_scale = q.shape[-1] ** (-0.5) head_size_og = q.size(2) @@ -958,7 +936,7 @@ class FlashAttnVarlenFunc(torch.autograd.Function): q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_varlen_forward( + out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( q, k, v, @@ -995,15 +973,13 @@ class FlashAttnVarlenFunc(torch.autograd.Function): @staticmethod def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ( - ctx.saved_tensors - ) + q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) head_size_og = dout.size(2) dout_padded = dout if head_size_og % 8 != 0: dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _flash_attn_varlen_backward( + _wrapped_flash_attn_varlen_backward( dout_padded, q, k, @@ -1030,25 +1006,7 @@ class FlashAttnVarlenFunc(torch.autograd.Function): dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension dk = dk[..., : dout.shape[-1]] dv = dv[..., : dout.shape[-1]] - return ( - dq, - dk, - dv, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None def flash_attn_qkvpacked_func( @@ -1106,7 +1064,7 @@ def flash_attn_qkvpacked_func( alibi_slopes, deterministic, return_attn_probs, - torch.is_grad_enabled(), + False if _XPU_AVAILABLE else torch.is_grad_enabled(), ) @@ -1184,7 +1142,7 @@ def flash_attn_kvpacked_func( alibi_slopes, deterministic, return_attn_probs, - torch.is_grad_enabled(), + False if _XPU_AVAILABLE else torch.is_grad_enabled(), ) @@ -1196,7 +1154,7 @@ def flash_attn_func( softmax_scale=None, causal=False, window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated + softcap=0.0, # 0.0 means deactivated alibi_slopes=None, deterministic=False, return_attn_probs=False, @@ -1261,7 +1219,7 @@ def flash_attn_func( alibi_slopes, deterministic, return_attn_probs, - torch.is_grad_enabled(), + False if _XPU_AVAILABLE else torch.is_grad_enabled(), ) @@ -1273,7 +1231,7 @@ def flash_attn_varlen_qkvpacked_func( softmax_scale=None, causal=False, window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated + softcap=0.0, # 0.0 means deactivated alibi_slopes=None, deterministic=False, return_attn_probs=False, @@ -1327,7 +1285,7 @@ def flash_attn_varlen_qkvpacked_func( alibi_slopes, deterministic, return_attn_probs, - torch.is_grad_enabled(), + False if _XPU_AVAILABLE else torch.is_grad_enabled(), ) @@ -1342,7 +1300,7 @@ def flash_attn_varlen_kvpacked_func( softmax_scale=None, causal=False, window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated + softcap=0.0, # 0.0 means deactivated alibi_slopes=None, deterministic=False, return_attn_probs=False, @@ -1419,7 +1377,7 @@ def flash_attn_varlen_kvpacked_func( alibi_slopes, deterministic, return_attn_probs, - torch.is_grad_enabled(), + False if _XPU_AVAILABLE else torch.is_grad_enabled(), ) @@ -1435,7 +1393,7 @@ def flash_attn_varlen_func( softmax_scale=None, causal=False, window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated + softcap=0.0, # 0.0 means deactivated alibi_slopes=None, deterministic=False, return_attn_probs=False, @@ -1513,7 +1471,7 @@ def flash_attn_varlen_func( deterministic, return_attn_probs, block_table, - False if q.device.type == "cpu" else torch.is_grad_enabled(), + False if _XPU_AVAILABLE or q.device.type == "cpu" else torch.is_grad_enabled(), ) @@ -1532,7 +1490,7 @@ def flash_attn_with_kvcache( softmax_scale=None, causal=False, window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated + softcap=0.0, # 0.0 means deactivated rotary_interleaved=True, alibi_slopes=None, num_splits=0, diff --git a/build/torch211-cxx11-cu126-x86_64-linux/layers/__init__.py b/build/torch210-cxx11-xpu20253-x86_64-linux/layers/__init__.py similarity index 100% rename from build/torch211-cxx11-cu126-x86_64-linux/layers/__init__.py rename to build/torch210-cxx11-xpu20253-x86_64-linux/layers/__init__.py diff --git a/build/torch211-cxx11-cu128-x86_64-linux/layers/patch_embed.py b/build/torch210-cxx11-xpu20253-x86_64-linux/layers/patch_embed.py similarity index 100% rename from build/torch211-cxx11-cu128-x86_64-linux/layers/patch_embed.py rename to build/torch210-cxx11-xpu20253-x86_64-linux/layers/patch_embed.py diff --git a/build/torch211-cxx11-cu128-x86_64-linux/layers/rotary.py b/build/torch210-cxx11-xpu20253-x86_64-linux/layers/rotary.py similarity index 100% rename from build/torch211-cxx11-cu128-x86_64-linux/layers/rotary.py rename to build/torch210-cxx11-xpu20253-x86_64-linux/layers/rotary.py diff --git a/build/torch210-cxx11-xpu20253-x86_64-linux/metadata.json b/build/torch210-cxx11-xpu20253-x86_64-linux/metadata.json index 19480d14e892636e0935cc1e767b18cdb8256b2f..9cf5deed9898dce769f4cc73913d3530b92a0bd8 100644 --- a/build/torch210-cxx11-xpu20253-x86_64-linux/metadata.json +++ b/build/torch210-cxx11-xpu20253-x86_64-linux/metadata.json @@ -1,10 +1,4 @@ { - "name": "flash-attn2", - "id": "_flash_attn2_xpu_a511a4c", "version": 1, - "license": "BSD-3-Clause", - "python-depends": [], - "backend": { - "type": "xpu" - } -} + "python-depends": [] +} \ No newline at end of file diff --git a/build/torch211-cxx11-cu126-x86_64-linux/ops/__init__.py b/build/torch210-cxx11-xpu20253-x86_64-linux/ops/__init__.py similarity index 100% rename from build/torch211-cxx11-cu126-x86_64-linux/ops/__init__.py rename to build/torch210-cxx11-xpu20253-x86_64-linux/ops/__init__.py diff --git a/build/torch211-cxx11-cu128-x86_64-linux/ops/activations.py b/build/torch210-cxx11-xpu20253-x86_64-linux/ops/activations.py similarity index 100% rename from build/torch211-cxx11-cu128-x86_64-linux/ops/activations.py rename to build/torch210-cxx11-xpu20253-x86_64-linux/ops/activations.py diff --git a/build/torch211-cxx11-cu128-x86_64-linux/ops/fused_dense.py b/build/torch210-cxx11-xpu20253-x86_64-linux/ops/fused_dense.py similarity index 100% rename from build/torch211-cxx11-cu128-x86_64-linux/ops/fused_dense.py rename to build/torch210-cxx11-xpu20253-x86_64-linux/ops/fused_dense.py diff --git a/build/torch211-cxx11-cu128-x86_64-linux/ops/layer_norm.py b/build/torch210-cxx11-xpu20253-x86_64-linux/ops/layer_norm.py similarity index 100% rename from build/torch211-cxx11-cu128-x86_64-linux/ops/layer_norm.py rename to build/torch210-cxx11-xpu20253-x86_64-linux/ops/layer_norm.py diff --git a/build/torch211-cxx11-cu128-x86_64-linux/ops/rms_norm.py b/build/torch210-cxx11-xpu20253-x86_64-linux/ops/rms_norm.py similarity index 100% rename from build/torch211-cxx11-cu128-x86_64-linux/ops/rms_norm.py rename to build/torch210-cxx11-xpu20253-x86_64-linux/ops/rms_norm.py diff --git a/build/torch211-cxx11-cu128-x86_64-linux/ops/triton/__init__.py b/build/torch210-cxx11-xpu20253-x86_64-linux/ops/triton/__init__.py similarity index 100% rename from build/torch211-cxx11-cu128-x86_64-linux/ops/triton/__init__.py rename to build/torch210-cxx11-xpu20253-x86_64-linux/ops/triton/__init__.py diff --git a/build/torch211-cxx11-cu128-x86_64-linux/ops/triton/cross_entropy.py b/build/torch210-cxx11-xpu20253-x86_64-linux/ops/triton/cross_entropy.py similarity index 100% rename from build/torch211-cxx11-cu128-x86_64-linux/ops/triton/cross_entropy.py rename to build/torch210-cxx11-xpu20253-x86_64-linux/ops/triton/cross_entropy.py diff --git a/build/torch211-cxx11-cu128-x86_64-linux/ops/triton/k_activations.py b/build/torch210-cxx11-xpu20253-x86_64-linux/ops/triton/k_activations.py similarity index 100% rename from build/torch211-cxx11-cu128-x86_64-linux/ops/triton/k_activations.py rename to build/torch210-cxx11-xpu20253-x86_64-linux/ops/triton/k_activations.py diff --git a/build/torch29-cxx11-cu130-aarch64-linux/ops/triton/layer_norm.py b/build/torch210-cxx11-xpu20253-x86_64-linux/ops/triton/layer_norm.py similarity index 100% rename from build/torch29-cxx11-cu130-aarch64-linux/ops/triton/layer_norm.py rename to build/torch210-cxx11-xpu20253-x86_64-linux/ops/triton/layer_norm.py diff --git a/build/torch211-cxx11-cu128-x86_64-linux/ops/triton/linear.py b/build/torch210-cxx11-xpu20253-x86_64-linux/ops/triton/linear.py similarity index 100% rename from build/torch211-cxx11-cu128-x86_64-linux/ops/triton/linear.py rename to build/torch210-cxx11-xpu20253-x86_64-linux/ops/triton/linear.py diff --git a/build/torch211-cxx11-cu128-x86_64-linux/ops/triton/mlp.py b/build/torch210-cxx11-xpu20253-x86_64-linux/ops/triton/mlp.py similarity index 100% rename from build/torch211-cxx11-cu128-x86_64-linux/ops/triton/mlp.py rename to build/torch210-cxx11-xpu20253-x86_64-linux/ops/triton/mlp.py diff --git a/build/torch211-cxx11-cu128-x86_64-linux/ops/triton/rotary.py b/build/torch210-cxx11-xpu20253-x86_64-linux/ops/triton/rotary.py similarity index 100% rename from build/torch211-cxx11-cu128-x86_64-linux/ops/triton/rotary.py rename to build/torch210-cxx11-xpu20253-x86_64-linux/ops/triton/rotary.py diff --git a/build/torch211-cxx11-cpu-x86_64-linux/__init__.py b/build/torch211-cxx11-cpu-x86_64-linux/__init__.py deleted file mode 100644 index ecc2f9d896b6c93f90b0a1499856dc0612177422..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cpu-x86_64-linux/__init__.py +++ /dev/null @@ -1,393 +0,0 @@ -from typing import Optional, List -import torch -from ._ops import ops as flash_attn_ops -from .flash_attn_interface import ( - flash_attn_func, - flash_attn_kvpacked_func, - flash_attn_qkvpacked_func, - flash_attn_varlen_func, - flash_attn_varlen_kvpacked_func, - flash_attn_varlen_qkvpacked_func, - flash_attn_with_kvcache, -) - - -def fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - return_softmax: bool = False, - gen: Optional[torch.Generator] = None, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention. - - Args: - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Optional output tensor, same shape as q - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - return_softmax: Whether to return softmax weights - gen: Optional random number generator - - Returns: - List of tensors: [output, softmax_lse, (softmax if return_softmax)] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.fwd( - q, - k, - v, - out, - alibi_slopes, - p_dropout, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - gen, - ) - - -def varlen_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - out: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - block_table: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - max_seqlen_q: int = 0, - max_seqlen_k: int = 0, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - zero_tensors: bool = False, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - return_softmax: bool = False, - gen: Optional[torch.Generator] = None, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention with variable sequence lengths. - - Args: - q: Query tensor of shape [total_q, num_heads, head_size] - k: Key tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - v: Value tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1] - cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1] - out: Optional output tensor of shape [total_q, num_heads, head_size] - seqused_k: Optional tensor specifying how many keys to use per batch element [batch_size] - leftpad_k: Optional left padding for keys of shape [batch_size] - block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq] - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - max_seqlen_q: Maximum sequence length for queries - max_seqlen_k: Maximum sequence length for keys - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - zero_tensors: Whether to zero tensors before computation - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - return_softmax: Whether to return softmax weights - gen: Optional random number generator - - Returns: - List of tensors: [output, softmax_lse, (softmax if return_softmax)] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.varlen_fwd( - q, - k, - v, - out, - cu_seqlens_q, - cu_seqlens_k, - seqused_k, - leftpad_k, - block_table, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - p_dropout, - softmax_scale, - zero_tensors, - is_causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - gen, - ) - - -def bwd( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - p_dropout: float = 0.0, - 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, - gen: Optional[torch.Generator] = None, - rng_state: Optional[torch.Tensor] = None, -) -> List[torch.Tensor]: - """ - Backward pass for multi-head attention. - - Args: - dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size] - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size] - softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q] - dq: Optional gradient tensor for queries, same shape as q - dk: Optional gradient tensor for keys, same shape as k - dv: Optional gradient tensor for values, same shape as v - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - deterministic: Whether to use deterministic algorithms - gen: Optional random number generator - rng_state: Optional RNG state from forward pass - - Returns: - List of tensors: [dq, dk, dv] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - alibi_slopes, - p_dropout, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - deterministic, - gen, - rng_state, - ) - - -def varlen_bwd( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - max_seqlen_q: int = 0, - max_seqlen_k: int = 0, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - zero_tensors: bool = False, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - deterministic: bool = False, - gen: Optional[torch.Generator] = None, - rng_state: Optional[torch.Tensor] = None, -) -> List[torch.Tensor]: - """ - Backward pass for multi-head attention with variable sequence lengths. - - Args: - dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size] - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size] - softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q] - cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1] - cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1] - dq: Optional gradient tensor for queries, same shape as q - dk: Optional gradient tensor for keys, same shape as k - dv: Optional gradient tensor for values, same shape as v - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - max_seqlen_q: Maximum sequence length for queries - max_seqlen_k: Maximum sequence length for keys - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - zero_tensors: Whether to zero tensors before computation - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - deterministic: Whether to use deterministic algorithms - gen: Optional random number generator - rng_state: Optional RNG state from forward pass - - Returns: - List of tensors: [dq, dk, dv] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.varlen_bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - p_dropout, - softmax_scale, - zero_tensors, - is_causal, - window_size_left, - window_size_right, - softcap, - deterministic, - gen, - rng_state, - ) - - -def fwd_kvcache( - q: torch.Tensor, - kcache: torch.Tensor, - vcache: torch.Tensor, - k: Optional[torch.Tensor] = None, - v: Optional[torch.Tensor] = None, - seqlens_k: Optional[torch.Tensor] = None, - rotary_cos: Optional[torch.Tensor] = None, - rotary_sin: Optional[torch.Tensor] = None, - cache_batch_idx: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - block_table: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - out: 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, - is_rotary_interleaved: bool = False, - num_splits: int = 1, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention with KV cache. - - Args: - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - kcache: Key cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - vcache: Value cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - k: Optional new keys tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size] - v: Optional new values tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size] - seqlens_k: Optional sequence lengths for keys of shape [batch_size] - rotary_cos: Optional rotary cosine tensor of shape [seqlen_ro, rotary_dim/2] - rotary_sin: Optional rotary sine tensor of shape [seqlen_ro, rotary_dim/2] - cache_batch_idx: Optional indices to index into the KV cache - leftpad_k: Optional left padding for keys of shape [batch_size] - block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq] - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - out: Optional output tensor, same shape as q - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - is_rotary_interleaved: Whether rotary embeddings are interleaved - num_splits: Number of splits for computation - - Returns: - List of tensors: [output, softmax_lse] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.fwd_kvcache( - q, - kcache, - vcache, - k, - v, - seqlens_k, - rotary_cos, - rotary_sin, - cache_batch_idx, - leftpad_k, - block_table, - alibi_slopes, - out, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - is_rotary_interleaved, - num_splits, - ) diff --git a/build/torch211-cxx11-cpu-x86_64-linux/_flash_attn2_cpu_f12afc9.abi3.so b/build/torch211-cxx11-cpu-x86_64-linux/_flash_attn2_cpu_f12afc9.abi3.so deleted file mode 100644 index 7f851338deb95e07d15b3888601b95aceebd6808..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cpu-x86_64-linux/_flash_attn2_cpu_f12afc9.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:a5fa36d84dbf0b9c749529e39547e9fc17fbb289640e715281c0beb7d82e012c -size 201552 diff --git a/build/torch211-cxx11-cpu-x86_64-linux/_ops.py b/build/torch211-cxx11-cpu-x86_64-linux/_ops.py deleted file mode 100644 index 90538036e3208c65483d0140dc37ea54c295f477..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cpu-x86_64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _flash_attn2_cpu_f12afc9 -ops = torch.ops._flash_attn2_cpu_f12afc9 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_flash_attn2_cpu_f12afc9::{op_name}" diff --git a/build/torch211-cxx11-cpu-x86_64-linux/bert_padding.py b/build/torch211-cxx11-cpu-x86_64-linux/bert_padding.py deleted file mode 100644 index 3c2d35159a014a9d03aabead9e52e009168696ea..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cpu-x86_64-linux/bert_padding.py +++ /dev/null @@ -1,218 +0,0 @@ -# Adapted from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/padding.py - -import torch -import torch.nn.functional as F -from einops import rearrange, repeat - - -class IndexFirstAxis(torch.autograd.Function): - @staticmethod - def forward(ctx, input, indices): - ctx.save_for_backward(indices) - assert input.ndim >= 2 - ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] - second_dim = other_shape.numel() - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - # return input[indices] - return torch.gather( - rearrange(input, "b ... -> b (...)"), 0, repeat(indices, "z -> z d", d=second_dim) - ).reshape(-1, *other_shape) - - @staticmethod - def backward(ctx, grad_output): - (indices,) = ctx.saved_tensors - assert grad_output.ndim >= 2 - other_shape = grad_output.shape[1:] - grad_output = rearrange(grad_output, "b ... -> b (...)") - grad_input = torch.zeros( - [ctx.first_axis_dim, grad_output.shape[1]], - device=grad_output.device, - dtype=grad_output.dtype, - ) - # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. - # grad_input[indices] = grad_output - grad_input.scatter_(0, repeat(indices, "z -> z d", d=grad_output.shape[1]), grad_output) - return grad_input.reshape(ctx.first_axis_dim, *other_shape), None - - -index_first_axis = IndexFirstAxis.apply - - -class IndexPutFirstAxis(torch.autograd.Function): - @staticmethod - def forward(ctx, values, indices, first_axis_dim): - ctx.save_for_backward(indices) - assert indices.ndim == 1 - assert values.ndim >= 2 - output = torch.zeros( - first_axis_dim, *values.shape[1:], device=values.device, dtype=values.dtype - ) - # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. - output[indices] = values - # output.scatter_(0, repeat(indices, 'z -> z d', d=values.shape[1]), values) - return output - - @staticmethod - def backward(ctx, grad_output): - (indices,) = ctx.saved_tensors - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - grad_values = grad_output[indices] - # grad_values = torch.gather(grad_output, 0, repeat(indices, 'z -> z d', d=grad_output.shape[1])) - return grad_values, None, None - - -index_put_first_axis = IndexPutFirstAxis.apply - - -class IndexFirstAxisResidual(torch.autograd.Function): - @staticmethod - def forward(ctx, input, indices): - ctx.save_for_backward(indices) - assert input.ndim >= 2 - ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] - second_dim = other_shape.numel() - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - output = input[indices] - # We don't want to reshape input (b ... -> b (...)) since it could change the channel_last - # memory format to channel_first. In other words, input might not be contiguous. - # If we don't detach, Pytorch complains about output being a view and is being modified inplace - return output, input.detach() - - @staticmethod - def backward(ctx, grad_output, grad_residual): - (indices,) = ctx.saved_tensors - assert grad_output.ndim >= 2 - other_shape = grad_output.shape[1:] - assert grad_residual.shape[1:] == other_shape - grad_input = grad_residual - # grad_input[indices] += grad_output - indices = indices.reshape(indices.shape[0], *((1,) * (grad_output.ndim - 1))) - indices = indices.expand_as(grad_output) - grad_input.scatter_add_(0, indices, grad_output) - return grad_input.reshape(ctx.first_axis_dim, *other_shape), None - - -index_first_axis_residual = IndexFirstAxisResidual.apply - - -def unpad_input(hidden_states, attention_mask, unused_mask=None): - """ - Arguments: - hidden_states: (batch, seqlen, ...) - attention_mask: (batch, seqlen), bool / int, 1 means valid and 0 means not valid. - unused_mask: (batch, seqlen), bool / int, 1 means the element is allocated but unused. - Return: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens selected in attention_mask + unused_mask. - indices: (total_nnz), the indices of masked tokens from the flattened input sequence. - cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. - max_seqlen_in_batch: int - seqused: (batch), returns the number of tokens selected in attention_mask + unused_mask. - """ - all_masks = (attention_mask + unused_mask) if unused_mask is not None else attention_mask - seqlens_in_batch = all_masks.sum(dim=-1, dtype=torch.int32) - used_seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) - indices = torch.nonzero(all_masks.flatten(), as_tuple=False).flatten() - max_seqlen_in_batch = seqlens_in_batch.max().item() - cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) - # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the - # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim - # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to - # index with integer indices. Moreover, torch's index is a bit slower than it needs to be, - # so we write custom forward and backward to make it a bit faster. - return ( - index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices), - indices, - cu_seqlens, - max_seqlen_in_batch, - used_seqlens_in_batch, - ) - - -def unpad_input_for_concatenated_sequences(hidden_states, attention_mask_in_length): - """ - Supports concatenating short samples in one sequence. The attention_mask_in_length is utilized to mask other short samples. It helps efficient training of variant lengths-based samples (e.g., the supervised fine-tuning task in large language model). - The motivation for this function is explained [here](https://github.com/Dao-AILab/flash-attention/issues/432#issuecomment-1668822286). - - For example, if batch = 3 and seqlen = 6, the attention_mask_in_length is: - ``` - [ - [2, 3, 0, 0, 0, 0], - [3, 2, 0, 0, 0, 0], - [6, 0, 0, 0, 0, 0] - ] - ``` - , which refers to the 3D-attention mask: - ``` - [ - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [0, 0, 1, 0, 0, 0], - [0, 0, 1, 1, 0, 0], - [0, 0, 1, 1, 1, 0], - [0, 0, 0, 0, 0, 1] - ], - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [0, 0, 0, 1, 0, 0], - [0, 0, 0, 1, 1, 0], - [0, 0, 0, 0, 0, 1] - ], - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [1, 1, 1, 1, 0, 0], - [1, 1, 1, 1, 1, 0], - [1, 1, 1, 1, 1, 1] - ] - ] - ```. - - Arguments: - hidden_states: (batch, seqlen, ...) - attention_mask_in_length: (batch, seqlen), int, a nonzero number (e.g., 1, 2, 3, etc.) means length of concatenated sequence in b-th batch, and 0 means none. - Return: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. - indices: (total_nnz), the indices of non-masked tokens from the flattened input sequence. - cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. - max_seqlen_in_batch: int - """ - length = attention_mask_in_length.sum(dim=-1) - seqlen = attention_mask_in_length.size(-1) - attention_mask_2d = torch.arange(seqlen, device=length.device, dtype=length.dtype).expand(len(length), seqlen) < length.unsqueeze(1) - real_indices_idx = torch.nonzero(attention_mask_in_length.flatten(), as_tuple=False).flatten() - seqlens_in_batch = attention_mask_in_length.flatten()[real_indices_idx] - indices = torch.nonzero(attention_mask_2d.flatten(), as_tuple=False).flatten() - max_seqlen_in_batch = seqlens_in_batch.max().item() - cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) - # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the - # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim - # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to - # index with integer indices. Moreover, torch's index is a bit slower than it needs to be, - # so we write custom forward and backward to make it a bit faster. - return ( - index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices), - indices, - cu_seqlens, - max_seqlen_in_batch, - ) - - -def pad_input(hidden_states, indices, batch, seqlen): - """ - Arguments: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. - indices: (total_nnz), the indices that represent the non-masked tokens of the original padded input sequence. - batch: int, batch size for the padded sequence. - seqlen: int, maximum sequence length for the padded sequence. - Return: - hidden_states: (batch, seqlen, ...) - """ - dim = hidden_states.shape[-1] - # output = torch.zeros((batch * seqlen), dim, device=hidden_states.device, dtype=hidden_states.dtype) - # output[indices] = hidden_states - output = index_put_first_axis(hidden_states, indices, batch * seqlen) - return rearrange(output, "(b s) ... -> b s ...", b=batch) diff --git a/build/torch211-cxx11-cpu-x86_64-linux/flash_attn2/__init__.py b/build/torch211-cxx11-cpu-x86_64-linux/flash_attn2/__init__.py deleted file mode 100644 index a9b2672c1cd85b74c1b3ded0fc0b2100e1aeac23..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cpu-x86_64-linux/flash_attn2/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -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-cpu-x86_64-linux/flash_attn_interface.py b/build/torch211-cxx11-cpu-x86_64-linux/flash_attn_interface.py deleted file mode 100644 index b014e569dc9812a9c840d1043f2cebe767b4d477..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cpu-x86_64-linux/flash_attn_interface.py +++ /dev/null @@ -1,1631 +0,0 @@ -# Copyright (c) 2023, Tri Dao. - -from typing import Optional, Sequence, Tuple, Union - -import torch -import torch.nn as nn -import os - -# # isort: off -# # We need to import the CUDA kernels after importing torch -# USE_TRITON_ROCM = os.getenv("FLASH_ATTENTION_TRITON_AMD_ENABLE", "FALSE") == "TRUE" -# if USE_TRITON_ROCM: -# from .flash_attn_triton_amd import interface_fa as flash_attn -# else: -# import flash_attn_2_cuda as flash_attn - - -from . import _ops -from ._ops import ops as flash_attn -from ._ops import add_op_namespace_prefix - -# # isort: on - -def maybe_contiguous(x): - return x.contiguous() if x is not None and x.stride(-1) != 1 else x - - -def _get_device(): - if torch.xpu.is_available(): - return "xpu" - elif torch.cuda.is_available(): - return "cuda" - else: - return "cpu" - - -def _get_block_size_n(device, head_dim, is_dropout, is_causal): - assert head_dim <= 256 - - if device.type == "xpu": - if head_dim <= 96: - return 64 - elif head_dim <= 128: - return 32 - elif head_dim <= 256: - return 64 - else: - return 32 - - # This should match the block sizes in the CUDA kernel - major, minor = torch.cuda.get_device_capability(device) - is_sm8x = major == 8 and minor > 0 # Only include sm86 and sm89, exclude sm80 (A100) - is_sm80 = major == 8 and minor == 0 - is_sm90 = major == 9 and minor == 0 - if head_dim <= 32: - return 128 - if head_dim <= 64: - return 128 if not is_dropout else 64 - elif head_dim <= 96: - return 64 - elif head_dim <= 128: - if is_sm8x: - return 64 if (not is_dropout and is_causal) else 32 - else: - return 64 if not is_dropout else 32 - elif head_dim <= 192: - return 64 - elif head_dim <= 224: - return 64 - elif head_dim <= 256: - return 64 - - -def round_multiple(x, m): - return (x + m - 1) // m * m - - -# torch.compile() support is only enabled for pytorch >= 2.4 -# The reason for this is that we are using the new custom_op and register_fake -# APIs, which support inplace modification of inputs in the function itself -if torch.__version__ >= "2.4.0": - _torch_custom_op_wrapper = torch.library.custom_op - _torch_register_fake_wrapper = torch.library.register_fake -else: - def noop_custom_op_wrapper(name, fn=None, /, *, mutates_args, device_types=None, schema=None): - def wrap(func): - return func - if fn is None: - return wrap - return fn - def noop_register_fake_wrapper(op, fn=None, /, *, lib=None, _stacklevel=1): - def wrap(func): - return func - if fn is None: - return wrap - return fn - _torch_custom_op_wrapper = noop_custom_op_wrapper - _torch_register_fake_wrapper = noop_register_fake_wrapper - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_forward"), mutates_args=(), device_types=_get_device()) -def _flash_attn_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - return_softmax: bool -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - out, softmax_lse, S_dmask, rng_state = flash_attn.fwd( - q, - k, - v, - None, - alibi_slopes, - dropout_p, - softmax_scale, - causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - None, - ) - return out, softmax_lse, S_dmask, rng_state - - -@_torch_register_fake_wrapper(add_op_namespace_prefix("_flash_attn_forward")) -def _flash_attn_forward_fake( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - return_softmax: bool -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - batch_size, seqlen_q, num_heads, head_size = q.shape - seqlen_k = k.shape[1] - out = torch.empty_like(q) - softmax_lse = torch.empty((batch_size, num_heads, seqlen_q), dtype=torch.float32, device=q.device, layout=q.layout) - p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) - if return_softmax: - p = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128), round_multiple(seqlen_k, 128)), dtype=q.dtype, device=q.device, layout=q.layout) - rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) - - return out, softmax_lse, p, rng_state - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_forward = _ops.ops._flash_attn_forward -else: - _wrapped_flash_attn_forward = _flash_attn_forward - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_varlen_forward"), mutates_args=(), device_types=_get_device()) -def _flash_attn_varlen_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - alibi_slopes: Optional[torch.Tensor] = None, - return_softmax: bool = False, - block_table: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - out, softmax_lse, S_dmask, rng_state = flash_attn.varlen_fwd( - q, - k, - v, - None, - cu_seqlens_q, - cu_seqlens_k, - seqused_k, - leftpad_k, - block_table, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - zero_tensors, - causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - None, - ) - # if out.isnan().any() or softmax_lse.isnan().any(): - # breakpoint() - return out, softmax_lse, S_dmask, rng_state - - -@_torch_register_fake_wrapper(add_op_namespace_prefix("_flash_attn_varlen_forward")) -def _flash_attn_varlen_forward_fake( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - alibi_slopes: Optional[torch.Tensor] = None, - return_softmax: bool = False, - block_table: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - paged_kv = block_table is not None - batch_size = cu_seqlens_q.numel() - 1 - total_q, num_heads, _ = q.shape - - out = torch.empty_like(q) - softmax_lse = torch.empty((num_heads, total_q), dtype=torch.float32, device=q.device, layout=q.layout) - p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) - seqlen_q_rounded = round_multiple(max_seqlen_q, 128) - seqlen_k_rounded = round_multiple(max_seqlen_k, 128) - if return_softmax: - p = torch.empty((batch_size, num_heads, seqlen_q_rounded, seqlen_k_rounded), dtype=q.dtype, device=q.device, layout=q.layout) - rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) - return out, softmax_lse, p, rng_state - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_varlen_forward = _ops.ops._flash_attn_varlen_forward -else: - _wrapped_flash_attn_varlen_forward = _flash_attn_varlen_forward - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_backward"), mutates_args=("dq", "dk", "dv"), device_types=_get_device()) -def _flash_attn_backward( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, -) -> 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)] - ( - dq, - dk, - dv, - softmax_d, - ) = flash_attn.bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - alibi_slopes, - dropout_p, - softmax_scale, - causal, - window_size_left, - window_size_right, - softcap, - deterministic, - None, - rng_state, - ) - return softmax_d - - -@_torch_register_fake_wrapper(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, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, -) -> torch.Tensor: - dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] - if dq is None: - dq = torch.empty_like(q) - if dk is None: - dk = torch.empty_like(k) - if dv is None: - dv = torch.empty_like(v) - batch_size, seqlen_q, num_heads, _ = q.shape - softmax_d = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128)), device=q.device, dtype=torch.float32) - - return softmax_d - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_backward = _ops.ops._flash_attn_backward -else: - _wrapped_flash_attn_backward = _flash_attn_backward - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_varlen_backward"), mutates_args=("dq", "dk", "dv"), device_types=_get_device()) -def _flash_attn_varlen_backward( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> 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)] - ( - dq, - dk, - dv, - softmax_d, - ) = flash_attn.varlen_bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - zero_tensors, - causal, - window_size_left, - window_size_right, - softcap, - deterministic, - None, - rng_state, - ) - # if dk.isnan().any() or dk.isnan().any() or dv.isnan().any() or softmax_d.isnan().any(): - # breakpoint() - return softmax_d - - -@_torch_register_fake_wrapper(add_op_namespace_prefix("_flash_attn_varlen_backward")) -def _flash_attn_varlen_backward_fake( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> torch.Tensor: - dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] - batch_size = cu_seqlens_q.numel() - 1 - total_q, num_heads, _ = q.shape - - if dq is None: - dq = torch.empty_like(q) - if dk is None: - dk = torch.empty_like(k) - if dv is None: - dv = torch.empty_like(v) - softmax_d = torch.empty((num_heads, total_q + 128 * batch_size), device=q.device, dtype=torch.float32) - - return softmax_d - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_varlen_backward = _ops.ops._flash_attn_varlen_backward -else: - _wrapped_flash_attn_varlen_backward = _flash_attn_varlen_backward - - -class FlashAttnQKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and qkv.requires_grad - if softmax_scale is None: - softmax_scale = qkv.shape[-1] ** (-0.5) - q, k, v = qkv[:, :, 0].detach(), qkv[:, :, 1].detach(), qkv[:, :, 2].detach() - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) - dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dqkv[:, :, 0], - dqkv[:, :, 1], - dqkv[:, :, 2], - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dqkv = dqkv[..., : dout.shape[-1]] # We could have padded the head dimension - return dqkv, None, None, None, None, None, None, None, None, None - - -class FlashAttnVarlenQKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - cu_seqlens, - max_seqlen, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and qkv.requires_grad - if softmax_scale is None: - softmax_scale = qkv.shape[-1] ** (-0.5) - q, k, v = qkv[:, 0].detach(), qkv[:, 1].detach(), qkv[:, 2].detach() - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens, - cu_seqlens, - max_seqlen, - max_seqlen, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=None, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens, rng_state) - ctx.dropout_p = dropout_p - ctx.max_seqlen = max_seqlen - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens, rng_state = ctx.saved_tensors - qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) - dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dqkv[:, 0], - dqkv[:, 1], - dqkv[:, 2], - cu_seqlens, - cu_seqlens, - ctx.max_seqlen, - ctx.max_seqlen, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - 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 FlashAttnKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - kv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, kv] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - k, v = kv[:, :, 0].detach(), kv[:, :, 1].detach() - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - dq = torch.empty_like(q) - kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) - dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dkv[:, :, 0], - dkv[:, :, 1], - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dkv = dkv[..., : dout.shape[-1]] - return dq, dkv, None, None, None, None, None, None, None, None, None - - -class FlashAttnVarlenKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, kv] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - k, v = kv[:, 0].detach(), kv[:, 1].detach() - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=None, - ) - if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state - ) - ctx.dropout_p = dropout_p - 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.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors - dq = torch.empty_like(q) - kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) - dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dkv[:, 0], - dkv[:, 1], - cu_seqlens_q, - cu_seqlens_k, - ctx.max_seqlen_q, - ctx.max_seqlen_k, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dkv = dkv[..., : dout.shape[-1]] - return dq, dkv, None, None, None, None, None, None, None, None, None, None, None, None, None - - -class FlashAttnFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - k, - v, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, k, v] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dk = dk[..., : dout.shape[-1]] - dv = dv[..., : dout.shape[-1]] - return dq, dk, dv, 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, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - block_table, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, k, v] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=block_table, - ) - if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state - ) - ctx.dropout_p = dropout_p - 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.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors - dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - ctx.max_seqlen_q, - ctx.max_seqlen_k, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dk = dk[..., : dout.shape[-1]] - dv = dv[..., : dout.shape[-1]] - return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None - - -def flash_attn_qkvpacked_func( - qkv, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # <=0.0 means deactivate - alibi_slopes=None, - deterministic=False, - 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, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_kvpacked_func( - q, - kv, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, -): - """dropout_p should be set to 0.0 during evaluation - If 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 K, V. - 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) - kv: (batch_size, seqlen, 2, 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. - softcap: float. Anything > 0 activates softcapping 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). - 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 FlashAttnKVPackedFunc.apply( - q, - kv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_func( - q, - k, - v, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - 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). - 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 FlashAttnFunc.apply( - q, - k, - v, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_qkvpacked_func( - qkv, - cu_seqlens, - max_seqlen, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - 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_varlen_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_varlen_kvpacked_func and flash_attn_varlen_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: (total, 3, nheads, headdim), where total = total number of tokens in the batch. - cu_seqlens: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into qkv. - max_seqlen: int. Maximum sequence length in the batch. - 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenQKVPackedFunc.apply( - qkv, - cu_seqlens, - max_seqlen, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_kvpacked_func( - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, -): - """dropout_p should be set to 0.0 during evaluation - If 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 K, V. - 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: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch. - kv: (total_k, 2, nheads_k, headdim), where total_k = total number of key tokens in the batch. - cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into q. - cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into kv. - max_seqlen_q: int. Maximum query sequence length in the batch. - max_seqlen_k: int. Maximum key sequence length in the batch. - 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 + 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenKVPackedFunc.apply( - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_func( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, - block_table=None, -): - """dropout_p should be set to 0.0 during evaluation - Supports multi-query and grouped-query attention (MQA/GQA) by passing in K, V 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: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch. - k: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. - v: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. - cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into q. - cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into kv. - max_seqlen_q: int. Maximum query sequence length in the batch. - max_seqlen_k: int. Maximum key sequence length in the batch. - 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 + 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenFunc.apply( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - block_table, - False if q.device.type == "cpu" else torch.is_grad_enabled(), - ) - - -def flash_attn_with_kvcache( - q, - k_cache, - v_cache, - k=None, - v=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, - block_table: Optional[torch.Tensor] = None, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - rotary_interleaved=True, - alibi_slopes=None, - num_splits=0, - 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 block_table, - or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_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) if there's no block_table, - or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_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). Similar to k. - 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. - block_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). - 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. - 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" - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-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) - cache_batch_idx = maybe_contiguous(cache_batch_idx) - block_table = maybe_contiguous(block_table) - out, softmax_lse = flash_attn.fwd_kvcache( - q, - k_cache, - v_cache, - k, - v, - cache_seqlens, - rotary_cos, - rotary_sin, - cache_batch_idx, - cache_leftpad, - block_table, - alibi_slopes, - None, - softmax_scale, - causal, - window_size[0], - window_size[1], - softcap, - rotary_interleaved, - num_splits, - ) - return (out, softmax_lse) if return_softmax_lse else out diff --git a/build/torch211-cxx11-cpu-x86_64-linux/metadata.json b/build/torch211-cxx11-cpu-x86_64-linux/metadata.json deleted file mode 100644 index 7b0e048efb86b18a5f62e93d0c49d1cc2bb76c4f..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cpu-x86_64-linux/metadata.json +++ /dev/null @@ -1,36 +0,0 @@ -{ - "name": "flash-attn2", - "id": "_flash_attn2_cpu_f12afc9", - "version": 1, - "license": "BSD-3-Clause", - "python-depends": [], - "backend": { - "type": "cpu" - }, - "digest": { - "algorithm": "sha256", - "files": { - "__init__.py": "Cd9MEdd/WUYnBmfiQpC4SyPqAxrTilTjvsvswNju4zc=", - "_flash_attn2_cpu_f12afc9.abi3.so": "pfo22E2/C5x0lSnjlUfp/Bf7solkDnFSgcC+t9guASw=", - "_ops.py": "Ru9fdaV75iST1v660HbpH/kA76qidjALxUmNhGcmOz0=", - "bert_padding.py": "gF1EmsdJ+HpQ86MRQ4VxDw+Sb/RVISdQALdNnoByHlw=", - "flash_attn2/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY=", - "flash_attn_interface.py": "W7a9bHR2ly6QL3aKqsBqxWA7M7ZWYcZPmtsPJu6rAfY=", - "layers/__init__.py": "47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU=", - "layers/patch_embed.py": "H58CgME/qSOPTZLOG08wFgrQS1j34pvNwMPrkTj3Ek4=", - "layers/rotary.py": "rFIdNrZQFNawwLD6d7t5LE2m+M93wwDZ8EuExw+gfD8=", - "ops/__init__.py": "47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU=", - "ops/activations.py": "t5lzNg1In8LP6bKeTnyeMizwqjv27JGbJ6ylPdGvZYg=", - "ops/fused_dense.py": "BqT8VTSCGdLJAjh6WBMkCRFYrdVodPeoElWBvH6AIzA=", - "ops/layer_norm.py": "zr7NXIm+2mtEynTp1CS0fbFGI2Mqdp41dY4AfDWF6EQ=", - "ops/rms_norm.py": "XEnihcj0a4aSz4LO55m5iKGVn4HKTeKN8TIyHjuDgxI=", - "ops/triton/__init__.py": "AbpHGcgLb+kRsJGnwFEktk7uzpZOCcBY74+YBdrKVGs=", - "ops/triton/cross_entropy.py": "y56nmuOtpSNxM7sDZcUSgXKIsIEayfuvrU8bAVog/nw=", - "ops/triton/k_activations.py": "+Z3vIyO4JkqBMipKsPvhzmxljtBdIhJCsl/M+/ESqBo=", - "ops/triton/layer_norm.py": "BBQdJiqopbIS12viUqzx3SvOQCNbDSAYLJXhpyMGyC8=", - "ops/triton/linear.py": "OtRvKz8xdpl+7v3q/ZTaS9fdBt9XrzMyapgRr50uBbM=", - "ops/triton/mlp.py": "uWA+D7HyuRUsOT+4nhi3vvCDNZnNEXwYcHMQ02UREfQ=", - "ops/triton/rotary.py": "0YSrlbep0lJmzLpTiANs0slTFeGcQQqR+el6uzbSLvI=" - } - } -} \ No newline at end of file diff --git a/build/torch211-cxx11-cpu-x86_64-linux/metadata.json.sigstore b/build/torch211-cxx11-cpu-x86_64-linux/metadata.json.sigstore deleted file mode 100644 index 114cee011f503c28d2cd49449e0272e3e4da8dad..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cpu-x86_64-linux/metadata.json.sigstore +++ /dev/null @@ -1 +0,0 @@ -{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json","verificationMaterial":{"certificate":{"rawBytes":"MIIHSTCCBtCgAwIBAgIUGiLMI0Ts4gy4DA6AWgzC3kimZWUwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwNjIzMTU0OTExWhcNMjYwNjIzMTU1OTExWjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEr/VpkuPYvUSaG/Cp6U17dBMq4aqb6rtpu6x4g8t/fN1J3NDEIXCix0R3hYFtPPqJOtN886j6pTnFqHw7odjb4aOCBe8wggXrMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUjJZlZNTQhdc2n8yuL3mIJY4vss0wHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoZjEyYWZjOTJjNDBiYmQ3NmQ5OTBkOTExZWJhYmYyOGU0OTBhN2JkZTATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoZjEyYWZjOTJjNDBiYmQ3NmQ5OTBkOTExZWJhYmYyOGU0OTBhN2JkZTAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoZjEyYWZjOTJjNDBiYmQ3NmQ5OTBkOTExZWJhYmYyOGU0OTBhN2JkZTAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKGYxMmFmYzkyYzQwYmJkNzZkOTkwZDkxMWViYWJmMjhlNDkwYTdiZGUwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMjgwMzYyNjg5MzUvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBiQYKKwYBBAHWeQIEAgR7BHkAdwB1AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABnvUr1osAAAQDAEYwRAIgM2W3kxrFuAobB17k6xeYeOCUUEV90XeJ5+ngmgeZEPACIAVOxfiCM1BgeqcDdzGBULq/IWtQudRy8E8WLdTgaaukMAoGCCqGSM49BAMDA2cAMGQCMBGZ27+Y69UYWyO6InnWDoHGNWeNgHaTyp+wimtERhbKfa8EXPaOc96oO+eY7RISIwIwP0RpaL6kij+Wv+IvZa3/UgYGvAW9mjQ+oWaKeDMmuBPTCQSZ1E2V0rD6y2X3X20j"},"tlogEntries":[{"logIndex":"1928479443","logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="},"kindVersion":{"kind":"hashedrekord","version":"0.0.1"},"integratedTime":"1782229751","inclusionPromise":{"signedEntryTimestamp":"MEQCIBi6fO98lx2YvYzUswcHqz+i9HHINRVwjTO5qgf9ivbcAiARfSFIOAmSacVm4wLssKZutwV1beCEx1IiTxJlhbr1FQ=="},"inclusionProof":{"logIndex":"1806575181","rootHash":"pt4yEy5w4PRpd50tykvO7WJFO1CsOs/PTfwGa6A3+Ns=","treeSize":"1806575188","hashes":["a4l2nIvFffiUbb5BDYaCY70nfjejwZ0Hc0eLLZLSFp4=","bcFMgWwx81q5fEJhgyFkBtPZFhBQkEoy9f0Mn3wj6Ng=","SlLYBytKXQtJ8ptdcSen7uAhNnGKmmyY/Zb1XTo2BSE=","TdY8I++ImZxbubLWS42twZR5/vsytPo0r1b0ykul7rc=","KOZkub6LCMNB6pWG1D9P15sOmzpCf5n1qwc365lHTjs=","zruZk0GGXsDZoTdXQuAXGby4NXFY3Xu2Sij9RLQjfgY=","V5gig8R2JeCm7rMID/aKGPW9/XTTywm3LyvVmOXyKWU=","FzVnFZVWkh/v9DZ4ulW796qT3F1KfQgNLEMT9xdl+5A=","BeJEZvGhIqI3T3bd4SKopQeAtyUUubTj1394a27605U=","TSzTqGkRADCpuZ5PwHaPV9gpzMTICOfnAYVjz2hJjQY=","sYnu6tOuM31ewt4jLBjWIwrC5Sd9suVDOXASLgbkP4M=","Y4xfXQrKMXTYkmle88lw9ofERyxOD79P+2S/XJkZfJ8=","V9c1wGJ9L8DpcbTK6ljbIoNaakhzwNhY/80E6RI15pM=","OMfDh4k9YD0E+t4sjqJsh+vg8D1ueo1cC07exyxmr6s=","+vHKm4wNwGJ46ceyZ8mePD7Vn1SediStNHCStRoTSWQ=","5HwWjuGNDisetbiWM0cWIXM65TIM3EJlqk13ntCQlcU=","NDIRlqxO0IaGr94q86MIU7VCAN3yEEgOHS4loya5lu4=","daxmZaajRpZV+JxHiOYZhJBiSKN5ucqjh2WnGbHhirw=","DOCeoSMovIvLExkhIvisow9AuNXgeWs4ECkyR6EcqYU="],"checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n1806575188\npt4yEy5w4PRpd50tykvO7WJFO1CsOs/PTfwGa6A3+Ns=\n\n— rekor.sigstore.dev wNI9ajBEAiB3dc+EqALJfiR7cVtkJ5WiGKid8S0Hqp5hi/+y1ZCuqwIgYloIyX3AFkoiiPAPYss8+NLpnCw1+v+fx4dmpRr+Vg4=\n"}},"canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiJkZjdlZDQyNzgzMmRjMWI1YThjMzY4YjQyMTcyZTIwMDYyMDNhZTgwY2JmMGMwYmE3MWFkZmMyY2Q3YmQyYWI2In19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJUUNQNWE0Rld0enlubk9ocG5lTTJnd1JOWDRWRnpiUWpsVmJzYTJPYXF4QXNRSWdHUVVEc3B6YUpsdHd1c3A0MFJFU25FYXlHdllDTHlsY3F5NGJyVG5relJVPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRWRU5EUW5SRFowRjNTVUpCWjBsVlIybE1UVWt3VkhNMFozazBSRUUyUVZkbmVrTXphMmx0V2xkVmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDVxU1hwTlZGVXdUMVJGZUZkb1kwNU5hbGwzVG1wSmVrMVVWVEZQVkVWNFYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZ5TDFad2EzVlFXWFpWVTJGSEwwTndObFV4TjJSQ1RYRTBZWEZpTm5KMGNIVTJlRFFLWnpoMEwyWk9NVW96VGtSRlNWaERhWGd3VWpOb1dVWjBVRkJ4U2s5MFRqZzRObW8yY0ZSdVJuRklkemR2WkdwaU5HRlBRMEpsT0hkbloxaHlUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlZxU2xwc0NscE9WRkZvWkdNeWJqaDVkVXd6YlVsS1dUUjJjM013ZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOWFha1Y1V1ZkYWFrOVVTbXBPUkVKcFdXMVJNMDV0VVRWUFZFSnJDazlVUlhoYVYwcG9XVzFaZVU5SFZUQlBWRUpvVGpKS2ExcFVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMXBxUlhsWlYxcHFUMVJLYWs1RVFtbFpiVkV6VG0xUk5VOVVRbXRQVkVWNFdsZEthRmx0V1hrS1QwZFZNRTlVUW1oT01rcHJXbFJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMXBxUlhsWlYxcHFUMVJLYWs1RVFtbFpiVkV6VG0xUk5VOVVRbXNLVDFSRmVGcFhTbWhaYlZsNVQwZFZNRTlVUW1oT01rcHJXbFJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwZFplRTF0Um0wS1dYcHJlVmw2VVhkWmJVcHJUbnBhYTA5VWEzZGFSR3Q0VFZkV2FWbFhTbTFOYW1oc1RrUnJkMWxVWkdsYVIxVjNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOYW1kM1RYcFplVTVxWnpWTmVsVjJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwVVZsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTTjBKSWEwRUtaSGRDTVVGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtNTJWWEl4YjNOQlFVRlJSQXBCUlZsM1VrRkpaMDB5VnpOcmVISkdkVUZ2WWtJeE4yczJlR1ZaWlU5RFZWVkZWamt3V0dWS05TdHVaMjFuWlZwRlVFRkRTVUZXVDNobWFVTk5NVUpuQ21WeFkwUmtla2RDVlV4eEwwbFhkRkYxWkZKNU9FVTRWMHhrVkdkaFlYVnJUVUZ2UjBORGNVZFRUVFE1UWtGTlJFRXlZMEZOUjFGRFRVSkhXakkzSzFrS05qbFZXVmQ1VHpaSmJtNVhSRzlJUjA1WFpVNW5TR0ZVZVhBcmQybHRkRVZTYUdKTFptRTRSVmhRWVU5ak9UWnZUeXRsV1RkU1NWTkpkMGwzVURCU2NBcGhURFpyYVdvclYzWXJTWFphWVRNdlZXZFpSM1pCVnpsdGFsRXJiMWRoUzJWRVRXMTFRbEJVUTFGVFdqRkZNbFl3Y2tRMmVUSllNMWd5TUdvS0xTMHRMUzFGVGtRZ1EwVlNWRWxHU1VOQlZFVXRMUzB0TFFvPSJ9fX19"}],"timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyzADAgEAMIICwgYJKoZIhvcNAQcCoIICszCCAq8CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQggGbM3/vDImGGEkyzvOfIKrWFEzxfWoWYmfUcUQOrlwwCFQDKjKtKOz9PN12X9Nns+S0EDR6F4BgPMjAyNjA2MjMxNTQ5MTFaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHcMIIB2AIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MDYyMzE1NDkxMVowLwYJKoZIhvcNAQkEMSIEIN8E4/MZ/6RirDTB2GTaxGkkoW71/HlatUK8blxrscCnMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRoMGYCMQCCQz5p7gHsi8NxSax4mDApMg/MoerlA6+IPMuONCeMdKfjRnRBMVYxoxxVI4IfCXcCMQDt72CaaymgSOLmejC8gu9tH8s7y4VF9CVrDUQmDuvu1DEoROrBzXSyKGblUYqGMuY="}]}},"messageSignature":{"messageDigest":{"algorithm":"SHA2_256","digest":"337UJ4MtwbWow2i0IXLiAGIDroDL8MC6ca38LNe9KrY="},"signature":"MEUCIQCP5a4FWtzynnOhpneM2gwRNX4VFzbQjlVbsa2OaqxAsQIgGQUDspzaJltwusp40RESnEayGvYCLylcqy4brTnkzRU="}} \ No newline at end of file diff --git a/build/torch211-cxx11-cpu-x86_64-linux/ops/triton/layer_norm.py b/build/torch211-cxx11-cpu-x86_64-linux/ops/triton/layer_norm.py deleted file mode 100644 index 5e6f0b3697b612e6a19031bbc55ac76ad233dc85..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cpu-x86_64-linux/ops/triton/layer_norm.py +++ /dev/null @@ -1,1254 +0,0 @@ -# Copyright (c) 2024, Tri Dao. -# Implement dropout + residual + layer_norm / rms_norm. - -# Based on the Triton LayerNorm tutorial: https://triton-lang.org/main/getting-started/tutorials/05-layer-norm.html -# For the backward pass, we keep weight_grad and bias_grad in registers and accumulate. -# This is faster for dimensions up to 8k, but after that it's much slower due to register spilling. -# The models we train have hidden dim up to 8k anyway (e.g. Llama 70B), so this is fine. - -import math -from typing import Optional, List - -import torch -import torch.nn.functional as F -from torch import Tensor - -import triton -import triton.language as tl - -from flash_attn.utils.torch import custom_fwd, custom_bwd -from flash_attn.utils.library import triton_op - -from ..._ops import add_op_namespace_prefix - - -def maybe_contiguous_lastdim(x): - return x.contiguous() if x is not None and x.stride(-1) != 1 else x - - -def maybe_contiguous(x): - return x.contiguous() if x is not None else None - - -def triton_autotune_configs(): - # Return configs with a valid warp count for the current device - configs = [] - # Maximum threads per block is architecture-dependent in theory, but in reality all are 1024 - max_threads_per_block = 1024 - # Default to warp size 32 if not defined by device - warp_size = getattr(torch.cuda.get_device_properties(torch.cuda.current_device()), "warp_size", 32) - # Autotune for warp counts which are powers of 2 and do not exceed thread per block limit - return [triton.Config({}, num_warps=warp_count) for warp_count in [1, 2, 4, 8, 16, 32] - if warp_count * warp_size <= max_threads_per_block] - # return [triton.Config({}, num_warps=8)] - - -def layer_norm_ref( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - zero_centered_weight=False, - dropout_mask=None, - dropout_mask1=None, - upcast=False, -): - dtype = x.dtype - if upcast: - x = x.float() - weight = weight.float() - bias = bias.float() if bias is not None else None - residual = residual.float() if residual is not None else residual - x1 = x1.float() if x1 is not None else None - weight1 = weight1.float() if weight1 is not None else None - bias1 = bias1.float() if bias1 is not None else None - if zero_centered_weight: - weight = weight + 1.0 - if weight1 is not None: - weight1 = weight1 + 1.0 - if x1 is not None: - assert rowscale is None, "rowscale is not supported with parallel LayerNorm" - if rowscale is not None: - x = x * rowscale[..., None] - if dropout_p > 0.0: - if dropout_mask is not None: - x = x.masked_fill(~dropout_mask, 0.0) / (1.0 - dropout_p) - else: - x = F.dropout(x, p=dropout_p) - if x1 is not None: - if dropout_mask1 is not None: - x1 = x1.masked_fill(~dropout_mask1, 0.0) / (1.0 - dropout_p) - else: - x1 = F.dropout(x1, p=dropout_p) - if x1 is not None: - x = x + x1 - if residual is not None: - x = (x + residual).to(x.dtype) - out = F.layer_norm(x.to(weight.dtype), x.shape[-1:], weight=weight, bias=bias, eps=eps).to( - dtype - ) - if weight1 is None: - return out if not prenorm else (out, x) - else: - out1 = F.layer_norm( - x.to(weight1.dtype), x.shape[-1:], weight=weight1, bias=bias1, eps=eps - ).to(dtype) - return (out, out1) if not prenorm else (out, out1, x) - - -def rms_norm_ref( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - zero_centered_weight=False, - dropout_mask=None, - dropout_mask1=None, - upcast=False, -): - dtype = x.dtype - if upcast: - x = x.float() - weight = weight.float() - bias = bias.float() if bias is not None else None - residual = residual.float() if residual is not None else residual - x1 = x1.float() if x1 is not None else None - weight1 = weight1.float() if weight1 is not None else None - bias1 = bias1.float() if bias1 is not None else None - if zero_centered_weight: - weight = weight + 1.0 - if weight1 is not None: - weight1 = weight1 + 1.0 - if x1 is not None: - assert rowscale is None, "rowscale is not supported with parallel LayerNorm" - if rowscale is not None: - x = x * rowscale[..., None] - if dropout_p > 0.0: - if dropout_mask is not None: - x = x.masked_fill(~dropout_mask, 0.0) / (1.0 - dropout_p) - else: - x = F.dropout(x, p=dropout_p) - if x1 is not None: - if dropout_mask1 is not None: - x1 = x1.masked_fill(~dropout_mask1, 0.0) / (1.0 - dropout_p) - else: - x1 = F.dropout(x1, p=dropout_p) - if x1 is not None: - x = x + x1 - if residual is not None: - x = (x + residual).to(x.dtype) - rstd = 1 / torch.sqrt((x.square()).mean(dim=-1, keepdim=True) + eps) - out = ((x * rstd * weight) + bias if bias is not None else (x * rstd * weight)).to(dtype) - if weight1 is None: - return out if not prenorm else (out, x) - else: - out1 = ((x * rstd * weight1) + bias1 if bias1 is not None else (x * rstd * weight1)).to( - dtype - ) - return (out, out1) if not prenorm else (out, out1, x) - - -@triton.autotune( - configs=triton_autotune_configs(), - key=["N", "HAS_RESIDUAL", "STORE_RESIDUAL_OUT", "IS_RMS_NORM", "HAS_BIAS", "HAS_X1", "HAS_W1", "HAS_B1"], -) -# torch compile doesn't like triton.heuristics, so we set these manually when calling the kernel -# @triton.heuristics({"HAS_BIAS": lambda args: args["B"] is not None}) -# @triton.heuristics({"HAS_RESIDUAL": lambda args: args["RESIDUAL"] is not None}) -# @triton.heuristics({"HAS_X1": lambda args: args["X1"] is not None}) -# @triton.heuristics({"HAS_W1": lambda args: args["W1"] is not None}) -# @triton.heuristics({"HAS_B1": lambda args: args["B1"] is not None}) -@triton.jit -def _layer_norm_fwd_1pass_kernel( - X, # pointer to the input - Y, # pointer to the output - W, # pointer to the weights - B, # pointer to the biases - RESIDUAL, # pointer to the residual - X1, - W1, - B1, - Y1, - RESIDUAL_OUT, # pointer to the residual - ROWSCALE, - SEEDS, # Dropout seeds for each row - DROPOUT_MASK, - DROPOUT_MASK1, - Mean, # pointer to the mean - Rstd, # pointer to the 1/std - stride_x_row, # how much to increase the pointer when moving by 1 row - stride_y_row, - stride_res_row, - stride_res_out_row, - stride_x1_row, - stride_y1_row, - M, # number of rows in X - N, # number of columns in X - eps, # epsilon to avoid division by zero - dropout_p, # Dropout probability - zero_centered_weight, # If true, add 1.0 to the weight - IS_RMS_NORM: tl.constexpr, - BLOCK_N: tl.constexpr, - HAS_RESIDUAL: tl.constexpr, - STORE_RESIDUAL_OUT: tl.constexpr, - HAS_BIAS: tl.constexpr, - HAS_DROPOUT: tl.constexpr, - STORE_DROPOUT_MASK: tl.constexpr, - HAS_ROWSCALE: tl.constexpr, - HAS_X1: tl.constexpr, - HAS_W1: tl.constexpr, - HAS_B1: tl.constexpr, -): - # Map the program id to the row of X and Y it should compute. - row = tl.program_id(0) - X += row * stride_x_row - Y += row * stride_y_row - if HAS_RESIDUAL: - RESIDUAL += row * stride_res_row - if STORE_RESIDUAL_OUT: - RESIDUAL_OUT += row * stride_res_out_row - if HAS_X1: - X1 += row * stride_x1_row - if HAS_W1: - Y1 += row * stride_y1_row - # Compute mean and variance - cols = tl.arange(0, BLOCK_N) - x = tl.load(X + cols, mask=cols < N, other=0.0).to(tl.float32) - if HAS_ROWSCALE: - rowscale = tl.load(ROWSCALE + row).to(tl.float32) - x *= rowscale - if HAS_DROPOUT: - # Compute dropout mask - # 7 rounds is good enough, and reduces register pressure - keep_mask = tl.rand(tl.load(SEEDS + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - x = tl.where(keep_mask, x / (1.0 - dropout_p), 0.0) - if STORE_DROPOUT_MASK: - tl.store(DROPOUT_MASK + row * N + cols, keep_mask, mask=cols < N) - if HAS_X1: - x1 = tl.load(X1 + cols, mask=cols < N, other=0.0).to(tl.float32) - if HAS_ROWSCALE: - rowscale = tl.load(ROWSCALE + M + row).to(tl.float32) - x1 *= rowscale - if HAS_DROPOUT: - # Compute dropout mask - # 7 rounds is good enough, and reduces register pressure - keep_mask = ( - tl.rand(tl.load(SEEDS + M + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - ) - x1 = tl.where(keep_mask, x1 / (1.0 - dropout_p), 0.0) - if STORE_DROPOUT_MASK: - tl.store(DROPOUT_MASK1 + row * N + cols, keep_mask, mask=cols < N) - x += x1 - if HAS_RESIDUAL: - residual = tl.load(RESIDUAL + cols, mask=cols < N, other=0.0).to(tl.float32) - x += residual - if STORE_RESIDUAL_OUT: - tl.store(RESIDUAL_OUT + cols, x, mask=cols < N) - if not IS_RMS_NORM: - mean = tl.sum(x, axis=0) / N - tl.store(Mean + row, mean) - xbar = tl.where(cols < N, x - mean, 0.0) - var = tl.sum(xbar * xbar, axis=0) / N - else: - xbar = tl.where(cols < N, x, 0.0) - var = tl.sum(xbar * xbar, axis=0) / N - rstd = 1 / tl.sqrt(var + eps) - tl.store(Rstd + row, rstd) - # Normalize and apply linear transformation - mask = cols < N - w = tl.load(W + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w += 1.0 - if HAS_BIAS: - b = tl.load(B + cols, mask=mask).to(tl.float32) - x_hat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd - y = x_hat * w + b if HAS_BIAS else x_hat * w - # Write output - tl.store(Y + cols, y, mask=mask) - if HAS_W1: - w1 = tl.load(W1 + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w1 += 1.0 - if HAS_B1: - b1 = tl.load(B1 + cols, mask=mask).to(tl.float32) - y1 = x_hat * w1 + b1 if HAS_B1 else x_hat * w1 - tl.store(Y1 + cols, y1, mask=mask) - - -def _layer_norm_fwd( - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - residual: Optional[Tensor] = None, - x1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - out_dtype: Optional[torch.dtype] = None, - residual_dtype: Optional[torch.dtype] = None, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - return_dropout_mask: bool = False, - out: Optional[Tensor] = None, - residual_out: Optional[Tensor] = None -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - # Need to wrap to handle the case where residual_out is a alias of x, which makes torch.library - # and torch.compile unhappy. Also allocate memory for out and residual_out if they are None - # so that _layer_norm_fwd_impl doesn't have to return them. - if out is None: - out = torch.empty_like(x, dtype=x.dtype if out_dtype is None else out_dtype) - if residual is not None: - residual_dtype = residual.dtype - if residual_out is None and ( - residual is not None - or (residual_dtype is not None and residual_dtype != x.dtype) - or dropout_p > 0.0 - or rowscale is not None - or x1 is not None - ): - residual_out = torch.empty_like( - x, dtype=residual_dtype if residual_dtype is not None else x.dtype - ) - else: - residual_out = None - y1, mean, rstd, seeds, dropout_mask, dropout_mask1 = _layer_norm_fwd_impl( - x, - weight, - bias, - eps, - out, - residual=residual, - x1=x1, - weight1=weight1, - bias1=bias1, - dropout_p=dropout_p, - rowscale=rowscale, - zero_centered_weight=zero_centered_weight, - is_rms_norm=is_rms_norm, - return_dropout_mask=return_dropout_mask, - residual_out=residual_out, - ) - # residual_out is None if residual is None and residual_dtype == input_dtype and dropout_p == 0.0 - if residual_out is None: - residual_out = x - return out, y1, mean, rstd, residual_out, seeds, dropout_mask, dropout_mask1 - - -# [2025-04-28] torch.library.triton_op ignores the schema argument, but here we need the schema -# since we're returning a tuple of tensors -@triton_op(add_op_namespace_prefix("layer_norm_fwd_impl"), mutates_args={"out", "residual_out"}, - schema="(Tensor x, Tensor weight, Tensor bias, float eps, Tensor(a!) out, Tensor? residual, Tensor? x1, Tensor? weight1, Tensor? bias1, float dropout_p, Tensor? rowscale, bool zero_centered_weight, bool is_rms_norm, bool return_dropout_mask, Tensor(a!)? residual_out) -> (Tensor y1, Tensor mean, Tensor rstd, Tensor seeds, Tensor dropout_mask, Tensor dropout_mask1)") -def _layer_norm_fwd_impl( - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - out: Tensor, - residual: Optional[Tensor] = None, - x1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - return_dropout_mask: bool = False, - residual_out: Optional[Tensor] = None -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - M, N = x.shape - assert x.stride(-1) == 1 - if residual is not None: - assert residual.stride(-1) == 1 - assert residual.shape == (M, N) - assert weight.shape == (N,) - assert weight.stride(-1) == 1 - if bias is not None: - assert bias.stride(-1) == 1 - assert bias.shape == (N,) - if x1 is not None: - assert x1.shape == x.shape - assert rowscale is None - assert x1.stride(-1) == 1 - if weight1 is not None: - assert weight1.shape == (N,) - assert weight1.stride(-1) == 1 - if bias1 is not None: - assert bias1.shape == (N,) - assert bias1.stride(-1) == 1 - if rowscale is not None: - assert rowscale.is_contiguous() - assert rowscale.shape == (M,) - assert out.shape == x.shape - assert out.stride(-1) == 1 - if residual_out is not None: - assert residual_out.shape == x.shape - assert residual_out.stride(-1) == 1 - if weight1 is not None: - y1 = torch.empty_like(out) - assert y1.stride(-1) == 1 - else: - y1 = None - mean = torch.empty((M,), dtype=torch.float32, device=x.device) if not is_rms_norm else None - rstd = torch.empty((M,), dtype=torch.float32, device=x.device) - if dropout_p > 0.0: - seeds = torch.randint( - 2**32, (M if x1 is None else 2 * M,), device=x.device, dtype=torch.int64 - ) - else: - seeds = None - if return_dropout_mask and dropout_p > 0.0: - dropout_mask = torch.empty(M, N, device=x.device, dtype=torch.bool) - if x1 is not None: - dropout_mask1 = torch.empty(M, N, device=x.device, dtype=torch.bool) - else: - dropout_mask1 = None - else: - dropout_mask, dropout_mask1 = None, None - # Less than 64KB per feature: enqueue fused kernel - MAX_FUSED_SIZE = 65536 // x.element_size() - BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) - if N > BLOCK_N: - raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") - with torch.cuda.device(x.device.index): - torch.library.wrap_triton(_layer_norm_fwd_1pass_kernel)[(M,)]( - x, - out, - weight, - bias, - residual, - x1, - weight1, - bias1, - y1, - residual_out, - rowscale, - seeds, - dropout_mask, - dropout_mask1, - mean, - rstd, - x.stride(0), - out.stride(0), - residual.stride(0) if residual is not None else 0, - residual_out.stride(0) if residual_out is not None else 0, - x1.stride(0) if x1 is not None else 0, - y1.stride(0) if y1 is not None else 0, - M, - N, - eps, - dropout_p, - # Passing bool make torch inductor very unhappy since it then tries to compare to int_max - int(zero_centered_weight), - is_rms_norm, - BLOCK_N, - residual is not None, - residual_out is not None, - bias is not None, - dropout_p > 0.0, - dropout_mask is not None, - rowscale is not None, - HAS_X1=x1 is not None, - HAS_W1=weight1 is not None, - HAS_B1=bias1 is not None, - ) - return y1, mean, rstd, seeds, dropout_mask, dropout_mask1 - - -@triton.autotune( - configs=triton_autotune_configs(), - key=["N", "HAS_DRESIDUAL", "STORE_DRESIDUAL", "IS_RMS_NORM", "HAS_BIAS", "HAS_DROPOUT"], -) -# torch compile doesn't like triton.heuristics, so we set these manually when calling the kernel -# @triton.heuristics({"HAS_BIAS": lambda args: args["B"] is not None}) -# @triton.heuristics({"HAS_DRESIDUAL": lambda args: args["DRESIDUAL"] is not None}) -# @triton.heuristics({"STORE_DRESIDUAL": lambda args: args["DRESIDUAL_IN"] is not None}) -# @triton.heuristics({"HAS_ROWSCALE": lambda args: args["ROWSCALE"] is not None}) -# @triton.heuristics({"HAS_DY1": lambda args: args["DY1"] is not None}) -# @triton.heuristics({"HAS_DX1": lambda args: args["DX1"] is not None}) -# @triton.heuristics({"HAS_B1": lambda args: args["DB1"] is not None}) -# @triton.heuristics({"RECOMPUTE_OUTPUT": lambda args: args["Y"] is not None}) -@triton.jit -def _layer_norm_bwd_kernel( - X, # pointer to the input - W, # pointer to the weights - B, # pointer to the biases - Y, # pointer to the output to be recomputed - DY, # pointer to the output gradient - DX, # pointer to the input gradient - DW, # pointer to the partial sum of weights gradient - DB, # pointer to the partial sum of biases gradient - DRESIDUAL, - W1, - DY1, - DX1, - DW1, - DB1, - DRESIDUAL_IN, - ROWSCALE, - SEEDS, - Mean, # pointer to the mean - Rstd, # pointer to the 1/std - stride_x_row, # how much to increase the pointer when moving by 1 row - stride_y_row, - stride_dy_row, - stride_dx_row, - stride_dres_row, - stride_dy1_row, - stride_dx1_row, - stride_dres_in_row, - M, # number of rows in X - N, # number of columns in X - eps, # epsilon to avoid division by zero - dropout_p, - zero_centered_weight, - rows_per_program, - IS_RMS_NORM: tl.constexpr, - BLOCK_N: tl.constexpr, - HAS_DRESIDUAL: tl.constexpr, - STORE_DRESIDUAL: tl.constexpr, - HAS_BIAS: tl.constexpr, - HAS_DROPOUT: tl.constexpr, - HAS_ROWSCALE: tl.constexpr, - HAS_DY1: tl.constexpr, - HAS_DX1: tl.constexpr, - HAS_B1: tl.constexpr, - RECOMPUTE_OUTPUT: tl.constexpr, -): - # Map the program id to the elements of X, DX, and DY it should compute. - row_block_id = tl.program_id(0) - row_start = row_block_id * rows_per_program - # Do not early exit if row_start >= M, because we need to write DW and DB - cols = tl.arange(0, BLOCK_N) - mask = cols < N - X += row_start * stride_x_row - if HAS_DRESIDUAL: - DRESIDUAL += row_start * stride_dres_row - if STORE_DRESIDUAL: - DRESIDUAL_IN += row_start * stride_dres_in_row - DY += row_start * stride_dy_row - DX += row_start * stride_dx_row - if HAS_DY1: - DY1 += row_start * stride_dy1_row - if HAS_DX1: - DX1 += row_start * stride_dx1_row - if RECOMPUTE_OUTPUT: - Y += row_start * stride_y_row - w = tl.load(W + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w += 1.0 - if RECOMPUTE_OUTPUT and HAS_BIAS: - b = tl.load(B + cols, mask=mask, other=0.0).to(tl.float32) - if HAS_DY1: - w1 = tl.load(W1 + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w1 += 1.0 - dw = tl.zeros((BLOCK_N,), dtype=tl.float32) - if HAS_BIAS: - db = tl.zeros((BLOCK_N,), dtype=tl.float32) - if HAS_DY1: - dw1 = tl.zeros((BLOCK_N,), dtype=tl.float32) - if HAS_B1: - db1 = tl.zeros((BLOCK_N,), dtype=tl.float32) - row_end = min((row_block_id + 1) * rows_per_program, M) - for row in range(row_start, row_end): - # Load data to SRAM - x = tl.load(X + cols, mask=mask, other=0).to(tl.float32) - dy = tl.load(DY + cols, mask=mask, other=0).to(tl.float32) - if HAS_DY1: - dy1 = tl.load(DY1 + cols, mask=mask, other=0).to(tl.float32) - if not IS_RMS_NORM: - mean = tl.load(Mean + row) - rstd = tl.load(Rstd + row) - # Compute dx - xhat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd - xhat = tl.where(mask, xhat, 0.0) - if RECOMPUTE_OUTPUT: - y = xhat * w + b if HAS_BIAS else xhat * w - tl.store(Y + cols, y, mask=mask) - wdy = w * dy - dw += dy * xhat - if HAS_BIAS: - db += dy - if HAS_DY1: - wdy += w1 * dy1 - dw1 += dy1 * xhat - if HAS_B1: - db1 += dy1 - if not IS_RMS_NORM: - c1 = tl.sum(xhat * wdy, axis=0) / N - c2 = tl.sum(wdy, axis=0) / N - dx = (wdy - (xhat * c1 + c2)) * rstd - else: - c1 = tl.sum(xhat * wdy, axis=0) / N - dx = (wdy - xhat * c1) * rstd - if HAS_DRESIDUAL: - dres = tl.load(DRESIDUAL + cols, mask=mask, other=0).to(tl.float32) - dx += dres - # Write dx - if STORE_DRESIDUAL: - tl.store(DRESIDUAL_IN + cols, dx, mask=mask) - if HAS_DX1: - if HAS_DROPOUT: - keep_mask = ( - tl.rand(tl.load(SEEDS + M + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - ) - dx1 = tl.where(keep_mask, dx / (1.0 - dropout_p), 0.0) - else: - dx1 = dx - tl.store(DX1 + cols, dx1, mask=mask) - if HAS_DROPOUT: - keep_mask = tl.rand(tl.load(SEEDS + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - dx = tl.where(keep_mask, dx / (1.0 - dropout_p), 0.0) - if HAS_ROWSCALE: - rowscale = tl.load(ROWSCALE + row).to(tl.float32) - dx *= rowscale - tl.store(DX + cols, dx, mask=mask) - - X += stride_x_row - if HAS_DRESIDUAL: - DRESIDUAL += stride_dres_row - if STORE_DRESIDUAL: - DRESIDUAL_IN += stride_dres_in_row - if RECOMPUTE_OUTPUT: - Y += stride_y_row - DY += stride_dy_row - DX += stride_dx_row - if HAS_DY1: - DY1 += stride_dy1_row - if HAS_DX1: - DX1 += stride_dx1_row - tl.store(DW + row_block_id * N + cols, dw, mask=mask) - if HAS_BIAS: - tl.store(DB + row_block_id * N + cols, db, mask=mask) - if HAS_DY1: - tl.store(DW1 + row_block_id * N + cols, dw1, mask=mask) - if HAS_B1: - tl.store(DB1 + row_block_id * N + cols, db1, mask=mask) - - -def _layer_norm_bwd( - dy: Tensor, - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - mean: Tensor, - rstd: Tensor, - dresidual: Optional[Tensor] = None, - dy1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - seeds: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - has_residual: bool = False, - has_x1: bool = False, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - x_dtype: Optional[torch.dtype] = None, - recompute_output: bool = False, -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - # Need to wrap to handle the case where dresidual_in or dx1 are aliases of x, - # which makes torch.library unhappy - dx, dw, db, dresidual_in, dx1, dw1, db1, y = _layer_norm_bwd_impl( - dy, - x, - weight, - bias, - eps, - mean, - rstd, - dresidual, - dy1, - weight1, - bias1, - seeds, - dropout_p, - rowscale, - has_residual, - has_x1, - zero_centered_weight, - is_rms_norm, - x_dtype=x_dtype, - recompute_output=recompute_output, - ) - # Don't need to compute dresidual_in separately in this case - if has_residual and dx.dtype == x.dtype and dropout_p == 0.0 and rowscale is None: - dresidual_in = dx - if has_x1 and dropout_p == 0.0: - dx1 = dx - return dx, dw, db, dresidual_in, dx1, dw1, db1, y - - - -@triton_op(add_op_namespace_prefix("layer_norm_bwd_impl"), mutates_args={}, - schema="(Tensor dy, Tensor x, Tensor weight, Tensor bias, float eps, Tensor mean, Tensor rstd, Tensor? dresidual, Tensor? dy1, Tensor? weight1, Tensor? bias1, Tensor? seeds, float dropout_p, Tensor? rowscale, bool has_residual, bool has_x1, bool zero_centered_weight, bool is_rms_norm, ScalarType? x_dtype, bool recompute_output) -> (Tensor dx, Tensor dw, Tensor db, Tensor dresidual_in, Tensor dx1, Tensor dw1, Tensor db1, Tensor y)", - allow_decomposition=False, # Don't let torch.compile trace inside - ) -def _layer_norm_bwd_impl( - dy: Tensor, - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - mean: Tensor, - rstd: Tensor, - dresidual: Optional[Tensor] = None, - dy1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - seeds: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - has_residual: bool = False, - has_x1: bool = False, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - x_dtype: Optional[torch.dtype] = None, - recompute_output: bool = False, -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - M, N = x.shape - assert x.stride(-1) == 1 - dy = maybe_contiguous_lastdim(dy) - assert dy.stride(-1) == 1 - assert dy.shape == (M, N) - if dresidual is not None: - dresidual = maybe_contiguous_lastdim(dresidual) - assert dresidual.stride(-1) == 1 - assert dresidual.shape == (M, N) - assert weight.shape == (N,) - assert weight.stride(-1) == 1 - if bias is not None: - assert bias.stride(-1) == 1 - assert bias.shape == (N,) - if dy1 is not None: - dy1 = maybe_contiguous_lastdim(dy1) - assert weight1 is not None - assert dy1.shape == dy.shape - assert dy1.stride(-1) == 1 - if weight1 is not None: - assert weight1.shape == (N,) - assert weight1.stride(-1) == 1 - if bias1 is not None: - assert bias1.shape == (N,) - assert bias1.stride(-1) == 1 - if seeds is not None: - assert seeds.is_contiguous() - assert seeds.shape == (M if not has_x1 else M * 2,) - if rowscale is not None: - assert rowscale.is_contiguous() - assert rowscale.shape == (M,) - # allocate output - dx = ( - torch.empty_like(x) - if x_dtype is None - else torch.empty(M, N, dtype=x_dtype, device=x.device) - ) - dresidual_in = ( - torch.empty_like(x) - if has_residual - and (dx.dtype != x.dtype or dropout_p > 0.0 or rowscale is not None or has_x1) - else None - ) - dx1 = torch.empty_like(dx) if (has_x1 and dropout_p > 0.0) else None - y = torch.empty(M, N, dtype=dy.dtype, device=dy.device) if recompute_output else None - if recompute_output: - assert weight1 is None, "recompute_output is not supported with parallel LayerNorm" - - # Less than 64KB per feature: enqueue fused kernel - MAX_FUSED_SIZE = 65536 // x.element_size() - BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) - if N > BLOCK_N: - raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") - # Increasing the multiple (e.g. 8) will allow more thread blocks to be launched and hide the - # latency of the gmem reads/writes, but will increase the time of summing up dw / db. - sm_count = torch.cuda.get_device_properties(x.device).multi_processor_count * 8 - _dw = torch.empty((sm_count, N), dtype=torch.float32, device=weight.device) - _db = ( - torch.empty((sm_count, N), dtype=torch.float32, device=bias.device) - if bias is not None - else None - ) - _dw1 = torch.empty_like(_dw) if weight1 is not None else None - _db1 = torch.empty_like(_db) if bias1 is not None else None - rows_per_program = math.ceil(M / sm_count) - grid = (sm_count,) - with torch.cuda.device(x.device.index): - torch.library.wrap_triton(_layer_norm_bwd_kernel)[grid]( - x, - weight, - bias, - y, - dy, - dx, - _dw, - _db, - dresidual, - weight1, - dy1, - dx1, - _dw1, - _db1, - dresidual_in, - rowscale, - seeds, - mean, - rstd, - x.stride(0), - 0 if not recompute_output else y.stride(0), - dy.stride(0), - dx.stride(0), - dresidual.stride(0) if dresidual is not None else 0, - dy1.stride(0) if dy1 is not None else 0, - dx1.stride(0) if dx1 is not None else 0, - dresidual_in.stride(0) if dresidual_in is not None else 0, - M, - N, - eps, - dropout_p, - # Passing bool make torch inductor very unhappy since it then tries to compare to int_max - int(zero_centered_weight), - rows_per_program, - is_rms_norm, - BLOCK_N, - dresidual is not None, - dresidual_in is not None, - bias is not None, - dropout_p > 0.0, - HAS_ROWSCALE=rowscale is not None, - HAS_DY1=dy1 is not None, - HAS_DX1=dx1 is not None, - HAS_B1=bias1 is not None, - RECOMPUTE_OUTPUT=y is not None, - ) - dw = _dw.sum(0).to(weight.dtype) - db = _db.sum(0).to(bias.dtype) if bias is not None else None - dw1 = _dw1.sum(0).to(weight1.dtype) if weight1 is not None else None - db1 = _db1.sum(0).to(bias1.dtype) if bias1 is not None else None - # dresidual_in and dx1 could be None, the wrapper will handle assigning them from dx - return dx, dw, db, dresidual_in, dx1, dw1, db1, y - - -class LayerNormFn(torch.autograd.Function): - - @staticmethod - def forward( - ctx, - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - residual_in_fp32=False, - zero_centered_weight=False, - is_rms_norm=False, - return_dropout_mask=False, - out_dtype=None, - out=None, - residual_out=None - ): - x_shape_og = x.shape - # reshape input data into 2D tensor - x = maybe_contiguous_lastdim(x.reshape(-1, x.shape[-1])) - if residual is not None: - assert residual.shape == x_shape_og - residual = maybe_contiguous_lastdim(residual.reshape(-1, residual.shape[-1])) - if x1 is not None: - assert x1.shape == x_shape_og - assert rowscale is None, "rowscale is not supported with parallel LayerNorm" - x1 = maybe_contiguous_lastdim(x1.reshape(-1, x1.shape[-1])) - weight = weight.contiguous() - bias = maybe_contiguous(bias) - weight1 = maybe_contiguous(weight1) - bias1 = maybe_contiguous(bias1) - if rowscale is not None: - rowscale = rowscale.reshape(-1).contiguous() - residual_dtype = ( - residual.dtype - if residual is not None - else (torch.float32 if residual_in_fp32 else None) - ) - if out is not None: - out = out.reshape(-1, out.shape[-1]) - if residual_out is not None: - residual_out = residual_out.reshape(-1, residual_out.shape[-1]) - y, y1, mean, rstd, residual_out, seeds, dropout_mask, dropout_mask1 = _layer_norm_fwd( - x, - weight, - bias, - eps, - residual, - x1, - weight1, - bias1, - dropout_p=dropout_p, - rowscale=rowscale, - out_dtype=out_dtype, - residual_dtype=residual_dtype, - zero_centered_weight=zero_centered_weight, - is_rms_norm=is_rms_norm, - return_dropout_mask=return_dropout_mask, - out=out, - residual_out=residual_out, - ) - ctx.save_for_backward( - residual_out, weight, bias, weight1, bias1, rowscale, seeds, mean, rstd - ) - ctx.x_shape_og = x_shape_og - ctx.eps = eps - ctx.dropout_p = dropout_p - ctx.is_rms_norm = is_rms_norm - ctx.has_residual = residual is not None - ctx.has_x1 = x1 is not None - ctx.prenorm = prenorm - ctx.x_dtype = x.dtype - ctx.zero_centered_weight = zero_centered_weight - y = y.reshape(x_shape_og) - y1 = y1.reshape(x_shape_og) if y1 is not None else None - residual_out = residual_out.reshape(x_shape_og) if residual_out is not None else None - dropout_mask = dropout_mask.reshape(x_shape_og) if dropout_mask is not None else None - dropout_mask1 = dropout_mask1.reshape(x_shape_og) if dropout_mask1 is not None else None - if not return_dropout_mask: - if weight1 is None: - return y if not prenorm else (y, residual_out) - else: - return (y, y1) if not prenorm else (y, y1, residual_out) - else: - if weight1 is None: - return ( - (y, dropout_mask, dropout_mask1) - if not prenorm - else (y, residual_out, dropout_mask, dropout_mask1) - ) - else: - return ( - (y, y1, dropout_mask, dropout_mask1) - if not prenorm - else (y, y1, residual_out, dropout_mask, dropout_mask1) - ) - - @staticmethod - def backward(ctx, dy, *args): - x, weight, bias, weight1, bias1, rowscale, seeds, mean, rstd = ctx.saved_tensors - dy = dy.reshape(-1, dy.shape[-1]) - if weight1 is not None: - dy1, args = args[0], args[1:] - dy1 = dy1.reshape(-1, dy1.shape[-1]) - assert dy1.shape == x.shape - else: - dy1 = None - if ctx.prenorm: - dresidual = args[0] - dresidual = dresidual.reshape(-1, dresidual.shape[-1]) - assert dresidual.shape == x.shape - else: - dresidual = None - dx, dw, db, dresidual_in, dx1, dw1, db1, _ = _layer_norm_bwd( - dy, - x, - weight, - bias, - ctx.eps, - mean, - rstd, - dresidual, - dy1, - weight1, - bias1, - seeds, - ctx.dropout_p, - rowscale, - ctx.has_residual, - ctx.has_x1, - ctx.zero_centered_weight, - ctx.is_rms_norm, - x_dtype=ctx.x_dtype, - recompute_output=False, - ) - return ( - dx.reshape(ctx.x_shape_og), - dw, - db, - dresidual_in.reshape(ctx.x_shape_og) if ctx.has_residual else None, - dx1.reshape(ctx.x_shape_og) if dx1 is not None else None, - dw1, - db1, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) - - -def layer_norm_fn( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - residual_in_fp32=False, - zero_centered_weight=False, - is_rms_norm=False, - return_dropout_mask=False, - out_dtype=None, - out=None, - residual_out=None -): - return LayerNormFn.apply( - x, - weight, - bias, - residual, - x1, - weight1, - bias1, - eps, - dropout_p, - rowscale, - prenorm, - residual_in_fp32, - zero_centered_weight, - is_rms_norm, - return_dropout_mask, - out_dtype, - out, - residual_out - ) - - -def rms_norm_fn( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - residual_in_fp32=False, - zero_centered_weight=False, - return_dropout_mask=False, - out_dtype=None, - out=None, - residual_out=None -): - return LayerNormFn.apply( - x, - weight, - bias, - residual, - x1, - weight1, - bias1, - eps, - dropout_p, - rowscale, - prenorm, - residual_in_fp32, - zero_centered_weight, - True, - return_dropout_mask, - out_dtype, - out, - residual_out - ) - - -class RMSNorm(torch.nn.Module): - - def __init__(self, hidden_size, eps=1e-5, dropout_p=0.0, zero_centered_weight=False, - device=None, dtype=None): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.eps = eps - if dropout_p > 0.0: - self.drop = torch.nn.Dropout(dropout_p) - else: - self.drop = None - self.zero_centered_weight = zero_centered_weight - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.register_parameter("bias", None) - self.reset_parameters() - - def reset_parameters(self): - if not self.zero_centered_weight: - torch.nn.init.ones_(self.weight) - else: - torch.nn.init.zeros_(self.weight) - - def forward(self, x, residual=None, prenorm=False, residual_in_fp32=False): - return rms_norm_fn( - x, - self.weight, - self.bias, - residual=residual, - eps=self.eps, - dropout_p=self.drop.p if self.drop is not None and self.training else 0.0, - prenorm=prenorm, - residual_in_fp32=residual_in_fp32, - zero_centered_weight=self.zero_centered_weight, - ) - - -class LayerNormLinearFn(torch.autograd.Function): - - @staticmethod - @custom_fwd - def forward( - ctx, - x, - norm_weight, - norm_bias, - linear_weight, - linear_bias, - residual=None, - eps=1e-6, - prenorm=False, - residual_in_fp32=False, - is_rms_norm=False, - ): - x_shape_og = x.shape - # reshape input data into 2D tensor - x = maybe_contiguous_lastdim(x.reshape(-1, x.shape[-1])) - if residual is not None: - assert residual.shape == x_shape_og - residual = maybe_contiguous_lastdim(residual.reshape(-1, residual.shape[-1])) - norm_weight = norm_weight.contiguous() - norm_bias = maybe_contiguous(norm_bias) - residual_dtype = ( - residual.dtype - if residual is not None - else (torch.float32 if residual_in_fp32 else None) - ) - y, _, mean, rstd, residual_out, *rest = _layer_norm_fwd( - x, - norm_weight, - norm_bias, - eps, - residual, - out_dtype=None if not torch.is_autocast_enabled() else torch.get_autocast_dtype("cuda"), - residual_dtype=residual_dtype, - is_rms_norm=is_rms_norm, - ) - y = y.reshape(x_shape_og) - dtype = torch.get_autocast_dtype("cuda") if torch.is_autocast_enabled() else y.dtype - linear_weight = linear_weight.to(dtype) - linear_bias = linear_bias.to(dtype) if linear_bias is not None else None - out = F.linear(y.to(linear_weight.dtype), linear_weight, linear_bias) - # We don't store y, will be recomputed in the backward pass to save memory - ctx.save_for_backward(residual_out, norm_weight, norm_bias, linear_weight, mean, rstd) - ctx.x_shape_og = x_shape_og - ctx.eps = eps - ctx.is_rms_norm = is_rms_norm - ctx.has_residual = residual is not None - ctx.prenorm = prenorm - ctx.x_dtype = x.dtype - ctx.linear_bias_is_none = linear_bias is None - return out if not prenorm else (out, residual_out.reshape(x_shape_og)) - - @staticmethod - @custom_bwd - def backward(ctx, dout, *args): - x, norm_weight, norm_bias, linear_weight, mean, rstd = ctx.saved_tensors - dout = dout.reshape(-1, dout.shape[-1]) - dy = F.linear(dout, linear_weight.t()) - dlinear_bias = None if ctx.linear_bias_is_none else dout.sum(0) - dy = maybe_contiguous_lastdim(dy) - assert dy.shape == x.shape - if ctx.prenorm: - dresidual = args[0] - dresidual = maybe_contiguous_lastdim(dresidual.reshape(-1, dresidual.shape[-1])) - assert dresidual.shape == x.shape - else: - dresidual = None - dx, dnorm_weight, dnorm_bias, dresidual_in, _, _, _, y = _layer_norm_bwd( - dy, - x, - norm_weight, - norm_bias, - ctx.eps, - mean, - rstd, - dresidual=dresidual, - has_residual=ctx.has_residual, - is_rms_norm=ctx.is_rms_norm, - x_dtype=ctx.x_dtype, - recompute_output=True, - ) - dlinear_weight = torch.einsum("bo,bi->oi", dout, y) - return ( - dx.reshape(ctx.x_shape_og), - dnorm_weight, - dnorm_bias, - dlinear_weight, - dlinear_bias, - dresidual_in.reshape(ctx.x_shape_og) if ctx.has_residual else None, - None, - None, - None, - None, - ) - - -def layer_norm_linear_fn( - x, - norm_weight, - norm_bias, - linear_weight, - linear_bias, - residual=None, - eps=1e-6, - prenorm=False, - residual_in_fp32=False, - is_rms_norm=False, -): - return LayerNormLinearFn.apply( - x, - norm_weight, - norm_bias, - linear_weight, - linear_bias, - residual, - eps, - prenorm, - residual_in_fp32, - is_rms_norm, - ) diff --git a/build/torch211-cxx11-cu126-aarch64-linux/__init__.py b/build/torch211-cxx11-cu126-aarch64-linux/__init__.py deleted file mode 100644 index ecc2f9d896b6c93f90b0a1499856dc0612177422..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu126-aarch64-linux/__init__.py +++ /dev/null @@ -1,393 +0,0 @@ -from typing import Optional, List -import torch -from ._ops import ops as flash_attn_ops -from .flash_attn_interface import ( - flash_attn_func, - flash_attn_kvpacked_func, - flash_attn_qkvpacked_func, - flash_attn_varlen_func, - flash_attn_varlen_kvpacked_func, - flash_attn_varlen_qkvpacked_func, - flash_attn_with_kvcache, -) - - -def fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - return_softmax: bool = False, - gen: Optional[torch.Generator] = None, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention. - - Args: - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Optional output tensor, same shape as q - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - return_softmax: Whether to return softmax weights - gen: Optional random number generator - - Returns: - List of tensors: [output, softmax_lse, (softmax if return_softmax)] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.fwd( - q, - k, - v, - out, - alibi_slopes, - p_dropout, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - gen, - ) - - -def varlen_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - out: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - block_table: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - max_seqlen_q: int = 0, - max_seqlen_k: int = 0, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - zero_tensors: bool = False, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - return_softmax: bool = False, - gen: Optional[torch.Generator] = None, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention with variable sequence lengths. - - Args: - q: Query tensor of shape [total_q, num_heads, head_size] - k: Key tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - v: Value tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1] - cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1] - out: Optional output tensor of shape [total_q, num_heads, head_size] - seqused_k: Optional tensor specifying how many keys to use per batch element [batch_size] - leftpad_k: Optional left padding for keys of shape [batch_size] - block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq] - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - max_seqlen_q: Maximum sequence length for queries - max_seqlen_k: Maximum sequence length for keys - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - zero_tensors: Whether to zero tensors before computation - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - return_softmax: Whether to return softmax weights - gen: Optional random number generator - - Returns: - List of tensors: [output, softmax_lse, (softmax if return_softmax)] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.varlen_fwd( - q, - k, - v, - out, - cu_seqlens_q, - cu_seqlens_k, - seqused_k, - leftpad_k, - block_table, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - p_dropout, - softmax_scale, - zero_tensors, - is_causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - gen, - ) - - -def bwd( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - p_dropout: float = 0.0, - 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, - gen: Optional[torch.Generator] = None, - rng_state: Optional[torch.Tensor] = None, -) -> List[torch.Tensor]: - """ - Backward pass for multi-head attention. - - Args: - dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size] - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size] - softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q] - dq: Optional gradient tensor for queries, same shape as q - dk: Optional gradient tensor for keys, same shape as k - dv: Optional gradient tensor for values, same shape as v - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - deterministic: Whether to use deterministic algorithms - gen: Optional random number generator - rng_state: Optional RNG state from forward pass - - Returns: - List of tensors: [dq, dk, dv] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - alibi_slopes, - p_dropout, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - deterministic, - gen, - rng_state, - ) - - -def varlen_bwd( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - max_seqlen_q: int = 0, - max_seqlen_k: int = 0, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - zero_tensors: bool = False, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - deterministic: bool = False, - gen: Optional[torch.Generator] = None, - rng_state: Optional[torch.Tensor] = None, -) -> List[torch.Tensor]: - """ - Backward pass for multi-head attention with variable sequence lengths. - - Args: - dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size] - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size] - softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q] - cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1] - cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1] - dq: Optional gradient tensor for queries, same shape as q - dk: Optional gradient tensor for keys, same shape as k - dv: Optional gradient tensor for values, same shape as v - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - max_seqlen_q: Maximum sequence length for queries - max_seqlen_k: Maximum sequence length for keys - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - zero_tensors: Whether to zero tensors before computation - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - deterministic: Whether to use deterministic algorithms - gen: Optional random number generator - rng_state: Optional RNG state from forward pass - - Returns: - List of tensors: [dq, dk, dv] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.varlen_bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - p_dropout, - softmax_scale, - zero_tensors, - is_causal, - window_size_left, - window_size_right, - softcap, - deterministic, - gen, - rng_state, - ) - - -def fwd_kvcache( - q: torch.Tensor, - kcache: torch.Tensor, - vcache: torch.Tensor, - k: Optional[torch.Tensor] = None, - v: Optional[torch.Tensor] = None, - seqlens_k: Optional[torch.Tensor] = None, - rotary_cos: Optional[torch.Tensor] = None, - rotary_sin: Optional[torch.Tensor] = None, - cache_batch_idx: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - block_table: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - out: 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, - is_rotary_interleaved: bool = False, - num_splits: int = 1, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention with KV cache. - - Args: - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - kcache: Key cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - vcache: Value cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - k: Optional new keys tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size] - v: Optional new values tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size] - seqlens_k: Optional sequence lengths for keys of shape [batch_size] - rotary_cos: Optional rotary cosine tensor of shape [seqlen_ro, rotary_dim/2] - rotary_sin: Optional rotary sine tensor of shape [seqlen_ro, rotary_dim/2] - cache_batch_idx: Optional indices to index into the KV cache - leftpad_k: Optional left padding for keys of shape [batch_size] - block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq] - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - out: Optional output tensor, same shape as q - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - is_rotary_interleaved: Whether rotary embeddings are interleaved - num_splits: Number of splits for computation - - Returns: - List of tensors: [output, softmax_lse] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.fwd_kvcache( - q, - kcache, - vcache, - k, - v, - seqlens_k, - rotary_cos, - rotary_sin, - cache_batch_idx, - leftpad_k, - block_table, - alibi_slopes, - out, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - is_rotary_interleaved, - num_splits, - ) diff --git a/build/torch211-cxx11-cu126-aarch64-linux/_flash_attn2_cuda_f12afc9.abi3.so b/build/torch211-cxx11-cu126-aarch64-linux/_flash_attn2_cuda_f12afc9.abi3.so deleted file mode 100644 index d2c6fa5ae11231ef44b4aea5b18a730b3f145db9..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu126-aarch64-linux/_flash_attn2_cuda_f12afc9.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:d9723af07fc19a183a32c259c203d6062379babf1070527da5674be1660a062d -size 446583112 diff --git a/build/torch211-cxx11-cu126-aarch64-linux/_ops.py b/build/torch211-cxx11-cu126-aarch64-linux/_ops.py deleted file mode 100644 index af5b731cb04b5e6ce54e862db0501f3683c3163f..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu126-aarch64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _flash_attn2_cuda_f12afc9 -ops = torch.ops._flash_attn2_cuda_f12afc9 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_flash_attn2_cuda_f12afc9::{op_name}" diff --git a/build/torch211-cxx11-cu126-aarch64-linux/bert_padding.py b/build/torch211-cxx11-cu126-aarch64-linux/bert_padding.py deleted file mode 100644 index 3c2d35159a014a9d03aabead9e52e009168696ea..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu126-aarch64-linux/bert_padding.py +++ /dev/null @@ -1,218 +0,0 @@ -# Adapted from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/padding.py - -import torch -import torch.nn.functional as F -from einops import rearrange, repeat - - -class IndexFirstAxis(torch.autograd.Function): - @staticmethod - def forward(ctx, input, indices): - ctx.save_for_backward(indices) - assert input.ndim >= 2 - ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] - second_dim = other_shape.numel() - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - # return input[indices] - return torch.gather( - rearrange(input, "b ... -> b (...)"), 0, repeat(indices, "z -> z d", d=second_dim) - ).reshape(-1, *other_shape) - - @staticmethod - def backward(ctx, grad_output): - (indices,) = ctx.saved_tensors - assert grad_output.ndim >= 2 - other_shape = grad_output.shape[1:] - grad_output = rearrange(grad_output, "b ... -> b (...)") - grad_input = torch.zeros( - [ctx.first_axis_dim, grad_output.shape[1]], - device=grad_output.device, - dtype=grad_output.dtype, - ) - # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. - # grad_input[indices] = grad_output - grad_input.scatter_(0, repeat(indices, "z -> z d", d=grad_output.shape[1]), grad_output) - return grad_input.reshape(ctx.first_axis_dim, *other_shape), None - - -index_first_axis = IndexFirstAxis.apply - - -class IndexPutFirstAxis(torch.autograd.Function): - @staticmethod - def forward(ctx, values, indices, first_axis_dim): - ctx.save_for_backward(indices) - assert indices.ndim == 1 - assert values.ndim >= 2 - output = torch.zeros( - first_axis_dim, *values.shape[1:], device=values.device, dtype=values.dtype - ) - # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. - output[indices] = values - # output.scatter_(0, repeat(indices, 'z -> z d', d=values.shape[1]), values) - return output - - @staticmethod - def backward(ctx, grad_output): - (indices,) = ctx.saved_tensors - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - grad_values = grad_output[indices] - # grad_values = torch.gather(grad_output, 0, repeat(indices, 'z -> z d', d=grad_output.shape[1])) - return grad_values, None, None - - -index_put_first_axis = IndexPutFirstAxis.apply - - -class IndexFirstAxisResidual(torch.autograd.Function): - @staticmethod - def forward(ctx, input, indices): - ctx.save_for_backward(indices) - assert input.ndim >= 2 - ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] - second_dim = other_shape.numel() - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - output = input[indices] - # We don't want to reshape input (b ... -> b (...)) since it could change the channel_last - # memory format to channel_first. In other words, input might not be contiguous. - # If we don't detach, Pytorch complains about output being a view and is being modified inplace - return output, input.detach() - - @staticmethod - def backward(ctx, grad_output, grad_residual): - (indices,) = ctx.saved_tensors - assert grad_output.ndim >= 2 - other_shape = grad_output.shape[1:] - assert grad_residual.shape[1:] == other_shape - grad_input = grad_residual - # grad_input[indices] += grad_output - indices = indices.reshape(indices.shape[0], *((1,) * (grad_output.ndim - 1))) - indices = indices.expand_as(grad_output) - grad_input.scatter_add_(0, indices, grad_output) - return grad_input.reshape(ctx.first_axis_dim, *other_shape), None - - -index_first_axis_residual = IndexFirstAxisResidual.apply - - -def unpad_input(hidden_states, attention_mask, unused_mask=None): - """ - Arguments: - hidden_states: (batch, seqlen, ...) - attention_mask: (batch, seqlen), bool / int, 1 means valid and 0 means not valid. - unused_mask: (batch, seqlen), bool / int, 1 means the element is allocated but unused. - Return: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens selected in attention_mask + unused_mask. - indices: (total_nnz), the indices of masked tokens from the flattened input sequence. - cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. - max_seqlen_in_batch: int - seqused: (batch), returns the number of tokens selected in attention_mask + unused_mask. - """ - all_masks = (attention_mask + unused_mask) if unused_mask is not None else attention_mask - seqlens_in_batch = all_masks.sum(dim=-1, dtype=torch.int32) - used_seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) - indices = torch.nonzero(all_masks.flatten(), as_tuple=False).flatten() - max_seqlen_in_batch = seqlens_in_batch.max().item() - cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) - # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the - # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim - # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to - # index with integer indices. Moreover, torch's index is a bit slower than it needs to be, - # so we write custom forward and backward to make it a bit faster. - return ( - index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices), - indices, - cu_seqlens, - max_seqlen_in_batch, - used_seqlens_in_batch, - ) - - -def unpad_input_for_concatenated_sequences(hidden_states, attention_mask_in_length): - """ - Supports concatenating short samples in one sequence. The attention_mask_in_length is utilized to mask other short samples. It helps efficient training of variant lengths-based samples (e.g., the supervised fine-tuning task in large language model). - The motivation for this function is explained [here](https://github.com/Dao-AILab/flash-attention/issues/432#issuecomment-1668822286). - - For example, if batch = 3 and seqlen = 6, the attention_mask_in_length is: - ``` - [ - [2, 3, 0, 0, 0, 0], - [3, 2, 0, 0, 0, 0], - [6, 0, 0, 0, 0, 0] - ] - ``` - , which refers to the 3D-attention mask: - ``` - [ - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [0, 0, 1, 0, 0, 0], - [0, 0, 1, 1, 0, 0], - [0, 0, 1, 1, 1, 0], - [0, 0, 0, 0, 0, 1] - ], - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [0, 0, 0, 1, 0, 0], - [0, 0, 0, 1, 1, 0], - [0, 0, 0, 0, 0, 1] - ], - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [1, 1, 1, 1, 0, 0], - [1, 1, 1, 1, 1, 0], - [1, 1, 1, 1, 1, 1] - ] - ] - ```. - - Arguments: - hidden_states: (batch, seqlen, ...) - attention_mask_in_length: (batch, seqlen), int, a nonzero number (e.g., 1, 2, 3, etc.) means length of concatenated sequence in b-th batch, and 0 means none. - Return: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. - indices: (total_nnz), the indices of non-masked tokens from the flattened input sequence. - cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. - max_seqlen_in_batch: int - """ - length = attention_mask_in_length.sum(dim=-1) - seqlen = attention_mask_in_length.size(-1) - attention_mask_2d = torch.arange(seqlen, device=length.device, dtype=length.dtype).expand(len(length), seqlen) < length.unsqueeze(1) - real_indices_idx = torch.nonzero(attention_mask_in_length.flatten(), as_tuple=False).flatten() - seqlens_in_batch = attention_mask_in_length.flatten()[real_indices_idx] - indices = torch.nonzero(attention_mask_2d.flatten(), as_tuple=False).flatten() - max_seqlen_in_batch = seqlens_in_batch.max().item() - cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) - # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the - # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim - # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to - # index with integer indices. Moreover, torch's index is a bit slower than it needs to be, - # so we write custom forward and backward to make it a bit faster. - return ( - index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices), - indices, - cu_seqlens, - max_seqlen_in_batch, - ) - - -def pad_input(hidden_states, indices, batch, seqlen): - """ - Arguments: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. - indices: (total_nnz), the indices that represent the non-masked tokens of the original padded input sequence. - batch: int, batch size for the padded sequence. - seqlen: int, maximum sequence length for the padded sequence. - Return: - hidden_states: (batch, seqlen, ...) - """ - dim = hidden_states.shape[-1] - # output = torch.zeros((batch * seqlen), dim, device=hidden_states.device, dtype=hidden_states.dtype) - # output[indices] = hidden_states - output = index_put_first_axis(hidden_states, indices, batch * seqlen) - return rearrange(output, "(b s) ... -> b s ...", b=batch) diff --git a/build/torch211-cxx11-cu126-aarch64-linux/flash_attn2/__init__.py b/build/torch211-cxx11-cu126-aarch64-linux/flash_attn2/__init__.py deleted file mode 100644 index a9b2672c1cd85b74c1b3ded0fc0b2100e1aeac23..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu126-aarch64-linux/flash_attn2/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -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-cu126-aarch64-linux/flash_attn_interface.py b/build/torch211-cxx11-cu126-aarch64-linux/flash_attn_interface.py deleted file mode 100644 index b014e569dc9812a9c840d1043f2cebe767b4d477..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu126-aarch64-linux/flash_attn_interface.py +++ /dev/null @@ -1,1631 +0,0 @@ -# Copyright (c) 2023, Tri Dao. - -from typing import Optional, Sequence, Tuple, Union - -import torch -import torch.nn as nn -import os - -# # isort: off -# # We need to import the CUDA kernels after importing torch -# USE_TRITON_ROCM = os.getenv("FLASH_ATTENTION_TRITON_AMD_ENABLE", "FALSE") == "TRUE" -# if USE_TRITON_ROCM: -# from .flash_attn_triton_amd import interface_fa as flash_attn -# else: -# import flash_attn_2_cuda as flash_attn - - -from . import _ops -from ._ops import ops as flash_attn -from ._ops import add_op_namespace_prefix - -# # isort: on - -def maybe_contiguous(x): - return x.contiguous() if x is not None and x.stride(-1) != 1 else x - - -def _get_device(): - if torch.xpu.is_available(): - return "xpu" - elif torch.cuda.is_available(): - return "cuda" - else: - return "cpu" - - -def _get_block_size_n(device, head_dim, is_dropout, is_causal): - assert head_dim <= 256 - - if device.type == "xpu": - if head_dim <= 96: - return 64 - elif head_dim <= 128: - return 32 - elif head_dim <= 256: - return 64 - else: - return 32 - - # This should match the block sizes in the CUDA kernel - major, minor = torch.cuda.get_device_capability(device) - is_sm8x = major == 8 and minor > 0 # Only include sm86 and sm89, exclude sm80 (A100) - is_sm80 = major == 8 and minor == 0 - is_sm90 = major == 9 and minor == 0 - if head_dim <= 32: - return 128 - if head_dim <= 64: - return 128 if not is_dropout else 64 - elif head_dim <= 96: - return 64 - elif head_dim <= 128: - if is_sm8x: - return 64 if (not is_dropout and is_causal) else 32 - else: - return 64 if not is_dropout else 32 - elif head_dim <= 192: - return 64 - elif head_dim <= 224: - return 64 - elif head_dim <= 256: - return 64 - - -def round_multiple(x, m): - return (x + m - 1) // m * m - - -# torch.compile() support is only enabled for pytorch >= 2.4 -# The reason for this is that we are using the new custom_op and register_fake -# APIs, which support inplace modification of inputs in the function itself -if torch.__version__ >= "2.4.0": - _torch_custom_op_wrapper = torch.library.custom_op - _torch_register_fake_wrapper = torch.library.register_fake -else: - def noop_custom_op_wrapper(name, fn=None, /, *, mutates_args, device_types=None, schema=None): - def wrap(func): - return func - if fn is None: - return wrap - return fn - def noop_register_fake_wrapper(op, fn=None, /, *, lib=None, _stacklevel=1): - def wrap(func): - return func - if fn is None: - return wrap - return fn - _torch_custom_op_wrapper = noop_custom_op_wrapper - _torch_register_fake_wrapper = noop_register_fake_wrapper - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_forward"), mutates_args=(), device_types=_get_device()) -def _flash_attn_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - return_softmax: bool -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - out, softmax_lse, S_dmask, rng_state = flash_attn.fwd( - q, - k, - v, - None, - alibi_slopes, - dropout_p, - softmax_scale, - causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - None, - ) - return out, softmax_lse, S_dmask, rng_state - - -@_torch_register_fake_wrapper(add_op_namespace_prefix("_flash_attn_forward")) -def _flash_attn_forward_fake( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - return_softmax: bool -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - batch_size, seqlen_q, num_heads, head_size = q.shape - seqlen_k = k.shape[1] - out = torch.empty_like(q) - softmax_lse = torch.empty((batch_size, num_heads, seqlen_q), dtype=torch.float32, device=q.device, layout=q.layout) - p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) - if return_softmax: - p = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128), round_multiple(seqlen_k, 128)), dtype=q.dtype, device=q.device, layout=q.layout) - rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) - - return out, softmax_lse, p, rng_state - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_forward = _ops.ops._flash_attn_forward -else: - _wrapped_flash_attn_forward = _flash_attn_forward - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_varlen_forward"), mutates_args=(), device_types=_get_device()) -def _flash_attn_varlen_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - alibi_slopes: Optional[torch.Tensor] = None, - return_softmax: bool = False, - block_table: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - out, softmax_lse, S_dmask, rng_state = flash_attn.varlen_fwd( - q, - k, - v, - None, - cu_seqlens_q, - cu_seqlens_k, - seqused_k, - leftpad_k, - block_table, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - zero_tensors, - causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - None, - ) - # if out.isnan().any() or softmax_lse.isnan().any(): - # breakpoint() - return out, softmax_lse, S_dmask, rng_state - - -@_torch_register_fake_wrapper(add_op_namespace_prefix("_flash_attn_varlen_forward")) -def _flash_attn_varlen_forward_fake( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - alibi_slopes: Optional[torch.Tensor] = None, - return_softmax: bool = False, - block_table: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - paged_kv = block_table is not None - batch_size = cu_seqlens_q.numel() - 1 - total_q, num_heads, _ = q.shape - - out = torch.empty_like(q) - softmax_lse = torch.empty((num_heads, total_q), dtype=torch.float32, device=q.device, layout=q.layout) - p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) - seqlen_q_rounded = round_multiple(max_seqlen_q, 128) - seqlen_k_rounded = round_multiple(max_seqlen_k, 128) - if return_softmax: - p = torch.empty((batch_size, num_heads, seqlen_q_rounded, seqlen_k_rounded), dtype=q.dtype, device=q.device, layout=q.layout) - rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) - return out, softmax_lse, p, rng_state - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_varlen_forward = _ops.ops._flash_attn_varlen_forward -else: - _wrapped_flash_attn_varlen_forward = _flash_attn_varlen_forward - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_backward"), mutates_args=("dq", "dk", "dv"), device_types=_get_device()) -def _flash_attn_backward( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, -) -> 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)] - ( - dq, - dk, - dv, - softmax_d, - ) = flash_attn.bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - alibi_slopes, - dropout_p, - softmax_scale, - causal, - window_size_left, - window_size_right, - softcap, - deterministic, - None, - rng_state, - ) - return softmax_d - - -@_torch_register_fake_wrapper(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, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, -) -> torch.Tensor: - dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] - if dq is None: - dq = torch.empty_like(q) - if dk is None: - dk = torch.empty_like(k) - if dv is None: - dv = torch.empty_like(v) - batch_size, seqlen_q, num_heads, _ = q.shape - softmax_d = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128)), device=q.device, dtype=torch.float32) - - return softmax_d - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_backward = _ops.ops._flash_attn_backward -else: - _wrapped_flash_attn_backward = _flash_attn_backward - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_varlen_backward"), mutates_args=("dq", "dk", "dv"), device_types=_get_device()) -def _flash_attn_varlen_backward( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> 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)] - ( - dq, - dk, - dv, - softmax_d, - ) = flash_attn.varlen_bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - zero_tensors, - causal, - window_size_left, - window_size_right, - softcap, - deterministic, - None, - rng_state, - ) - # if dk.isnan().any() or dk.isnan().any() or dv.isnan().any() or softmax_d.isnan().any(): - # breakpoint() - return softmax_d - - -@_torch_register_fake_wrapper(add_op_namespace_prefix("_flash_attn_varlen_backward")) -def _flash_attn_varlen_backward_fake( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> torch.Tensor: - dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] - batch_size = cu_seqlens_q.numel() - 1 - total_q, num_heads, _ = q.shape - - if dq is None: - dq = torch.empty_like(q) - if dk is None: - dk = torch.empty_like(k) - if dv is None: - dv = torch.empty_like(v) - softmax_d = torch.empty((num_heads, total_q + 128 * batch_size), device=q.device, dtype=torch.float32) - - return softmax_d - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_varlen_backward = _ops.ops._flash_attn_varlen_backward -else: - _wrapped_flash_attn_varlen_backward = _flash_attn_varlen_backward - - -class FlashAttnQKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and qkv.requires_grad - if softmax_scale is None: - softmax_scale = qkv.shape[-1] ** (-0.5) - q, k, v = qkv[:, :, 0].detach(), qkv[:, :, 1].detach(), qkv[:, :, 2].detach() - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) - dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dqkv[:, :, 0], - dqkv[:, :, 1], - dqkv[:, :, 2], - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dqkv = dqkv[..., : dout.shape[-1]] # We could have padded the head dimension - return dqkv, None, None, None, None, None, None, None, None, None - - -class FlashAttnVarlenQKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - cu_seqlens, - max_seqlen, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and qkv.requires_grad - if softmax_scale is None: - softmax_scale = qkv.shape[-1] ** (-0.5) - q, k, v = qkv[:, 0].detach(), qkv[:, 1].detach(), qkv[:, 2].detach() - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens, - cu_seqlens, - max_seqlen, - max_seqlen, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=None, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens, rng_state) - ctx.dropout_p = dropout_p - ctx.max_seqlen = max_seqlen - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens, rng_state = ctx.saved_tensors - qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) - dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dqkv[:, 0], - dqkv[:, 1], - dqkv[:, 2], - cu_seqlens, - cu_seqlens, - ctx.max_seqlen, - ctx.max_seqlen, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - 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 FlashAttnKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - kv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, kv] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - k, v = kv[:, :, 0].detach(), kv[:, :, 1].detach() - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - dq = torch.empty_like(q) - kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) - dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dkv[:, :, 0], - dkv[:, :, 1], - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dkv = dkv[..., : dout.shape[-1]] - return dq, dkv, None, None, None, None, None, None, None, None, None - - -class FlashAttnVarlenKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, kv] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - k, v = kv[:, 0].detach(), kv[:, 1].detach() - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=None, - ) - if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state - ) - ctx.dropout_p = dropout_p - 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.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors - dq = torch.empty_like(q) - kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) - dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dkv[:, 0], - dkv[:, 1], - cu_seqlens_q, - cu_seqlens_k, - ctx.max_seqlen_q, - ctx.max_seqlen_k, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dkv = dkv[..., : dout.shape[-1]] - return dq, dkv, None, None, None, None, None, None, None, None, None, None, None, None, None - - -class FlashAttnFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - k, - v, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, k, v] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dk = dk[..., : dout.shape[-1]] - dv = dv[..., : dout.shape[-1]] - return dq, dk, dv, 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, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - block_table, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, k, v] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=block_table, - ) - if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state - ) - ctx.dropout_p = dropout_p - 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.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors - dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - ctx.max_seqlen_q, - ctx.max_seqlen_k, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dk = dk[..., : dout.shape[-1]] - dv = dv[..., : dout.shape[-1]] - return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None - - -def flash_attn_qkvpacked_func( - qkv, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # <=0.0 means deactivate - alibi_slopes=None, - deterministic=False, - 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, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_kvpacked_func( - q, - kv, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, -): - """dropout_p should be set to 0.0 during evaluation - If 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 K, V. - 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) - kv: (batch_size, seqlen, 2, 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. - softcap: float. Anything > 0 activates softcapping 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). - 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 FlashAttnKVPackedFunc.apply( - q, - kv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_func( - q, - k, - v, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - 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). - 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 FlashAttnFunc.apply( - q, - k, - v, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_qkvpacked_func( - qkv, - cu_seqlens, - max_seqlen, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - 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_varlen_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_varlen_kvpacked_func and flash_attn_varlen_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: (total, 3, nheads, headdim), where total = total number of tokens in the batch. - cu_seqlens: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into qkv. - max_seqlen: int. Maximum sequence length in the batch. - 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenQKVPackedFunc.apply( - qkv, - cu_seqlens, - max_seqlen, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_kvpacked_func( - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, -): - """dropout_p should be set to 0.0 during evaluation - If 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 K, V. - 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: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch. - kv: (total_k, 2, nheads_k, headdim), where total_k = total number of key tokens in the batch. - cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into q. - cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into kv. - max_seqlen_q: int. Maximum query sequence length in the batch. - max_seqlen_k: int. Maximum key sequence length in the batch. - 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 + 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenKVPackedFunc.apply( - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_func( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, - block_table=None, -): - """dropout_p should be set to 0.0 during evaluation - Supports multi-query and grouped-query attention (MQA/GQA) by passing in K, V 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: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch. - k: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. - v: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. - cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into q. - cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into kv. - max_seqlen_q: int. Maximum query sequence length in the batch. - max_seqlen_k: int. Maximum key sequence length in the batch. - 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 + 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenFunc.apply( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - block_table, - False if q.device.type == "cpu" else torch.is_grad_enabled(), - ) - - -def flash_attn_with_kvcache( - q, - k_cache, - v_cache, - k=None, - v=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, - block_table: Optional[torch.Tensor] = None, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - rotary_interleaved=True, - alibi_slopes=None, - num_splits=0, - 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 block_table, - or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_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) if there's no block_table, - or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_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). Similar to k. - 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. - block_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). - 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. - 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" - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-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) - cache_batch_idx = maybe_contiguous(cache_batch_idx) - block_table = maybe_contiguous(block_table) - out, softmax_lse = flash_attn.fwd_kvcache( - q, - k_cache, - v_cache, - k, - v, - cache_seqlens, - rotary_cos, - rotary_sin, - cache_batch_idx, - cache_leftpad, - block_table, - alibi_slopes, - None, - softmax_scale, - causal, - window_size[0], - window_size[1], - softcap, - rotary_interleaved, - num_splits, - ) - return (out, softmax_lse) if return_softmax_lse else out diff --git a/build/torch211-cxx11-cu126-aarch64-linux/metadata.json b/build/torch211-cxx11-cu126-aarch64-linux/metadata.json deleted file mode 100644 index 8c16f1144fe08bf19da65ae6a12ac01f8acdc065..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu126-aarch64-linux/metadata.json +++ /dev/null @@ -1,40 +0,0 @@ -{ - "name": "flash-attn2", - "id": "_flash_attn2_cuda_f12afc9", - "version": 1, - "license": "BSD-3-Clause", - "python-depends": [], - "backend": { - "type": "cuda", - "archs": [ - "8.0", - "9.0" - ] - }, - "digest": { - "algorithm": "sha256", - "files": { - "__init__.py": "Cd9MEdd/WUYnBmfiQpC4SyPqAxrTilTjvsvswNju4zc=", - "_flash_attn2_cuda_f12afc9.abi3.so": "2XI68H/Bmhg6MsJZwgPWBiN5ur8QcFJ9pWdL4WYKBi0=", - "_ops.py": "GQbmTTeSOU+7Q2nI68Hz+1RvEogyfEFccx1ubUYw+2s=", - "bert_padding.py": "gF1EmsdJ+HpQ86MRQ4VxDw+Sb/RVISdQALdNnoByHlw=", - "flash_attn2/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY=", - "flash_attn_interface.py": "W7a9bHR2ly6QL3aKqsBqxWA7M7ZWYcZPmtsPJu6rAfY=", - "layers/__init__.py": "47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU=", - "layers/patch_embed.py": "H58CgME/qSOPTZLOG08wFgrQS1j34pvNwMPrkTj3Ek4=", - "layers/rotary.py": "rFIdNrZQFNawwLD6d7t5LE2m+M93wwDZ8EuExw+gfD8=", - "ops/__init__.py": "47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU=", - "ops/activations.py": "t5lzNg1In8LP6bKeTnyeMizwqjv27JGbJ6ylPdGvZYg=", - "ops/fused_dense.py": "BqT8VTSCGdLJAjh6WBMkCRFYrdVodPeoElWBvH6AIzA=", - "ops/layer_norm.py": "zr7NXIm+2mtEynTp1CS0fbFGI2Mqdp41dY4AfDWF6EQ=", - "ops/rms_norm.py": "XEnihcj0a4aSz4LO55m5iKGVn4HKTeKN8TIyHjuDgxI=", - "ops/triton/__init__.py": "AbpHGcgLb+kRsJGnwFEktk7uzpZOCcBY74+YBdrKVGs=", - "ops/triton/cross_entropy.py": "y56nmuOtpSNxM7sDZcUSgXKIsIEayfuvrU8bAVog/nw=", - "ops/triton/k_activations.py": "+Z3vIyO4JkqBMipKsPvhzmxljtBdIhJCsl/M+/ESqBo=", - "ops/triton/layer_norm.py": "BBQdJiqopbIS12viUqzx3SvOQCNbDSAYLJXhpyMGyC8=", - "ops/triton/linear.py": "OtRvKz8xdpl+7v3q/ZTaS9fdBt9XrzMyapgRr50uBbM=", - "ops/triton/mlp.py": "uWA+D7HyuRUsOT+4nhi3vvCDNZnNEXwYcHMQ02UREfQ=", - "ops/triton/rotary.py": "0YSrlbep0lJmzLpTiANs0slTFeGcQQqR+el6uzbSLvI=" - } - } -} \ No newline at end of file diff --git a/build/torch211-cxx11-cu126-aarch64-linux/metadata.json.sigstore b/build/torch211-cxx11-cu126-aarch64-linux/metadata.json.sigstore deleted file mode 100644 index 72df12fe8ab0b8752c33a9fe39bad394319d2882..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu126-aarch64-linux/metadata.json.sigstore +++ /dev/null @@ -1 +0,0 @@ -{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHSzCCBtGgAwIBAgIUJlLjauFMKjms/QQ19SyEeSB1jkgwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwNjIzMjM0MDQxWhcNMjYwNjIzMjM1MDQxWjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEuFBckUQX2LJelzPlwWFzy0lnA9zDRmyqtLy/fUMQKRyLtLaMNSOoo5gdSecr8DFIYHoqlKjnhnTs9MvVO98oxqOCBfAwggXsMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUwPqgjtdu0npnzjDDzPfz6h3Y5C0wHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoZjEyYWZjOTJjNDBiYmQ3NmQ5OTBkOTExZWJhYmYyOGU0OTBhN2JkZTATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoZjEyYWZjOTJjNDBiYmQ3NmQ5OTBkOTExZWJhYmYyOGU0OTBhN2JkZTAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoZjEyYWZjOTJjNDBiYmQ3NmQ5OTBkOTExZWJhYmYyOGU0OTBhN2JkZTAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKGYxMmFmYzkyYzQwYmJkNzZkOTkwZDkxMWViYWJmMjhlNDkwYTdiZGUwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMjgwMzYyNjg5MzUvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBigYKKwYBBAHWeQIEAgR8BHoAeAB2AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABnvbbgoIAAAQDAEcwRQIhAOYOWmPPwM3fTrUeK0fxDUmkb6ixB+uE4EGIlS5oKBBMAiBL4OdVVh/aGGzoxr0BPO6DZdT6CS9eBJP8BnSUdTt+PTAKBggqhkjOPQQDAwNoADBlAjEAi+oJ2fgcTYEL8rz4gzQANH2pv9CHpMVJWohhcTjLafUyEHst/8OmuQEQtfgbK7yiAjAbDPjT+u4QwTND5V3HRwqmC/GhWkLZ7qV1E/28HumuMgT0DDgg5kTILc6HUfe/lJw="}, "tlogEntries":[{"logIndex":"1933039176", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1782258041", "inclusionPromise":{"signedEntryTimestamp":"MEQCIEGqYhETWYlrrpHn8IOTEiLMxfAAe87qzmifEA4avbjoAiAMgtovjtJbEVhUdpjCU+Lt7Rm4TWGoNc6pBpDuQfnFnA=="}, "inclusionProof":{"logIndex":"1811134914", "rootHash":"Mo7p6GAE5vbthsKoVzCojKgL3apqvnCxbAuZzbRuPw8=", "treeSize":"1811134952", "hashes":["33KfZUwDSwH+Qgiyxfv5ooCqUlniFxHBVg1gaHhz7MA=", "LFPGcrnZg0FRabKIK8AD7gW4T7PCbHMY9uXkE83Iy3s=", "7YWKL24QXf8cYhvp0nYHsCnSjY02E5DgshMTGOU39Us=", "ubpQsIzVu9AUVq79n2HxbZuD+7M13ePZ4evoMULnbtU=", "4E/ydT4+W1Xgc4SLmGHd0eycRru2QL2uAqD6hMSqQl8=", "F8ujUbuv7OJbS3b2sw6l5a6cXYBPLmnOWNsaZpOeGuQ=", "e7/f+3LhBuOUmGH8mumO7KvVYW9S3qvnmbhfgUHyhUE=", "mwifXGJ4nqgB69Yt0Pem54qGmK8yLlI1ACFo38BZRc4=", "LzjqkPhZa4hiR0fFtJfuePCkK1oRLVJSKEzhGvUZzIg=", "WD8KwmQq5R8K69koYj5k4K7EKo2VF3LqJcx2xoXcd7o=", "99JTcUz0aJeZZocenAppVOLl8v91qMn0v7CmEjq9lws=", "OSDrt6jTKV6LYAkaG63p2BFz4SmbNmGb3uerhF5pOpQ=", "yehxDZkBWOGIzNDb93qxmQa6VBVo7kacVSoHyXMm3ZI=", "oUh+N71k3jWcTmlDDB6Esa5uqlCnbmzxTmDdut2A16M=", "2BsqkdK2omLvu+CpsFjpwqGi6m22H/7zcwCVkPwslM8=", "S8Oe9ZicoJqbVuL1I0dIdTgVpLc82Cm2CMOcV83sKVs=", "E1O2ilWlpFRLWGdbmJUMDZ2qbGu/eDZ12m6YjQ9/JPA=", "nyv2TDlfkz0XfRwjY98zsbUbOkBZ+Tgz7PML0c2uN14=", "OMfDh4k9YD0E+t4sjqJsh+vg8D1ueo1cC07exyxmr6s=", "+vHKm4wNwGJ46ceyZ8mePD7Vn1SediStNHCStRoTSWQ=", "5HwWjuGNDisetbiWM0cWIXM65TIM3EJlqk13ntCQlcU=", "NDIRlqxO0IaGr94q86MIU7VCAN3yEEgOHS4loya5lu4=", "daxmZaajRpZV+JxHiOYZhJBiSKN5ucqjh2WnGbHhirw=", "DOCeoSMovIvLExkhIvisow9AuNXgeWs4ECkyR6EcqYU="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n1811134952\nMo7p6GAE5vbthsKoVzCojKgL3apqvnCxbAuZzbRuPw8=\n\n— rekor.sigstore.dev wNI9ajBGAiEA+clLWbvNjUhVrnaMW3o3Dio+RWwnsX0m+7Hv/anh+dYCIQDduIUH8vyMexwCTOn/5hYtQXSCH4pxmtsJG4Aem+4cMA==\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiJkYWI1MzY2ZDY0MDYyODZhOGJhNzMzNDFmYmZkYzBlNzYzNDBhNDM2ODAwNjNlY2E0ZTBjODQ4ZjFjZWI5YjIwIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FWUNJUURnSHdUS3VRVm1OV0grUHZZMEdLZWR4T3R1WGRQTHVYQnFWam1Qb2R1ZmNBSWhBSVdzK2xhcENuSmhrMmczcmpEWEhFd3B6RkdpWTk4eXpSNk1xYThxZW05QSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRla05EUW5SSFowRjNTVUpCWjBsVlNteE1hbUYxUmsxTGFtMXpMMUZSTVRsVGVVVmxVMEl4YW10bmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDVxU1hwTmFrMHdUVVJSZUZkb1kwNU5hbGwzVG1wSmVrMXFUVEZOUkZGNFYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVYxUmtKamExVlJXREpNU21Wc2VsQnNkMWRHZW5rd2JHNUJPWHBFVW0xNWNYUk1lUzhLWmxWTlVVdFNlVXgwVEdGTlRsTlBiMjgxWjJSVFpXTnlPRVJHU1ZsSWIzRnNTMnB1YUc1VWN6bE5kbFpQT1RodmVIRlBRMEptUVhkbloxaHpUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlYzVUhGbkNtcDBaSFV3Ym5CdWVtcEVSSHBRWm5vMmFETlpOVU13ZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOWFha1Y1V1ZkYWFrOVVTbXBPUkVKcFdXMVJNMDV0VVRWUFZFSnJDazlVUlhoYVYwcG9XVzFaZVU5SFZUQlBWRUpvVGpKS2ExcFVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMXBxUlhsWlYxcHFUMVJLYWs1RVFtbFpiVkV6VG0xUk5VOVVRbXRQVkVWNFdsZEthRmx0V1hrS1QwZFZNRTlVUW1oT01rcHJXbFJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMXBxUlhsWlYxcHFUMVJLYWs1RVFtbFpiVkV6VG0xUk5VOVVRbXNLVDFSRmVGcFhTbWhaYlZsNVQwZFZNRTlVUW1oT01rcHJXbFJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwZFplRTF0Um0wS1dYcHJlVmw2VVhkWmJVcHJUbnBhYTA5VWEzZGFSR3Q0VFZkV2FWbFhTbTFOYW1oc1RrUnJkMWxVWkdsYVIxVjNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOYW1kM1RYcFplVTVxWnpWTmVsVjJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwWjFsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTT0VKSWIwRUtaVUZDTWtGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtNTJZbUpuYjBsQlFVRlJSQXBCUldOM1VsRkphRUZQV1U5WGJWQlFkMDB6WmxSeVZXVkxNR1o0UkZWdGEySTJhWGhDSzNWRk5FVkhTV3hUTlc5TFFrSk5RV2xDVERSUFpGWldhQzloQ2tkSGVtOTRjakJDVUU4MlJGcGtWRFpEVXpsbFFrcFFPRUp1VTFWa1ZIUXJVRlJCUzBKblozRm9hMnBQVUZGUlJFRjNUbTlCUkVKc1FXcEZRV2tyYjBvS01tWm5ZMVJaUlV3NGNubzBaM3BSUVU1SU1uQjJPVU5JY0UxV1NsZHZhR2hqVkdwTVlXWlZlVVZJYzNRdk9FOXRkVkZGVVhSbVoySkxOM2xwUVdwQllncEVVR3BVSzNVMFVYZFVUa1ExVmpOSVVuZHhiVU12UjJoWGEweGFOM0ZXTVVVdk1qaElkVzExVFdkVU1FUkVaMmMxYTFSSlRHTTJTRlZtWlM5c1NuYzlDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyTADAgEAMIICwAYJKoZIhvcNAQcCoIICsTCCAq0CAQMxDTALBglghkgBZQMEAgEwgbcGCyqGSIb3DQEJEAEEoIGnBIGkMIGhAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQg9GmRsc7E/n2MZXpdNiKJH1g7fMKt3HnX/bRiwtWB1RICFBI4G3Zn7paAZiG59w49otud1EzRGA8yMDI2MDYyMzIzNDA0MVowAwIBAaAypDAwLjEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MRUwEwYDVQQDEwxzaWdzdG9yZS10c2GgADGCAdswggHXAgEBMFEwOTEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MSAwHgYDVQQDExdzaWdzdG9yZS10c2Etc2VsZnNpZ25lZAIUOhNULwyQYe68wUMvy4qOiyojiwwwCwYJYIZIAWUDBAIBoIH8MBoGCSqGSIb3DQEJAzENBgsqhkiG9w0BCRABBDAcBgkqhkiG9w0BCQUxDxcNMjYwNjIzMjM0MDQxWjAvBgkqhkiG9w0BCQQxIgQgbL7/0KTcQMV0HFrr0g1ZXRfTRbOBYkHaJ800m2oE/B4wgY4GCyqGSIb3DQEJEAIvMX8wfTB7MHkEIIX5J7wHq2LKw7RDVsEO/IGyxog/2nq55thw2dE6zQW3MFUwPaQ7MDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAoGCCqGSM49BAMCBGcwZQIxALJWmLgpRciQGU60j+M2tLuvy2r6jaGq8x4O6NckQNL8gBwGu3sXtS6q2Ia07WKSDQIwJ7SrctC3Fcjdm7XQWZTWpF0xxdRlXPgSL2Eicl/UeUPSM928SJaFxIbibPFUnhA9"}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"2rU2bWQGKGqLpzNB+/3A52NApDaABj7KTgyEjxzrmyA="}, "signature":"MEYCIQDgHwTKuQVmNWH+PvY0GKedxOtuXdPLuXBqVjmPodufcAIhAIWs+lapCnJhk2g3rjDXHEwpzFGiY98yzR6Mqa8qem9A"}} \ No newline at end of file diff --git a/build/torch211-cxx11-cu126-aarch64-linux/ops/triton/layer_norm.py b/build/torch211-cxx11-cu126-aarch64-linux/ops/triton/layer_norm.py deleted file mode 100644 index 5e6f0b3697b612e6a19031bbc55ac76ad233dc85..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu126-aarch64-linux/ops/triton/layer_norm.py +++ /dev/null @@ -1,1254 +0,0 @@ -# Copyright (c) 2024, Tri Dao. -# Implement dropout + residual + layer_norm / rms_norm. - -# Based on the Triton LayerNorm tutorial: https://triton-lang.org/main/getting-started/tutorials/05-layer-norm.html -# For the backward pass, we keep weight_grad and bias_grad in registers and accumulate. -# This is faster for dimensions up to 8k, but after that it's much slower due to register spilling. -# The models we train have hidden dim up to 8k anyway (e.g. Llama 70B), so this is fine. - -import math -from typing import Optional, List - -import torch -import torch.nn.functional as F -from torch import Tensor - -import triton -import triton.language as tl - -from flash_attn.utils.torch import custom_fwd, custom_bwd -from flash_attn.utils.library import triton_op - -from ..._ops import add_op_namespace_prefix - - -def maybe_contiguous_lastdim(x): - return x.contiguous() if x is not None and x.stride(-1) != 1 else x - - -def maybe_contiguous(x): - return x.contiguous() if x is not None else None - - -def triton_autotune_configs(): - # Return configs with a valid warp count for the current device - configs = [] - # Maximum threads per block is architecture-dependent in theory, but in reality all are 1024 - max_threads_per_block = 1024 - # Default to warp size 32 if not defined by device - warp_size = getattr(torch.cuda.get_device_properties(torch.cuda.current_device()), "warp_size", 32) - # Autotune for warp counts which are powers of 2 and do not exceed thread per block limit - return [triton.Config({}, num_warps=warp_count) for warp_count in [1, 2, 4, 8, 16, 32] - if warp_count * warp_size <= max_threads_per_block] - # return [triton.Config({}, num_warps=8)] - - -def layer_norm_ref( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - zero_centered_weight=False, - dropout_mask=None, - dropout_mask1=None, - upcast=False, -): - dtype = x.dtype - if upcast: - x = x.float() - weight = weight.float() - bias = bias.float() if bias is not None else None - residual = residual.float() if residual is not None else residual - x1 = x1.float() if x1 is not None else None - weight1 = weight1.float() if weight1 is not None else None - bias1 = bias1.float() if bias1 is not None else None - if zero_centered_weight: - weight = weight + 1.0 - if weight1 is not None: - weight1 = weight1 + 1.0 - if x1 is not None: - assert rowscale is None, "rowscale is not supported with parallel LayerNorm" - if rowscale is not None: - x = x * rowscale[..., None] - if dropout_p > 0.0: - if dropout_mask is not None: - x = x.masked_fill(~dropout_mask, 0.0) / (1.0 - dropout_p) - else: - x = F.dropout(x, p=dropout_p) - if x1 is not None: - if dropout_mask1 is not None: - x1 = x1.masked_fill(~dropout_mask1, 0.0) / (1.0 - dropout_p) - else: - x1 = F.dropout(x1, p=dropout_p) - if x1 is not None: - x = x + x1 - if residual is not None: - x = (x + residual).to(x.dtype) - out = F.layer_norm(x.to(weight.dtype), x.shape[-1:], weight=weight, bias=bias, eps=eps).to( - dtype - ) - if weight1 is None: - return out if not prenorm else (out, x) - else: - out1 = F.layer_norm( - x.to(weight1.dtype), x.shape[-1:], weight=weight1, bias=bias1, eps=eps - ).to(dtype) - return (out, out1) if not prenorm else (out, out1, x) - - -def rms_norm_ref( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - zero_centered_weight=False, - dropout_mask=None, - dropout_mask1=None, - upcast=False, -): - dtype = x.dtype - if upcast: - x = x.float() - weight = weight.float() - bias = bias.float() if bias is not None else None - residual = residual.float() if residual is not None else residual - x1 = x1.float() if x1 is not None else None - weight1 = weight1.float() if weight1 is not None else None - bias1 = bias1.float() if bias1 is not None else None - if zero_centered_weight: - weight = weight + 1.0 - if weight1 is not None: - weight1 = weight1 + 1.0 - if x1 is not None: - assert rowscale is None, "rowscale is not supported with parallel LayerNorm" - if rowscale is not None: - x = x * rowscale[..., None] - if dropout_p > 0.0: - if dropout_mask is not None: - x = x.masked_fill(~dropout_mask, 0.0) / (1.0 - dropout_p) - else: - x = F.dropout(x, p=dropout_p) - if x1 is not None: - if dropout_mask1 is not None: - x1 = x1.masked_fill(~dropout_mask1, 0.0) / (1.0 - dropout_p) - else: - x1 = F.dropout(x1, p=dropout_p) - if x1 is not None: - x = x + x1 - if residual is not None: - x = (x + residual).to(x.dtype) - rstd = 1 / torch.sqrt((x.square()).mean(dim=-1, keepdim=True) + eps) - out = ((x * rstd * weight) + bias if bias is not None else (x * rstd * weight)).to(dtype) - if weight1 is None: - return out if not prenorm else (out, x) - else: - out1 = ((x * rstd * weight1) + bias1 if bias1 is not None else (x * rstd * weight1)).to( - dtype - ) - return (out, out1) if not prenorm else (out, out1, x) - - -@triton.autotune( - configs=triton_autotune_configs(), - key=["N", "HAS_RESIDUAL", "STORE_RESIDUAL_OUT", "IS_RMS_NORM", "HAS_BIAS", "HAS_X1", "HAS_W1", "HAS_B1"], -) -# torch compile doesn't like triton.heuristics, so we set these manually when calling the kernel -# @triton.heuristics({"HAS_BIAS": lambda args: args["B"] is not None}) -# @triton.heuristics({"HAS_RESIDUAL": lambda args: args["RESIDUAL"] is not None}) -# @triton.heuristics({"HAS_X1": lambda args: args["X1"] is not None}) -# @triton.heuristics({"HAS_W1": lambda args: args["W1"] is not None}) -# @triton.heuristics({"HAS_B1": lambda args: args["B1"] is not None}) -@triton.jit -def _layer_norm_fwd_1pass_kernel( - X, # pointer to the input - Y, # pointer to the output - W, # pointer to the weights - B, # pointer to the biases - RESIDUAL, # pointer to the residual - X1, - W1, - B1, - Y1, - RESIDUAL_OUT, # pointer to the residual - ROWSCALE, - SEEDS, # Dropout seeds for each row - DROPOUT_MASK, - DROPOUT_MASK1, - Mean, # pointer to the mean - Rstd, # pointer to the 1/std - stride_x_row, # how much to increase the pointer when moving by 1 row - stride_y_row, - stride_res_row, - stride_res_out_row, - stride_x1_row, - stride_y1_row, - M, # number of rows in X - N, # number of columns in X - eps, # epsilon to avoid division by zero - dropout_p, # Dropout probability - zero_centered_weight, # If true, add 1.0 to the weight - IS_RMS_NORM: tl.constexpr, - BLOCK_N: tl.constexpr, - HAS_RESIDUAL: tl.constexpr, - STORE_RESIDUAL_OUT: tl.constexpr, - HAS_BIAS: tl.constexpr, - HAS_DROPOUT: tl.constexpr, - STORE_DROPOUT_MASK: tl.constexpr, - HAS_ROWSCALE: tl.constexpr, - HAS_X1: tl.constexpr, - HAS_W1: tl.constexpr, - HAS_B1: tl.constexpr, -): - # Map the program id to the row of X and Y it should compute. - row = tl.program_id(0) - X += row * stride_x_row - Y += row * stride_y_row - if HAS_RESIDUAL: - RESIDUAL += row * stride_res_row - if STORE_RESIDUAL_OUT: - RESIDUAL_OUT += row * stride_res_out_row - if HAS_X1: - X1 += row * stride_x1_row - if HAS_W1: - Y1 += row * stride_y1_row - # Compute mean and variance - cols = tl.arange(0, BLOCK_N) - x = tl.load(X + cols, mask=cols < N, other=0.0).to(tl.float32) - if HAS_ROWSCALE: - rowscale = tl.load(ROWSCALE + row).to(tl.float32) - x *= rowscale - if HAS_DROPOUT: - # Compute dropout mask - # 7 rounds is good enough, and reduces register pressure - keep_mask = tl.rand(tl.load(SEEDS + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - x = tl.where(keep_mask, x / (1.0 - dropout_p), 0.0) - if STORE_DROPOUT_MASK: - tl.store(DROPOUT_MASK + row * N + cols, keep_mask, mask=cols < N) - if HAS_X1: - x1 = tl.load(X1 + cols, mask=cols < N, other=0.0).to(tl.float32) - if HAS_ROWSCALE: - rowscale = tl.load(ROWSCALE + M + row).to(tl.float32) - x1 *= rowscale - if HAS_DROPOUT: - # Compute dropout mask - # 7 rounds is good enough, and reduces register pressure - keep_mask = ( - tl.rand(tl.load(SEEDS + M + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - ) - x1 = tl.where(keep_mask, x1 / (1.0 - dropout_p), 0.0) - if STORE_DROPOUT_MASK: - tl.store(DROPOUT_MASK1 + row * N + cols, keep_mask, mask=cols < N) - x += x1 - if HAS_RESIDUAL: - residual = tl.load(RESIDUAL + cols, mask=cols < N, other=0.0).to(tl.float32) - x += residual - if STORE_RESIDUAL_OUT: - tl.store(RESIDUAL_OUT + cols, x, mask=cols < N) - if not IS_RMS_NORM: - mean = tl.sum(x, axis=0) / N - tl.store(Mean + row, mean) - xbar = tl.where(cols < N, x - mean, 0.0) - var = tl.sum(xbar * xbar, axis=0) / N - else: - xbar = tl.where(cols < N, x, 0.0) - var = tl.sum(xbar * xbar, axis=0) / N - rstd = 1 / tl.sqrt(var + eps) - tl.store(Rstd + row, rstd) - # Normalize and apply linear transformation - mask = cols < N - w = tl.load(W + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w += 1.0 - if HAS_BIAS: - b = tl.load(B + cols, mask=mask).to(tl.float32) - x_hat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd - y = x_hat * w + b if HAS_BIAS else x_hat * w - # Write output - tl.store(Y + cols, y, mask=mask) - if HAS_W1: - w1 = tl.load(W1 + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w1 += 1.0 - if HAS_B1: - b1 = tl.load(B1 + cols, mask=mask).to(tl.float32) - y1 = x_hat * w1 + b1 if HAS_B1 else x_hat * w1 - tl.store(Y1 + cols, y1, mask=mask) - - -def _layer_norm_fwd( - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - residual: Optional[Tensor] = None, - x1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - out_dtype: Optional[torch.dtype] = None, - residual_dtype: Optional[torch.dtype] = None, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - return_dropout_mask: bool = False, - out: Optional[Tensor] = None, - residual_out: Optional[Tensor] = None -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - # Need to wrap to handle the case where residual_out is a alias of x, which makes torch.library - # and torch.compile unhappy. Also allocate memory for out and residual_out if they are None - # so that _layer_norm_fwd_impl doesn't have to return them. - if out is None: - out = torch.empty_like(x, dtype=x.dtype if out_dtype is None else out_dtype) - if residual is not None: - residual_dtype = residual.dtype - if residual_out is None and ( - residual is not None - or (residual_dtype is not None and residual_dtype != x.dtype) - or dropout_p > 0.0 - or rowscale is not None - or x1 is not None - ): - residual_out = torch.empty_like( - x, dtype=residual_dtype if residual_dtype is not None else x.dtype - ) - else: - residual_out = None - y1, mean, rstd, seeds, dropout_mask, dropout_mask1 = _layer_norm_fwd_impl( - x, - weight, - bias, - eps, - out, - residual=residual, - x1=x1, - weight1=weight1, - bias1=bias1, - dropout_p=dropout_p, - rowscale=rowscale, - zero_centered_weight=zero_centered_weight, - is_rms_norm=is_rms_norm, - return_dropout_mask=return_dropout_mask, - residual_out=residual_out, - ) - # residual_out is None if residual is None and residual_dtype == input_dtype and dropout_p == 0.0 - if residual_out is None: - residual_out = x - return out, y1, mean, rstd, residual_out, seeds, dropout_mask, dropout_mask1 - - -# [2025-04-28] torch.library.triton_op ignores the schema argument, but here we need the schema -# since we're returning a tuple of tensors -@triton_op(add_op_namespace_prefix("layer_norm_fwd_impl"), mutates_args={"out", "residual_out"}, - schema="(Tensor x, Tensor weight, Tensor bias, float eps, Tensor(a!) out, Tensor? residual, Tensor? x1, Tensor? weight1, Tensor? bias1, float dropout_p, Tensor? rowscale, bool zero_centered_weight, bool is_rms_norm, bool return_dropout_mask, Tensor(a!)? residual_out) -> (Tensor y1, Tensor mean, Tensor rstd, Tensor seeds, Tensor dropout_mask, Tensor dropout_mask1)") -def _layer_norm_fwd_impl( - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - out: Tensor, - residual: Optional[Tensor] = None, - x1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - return_dropout_mask: bool = False, - residual_out: Optional[Tensor] = None -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - M, N = x.shape - assert x.stride(-1) == 1 - if residual is not None: - assert residual.stride(-1) == 1 - assert residual.shape == (M, N) - assert weight.shape == (N,) - assert weight.stride(-1) == 1 - if bias is not None: - assert bias.stride(-1) == 1 - assert bias.shape == (N,) - if x1 is not None: - assert x1.shape == x.shape - assert rowscale is None - assert x1.stride(-1) == 1 - if weight1 is not None: - assert weight1.shape == (N,) - assert weight1.stride(-1) == 1 - if bias1 is not None: - assert bias1.shape == (N,) - assert bias1.stride(-1) == 1 - if rowscale is not None: - assert rowscale.is_contiguous() - assert rowscale.shape == (M,) - assert out.shape == x.shape - assert out.stride(-1) == 1 - if residual_out is not None: - assert residual_out.shape == x.shape - assert residual_out.stride(-1) == 1 - if weight1 is not None: - y1 = torch.empty_like(out) - assert y1.stride(-1) == 1 - else: - y1 = None - mean = torch.empty((M,), dtype=torch.float32, device=x.device) if not is_rms_norm else None - rstd = torch.empty((M,), dtype=torch.float32, device=x.device) - if dropout_p > 0.0: - seeds = torch.randint( - 2**32, (M if x1 is None else 2 * M,), device=x.device, dtype=torch.int64 - ) - else: - seeds = None - if return_dropout_mask and dropout_p > 0.0: - dropout_mask = torch.empty(M, N, device=x.device, dtype=torch.bool) - if x1 is not None: - dropout_mask1 = torch.empty(M, N, device=x.device, dtype=torch.bool) - else: - dropout_mask1 = None - else: - dropout_mask, dropout_mask1 = None, None - # Less than 64KB per feature: enqueue fused kernel - MAX_FUSED_SIZE = 65536 // x.element_size() - BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) - if N > BLOCK_N: - raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") - with torch.cuda.device(x.device.index): - torch.library.wrap_triton(_layer_norm_fwd_1pass_kernel)[(M,)]( - x, - out, - weight, - bias, - residual, - x1, - weight1, - bias1, - y1, - residual_out, - rowscale, - seeds, - dropout_mask, - dropout_mask1, - mean, - rstd, - x.stride(0), - out.stride(0), - residual.stride(0) if residual is not None else 0, - residual_out.stride(0) if residual_out is not None else 0, - x1.stride(0) if x1 is not None else 0, - y1.stride(0) if y1 is not None else 0, - M, - N, - eps, - dropout_p, - # Passing bool make torch inductor very unhappy since it then tries to compare to int_max - int(zero_centered_weight), - is_rms_norm, - BLOCK_N, - residual is not None, - residual_out is not None, - bias is not None, - dropout_p > 0.0, - dropout_mask is not None, - rowscale is not None, - HAS_X1=x1 is not None, - HAS_W1=weight1 is not None, - HAS_B1=bias1 is not None, - ) - return y1, mean, rstd, seeds, dropout_mask, dropout_mask1 - - -@triton.autotune( - configs=triton_autotune_configs(), - key=["N", "HAS_DRESIDUAL", "STORE_DRESIDUAL", "IS_RMS_NORM", "HAS_BIAS", "HAS_DROPOUT"], -) -# torch compile doesn't like triton.heuristics, so we set these manually when calling the kernel -# @triton.heuristics({"HAS_BIAS": lambda args: args["B"] is not None}) -# @triton.heuristics({"HAS_DRESIDUAL": lambda args: args["DRESIDUAL"] is not None}) -# @triton.heuristics({"STORE_DRESIDUAL": lambda args: args["DRESIDUAL_IN"] is not None}) -# @triton.heuristics({"HAS_ROWSCALE": lambda args: args["ROWSCALE"] is not None}) -# @triton.heuristics({"HAS_DY1": lambda args: args["DY1"] is not None}) -# @triton.heuristics({"HAS_DX1": lambda args: args["DX1"] is not None}) -# @triton.heuristics({"HAS_B1": lambda args: args["DB1"] is not None}) -# @triton.heuristics({"RECOMPUTE_OUTPUT": lambda args: args["Y"] is not None}) -@triton.jit -def _layer_norm_bwd_kernel( - X, # pointer to the input - W, # pointer to the weights - B, # pointer to the biases - Y, # pointer to the output to be recomputed - DY, # pointer to the output gradient - DX, # pointer to the input gradient - DW, # pointer to the partial sum of weights gradient - DB, # pointer to the partial sum of biases gradient - DRESIDUAL, - W1, - DY1, - DX1, - DW1, - DB1, - DRESIDUAL_IN, - ROWSCALE, - SEEDS, - Mean, # pointer to the mean - Rstd, # pointer to the 1/std - stride_x_row, # how much to increase the pointer when moving by 1 row - stride_y_row, - stride_dy_row, - stride_dx_row, - stride_dres_row, - stride_dy1_row, - stride_dx1_row, - stride_dres_in_row, - M, # number of rows in X - N, # number of columns in X - eps, # epsilon to avoid division by zero - dropout_p, - zero_centered_weight, - rows_per_program, - IS_RMS_NORM: tl.constexpr, - BLOCK_N: tl.constexpr, - HAS_DRESIDUAL: tl.constexpr, - STORE_DRESIDUAL: tl.constexpr, - HAS_BIAS: tl.constexpr, - HAS_DROPOUT: tl.constexpr, - HAS_ROWSCALE: tl.constexpr, - HAS_DY1: tl.constexpr, - HAS_DX1: tl.constexpr, - HAS_B1: tl.constexpr, - RECOMPUTE_OUTPUT: tl.constexpr, -): - # Map the program id to the elements of X, DX, and DY it should compute. - row_block_id = tl.program_id(0) - row_start = row_block_id * rows_per_program - # Do not early exit if row_start >= M, because we need to write DW and DB - cols = tl.arange(0, BLOCK_N) - mask = cols < N - X += row_start * stride_x_row - if HAS_DRESIDUAL: - DRESIDUAL += row_start * stride_dres_row - if STORE_DRESIDUAL: - DRESIDUAL_IN += row_start * stride_dres_in_row - DY += row_start * stride_dy_row - DX += row_start * stride_dx_row - if HAS_DY1: - DY1 += row_start * stride_dy1_row - if HAS_DX1: - DX1 += row_start * stride_dx1_row - if RECOMPUTE_OUTPUT: - Y += row_start * stride_y_row - w = tl.load(W + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w += 1.0 - if RECOMPUTE_OUTPUT and HAS_BIAS: - b = tl.load(B + cols, mask=mask, other=0.0).to(tl.float32) - if HAS_DY1: - w1 = tl.load(W1 + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w1 += 1.0 - dw = tl.zeros((BLOCK_N,), dtype=tl.float32) - if HAS_BIAS: - db = tl.zeros((BLOCK_N,), dtype=tl.float32) - if HAS_DY1: - dw1 = tl.zeros((BLOCK_N,), dtype=tl.float32) - if HAS_B1: - db1 = tl.zeros((BLOCK_N,), dtype=tl.float32) - row_end = min((row_block_id + 1) * rows_per_program, M) - for row in range(row_start, row_end): - # Load data to SRAM - x = tl.load(X + cols, mask=mask, other=0).to(tl.float32) - dy = tl.load(DY + cols, mask=mask, other=0).to(tl.float32) - if HAS_DY1: - dy1 = tl.load(DY1 + cols, mask=mask, other=0).to(tl.float32) - if not IS_RMS_NORM: - mean = tl.load(Mean + row) - rstd = tl.load(Rstd + row) - # Compute dx - xhat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd - xhat = tl.where(mask, xhat, 0.0) - if RECOMPUTE_OUTPUT: - y = xhat * w + b if HAS_BIAS else xhat * w - tl.store(Y + cols, y, mask=mask) - wdy = w * dy - dw += dy * xhat - if HAS_BIAS: - db += dy - if HAS_DY1: - wdy += w1 * dy1 - dw1 += dy1 * xhat - if HAS_B1: - db1 += dy1 - if not IS_RMS_NORM: - c1 = tl.sum(xhat * wdy, axis=0) / N - c2 = tl.sum(wdy, axis=0) / N - dx = (wdy - (xhat * c1 + c2)) * rstd - else: - c1 = tl.sum(xhat * wdy, axis=0) / N - dx = (wdy - xhat * c1) * rstd - if HAS_DRESIDUAL: - dres = tl.load(DRESIDUAL + cols, mask=mask, other=0).to(tl.float32) - dx += dres - # Write dx - if STORE_DRESIDUAL: - tl.store(DRESIDUAL_IN + cols, dx, mask=mask) - if HAS_DX1: - if HAS_DROPOUT: - keep_mask = ( - tl.rand(tl.load(SEEDS + M + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - ) - dx1 = tl.where(keep_mask, dx / (1.0 - dropout_p), 0.0) - else: - dx1 = dx - tl.store(DX1 + cols, dx1, mask=mask) - if HAS_DROPOUT: - keep_mask = tl.rand(tl.load(SEEDS + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - dx = tl.where(keep_mask, dx / (1.0 - dropout_p), 0.0) - if HAS_ROWSCALE: - rowscale = tl.load(ROWSCALE + row).to(tl.float32) - dx *= rowscale - tl.store(DX + cols, dx, mask=mask) - - X += stride_x_row - if HAS_DRESIDUAL: - DRESIDUAL += stride_dres_row - if STORE_DRESIDUAL: - DRESIDUAL_IN += stride_dres_in_row - if RECOMPUTE_OUTPUT: - Y += stride_y_row - DY += stride_dy_row - DX += stride_dx_row - if HAS_DY1: - DY1 += stride_dy1_row - if HAS_DX1: - DX1 += stride_dx1_row - tl.store(DW + row_block_id * N + cols, dw, mask=mask) - if HAS_BIAS: - tl.store(DB + row_block_id * N + cols, db, mask=mask) - if HAS_DY1: - tl.store(DW1 + row_block_id * N + cols, dw1, mask=mask) - if HAS_B1: - tl.store(DB1 + row_block_id * N + cols, db1, mask=mask) - - -def _layer_norm_bwd( - dy: Tensor, - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - mean: Tensor, - rstd: Tensor, - dresidual: Optional[Tensor] = None, - dy1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - seeds: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - has_residual: bool = False, - has_x1: bool = False, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - x_dtype: Optional[torch.dtype] = None, - recompute_output: bool = False, -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - # Need to wrap to handle the case where dresidual_in or dx1 are aliases of x, - # which makes torch.library unhappy - dx, dw, db, dresidual_in, dx1, dw1, db1, y = _layer_norm_bwd_impl( - dy, - x, - weight, - bias, - eps, - mean, - rstd, - dresidual, - dy1, - weight1, - bias1, - seeds, - dropout_p, - rowscale, - has_residual, - has_x1, - zero_centered_weight, - is_rms_norm, - x_dtype=x_dtype, - recompute_output=recompute_output, - ) - # Don't need to compute dresidual_in separately in this case - if has_residual and dx.dtype == x.dtype and dropout_p == 0.0 and rowscale is None: - dresidual_in = dx - if has_x1 and dropout_p == 0.0: - dx1 = dx - return dx, dw, db, dresidual_in, dx1, dw1, db1, y - - - -@triton_op(add_op_namespace_prefix("layer_norm_bwd_impl"), mutates_args={}, - schema="(Tensor dy, Tensor x, Tensor weight, Tensor bias, float eps, Tensor mean, Tensor rstd, Tensor? dresidual, Tensor? dy1, Tensor? weight1, Tensor? bias1, Tensor? seeds, float dropout_p, Tensor? rowscale, bool has_residual, bool has_x1, bool zero_centered_weight, bool is_rms_norm, ScalarType? x_dtype, bool recompute_output) -> (Tensor dx, Tensor dw, Tensor db, Tensor dresidual_in, Tensor dx1, Tensor dw1, Tensor db1, Tensor y)", - allow_decomposition=False, # Don't let torch.compile trace inside - ) -def _layer_norm_bwd_impl( - dy: Tensor, - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - mean: Tensor, - rstd: Tensor, - dresidual: Optional[Tensor] = None, - dy1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - seeds: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - has_residual: bool = False, - has_x1: bool = False, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - x_dtype: Optional[torch.dtype] = None, - recompute_output: bool = False, -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - M, N = x.shape - assert x.stride(-1) == 1 - dy = maybe_contiguous_lastdim(dy) - assert dy.stride(-1) == 1 - assert dy.shape == (M, N) - if dresidual is not None: - dresidual = maybe_contiguous_lastdim(dresidual) - assert dresidual.stride(-1) == 1 - assert dresidual.shape == (M, N) - assert weight.shape == (N,) - assert weight.stride(-1) == 1 - if bias is not None: - assert bias.stride(-1) == 1 - assert bias.shape == (N,) - if dy1 is not None: - dy1 = maybe_contiguous_lastdim(dy1) - assert weight1 is not None - assert dy1.shape == dy.shape - assert dy1.stride(-1) == 1 - if weight1 is not None: - assert weight1.shape == (N,) - assert weight1.stride(-1) == 1 - if bias1 is not None: - assert bias1.shape == (N,) - assert bias1.stride(-1) == 1 - if seeds is not None: - assert seeds.is_contiguous() - assert seeds.shape == (M if not has_x1 else M * 2,) - if rowscale is not None: - assert rowscale.is_contiguous() - assert rowscale.shape == (M,) - # allocate output - dx = ( - torch.empty_like(x) - if x_dtype is None - else torch.empty(M, N, dtype=x_dtype, device=x.device) - ) - dresidual_in = ( - torch.empty_like(x) - if has_residual - and (dx.dtype != x.dtype or dropout_p > 0.0 or rowscale is not None or has_x1) - else None - ) - dx1 = torch.empty_like(dx) if (has_x1 and dropout_p > 0.0) else None - y = torch.empty(M, N, dtype=dy.dtype, device=dy.device) if recompute_output else None - if recompute_output: - assert weight1 is None, "recompute_output is not supported with parallel LayerNorm" - - # Less than 64KB per feature: enqueue fused kernel - MAX_FUSED_SIZE = 65536 // x.element_size() - BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) - if N > BLOCK_N: - raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") - # Increasing the multiple (e.g. 8) will allow more thread blocks to be launched and hide the - # latency of the gmem reads/writes, but will increase the time of summing up dw / db. - sm_count = torch.cuda.get_device_properties(x.device).multi_processor_count * 8 - _dw = torch.empty((sm_count, N), dtype=torch.float32, device=weight.device) - _db = ( - torch.empty((sm_count, N), dtype=torch.float32, device=bias.device) - if bias is not None - else None - ) - _dw1 = torch.empty_like(_dw) if weight1 is not None else None - _db1 = torch.empty_like(_db) if bias1 is not None else None - rows_per_program = math.ceil(M / sm_count) - grid = (sm_count,) - with torch.cuda.device(x.device.index): - torch.library.wrap_triton(_layer_norm_bwd_kernel)[grid]( - x, - weight, - bias, - y, - dy, - dx, - _dw, - _db, - dresidual, - weight1, - dy1, - dx1, - _dw1, - _db1, - dresidual_in, - rowscale, - seeds, - mean, - rstd, - x.stride(0), - 0 if not recompute_output else y.stride(0), - dy.stride(0), - dx.stride(0), - dresidual.stride(0) if dresidual is not None else 0, - dy1.stride(0) if dy1 is not None else 0, - dx1.stride(0) if dx1 is not None else 0, - dresidual_in.stride(0) if dresidual_in is not None else 0, - M, - N, - eps, - dropout_p, - # Passing bool make torch inductor very unhappy since it then tries to compare to int_max - int(zero_centered_weight), - rows_per_program, - is_rms_norm, - BLOCK_N, - dresidual is not None, - dresidual_in is not None, - bias is not None, - dropout_p > 0.0, - HAS_ROWSCALE=rowscale is not None, - HAS_DY1=dy1 is not None, - HAS_DX1=dx1 is not None, - HAS_B1=bias1 is not None, - RECOMPUTE_OUTPUT=y is not None, - ) - dw = _dw.sum(0).to(weight.dtype) - db = _db.sum(0).to(bias.dtype) if bias is not None else None - dw1 = _dw1.sum(0).to(weight1.dtype) if weight1 is not None else None - db1 = _db1.sum(0).to(bias1.dtype) if bias1 is not None else None - # dresidual_in and dx1 could be None, the wrapper will handle assigning them from dx - return dx, dw, db, dresidual_in, dx1, dw1, db1, y - - -class LayerNormFn(torch.autograd.Function): - - @staticmethod - def forward( - ctx, - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - residual_in_fp32=False, - zero_centered_weight=False, - is_rms_norm=False, - return_dropout_mask=False, - out_dtype=None, - out=None, - residual_out=None - ): - x_shape_og = x.shape - # reshape input data into 2D tensor - x = maybe_contiguous_lastdim(x.reshape(-1, x.shape[-1])) - if residual is not None: - assert residual.shape == x_shape_og - residual = maybe_contiguous_lastdim(residual.reshape(-1, residual.shape[-1])) - if x1 is not None: - assert x1.shape == x_shape_og - assert rowscale is None, "rowscale is not supported with parallel LayerNorm" - x1 = maybe_contiguous_lastdim(x1.reshape(-1, x1.shape[-1])) - weight = weight.contiguous() - bias = maybe_contiguous(bias) - weight1 = maybe_contiguous(weight1) - bias1 = maybe_contiguous(bias1) - if rowscale is not None: - rowscale = rowscale.reshape(-1).contiguous() - residual_dtype = ( - residual.dtype - if residual is not None - else (torch.float32 if residual_in_fp32 else None) - ) - if out is not None: - out = out.reshape(-1, out.shape[-1]) - if residual_out is not None: - residual_out = residual_out.reshape(-1, residual_out.shape[-1]) - y, y1, mean, rstd, residual_out, seeds, dropout_mask, dropout_mask1 = _layer_norm_fwd( - x, - weight, - bias, - eps, - residual, - x1, - weight1, - bias1, - dropout_p=dropout_p, - rowscale=rowscale, - out_dtype=out_dtype, - residual_dtype=residual_dtype, - zero_centered_weight=zero_centered_weight, - is_rms_norm=is_rms_norm, - return_dropout_mask=return_dropout_mask, - out=out, - residual_out=residual_out, - ) - ctx.save_for_backward( - residual_out, weight, bias, weight1, bias1, rowscale, seeds, mean, rstd - ) - ctx.x_shape_og = x_shape_og - ctx.eps = eps - ctx.dropout_p = dropout_p - ctx.is_rms_norm = is_rms_norm - ctx.has_residual = residual is not None - ctx.has_x1 = x1 is not None - ctx.prenorm = prenorm - ctx.x_dtype = x.dtype - ctx.zero_centered_weight = zero_centered_weight - y = y.reshape(x_shape_og) - y1 = y1.reshape(x_shape_og) if y1 is not None else None - residual_out = residual_out.reshape(x_shape_og) if residual_out is not None else None - dropout_mask = dropout_mask.reshape(x_shape_og) if dropout_mask is not None else None - dropout_mask1 = dropout_mask1.reshape(x_shape_og) if dropout_mask1 is not None else None - if not return_dropout_mask: - if weight1 is None: - return y if not prenorm else (y, residual_out) - else: - return (y, y1) if not prenorm else (y, y1, residual_out) - else: - if weight1 is None: - return ( - (y, dropout_mask, dropout_mask1) - if not prenorm - else (y, residual_out, dropout_mask, dropout_mask1) - ) - else: - return ( - (y, y1, dropout_mask, dropout_mask1) - if not prenorm - else (y, y1, residual_out, dropout_mask, dropout_mask1) - ) - - @staticmethod - def backward(ctx, dy, *args): - x, weight, bias, weight1, bias1, rowscale, seeds, mean, rstd = ctx.saved_tensors - dy = dy.reshape(-1, dy.shape[-1]) - if weight1 is not None: - dy1, args = args[0], args[1:] - dy1 = dy1.reshape(-1, dy1.shape[-1]) - assert dy1.shape == x.shape - else: - dy1 = None - if ctx.prenorm: - dresidual = args[0] - dresidual = dresidual.reshape(-1, dresidual.shape[-1]) - assert dresidual.shape == x.shape - else: - dresidual = None - dx, dw, db, dresidual_in, dx1, dw1, db1, _ = _layer_norm_bwd( - dy, - x, - weight, - bias, - ctx.eps, - mean, - rstd, - dresidual, - dy1, - weight1, - bias1, - seeds, - ctx.dropout_p, - rowscale, - ctx.has_residual, - ctx.has_x1, - ctx.zero_centered_weight, - ctx.is_rms_norm, - x_dtype=ctx.x_dtype, - recompute_output=False, - ) - return ( - dx.reshape(ctx.x_shape_og), - dw, - db, - dresidual_in.reshape(ctx.x_shape_og) if ctx.has_residual else None, - dx1.reshape(ctx.x_shape_og) if dx1 is not None else None, - dw1, - db1, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) - - -def layer_norm_fn( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - residual_in_fp32=False, - zero_centered_weight=False, - is_rms_norm=False, - return_dropout_mask=False, - out_dtype=None, - out=None, - residual_out=None -): - return LayerNormFn.apply( - x, - weight, - bias, - residual, - x1, - weight1, - bias1, - eps, - dropout_p, - rowscale, - prenorm, - residual_in_fp32, - zero_centered_weight, - is_rms_norm, - return_dropout_mask, - out_dtype, - out, - residual_out - ) - - -def rms_norm_fn( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - residual_in_fp32=False, - zero_centered_weight=False, - return_dropout_mask=False, - out_dtype=None, - out=None, - residual_out=None -): - return LayerNormFn.apply( - x, - weight, - bias, - residual, - x1, - weight1, - bias1, - eps, - dropout_p, - rowscale, - prenorm, - residual_in_fp32, - zero_centered_weight, - True, - return_dropout_mask, - out_dtype, - out, - residual_out - ) - - -class RMSNorm(torch.nn.Module): - - def __init__(self, hidden_size, eps=1e-5, dropout_p=0.0, zero_centered_weight=False, - device=None, dtype=None): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.eps = eps - if dropout_p > 0.0: - self.drop = torch.nn.Dropout(dropout_p) - else: - self.drop = None - self.zero_centered_weight = zero_centered_weight - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.register_parameter("bias", None) - self.reset_parameters() - - def reset_parameters(self): - if not self.zero_centered_weight: - torch.nn.init.ones_(self.weight) - else: - torch.nn.init.zeros_(self.weight) - - def forward(self, x, residual=None, prenorm=False, residual_in_fp32=False): - return rms_norm_fn( - x, - self.weight, - self.bias, - residual=residual, - eps=self.eps, - dropout_p=self.drop.p if self.drop is not None and self.training else 0.0, - prenorm=prenorm, - residual_in_fp32=residual_in_fp32, - zero_centered_weight=self.zero_centered_weight, - ) - - -class LayerNormLinearFn(torch.autograd.Function): - - @staticmethod - @custom_fwd - def forward( - ctx, - x, - norm_weight, - norm_bias, - linear_weight, - linear_bias, - residual=None, - eps=1e-6, - prenorm=False, - residual_in_fp32=False, - is_rms_norm=False, - ): - x_shape_og = x.shape - # reshape input data into 2D tensor - x = maybe_contiguous_lastdim(x.reshape(-1, x.shape[-1])) - if residual is not None: - assert residual.shape == x_shape_og - residual = maybe_contiguous_lastdim(residual.reshape(-1, residual.shape[-1])) - norm_weight = norm_weight.contiguous() - norm_bias = maybe_contiguous(norm_bias) - residual_dtype = ( - residual.dtype - if residual is not None - else (torch.float32 if residual_in_fp32 else None) - ) - y, _, mean, rstd, residual_out, *rest = _layer_norm_fwd( - x, - norm_weight, - norm_bias, - eps, - residual, - out_dtype=None if not torch.is_autocast_enabled() else torch.get_autocast_dtype("cuda"), - residual_dtype=residual_dtype, - is_rms_norm=is_rms_norm, - ) - y = y.reshape(x_shape_og) - dtype = torch.get_autocast_dtype("cuda") if torch.is_autocast_enabled() else y.dtype - linear_weight = linear_weight.to(dtype) - linear_bias = linear_bias.to(dtype) if linear_bias is not None else None - out = F.linear(y.to(linear_weight.dtype), linear_weight, linear_bias) - # We don't store y, will be recomputed in the backward pass to save memory - ctx.save_for_backward(residual_out, norm_weight, norm_bias, linear_weight, mean, rstd) - ctx.x_shape_og = x_shape_og - ctx.eps = eps - ctx.is_rms_norm = is_rms_norm - ctx.has_residual = residual is not None - ctx.prenorm = prenorm - ctx.x_dtype = x.dtype - ctx.linear_bias_is_none = linear_bias is None - return out if not prenorm else (out, residual_out.reshape(x_shape_og)) - - @staticmethod - @custom_bwd - def backward(ctx, dout, *args): - x, norm_weight, norm_bias, linear_weight, mean, rstd = ctx.saved_tensors - dout = dout.reshape(-1, dout.shape[-1]) - dy = F.linear(dout, linear_weight.t()) - dlinear_bias = None if ctx.linear_bias_is_none else dout.sum(0) - dy = maybe_contiguous_lastdim(dy) - assert dy.shape == x.shape - if ctx.prenorm: - dresidual = args[0] - dresidual = maybe_contiguous_lastdim(dresidual.reshape(-1, dresidual.shape[-1])) - assert dresidual.shape == x.shape - else: - dresidual = None - dx, dnorm_weight, dnorm_bias, dresidual_in, _, _, _, y = _layer_norm_bwd( - dy, - x, - norm_weight, - norm_bias, - ctx.eps, - mean, - rstd, - dresidual=dresidual, - has_residual=ctx.has_residual, - is_rms_norm=ctx.is_rms_norm, - x_dtype=ctx.x_dtype, - recompute_output=True, - ) - dlinear_weight = torch.einsum("bo,bi->oi", dout, y) - return ( - dx.reshape(ctx.x_shape_og), - dnorm_weight, - dnorm_bias, - dlinear_weight, - dlinear_bias, - dresidual_in.reshape(ctx.x_shape_og) if ctx.has_residual else None, - None, - None, - None, - None, - ) - - -def layer_norm_linear_fn( - x, - norm_weight, - norm_bias, - linear_weight, - linear_bias, - residual=None, - eps=1e-6, - prenorm=False, - residual_in_fp32=False, - is_rms_norm=False, -): - return LayerNormLinearFn.apply( - x, - norm_weight, - norm_bias, - linear_weight, - linear_bias, - residual, - eps, - prenorm, - residual_in_fp32, - is_rms_norm, - ) diff --git a/build/torch211-cxx11-cu126-x86_64-linux/__init__.py b/build/torch211-cxx11-cu126-x86_64-linux/__init__.py deleted file mode 100644 index ecc2f9d896b6c93f90b0a1499856dc0612177422..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu126-x86_64-linux/__init__.py +++ /dev/null @@ -1,393 +0,0 @@ -from typing import Optional, List -import torch -from ._ops import ops as flash_attn_ops -from .flash_attn_interface import ( - flash_attn_func, - flash_attn_kvpacked_func, - flash_attn_qkvpacked_func, - flash_attn_varlen_func, - flash_attn_varlen_kvpacked_func, - flash_attn_varlen_qkvpacked_func, - flash_attn_with_kvcache, -) - - -def fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - return_softmax: bool = False, - gen: Optional[torch.Generator] = None, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention. - - Args: - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Optional output tensor, same shape as q - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - return_softmax: Whether to return softmax weights - gen: Optional random number generator - - Returns: - List of tensors: [output, softmax_lse, (softmax if return_softmax)] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.fwd( - q, - k, - v, - out, - alibi_slopes, - p_dropout, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - gen, - ) - - -def varlen_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - out: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - block_table: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - max_seqlen_q: int = 0, - max_seqlen_k: int = 0, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - zero_tensors: bool = False, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - return_softmax: bool = False, - gen: Optional[torch.Generator] = None, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention with variable sequence lengths. - - Args: - q: Query tensor of shape [total_q, num_heads, head_size] - k: Key tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - v: Value tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1] - cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1] - out: Optional output tensor of shape [total_q, num_heads, head_size] - seqused_k: Optional tensor specifying how many keys to use per batch element [batch_size] - leftpad_k: Optional left padding for keys of shape [batch_size] - block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq] - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - max_seqlen_q: Maximum sequence length for queries - max_seqlen_k: Maximum sequence length for keys - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - zero_tensors: Whether to zero tensors before computation - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - return_softmax: Whether to return softmax weights - gen: Optional random number generator - - Returns: - List of tensors: [output, softmax_lse, (softmax if return_softmax)] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.varlen_fwd( - q, - k, - v, - out, - cu_seqlens_q, - cu_seqlens_k, - seqused_k, - leftpad_k, - block_table, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - p_dropout, - softmax_scale, - zero_tensors, - is_causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - gen, - ) - - -def bwd( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - p_dropout: float = 0.0, - 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, - gen: Optional[torch.Generator] = None, - rng_state: Optional[torch.Tensor] = None, -) -> List[torch.Tensor]: - """ - Backward pass for multi-head attention. - - Args: - dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size] - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size] - softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q] - dq: Optional gradient tensor for queries, same shape as q - dk: Optional gradient tensor for keys, same shape as k - dv: Optional gradient tensor for values, same shape as v - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - deterministic: Whether to use deterministic algorithms - gen: Optional random number generator - rng_state: Optional RNG state from forward pass - - Returns: - List of tensors: [dq, dk, dv] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - alibi_slopes, - p_dropout, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - deterministic, - gen, - rng_state, - ) - - -def varlen_bwd( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - max_seqlen_q: int = 0, - max_seqlen_k: int = 0, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - zero_tensors: bool = False, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - deterministic: bool = False, - gen: Optional[torch.Generator] = None, - rng_state: Optional[torch.Tensor] = None, -) -> List[torch.Tensor]: - """ - Backward pass for multi-head attention with variable sequence lengths. - - Args: - dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size] - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size] - softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q] - cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1] - cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1] - dq: Optional gradient tensor for queries, same shape as q - dk: Optional gradient tensor for keys, same shape as k - dv: Optional gradient tensor for values, same shape as v - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - max_seqlen_q: Maximum sequence length for queries - max_seqlen_k: Maximum sequence length for keys - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - zero_tensors: Whether to zero tensors before computation - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - deterministic: Whether to use deterministic algorithms - gen: Optional random number generator - rng_state: Optional RNG state from forward pass - - Returns: - List of tensors: [dq, dk, dv] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.varlen_bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - p_dropout, - softmax_scale, - zero_tensors, - is_causal, - window_size_left, - window_size_right, - softcap, - deterministic, - gen, - rng_state, - ) - - -def fwd_kvcache( - q: torch.Tensor, - kcache: torch.Tensor, - vcache: torch.Tensor, - k: Optional[torch.Tensor] = None, - v: Optional[torch.Tensor] = None, - seqlens_k: Optional[torch.Tensor] = None, - rotary_cos: Optional[torch.Tensor] = None, - rotary_sin: Optional[torch.Tensor] = None, - cache_batch_idx: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - block_table: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - out: 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, - is_rotary_interleaved: bool = False, - num_splits: int = 1, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention with KV cache. - - Args: - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - kcache: Key cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - vcache: Value cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - k: Optional new keys tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size] - v: Optional new values tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size] - seqlens_k: Optional sequence lengths for keys of shape [batch_size] - rotary_cos: Optional rotary cosine tensor of shape [seqlen_ro, rotary_dim/2] - rotary_sin: Optional rotary sine tensor of shape [seqlen_ro, rotary_dim/2] - cache_batch_idx: Optional indices to index into the KV cache - leftpad_k: Optional left padding for keys of shape [batch_size] - block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq] - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - out: Optional output tensor, same shape as q - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - is_rotary_interleaved: Whether rotary embeddings are interleaved - num_splits: Number of splits for computation - - Returns: - List of tensors: [output, softmax_lse] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.fwd_kvcache( - q, - kcache, - vcache, - k, - v, - seqlens_k, - rotary_cos, - rotary_sin, - cache_batch_idx, - leftpad_k, - block_table, - alibi_slopes, - out, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - is_rotary_interleaved, - num_splits, - ) diff --git a/build/torch211-cxx11-cu126-x86_64-linux/_flash_attn2_cuda_f12afc9.abi3.so b/build/torch211-cxx11-cu126-x86_64-linux/_flash_attn2_cuda_f12afc9.abi3.so deleted file mode 100644 index 2b110c91cc80fa162f569eb8508e97e927258b1d..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu126-x86_64-linux/_flash_attn2_cuda_f12afc9.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:145438f01c40b69a1ad1fc6443fb43dd0f145551a69d8f62ad2cd8330562e7b0 -size 446828600 diff --git a/build/torch211-cxx11-cu126-x86_64-linux/_ops.py b/build/torch211-cxx11-cu126-x86_64-linux/_ops.py deleted file mode 100644 index af5b731cb04b5e6ce54e862db0501f3683c3163f..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu126-x86_64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _flash_attn2_cuda_f12afc9 -ops = torch.ops._flash_attn2_cuda_f12afc9 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_flash_attn2_cuda_f12afc9::{op_name}" diff --git a/build/torch211-cxx11-cu126-x86_64-linux/bert_padding.py b/build/torch211-cxx11-cu126-x86_64-linux/bert_padding.py deleted file mode 100644 index 3c2d35159a014a9d03aabead9e52e009168696ea..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu126-x86_64-linux/bert_padding.py +++ /dev/null @@ -1,218 +0,0 @@ -# Adapted from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/padding.py - -import torch -import torch.nn.functional as F -from einops import rearrange, repeat - - -class IndexFirstAxis(torch.autograd.Function): - @staticmethod - def forward(ctx, input, indices): - ctx.save_for_backward(indices) - assert input.ndim >= 2 - ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] - second_dim = other_shape.numel() - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - # return input[indices] - return torch.gather( - rearrange(input, "b ... -> b (...)"), 0, repeat(indices, "z -> z d", d=second_dim) - ).reshape(-1, *other_shape) - - @staticmethod - def backward(ctx, grad_output): - (indices,) = ctx.saved_tensors - assert grad_output.ndim >= 2 - other_shape = grad_output.shape[1:] - grad_output = rearrange(grad_output, "b ... -> b (...)") - grad_input = torch.zeros( - [ctx.first_axis_dim, grad_output.shape[1]], - device=grad_output.device, - dtype=grad_output.dtype, - ) - # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. - # grad_input[indices] = grad_output - grad_input.scatter_(0, repeat(indices, "z -> z d", d=grad_output.shape[1]), grad_output) - return grad_input.reshape(ctx.first_axis_dim, *other_shape), None - - -index_first_axis = IndexFirstAxis.apply - - -class IndexPutFirstAxis(torch.autograd.Function): - @staticmethod - def forward(ctx, values, indices, first_axis_dim): - ctx.save_for_backward(indices) - assert indices.ndim == 1 - assert values.ndim >= 2 - output = torch.zeros( - first_axis_dim, *values.shape[1:], device=values.device, dtype=values.dtype - ) - # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. - output[indices] = values - # output.scatter_(0, repeat(indices, 'z -> z d', d=values.shape[1]), values) - return output - - @staticmethod - def backward(ctx, grad_output): - (indices,) = ctx.saved_tensors - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - grad_values = grad_output[indices] - # grad_values = torch.gather(grad_output, 0, repeat(indices, 'z -> z d', d=grad_output.shape[1])) - return grad_values, None, None - - -index_put_first_axis = IndexPutFirstAxis.apply - - -class IndexFirstAxisResidual(torch.autograd.Function): - @staticmethod - def forward(ctx, input, indices): - ctx.save_for_backward(indices) - assert input.ndim >= 2 - ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] - second_dim = other_shape.numel() - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - output = input[indices] - # We don't want to reshape input (b ... -> b (...)) since it could change the channel_last - # memory format to channel_first. In other words, input might not be contiguous. - # If we don't detach, Pytorch complains about output being a view and is being modified inplace - return output, input.detach() - - @staticmethod - def backward(ctx, grad_output, grad_residual): - (indices,) = ctx.saved_tensors - assert grad_output.ndim >= 2 - other_shape = grad_output.shape[1:] - assert grad_residual.shape[1:] == other_shape - grad_input = grad_residual - # grad_input[indices] += grad_output - indices = indices.reshape(indices.shape[0], *((1,) * (grad_output.ndim - 1))) - indices = indices.expand_as(grad_output) - grad_input.scatter_add_(0, indices, grad_output) - return grad_input.reshape(ctx.first_axis_dim, *other_shape), None - - -index_first_axis_residual = IndexFirstAxisResidual.apply - - -def unpad_input(hidden_states, attention_mask, unused_mask=None): - """ - Arguments: - hidden_states: (batch, seqlen, ...) - attention_mask: (batch, seqlen), bool / int, 1 means valid and 0 means not valid. - unused_mask: (batch, seqlen), bool / int, 1 means the element is allocated but unused. - Return: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens selected in attention_mask + unused_mask. - indices: (total_nnz), the indices of masked tokens from the flattened input sequence. - cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. - max_seqlen_in_batch: int - seqused: (batch), returns the number of tokens selected in attention_mask + unused_mask. - """ - all_masks = (attention_mask + unused_mask) if unused_mask is not None else attention_mask - seqlens_in_batch = all_masks.sum(dim=-1, dtype=torch.int32) - used_seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) - indices = torch.nonzero(all_masks.flatten(), as_tuple=False).flatten() - max_seqlen_in_batch = seqlens_in_batch.max().item() - cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) - # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the - # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim - # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to - # index with integer indices. Moreover, torch's index is a bit slower than it needs to be, - # so we write custom forward and backward to make it a bit faster. - return ( - index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices), - indices, - cu_seqlens, - max_seqlen_in_batch, - used_seqlens_in_batch, - ) - - -def unpad_input_for_concatenated_sequences(hidden_states, attention_mask_in_length): - """ - Supports concatenating short samples in one sequence. The attention_mask_in_length is utilized to mask other short samples. It helps efficient training of variant lengths-based samples (e.g., the supervised fine-tuning task in large language model). - The motivation for this function is explained [here](https://github.com/Dao-AILab/flash-attention/issues/432#issuecomment-1668822286). - - For example, if batch = 3 and seqlen = 6, the attention_mask_in_length is: - ``` - [ - [2, 3, 0, 0, 0, 0], - [3, 2, 0, 0, 0, 0], - [6, 0, 0, 0, 0, 0] - ] - ``` - , which refers to the 3D-attention mask: - ``` - [ - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [0, 0, 1, 0, 0, 0], - [0, 0, 1, 1, 0, 0], - [0, 0, 1, 1, 1, 0], - [0, 0, 0, 0, 0, 1] - ], - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [0, 0, 0, 1, 0, 0], - [0, 0, 0, 1, 1, 0], - [0, 0, 0, 0, 0, 1] - ], - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [1, 1, 1, 1, 0, 0], - [1, 1, 1, 1, 1, 0], - [1, 1, 1, 1, 1, 1] - ] - ] - ```. - - Arguments: - hidden_states: (batch, seqlen, ...) - attention_mask_in_length: (batch, seqlen), int, a nonzero number (e.g., 1, 2, 3, etc.) means length of concatenated sequence in b-th batch, and 0 means none. - Return: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. - indices: (total_nnz), the indices of non-masked tokens from the flattened input sequence. - cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. - max_seqlen_in_batch: int - """ - length = attention_mask_in_length.sum(dim=-1) - seqlen = attention_mask_in_length.size(-1) - attention_mask_2d = torch.arange(seqlen, device=length.device, dtype=length.dtype).expand(len(length), seqlen) < length.unsqueeze(1) - real_indices_idx = torch.nonzero(attention_mask_in_length.flatten(), as_tuple=False).flatten() - seqlens_in_batch = attention_mask_in_length.flatten()[real_indices_idx] - indices = torch.nonzero(attention_mask_2d.flatten(), as_tuple=False).flatten() - max_seqlen_in_batch = seqlens_in_batch.max().item() - cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) - # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the - # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim - # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to - # index with integer indices. Moreover, torch's index is a bit slower than it needs to be, - # so we write custom forward and backward to make it a bit faster. - return ( - index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices), - indices, - cu_seqlens, - max_seqlen_in_batch, - ) - - -def pad_input(hidden_states, indices, batch, seqlen): - """ - Arguments: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. - indices: (total_nnz), the indices that represent the non-masked tokens of the original padded input sequence. - batch: int, batch size for the padded sequence. - seqlen: int, maximum sequence length for the padded sequence. - Return: - hidden_states: (batch, seqlen, ...) - """ - dim = hidden_states.shape[-1] - # output = torch.zeros((batch * seqlen), dim, device=hidden_states.device, dtype=hidden_states.dtype) - # output[indices] = hidden_states - output = index_put_first_axis(hidden_states, indices, batch * seqlen) - return rearrange(output, "(b s) ... -> b s ...", b=batch) diff --git a/build/torch211-cxx11-cu126-x86_64-linux/flash_attn2/__init__.py b/build/torch211-cxx11-cu126-x86_64-linux/flash_attn2/__init__.py deleted file mode 100644 index a9b2672c1cd85b74c1b3ded0fc0b2100e1aeac23..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu126-x86_64-linux/flash_attn2/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -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-cu126-x86_64-linux/flash_attn_interface.py b/build/torch211-cxx11-cu126-x86_64-linux/flash_attn_interface.py deleted file mode 100644 index b014e569dc9812a9c840d1043f2cebe767b4d477..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu126-x86_64-linux/flash_attn_interface.py +++ /dev/null @@ -1,1631 +0,0 @@ -# Copyright (c) 2023, Tri Dao. - -from typing import Optional, Sequence, Tuple, Union - -import torch -import torch.nn as nn -import os - -# # isort: off -# # We need to import the CUDA kernels after importing torch -# USE_TRITON_ROCM = os.getenv("FLASH_ATTENTION_TRITON_AMD_ENABLE", "FALSE") == "TRUE" -# if USE_TRITON_ROCM: -# from .flash_attn_triton_amd import interface_fa as flash_attn -# else: -# import flash_attn_2_cuda as flash_attn - - -from . import _ops -from ._ops import ops as flash_attn -from ._ops import add_op_namespace_prefix - -# # isort: on - -def maybe_contiguous(x): - return x.contiguous() if x is not None and x.stride(-1) != 1 else x - - -def _get_device(): - if torch.xpu.is_available(): - return "xpu" - elif torch.cuda.is_available(): - return "cuda" - else: - return "cpu" - - -def _get_block_size_n(device, head_dim, is_dropout, is_causal): - assert head_dim <= 256 - - if device.type == "xpu": - if head_dim <= 96: - return 64 - elif head_dim <= 128: - return 32 - elif head_dim <= 256: - return 64 - else: - return 32 - - # This should match the block sizes in the CUDA kernel - major, minor = torch.cuda.get_device_capability(device) - is_sm8x = major == 8 and minor > 0 # Only include sm86 and sm89, exclude sm80 (A100) - is_sm80 = major == 8 and minor == 0 - is_sm90 = major == 9 and minor == 0 - if head_dim <= 32: - return 128 - if head_dim <= 64: - return 128 if not is_dropout else 64 - elif head_dim <= 96: - return 64 - elif head_dim <= 128: - if is_sm8x: - return 64 if (not is_dropout and is_causal) else 32 - else: - return 64 if not is_dropout else 32 - elif head_dim <= 192: - return 64 - elif head_dim <= 224: - return 64 - elif head_dim <= 256: - return 64 - - -def round_multiple(x, m): - return (x + m - 1) // m * m - - -# torch.compile() support is only enabled for pytorch >= 2.4 -# The reason for this is that we are using the new custom_op and register_fake -# APIs, which support inplace modification of inputs in the function itself -if torch.__version__ >= "2.4.0": - _torch_custom_op_wrapper = torch.library.custom_op - _torch_register_fake_wrapper = torch.library.register_fake -else: - def noop_custom_op_wrapper(name, fn=None, /, *, mutates_args, device_types=None, schema=None): - def wrap(func): - return func - if fn is None: - return wrap - return fn - def noop_register_fake_wrapper(op, fn=None, /, *, lib=None, _stacklevel=1): - def wrap(func): - return func - if fn is None: - return wrap - return fn - _torch_custom_op_wrapper = noop_custom_op_wrapper - _torch_register_fake_wrapper = noop_register_fake_wrapper - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_forward"), mutates_args=(), device_types=_get_device()) -def _flash_attn_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - return_softmax: bool -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - out, softmax_lse, S_dmask, rng_state = flash_attn.fwd( - q, - k, - v, - None, - alibi_slopes, - dropout_p, - softmax_scale, - causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - None, - ) - return out, softmax_lse, S_dmask, rng_state - - -@_torch_register_fake_wrapper(add_op_namespace_prefix("_flash_attn_forward")) -def _flash_attn_forward_fake( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - return_softmax: bool -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - batch_size, seqlen_q, num_heads, head_size = q.shape - seqlen_k = k.shape[1] - out = torch.empty_like(q) - softmax_lse = torch.empty((batch_size, num_heads, seqlen_q), dtype=torch.float32, device=q.device, layout=q.layout) - p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) - if return_softmax: - p = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128), round_multiple(seqlen_k, 128)), dtype=q.dtype, device=q.device, layout=q.layout) - rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) - - return out, softmax_lse, p, rng_state - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_forward = _ops.ops._flash_attn_forward -else: - _wrapped_flash_attn_forward = _flash_attn_forward - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_varlen_forward"), mutates_args=(), device_types=_get_device()) -def _flash_attn_varlen_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - alibi_slopes: Optional[torch.Tensor] = None, - return_softmax: bool = False, - block_table: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - out, softmax_lse, S_dmask, rng_state = flash_attn.varlen_fwd( - q, - k, - v, - None, - cu_seqlens_q, - cu_seqlens_k, - seqused_k, - leftpad_k, - block_table, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - zero_tensors, - causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - None, - ) - # if out.isnan().any() or softmax_lse.isnan().any(): - # breakpoint() - return out, softmax_lse, S_dmask, rng_state - - -@_torch_register_fake_wrapper(add_op_namespace_prefix("_flash_attn_varlen_forward")) -def _flash_attn_varlen_forward_fake( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - alibi_slopes: Optional[torch.Tensor] = None, - return_softmax: bool = False, - block_table: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - paged_kv = block_table is not None - batch_size = cu_seqlens_q.numel() - 1 - total_q, num_heads, _ = q.shape - - out = torch.empty_like(q) - softmax_lse = torch.empty((num_heads, total_q), dtype=torch.float32, device=q.device, layout=q.layout) - p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) - seqlen_q_rounded = round_multiple(max_seqlen_q, 128) - seqlen_k_rounded = round_multiple(max_seqlen_k, 128) - if return_softmax: - p = torch.empty((batch_size, num_heads, seqlen_q_rounded, seqlen_k_rounded), dtype=q.dtype, device=q.device, layout=q.layout) - rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) - return out, softmax_lse, p, rng_state - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_varlen_forward = _ops.ops._flash_attn_varlen_forward -else: - _wrapped_flash_attn_varlen_forward = _flash_attn_varlen_forward - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_backward"), mutates_args=("dq", "dk", "dv"), device_types=_get_device()) -def _flash_attn_backward( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, -) -> 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)] - ( - dq, - dk, - dv, - softmax_d, - ) = flash_attn.bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - alibi_slopes, - dropout_p, - softmax_scale, - causal, - window_size_left, - window_size_right, - softcap, - deterministic, - None, - rng_state, - ) - return softmax_d - - -@_torch_register_fake_wrapper(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, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, -) -> torch.Tensor: - dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] - if dq is None: - dq = torch.empty_like(q) - if dk is None: - dk = torch.empty_like(k) - if dv is None: - dv = torch.empty_like(v) - batch_size, seqlen_q, num_heads, _ = q.shape - softmax_d = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128)), device=q.device, dtype=torch.float32) - - return softmax_d - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_backward = _ops.ops._flash_attn_backward -else: - _wrapped_flash_attn_backward = _flash_attn_backward - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_varlen_backward"), mutates_args=("dq", "dk", "dv"), device_types=_get_device()) -def _flash_attn_varlen_backward( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> 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)] - ( - dq, - dk, - dv, - softmax_d, - ) = flash_attn.varlen_bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - zero_tensors, - causal, - window_size_left, - window_size_right, - softcap, - deterministic, - None, - rng_state, - ) - # if dk.isnan().any() or dk.isnan().any() or dv.isnan().any() or softmax_d.isnan().any(): - # breakpoint() - return softmax_d - - -@_torch_register_fake_wrapper(add_op_namespace_prefix("_flash_attn_varlen_backward")) -def _flash_attn_varlen_backward_fake( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> torch.Tensor: - dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] - batch_size = cu_seqlens_q.numel() - 1 - total_q, num_heads, _ = q.shape - - if dq is None: - dq = torch.empty_like(q) - if dk is None: - dk = torch.empty_like(k) - if dv is None: - dv = torch.empty_like(v) - softmax_d = torch.empty((num_heads, total_q + 128 * batch_size), device=q.device, dtype=torch.float32) - - return softmax_d - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_varlen_backward = _ops.ops._flash_attn_varlen_backward -else: - _wrapped_flash_attn_varlen_backward = _flash_attn_varlen_backward - - -class FlashAttnQKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and qkv.requires_grad - if softmax_scale is None: - softmax_scale = qkv.shape[-1] ** (-0.5) - q, k, v = qkv[:, :, 0].detach(), qkv[:, :, 1].detach(), qkv[:, :, 2].detach() - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) - dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dqkv[:, :, 0], - dqkv[:, :, 1], - dqkv[:, :, 2], - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dqkv = dqkv[..., : dout.shape[-1]] # We could have padded the head dimension - return dqkv, None, None, None, None, None, None, None, None, None - - -class FlashAttnVarlenQKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - cu_seqlens, - max_seqlen, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and qkv.requires_grad - if softmax_scale is None: - softmax_scale = qkv.shape[-1] ** (-0.5) - q, k, v = qkv[:, 0].detach(), qkv[:, 1].detach(), qkv[:, 2].detach() - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens, - cu_seqlens, - max_seqlen, - max_seqlen, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=None, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens, rng_state) - ctx.dropout_p = dropout_p - ctx.max_seqlen = max_seqlen - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens, rng_state = ctx.saved_tensors - qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) - dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dqkv[:, 0], - dqkv[:, 1], - dqkv[:, 2], - cu_seqlens, - cu_seqlens, - ctx.max_seqlen, - ctx.max_seqlen, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - 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 FlashAttnKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - kv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, kv] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - k, v = kv[:, :, 0].detach(), kv[:, :, 1].detach() - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - dq = torch.empty_like(q) - kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) - dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dkv[:, :, 0], - dkv[:, :, 1], - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dkv = dkv[..., : dout.shape[-1]] - return dq, dkv, None, None, None, None, None, None, None, None, None - - -class FlashAttnVarlenKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, kv] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - k, v = kv[:, 0].detach(), kv[:, 1].detach() - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=None, - ) - if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state - ) - ctx.dropout_p = dropout_p - 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.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors - dq = torch.empty_like(q) - kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) - dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dkv[:, 0], - dkv[:, 1], - cu_seqlens_q, - cu_seqlens_k, - ctx.max_seqlen_q, - ctx.max_seqlen_k, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dkv = dkv[..., : dout.shape[-1]] - return dq, dkv, None, None, None, None, None, None, None, None, None, None, None, None, None - - -class FlashAttnFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - k, - v, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, k, v] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dk = dk[..., : dout.shape[-1]] - dv = dv[..., : dout.shape[-1]] - return dq, dk, dv, 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, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - block_table, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, k, v] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=block_table, - ) - if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state - ) - ctx.dropout_p = dropout_p - 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.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors - dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - ctx.max_seqlen_q, - ctx.max_seqlen_k, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dk = dk[..., : dout.shape[-1]] - dv = dv[..., : dout.shape[-1]] - return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None - - -def flash_attn_qkvpacked_func( - qkv, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # <=0.0 means deactivate - alibi_slopes=None, - deterministic=False, - 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, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_kvpacked_func( - q, - kv, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, -): - """dropout_p should be set to 0.0 during evaluation - If 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 K, V. - 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) - kv: (batch_size, seqlen, 2, 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. - softcap: float. Anything > 0 activates softcapping 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). - 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 FlashAttnKVPackedFunc.apply( - q, - kv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_func( - q, - k, - v, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - 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). - 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 FlashAttnFunc.apply( - q, - k, - v, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_qkvpacked_func( - qkv, - cu_seqlens, - max_seqlen, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - 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_varlen_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_varlen_kvpacked_func and flash_attn_varlen_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: (total, 3, nheads, headdim), where total = total number of tokens in the batch. - cu_seqlens: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into qkv. - max_seqlen: int. Maximum sequence length in the batch. - 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenQKVPackedFunc.apply( - qkv, - cu_seqlens, - max_seqlen, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_kvpacked_func( - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, -): - """dropout_p should be set to 0.0 during evaluation - If 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 K, V. - 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: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch. - kv: (total_k, 2, nheads_k, headdim), where total_k = total number of key tokens in the batch. - cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into q. - cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into kv. - max_seqlen_q: int. Maximum query sequence length in the batch. - max_seqlen_k: int. Maximum key sequence length in the batch. - 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 + 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenKVPackedFunc.apply( - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_func( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, - block_table=None, -): - """dropout_p should be set to 0.0 during evaluation - Supports multi-query and grouped-query attention (MQA/GQA) by passing in K, V 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: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch. - k: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. - v: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. - cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into q. - cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into kv. - max_seqlen_q: int. Maximum query sequence length in the batch. - max_seqlen_k: int. Maximum key sequence length in the batch. - 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 + 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenFunc.apply( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - block_table, - False if q.device.type == "cpu" else torch.is_grad_enabled(), - ) - - -def flash_attn_with_kvcache( - q, - k_cache, - v_cache, - k=None, - v=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, - block_table: Optional[torch.Tensor] = None, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - rotary_interleaved=True, - alibi_slopes=None, - num_splits=0, - 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 block_table, - or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_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) if there's no block_table, - or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_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). Similar to k. - 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. - block_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). - 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. - 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" - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-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) - cache_batch_idx = maybe_contiguous(cache_batch_idx) - block_table = maybe_contiguous(block_table) - out, softmax_lse = flash_attn.fwd_kvcache( - q, - k_cache, - v_cache, - k, - v, - cache_seqlens, - rotary_cos, - rotary_sin, - cache_batch_idx, - cache_leftpad, - block_table, - alibi_slopes, - None, - softmax_scale, - causal, - window_size[0], - window_size[1], - softcap, - rotary_interleaved, - num_splits, - ) - return (out, softmax_lse) if return_softmax_lse else out diff --git a/build/torch211-cxx11-cu126-x86_64-linux/metadata.json b/build/torch211-cxx11-cu126-x86_64-linux/metadata.json deleted file mode 100644 index ae7ea0f6e36103b24761d1e9665f73f0e75e6b68..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu126-x86_64-linux/metadata.json +++ /dev/null @@ -1,40 +0,0 @@ -{ - "name": "flash-attn2", - "id": "_flash_attn2_cuda_f12afc9", - "version": 1, - "license": "BSD-3-Clause", - "python-depends": [], - "backend": { - "type": "cuda", - "archs": [ - "8.0", - "9.0" - ] - }, - "digest": { - "algorithm": "sha256", - "files": { - "__init__.py": "Cd9MEdd/WUYnBmfiQpC4SyPqAxrTilTjvsvswNju4zc=", - "_flash_attn2_cuda_f12afc9.abi3.so": "FFQ48BxAtpoa0fxkQ/tD3Q8UVVGmnY9irSzYMwVi57A=", - "_ops.py": "GQbmTTeSOU+7Q2nI68Hz+1RvEogyfEFccx1ubUYw+2s=", - "bert_padding.py": "gF1EmsdJ+HpQ86MRQ4VxDw+Sb/RVISdQALdNnoByHlw=", - "flash_attn2/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY=", - "flash_attn_interface.py": "W7a9bHR2ly6QL3aKqsBqxWA7M7ZWYcZPmtsPJu6rAfY=", - "layers/__init__.py": "47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU=", - "layers/patch_embed.py": "H58CgME/qSOPTZLOG08wFgrQS1j34pvNwMPrkTj3Ek4=", - "layers/rotary.py": "rFIdNrZQFNawwLD6d7t5LE2m+M93wwDZ8EuExw+gfD8=", - "ops/__init__.py": "47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU=", - "ops/activations.py": "t5lzNg1In8LP6bKeTnyeMizwqjv27JGbJ6ylPdGvZYg=", - "ops/fused_dense.py": "BqT8VTSCGdLJAjh6WBMkCRFYrdVodPeoElWBvH6AIzA=", - "ops/layer_norm.py": "zr7NXIm+2mtEynTp1CS0fbFGI2Mqdp41dY4AfDWF6EQ=", - "ops/rms_norm.py": "XEnihcj0a4aSz4LO55m5iKGVn4HKTeKN8TIyHjuDgxI=", - "ops/triton/__init__.py": "AbpHGcgLb+kRsJGnwFEktk7uzpZOCcBY74+YBdrKVGs=", - "ops/triton/cross_entropy.py": "y56nmuOtpSNxM7sDZcUSgXKIsIEayfuvrU8bAVog/nw=", - "ops/triton/k_activations.py": "+Z3vIyO4JkqBMipKsPvhzmxljtBdIhJCsl/M+/ESqBo=", - "ops/triton/layer_norm.py": "BBQdJiqopbIS12viUqzx3SvOQCNbDSAYLJXhpyMGyC8=", - "ops/triton/linear.py": "OtRvKz8xdpl+7v3q/ZTaS9fdBt9XrzMyapgRr50uBbM=", - "ops/triton/mlp.py": "uWA+D7HyuRUsOT+4nhi3vvCDNZnNEXwYcHMQ02UREfQ=", - "ops/triton/rotary.py": "0YSrlbep0lJmzLpTiANs0slTFeGcQQqR+el6uzbSLvI=" - } - } -} \ No newline at end of file diff --git a/build/torch211-cxx11-cu126-x86_64-linux/metadata.json.sigstore b/build/torch211-cxx11-cu126-x86_64-linux/metadata.json.sigstore deleted file mode 100644 index 5da946757f9a555be05accaf654fe333c129b54c..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu126-x86_64-linux/metadata.json.sigstore +++ /dev/null @@ -1 +0,0 @@ -{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json","verificationMaterial":{"certificate":{"rawBytes":"MIIHSzCCBtGgAwIBAgIUaQWJufQzcYs3v3N9gu34+DhDyH8wCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwNjI0MDEwMTIwWhcNMjYwNjI0MDExMTIwWjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEVMLzVqV3ivFoz+15m/x4yzRM2bwMZqAQzyaK6KYrsXD+/tqLB9Z5fMMN4mHPiOpq7mSVTOQcMymwwhI+WfSySaOCBfAwggXsMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQU2ruSHlw0MRMKYNKFavCe53H3ACswHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoZjEyYWZjOTJjNDBiYmQ3NmQ5OTBkOTExZWJhYmYyOGU0OTBhN2JkZTATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoZjEyYWZjOTJjNDBiYmQ3NmQ5OTBkOTExZWJhYmYyOGU0OTBhN2JkZTAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoZjEyYWZjOTJjNDBiYmQ3NmQ5OTBkOTExZWJhYmYyOGU0OTBhN2JkZTAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKGYxMmFmYzkyYzQwYmJkNzZkOTkwZDkxMWViYWJmMjhlNDkwYTdiZGUwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMjgwMzYyNjg5MzUvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBigYKKwYBBAHWeQIEAgR8BHoAeAB2AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABnvclV1AAAAQDAEcwRQIgIj88pdI0FaCPtiWCiQQwFxhCbgTPmhXqy9uBltyn+2YCIQDephoDPXIZLuszV0yccHUUbpkalfhG7f3+AOwjASML/TAKBggqhkjOPQQDAwNoADBlAjEAp6YKROZCJ60UC8Qwwrt9ekVudi2f7PNNCo3kXXpwmfyVt6PG5Gvbo5T55uNN9e+yAjATx4ooXodX56XK+d7KEcS92EP3bRtE5iUbnZzQ85Hg8FmJpNxDTzY2GRSagGhcwJ0="},"tlogEntries":[{"logIndex":"1933771847","logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="},"kindVersion":{"kind":"hashedrekord","version":"0.0.1"},"integratedTime":"1782262880","inclusionPromise":{"signedEntryTimestamp":"MEUCIEKR9pavBq0szL6jciHAgFHy6xWR01yhImUi8f/2fIeCAiEA2EGnAa616S19fwPkaBYetng9jjZ0duxVJA5sTn9zlI0="},"inclusionProof":{"logIndex":"1811867585","rootHash":"A7LwdeN/PdsoLGv8ygMMX8PeZHmMA/VqEInF0jKDknc=","treeSize":"1811867598","hashes":["J86cPjNstsGF7OR5TieCq58R7jvyRJlKx9Tt2IZnu/M=","6gK/ULS5tWhCl6wk3C41yokjvZ3+Lcl3RdqEsMWzf9c=","Id76JvKRbMe9+PkS70G0Y6XMOkQsyevj3o0veOVK808=","xCnh984Dnwub6LF3Qo3kuVr8wu7UNtHwZYs4UCquews=","RZ1nRLsOnSpgIrAv4uhDsq98VpFShdW8MN+R/JeLIaA=","/2NUs/oaWZiEqet9OIHRFeZyehq7AAMr58s1kVJ2Dpk=","uIoOwNpXy89eRCV104FzEaFwVnFw0a20emg3UDglyx0=","2dEHLecK0QsJnHeF6+nozowEff01gRId5Wiy55DLFwc=","n8rRW5qWPQUx8euARGGwdNeIfcl6p6rcgzXRG0/79n0=","mzp9m+2At7GZdgfx0m7RwtO9YzltCZ6UhZBy4p+lUdM=","ipWWLDkqay5JAYe6P4Nrqsck/1C0C5uwqT9QbFsiY1E=","COKvK4TO+dkpqUDGbTbSKewajXzAonx1B1YaRb/ymLw=","dQlR+DXGDKXm5V1fupxaMSOwNZ9F1MybXyyxB96yj80=","JKJdxSZwKPO6QFzLTxHCrUyI5sDU/CwY57zDlA9ni6Q=","OgfnQCh7xr8xUTErIw//YAKaKMDwRJWiThvYExQvLJM=","S8Oe9ZicoJqbVuL1I0dIdTgVpLc82Cm2CMOcV83sKVs=","E1O2ilWlpFRLWGdbmJUMDZ2qbGu/eDZ12m6YjQ9/JPA=","nyv2TDlfkz0XfRwjY98zsbUbOkBZ+Tgz7PML0c2uN14=","OMfDh4k9YD0E+t4sjqJsh+vg8D1ueo1cC07exyxmr6s=","+vHKm4wNwGJ46ceyZ8mePD7Vn1SediStNHCStRoTSWQ=","5HwWjuGNDisetbiWM0cWIXM65TIM3EJlqk13ntCQlcU=","NDIRlqxO0IaGr94q86MIU7VCAN3yEEgOHS4loya5lu4=","daxmZaajRpZV+JxHiOYZhJBiSKN5ucqjh2WnGbHhirw=","DOCeoSMovIvLExkhIvisow9AuNXgeWs4ECkyR6EcqYU="],"checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n1811867598\nA7LwdeN/PdsoLGv8ygMMX8PeZHmMA/VqEInF0jKDknc=\n\n— rekor.sigstore.dev wNI9ajBEAiBhOvac3BBrbBZqQk8cKfJFZ4pYFh1Znd+VxGZdm8B8hgIgLzhyoM5WfjTLWC9Oi3fdeXA9eJdrk+DrFr7G+dXdsuM=\n"}},"canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiJlY2JkNGMxMDE3ZjQwMGI5ZTdkOTExZTkxMjZkNDg1ZTNhMzNlZjQ4MzQxM2NlMDJmYTE2YWQ3ZTE2NWU4NzBiIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJRkd3MFlIQkZsOG9OSVZUVVVQQ1lIRG14WDNYclBIb3BzVFBRZVZHaklDeEFpRUE0VmJMdFVuVUZ3Y2tYR3BZbHV5ZVdwOFZHR2VQNUFKLzJ1cll2ekowZFlrPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRla05EUW5SSFowRjNTVUpCWjBsVllWRlhTblZtVVhwaldYTXpkak5PT1dkMU16UXJSR2hFZVVnNGQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDVxU1RCTlJFVjNUVlJKZDFkb1kwNU5hbGwzVG1wSk1FMUVSWGhOVkVsM1YycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZXVFV4NlZuRldNMmwyUm05Nkt6RTFiUzk0TkhsNlVrMHlZbmROV25GQlVYcDVZVXNLTmt0WmNuTllSQ3N2ZEhGTVFqbGFOV1pOVFU0MGJVaFFhVTl3Y1RkdFUxWlVUMUZqVFhsdGQzZG9TU3RYWmxONVUyRlBRMEptUVhkbloxaHpUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlV5Y25WVENraHNkekJOVWsxTFdVNUxSbUYyUTJVMU0wZ3pRVU56ZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOWFha1Y1V1ZkYWFrOVVTbXBPUkVKcFdXMVJNMDV0VVRWUFZFSnJDazlVUlhoYVYwcG9XVzFaZVU5SFZUQlBWRUpvVGpKS2ExcFVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMXBxUlhsWlYxcHFUMVJLYWs1RVFtbFpiVkV6VG0xUk5VOVVRbXRQVkVWNFdsZEthRmx0V1hrS1QwZFZNRTlVUW1oT01rcHJXbFJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMXBxUlhsWlYxcHFUMVJLYWs1RVFtbFpiVkV6VG0xUk5VOVVRbXNLVDFSRmVGcFhTbWhaYlZsNVQwZFZNRTlVUW1oT01rcHJXbFJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwZFplRTF0Um0wS1dYcHJlVmw2VVhkWmJVcHJUbnBhYTA5VWEzZGFSR3Q0VFZkV2FWbFhTbTFOYW1oc1RrUnJkMWxVWkdsYVIxVjNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOYW1kM1RYcFplVTVxWnpWTmVsVjJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwWjFsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTT0VKSWIwRUtaVUZDTWtGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtNTJZMnhXTVVGQlFVRlJSQXBCUldOM1VsRkpaMGxxT0Rod1pFa3dSbUZEVUhScFYwTnBVVkYzUm5ob1EySm5WRkJ0YUZoeGVUbDFRbXgwZVc0ck1sbERTVkZFWlhCb2IwUlFXRWxhQ2t4MWMzcFdNSGxqWTBoVlZXSndhMkZzWm1oSE4yWXpLMEZQZDJwQlUwMU1MMVJCUzBKblozRm9hMnBQVUZGUlJFRjNUbTlCUkVKc1FXcEZRWEEyV1VzS1VrOWFRMG8yTUZWRE9GRjNkM0owT1dWclZuVmthVEptTjFCT1RrTnZNMnRZV0hCM2JXWjVWblEyVUVjMVIzWmlielZVTlRWMVRrNDVaU3Q1UVdwQlZBcDRORzl2V0c5a1dEVTJXRXNyWkRkTFJXTlRPVEpGVUROaVVuUkZOV2xWWW01YWVsRTROVWhuT0VadFNuQk9lRVJVZWxreVIxSlRZV2RIYUdOM1NqQTlDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}],"timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyTADAgEAMIICwAYJKoZIhvcNAQcCoIICsTCCAq0CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgpxHZf6eNQMk48zw9nJ0ZulPZE40KFR4jbAzyI9MN1jICFQChSHT94Auyz5tLLH8fMFx+K14BQRgPMjAyNjA2MjQwMTAxMjBaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHaMIIB1gIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MDYyNDAxMDEyMFowLwYJKoZIhvcNAQkEMSIEIOAUa9KDcyYkFvqcFX194azllKRCUd5dh6hKe1YjojCBMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRmMGQCMB8gY2Y2AVs2uJtBr3tEfo9p71wZ5+LM/rOK5eUkw4cIUejfYOWDbCXSzIK+ZsibSQIwcTGqwgeijWmc8OeQHsUaj64B0aVTqWHeV3JsyvZjhR8m04nj6mjZG8uWVBsN25Jg"}]}},"messageSignature":{"messageDigest":{"algorithm":"SHA2_256","digest":"7L1MEBf0ALnn2RHpEm1IXjoz70g0E84C+hatfhZehws="},"signature":"MEUCIFGw0YHBFl8oNIVTUUPCYHDmxX3XrPHopsTPQeVGjICxAiEA4VbLtUnUFwckXGpYluyeWp8VGGeP5AJ/2urYvzJ0dYk="}} \ No newline at end of file diff --git a/build/torch211-cxx11-cu126-x86_64-linux/ops/triton/layer_norm.py b/build/torch211-cxx11-cu126-x86_64-linux/ops/triton/layer_norm.py deleted file mode 100644 index 5e6f0b3697b612e6a19031bbc55ac76ad233dc85..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu126-x86_64-linux/ops/triton/layer_norm.py +++ /dev/null @@ -1,1254 +0,0 @@ -# Copyright (c) 2024, Tri Dao. -# Implement dropout + residual + layer_norm / rms_norm. - -# Based on the Triton LayerNorm tutorial: https://triton-lang.org/main/getting-started/tutorials/05-layer-norm.html -# For the backward pass, we keep weight_grad and bias_grad in registers and accumulate. -# This is faster for dimensions up to 8k, but after that it's much slower due to register spilling. -# The models we train have hidden dim up to 8k anyway (e.g. Llama 70B), so this is fine. - -import math -from typing import Optional, List - -import torch -import torch.nn.functional as F -from torch import Tensor - -import triton -import triton.language as tl - -from flash_attn.utils.torch import custom_fwd, custom_bwd -from flash_attn.utils.library import triton_op - -from ..._ops import add_op_namespace_prefix - - -def maybe_contiguous_lastdim(x): - return x.contiguous() if x is not None and x.stride(-1) != 1 else x - - -def maybe_contiguous(x): - return x.contiguous() if x is not None else None - - -def triton_autotune_configs(): - # Return configs with a valid warp count for the current device - configs = [] - # Maximum threads per block is architecture-dependent in theory, but in reality all are 1024 - max_threads_per_block = 1024 - # Default to warp size 32 if not defined by device - warp_size = getattr(torch.cuda.get_device_properties(torch.cuda.current_device()), "warp_size", 32) - # Autotune for warp counts which are powers of 2 and do not exceed thread per block limit - return [triton.Config({}, num_warps=warp_count) for warp_count in [1, 2, 4, 8, 16, 32] - if warp_count * warp_size <= max_threads_per_block] - # return [triton.Config({}, num_warps=8)] - - -def layer_norm_ref( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - zero_centered_weight=False, - dropout_mask=None, - dropout_mask1=None, - upcast=False, -): - dtype = x.dtype - if upcast: - x = x.float() - weight = weight.float() - bias = bias.float() if bias is not None else None - residual = residual.float() if residual is not None else residual - x1 = x1.float() if x1 is not None else None - weight1 = weight1.float() if weight1 is not None else None - bias1 = bias1.float() if bias1 is not None else None - if zero_centered_weight: - weight = weight + 1.0 - if weight1 is not None: - weight1 = weight1 + 1.0 - if x1 is not None: - assert rowscale is None, "rowscale is not supported with parallel LayerNorm" - if rowscale is not None: - x = x * rowscale[..., None] - if dropout_p > 0.0: - if dropout_mask is not None: - x = x.masked_fill(~dropout_mask, 0.0) / (1.0 - dropout_p) - else: - x = F.dropout(x, p=dropout_p) - if x1 is not None: - if dropout_mask1 is not None: - x1 = x1.masked_fill(~dropout_mask1, 0.0) / (1.0 - dropout_p) - else: - x1 = F.dropout(x1, p=dropout_p) - if x1 is not None: - x = x + x1 - if residual is not None: - x = (x + residual).to(x.dtype) - out = F.layer_norm(x.to(weight.dtype), x.shape[-1:], weight=weight, bias=bias, eps=eps).to( - dtype - ) - if weight1 is None: - return out if not prenorm else (out, x) - else: - out1 = F.layer_norm( - x.to(weight1.dtype), x.shape[-1:], weight=weight1, bias=bias1, eps=eps - ).to(dtype) - return (out, out1) if not prenorm else (out, out1, x) - - -def rms_norm_ref( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - zero_centered_weight=False, - dropout_mask=None, - dropout_mask1=None, - upcast=False, -): - dtype = x.dtype - if upcast: - x = x.float() - weight = weight.float() - bias = bias.float() if bias is not None else None - residual = residual.float() if residual is not None else residual - x1 = x1.float() if x1 is not None else None - weight1 = weight1.float() if weight1 is not None else None - bias1 = bias1.float() if bias1 is not None else None - if zero_centered_weight: - weight = weight + 1.0 - if weight1 is not None: - weight1 = weight1 + 1.0 - if x1 is not None: - assert rowscale is None, "rowscale is not supported with parallel LayerNorm" - if rowscale is not None: - x = x * rowscale[..., None] - if dropout_p > 0.0: - if dropout_mask is not None: - x = x.masked_fill(~dropout_mask, 0.0) / (1.0 - dropout_p) - else: - x = F.dropout(x, p=dropout_p) - if x1 is not None: - if dropout_mask1 is not None: - x1 = x1.masked_fill(~dropout_mask1, 0.0) / (1.0 - dropout_p) - else: - x1 = F.dropout(x1, p=dropout_p) - if x1 is not None: - x = x + x1 - if residual is not None: - x = (x + residual).to(x.dtype) - rstd = 1 / torch.sqrt((x.square()).mean(dim=-1, keepdim=True) + eps) - out = ((x * rstd * weight) + bias if bias is not None else (x * rstd * weight)).to(dtype) - if weight1 is None: - return out if not prenorm else (out, x) - else: - out1 = ((x * rstd * weight1) + bias1 if bias1 is not None else (x * rstd * weight1)).to( - dtype - ) - return (out, out1) if not prenorm else (out, out1, x) - - -@triton.autotune( - configs=triton_autotune_configs(), - key=["N", "HAS_RESIDUAL", "STORE_RESIDUAL_OUT", "IS_RMS_NORM", "HAS_BIAS", "HAS_X1", "HAS_W1", "HAS_B1"], -) -# torch compile doesn't like triton.heuristics, so we set these manually when calling the kernel -# @triton.heuristics({"HAS_BIAS": lambda args: args["B"] is not None}) -# @triton.heuristics({"HAS_RESIDUAL": lambda args: args["RESIDUAL"] is not None}) -# @triton.heuristics({"HAS_X1": lambda args: args["X1"] is not None}) -# @triton.heuristics({"HAS_W1": lambda args: args["W1"] is not None}) -# @triton.heuristics({"HAS_B1": lambda args: args["B1"] is not None}) -@triton.jit -def _layer_norm_fwd_1pass_kernel( - X, # pointer to the input - Y, # pointer to the output - W, # pointer to the weights - B, # pointer to the biases - RESIDUAL, # pointer to the residual - X1, - W1, - B1, - Y1, - RESIDUAL_OUT, # pointer to the residual - ROWSCALE, - SEEDS, # Dropout seeds for each row - DROPOUT_MASK, - DROPOUT_MASK1, - Mean, # pointer to the mean - Rstd, # pointer to the 1/std - stride_x_row, # how much to increase the pointer when moving by 1 row - stride_y_row, - stride_res_row, - stride_res_out_row, - stride_x1_row, - stride_y1_row, - M, # number of rows in X - N, # number of columns in X - eps, # epsilon to avoid division by zero - dropout_p, # Dropout probability - zero_centered_weight, # If true, add 1.0 to the weight - IS_RMS_NORM: tl.constexpr, - BLOCK_N: tl.constexpr, - HAS_RESIDUAL: tl.constexpr, - STORE_RESIDUAL_OUT: tl.constexpr, - HAS_BIAS: tl.constexpr, - HAS_DROPOUT: tl.constexpr, - STORE_DROPOUT_MASK: tl.constexpr, - HAS_ROWSCALE: tl.constexpr, - HAS_X1: tl.constexpr, - HAS_W1: tl.constexpr, - HAS_B1: tl.constexpr, -): - # Map the program id to the row of X and Y it should compute. - row = tl.program_id(0) - X += row * stride_x_row - Y += row * stride_y_row - if HAS_RESIDUAL: - RESIDUAL += row * stride_res_row - if STORE_RESIDUAL_OUT: - RESIDUAL_OUT += row * stride_res_out_row - if HAS_X1: - X1 += row * stride_x1_row - if HAS_W1: - Y1 += row * stride_y1_row - # Compute mean and variance - cols = tl.arange(0, BLOCK_N) - x = tl.load(X + cols, mask=cols < N, other=0.0).to(tl.float32) - if HAS_ROWSCALE: - rowscale = tl.load(ROWSCALE + row).to(tl.float32) - x *= rowscale - if HAS_DROPOUT: - # Compute dropout mask - # 7 rounds is good enough, and reduces register pressure - keep_mask = tl.rand(tl.load(SEEDS + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - x = tl.where(keep_mask, x / (1.0 - dropout_p), 0.0) - if STORE_DROPOUT_MASK: - tl.store(DROPOUT_MASK + row * N + cols, keep_mask, mask=cols < N) - if HAS_X1: - x1 = tl.load(X1 + cols, mask=cols < N, other=0.0).to(tl.float32) - if HAS_ROWSCALE: - rowscale = tl.load(ROWSCALE + M + row).to(tl.float32) - x1 *= rowscale - if HAS_DROPOUT: - # Compute dropout mask - # 7 rounds is good enough, and reduces register pressure - keep_mask = ( - tl.rand(tl.load(SEEDS + M + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - ) - x1 = tl.where(keep_mask, x1 / (1.0 - dropout_p), 0.0) - if STORE_DROPOUT_MASK: - tl.store(DROPOUT_MASK1 + row * N + cols, keep_mask, mask=cols < N) - x += x1 - if HAS_RESIDUAL: - residual = tl.load(RESIDUAL + cols, mask=cols < N, other=0.0).to(tl.float32) - x += residual - if STORE_RESIDUAL_OUT: - tl.store(RESIDUAL_OUT + cols, x, mask=cols < N) - if not IS_RMS_NORM: - mean = tl.sum(x, axis=0) / N - tl.store(Mean + row, mean) - xbar = tl.where(cols < N, x - mean, 0.0) - var = tl.sum(xbar * xbar, axis=0) / N - else: - xbar = tl.where(cols < N, x, 0.0) - var = tl.sum(xbar * xbar, axis=0) / N - rstd = 1 / tl.sqrt(var + eps) - tl.store(Rstd + row, rstd) - # Normalize and apply linear transformation - mask = cols < N - w = tl.load(W + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w += 1.0 - if HAS_BIAS: - b = tl.load(B + cols, mask=mask).to(tl.float32) - x_hat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd - y = x_hat * w + b if HAS_BIAS else x_hat * w - # Write output - tl.store(Y + cols, y, mask=mask) - if HAS_W1: - w1 = tl.load(W1 + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w1 += 1.0 - if HAS_B1: - b1 = tl.load(B1 + cols, mask=mask).to(tl.float32) - y1 = x_hat * w1 + b1 if HAS_B1 else x_hat * w1 - tl.store(Y1 + cols, y1, mask=mask) - - -def _layer_norm_fwd( - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - residual: Optional[Tensor] = None, - x1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - out_dtype: Optional[torch.dtype] = None, - residual_dtype: Optional[torch.dtype] = None, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - return_dropout_mask: bool = False, - out: Optional[Tensor] = None, - residual_out: Optional[Tensor] = None -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - # Need to wrap to handle the case where residual_out is a alias of x, which makes torch.library - # and torch.compile unhappy. Also allocate memory for out and residual_out if they are None - # so that _layer_norm_fwd_impl doesn't have to return them. - if out is None: - out = torch.empty_like(x, dtype=x.dtype if out_dtype is None else out_dtype) - if residual is not None: - residual_dtype = residual.dtype - if residual_out is None and ( - residual is not None - or (residual_dtype is not None and residual_dtype != x.dtype) - or dropout_p > 0.0 - or rowscale is not None - or x1 is not None - ): - residual_out = torch.empty_like( - x, dtype=residual_dtype if residual_dtype is not None else x.dtype - ) - else: - residual_out = None - y1, mean, rstd, seeds, dropout_mask, dropout_mask1 = _layer_norm_fwd_impl( - x, - weight, - bias, - eps, - out, - residual=residual, - x1=x1, - weight1=weight1, - bias1=bias1, - dropout_p=dropout_p, - rowscale=rowscale, - zero_centered_weight=zero_centered_weight, - is_rms_norm=is_rms_norm, - return_dropout_mask=return_dropout_mask, - residual_out=residual_out, - ) - # residual_out is None if residual is None and residual_dtype == input_dtype and dropout_p == 0.0 - if residual_out is None: - residual_out = x - return out, y1, mean, rstd, residual_out, seeds, dropout_mask, dropout_mask1 - - -# [2025-04-28] torch.library.triton_op ignores the schema argument, but here we need the schema -# since we're returning a tuple of tensors -@triton_op(add_op_namespace_prefix("layer_norm_fwd_impl"), mutates_args={"out", "residual_out"}, - schema="(Tensor x, Tensor weight, Tensor bias, float eps, Tensor(a!) out, Tensor? residual, Tensor? x1, Tensor? weight1, Tensor? bias1, float dropout_p, Tensor? rowscale, bool zero_centered_weight, bool is_rms_norm, bool return_dropout_mask, Tensor(a!)? residual_out) -> (Tensor y1, Tensor mean, Tensor rstd, Tensor seeds, Tensor dropout_mask, Tensor dropout_mask1)") -def _layer_norm_fwd_impl( - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - out: Tensor, - residual: Optional[Tensor] = None, - x1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - return_dropout_mask: bool = False, - residual_out: Optional[Tensor] = None -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - M, N = x.shape - assert x.stride(-1) == 1 - if residual is not None: - assert residual.stride(-1) == 1 - assert residual.shape == (M, N) - assert weight.shape == (N,) - assert weight.stride(-1) == 1 - if bias is not None: - assert bias.stride(-1) == 1 - assert bias.shape == (N,) - if x1 is not None: - assert x1.shape == x.shape - assert rowscale is None - assert x1.stride(-1) == 1 - if weight1 is not None: - assert weight1.shape == (N,) - assert weight1.stride(-1) == 1 - if bias1 is not None: - assert bias1.shape == (N,) - assert bias1.stride(-1) == 1 - if rowscale is not None: - assert rowscale.is_contiguous() - assert rowscale.shape == (M,) - assert out.shape == x.shape - assert out.stride(-1) == 1 - if residual_out is not None: - assert residual_out.shape == x.shape - assert residual_out.stride(-1) == 1 - if weight1 is not None: - y1 = torch.empty_like(out) - assert y1.stride(-1) == 1 - else: - y1 = None - mean = torch.empty((M,), dtype=torch.float32, device=x.device) if not is_rms_norm else None - rstd = torch.empty((M,), dtype=torch.float32, device=x.device) - if dropout_p > 0.0: - seeds = torch.randint( - 2**32, (M if x1 is None else 2 * M,), device=x.device, dtype=torch.int64 - ) - else: - seeds = None - if return_dropout_mask and dropout_p > 0.0: - dropout_mask = torch.empty(M, N, device=x.device, dtype=torch.bool) - if x1 is not None: - dropout_mask1 = torch.empty(M, N, device=x.device, dtype=torch.bool) - else: - dropout_mask1 = None - else: - dropout_mask, dropout_mask1 = None, None - # Less than 64KB per feature: enqueue fused kernel - MAX_FUSED_SIZE = 65536 // x.element_size() - BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) - if N > BLOCK_N: - raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") - with torch.cuda.device(x.device.index): - torch.library.wrap_triton(_layer_norm_fwd_1pass_kernel)[(M,)]( - x, - out, - weight, - bias, - residual, - x1, - weight1, - bias1, - y1, - residual_out, - rowscale, - seeds, - dropout_mask, - dropout_mask1, - mean, - rstd, - x.stride(0), - out.stride(0), - residual.stride(0) if residual is not None else 0, - residual_out.stride(0) if residual_out is not None else 0, - x1.stride(0) if x1 is not None else 0, - y1.stride(0) if y1 is not None else 0, - M, - N, - eps, - dropout_p, - # Passing bool make torch inductor very unhappy since it then tries to compare to int_max - int(zero_centered_weight), - is_rms_norm, - BLOCK_N, - residual is not None, - residual_out is not None, - bias is not None, - dropout_p > 0.0, - dropout_mask is not None, - rowscale is not None, - HAS_X1=x1 is not None, - HAS_W1=weight1 is not None, - HAS_B1=bias1 is not None, - ) - return y1, mean, rstd, seeds, dropout_mask, dropout_mask1 - - -@triton.autotune( - configs=triton_autotune_configs(), - key=["N", "HAS_DRESIDUAL", "STORE_DRESIDUAL", "IS_RMS_NORM", "HAS_BIAS", "HAS_DROPOUT"], -) -# torch compile doesn't like triton.heuristics, so we set these manually when calling the kernel -# @triton.heuristics({"HAS_BIAS": lambda args: args["B"] is not None}) -# @triton.heuristics({"HAS_DRESIDUAL": lambda args: args["DRESIDUAL"] is not None}) -# @triton.heuristics({"STORE_DRESIDUAL": lambda args: args["DRESIDUAL_IN"] is not None}) -# @triton.heuristics({"HAS_ROWSCALE": lambda args: args["ROWSCALE"] is not None}) -# @triton.heuristics({"HAS_DY1": lambda args: args["DY1"] is not None}) -# @triton.heuristics({"HAS_DX1": lambda args: args["DX1"] is not None}) -# @triton.heuristics({"HAS_B1": lambda args: args["DB1"] is not None}) -# @triton.heuristics({"RECOMPUTE_OUTPUT": lambda args: args["Y"] is not None}) -@triton.jit -def _layer_norm_bwd_kernel( - X, # pointer to the input - W, # pointer to the weights - B, # pointer to the biases - Y, # pointer to the output to be recomputed - DY, # pointer to the output gradient - DX, # pointer to the input gradient - DW, # pointer to the partial sum of weights gradient - DB, # pointer to the partial sum of biases gradient - DRESIDUAL, - W1, - DY1, - DX1, - DW1, - DB1, - DRESIDUAL_IN, - ROWSCALE, - SEEDS, - Mean, # pointer to the mean - Rstd, # pointer to the 1/std - stride_x_row, # how much to increase the pointer when moving by 1 row - stride_y_row, - stride_dy_row, - stride_dx_row, - stride_dres_row, - stride_dy1_row, - stride_dx1_row, - stride_dres_in_row, - M, # number of rows in X - N, # number of columns in X - eps, # epsilon to avoid division by zero - dropout_p, - zero_centered_weight, - rows_per_program, - IS_RMS_NORM: tl.constexpr, - BLOCK_N: tl.constexpr, - HAS_DRESIDUAL: tl.constexpr, - STORE_DRESIDUAL: tl.constexpr, - HAS_BIAS: tl.constexpr, - HAS_DROPOUT: tl.constexpr, - HAS_ROWSCALE: tl.constexpr, - HAS_DY1: tl.constexpr, - HAS_DX1: tl.constexpr, - HAS_B1: tl.constexpr, - RECOMPUTE_OUTPUT: tl.constexpr, -): - # Map the program id to the elements of X, DX, and DY it should compute. - row_block_id = tl.program_id(0) - row_start = row_block_id * rows_per_program - # Do not early exit if row_start >= M, because we need to write DW and DB - cols = tl.arange(0, BLOCK_N) - mask = cols < N - X += row_start * stride_x_row - if HAS_DRESIDUAL: - DRESIDUAL += row_start * stride_dres_row - if STORE_DRESIDUAL: - DRESIDUAL_IN += row_start * stride_dres_in_row - DY += row_start * stride_dy_row - DX += row_start * stride_dx_row - if HAS_DY1: - DY1 += row_start * stride_dy1_row - if HAS_DX1: - DX1 += row_start * stride_dx1_row - if RECOMPUTE_OUTPUT: - Y += row_start * stride_y_row - w = tl.load(W + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w += 1.0 - if RECOMPUTE_OUTPUT and HAS_BIAS: - b = tl.load(B + cols, mask=mask, other=0.0).to(tl.float32) - if HAS_DY1: - w1 = tl.load(W1 + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w1 += 1.0 - dw = tl.zeros((BLOCK_N,), dtype=tl.float32) - if HAS_BIAS: - db = tl.zeros((BLOCK_N,), dtype=tl.float32) - if HAS_DY1: - dw1 = tl.zeros((BLOCK_N,), dtype=tl.float32) - if HAS_B1: - db1 = tl.zeros((BLOCK_N,), dtype=tl.float32) - row_end = min((row_block_id + 1) * rows_per_program, M) - for row in range(row_start, row_end): - # Load data to SRAM - x = tl.load(X + cols, mask=mask, other=0).to(tl.float32) - dy = tl.load(DY + cols, mask=mask, other=0).to(tl.float32) - if HAS_DY1: - dy1 = tl.load(DY1 + cols, mask=mask, other=0).to(tl.float32) - if not IS_RMS_NORM: - mean = tl.load(Mean + row) - rstd = tl.load(Rstd + row) - # Compute dx - xhat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd - xhat = tl.where(mask, xhat, 0.0) - if RECOMPUTE_OUTPUT: - y = xhat * w + b if HAS_BIAS else xhat * w - tl.store(Y + cols, y, mask=mask) - wdy = w * dy - dw += dy * xhat - if HAS_BIAS: - db += dy - if HAS_DY1: - wdy += w1 * dy1 - dw1 += dy1 * xhat - if HAS_B1: - db1 += dy1 - if not IS_RMS_NORM: - c1 = tl.sum(xhat * wdy, axis=0) / N - c2 = tl.sum(wdy, axis=0) / N - dx = (wdy - (xhat * c1 + c2)) * rstd - else: - c1 = tl.sum(xhat * wdy, axis=0) / N - dx = (wdy - xhat * c1) * rstd - if HAS_DRESIDUAL: - dres = tl.load(DRESIDUAL + cols, mask=mask, other=0).to(tl.float32) - dx += dres - # Write dx - if STORE_DRESIDUAL: - tl.store(DRESIDUAL_IN + cols, dx, mask=mask) - if HAS_DX1: - if HAS_DROPOUT: - keep_mask = ( - tl.rand(tl.load(SEEDS + M + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - ) - dx1 = tl.where(keep_mask, dx / (1.0 - dropout_p), 0.0) - else: - dx1 = dx - tl.store(DX1 + cols, dx1, mask=mask) - if HAS_DROPOUT: - keep_mask = tl.rand(tl.load(SEEDS + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - dx = tl.where(keep_mask, dx / (1.0 - dropout_p), 0.0) - if HAS_ROWSCALE: - rowscale = tl.load(ROWSCALE + row).to(tl.float32) - dx *= rowscale - tl.store(DX + cols, dx, mask=mask) - - X += stride_x_row - if HAS_DRESIDUAL: - DRESIDUAL += stride_dres_row - if STORE_DRESIDUAL: - DRESIDUAL_IN += stride_dres_in_row - if RECOMPUTE_OUTPUT: - Y += stride_y_row - DY += stride_dy_row - DX += stride_dx_row - if HAS_DY1: - DY1 += stride_dy1_row - if HAS_DX1: - DX1 += stride_dx1_row - tl.store(DW + row_block_id * N + cols, dw, mask=mask) - if HAS_BIAS: - tl.store(DB + row_block_id * N + cols, db, mask=mask) - if HAS_DY1: - tl.store(DW1 + row_block_id * N + cols, dw1, mask=mask) - if HAS_B1: - tl.store(DB1 + row_block_id * N + cols, db1, mask=mask) - - -def _layer_norm_bwd( - dy: Tensor, - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - mean: Tensor, - rstd: Tensor, - dresidual: Optional[Tensor] = None, - dy1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - seeds: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - has_residual: bool = False, - has_x1: bool = False, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - x_dtype: Optional[torch.dtype] = None, - recompute_output: bool = False, -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - # Need to wrap to handle the case where dresidual_in or dx1 are aliases of x, - # which makes torch.library unhappy - dx, dw, db, dresidual_in, dx1, dw1, db1, y = _layer_norm_bwd_impl( - dy, - x, - weight, - bias, - eps, - mean, - rstd, - dresidual, - dy1, - weight1, - bias1, - seeds, - dropout_p, - rowscale, - has_residual, - has_x1, - zero_centered_weight, - is_rms_norm, - x_dtype=x_dtype, - recompute_output=recompute_output, - ) - # Don't need to compute dresidual_in separately in this case - if has_residual and dx.dtype == x.dtype and dropout_p == 0.0 and rowscale is None: - dresidual_in = dx - if has_x1 and dropout_p == 0.0: - dx1 = dx - return dx, dw, db, dresidual_in, dx1, dw1, db1, y - - - -@triton_op(add_op_namespace_prefix("layer_norm_bwd_impl"), mutates_args={}, - schema="(Tensor dy, Tensor x, Tensor weight, Tensor bias, float eps, Tensor mean, Tensor rstd, Tensor? dresidual, Tensor? dy1, Tensor? weight1, Tensor? bias1, Tensor? seeds, float dropout_p, Tensor? rowscale, bool has_residual, bool has_x1, bool zero_centered_weight, bool is_rms_norm, ScalarType? x_dtype, bool recompute_output) -> (Tensor dx, Tensor dw, Tensor db, Tensor dresidual_in, Tensor dx1, Tensor dw1, Tensor db1, Tensor y)", - allow_decomposition=False, # Don't let torch.compile trace inside - ) -def _layer_norm_bwd_impl( - dy: Tensor, - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - mean: Tensor, - rstd: Tensor, - dresidual: Optional[Tensor] = None, - dy1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - seeds: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - has_residual: bool = False, - has_x1: bool = False, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - x_dtype: Optional[torch.dtype] = None, - recompute_output: bool = False, -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - M, N = x.shape - assert x.stride(-1) == 1 - dy = maybe_contiguous_lastdim(dy) - assert dy.stride(-1) == 1 - assert dy.shape == (M, N) - if dresidual is not None: - dresidual = maybe_contiguous_lastdim(dresidual) - assert dresidual.stride(-1) == 1 - assert dresidual.shape == (M, N) - assert weight.shape == (N,) - assert weight.stride(-1) == 1 - if bias is not None: - assert bias.stride(-1) == 1 - assert bias.shape == (N,) - if dy1 is not None: - dy1 = maybe_contiguous_lastdim(dy1) - assert weight1 is not None - assert dy1.shape == dy.shape - assert dy1.stride(-1) == 1 - if weight1 is not None: - assert weight1.shape == (N,) - assert weight1.stride(-1) == 1 - if bias1 is not None: - assert bias1.shape == (N,) - assert bias1.stride(-1) == 1 - if seeds is not None: - assert seeds.is_contiguous() - assert seeds.shape == (M if not has_x1 else M * 2,) - if rowscale is not None: - assert rowscale.is_contiguous() - assert rowscale.shape == (M,) - # allocate output - dx = ( - torch.empty_like(x) - if x_dtype is None - else torch.empty(M, N, dtype=x_dtype, device=x.device) - ) - dresidual_in = ( - torch.empty_like(x) - if has_residual - and (dx.dtype != x.dtype or dropout_p > 0.0 or rowscale is not None or has_x1) - else None - ) - dx1 = torch.empty_like(dx) if (has_x1 and dropout_p > 0.0) else None - y = torch.empty(M, N, dtype=dy.dtype, device=dy.device) if recompute_output else None - if recompute_output: - assert weight1 is None, "recompute_output is not supported with parallel LayerNorm" - - # Less than 64KB per feature: enqueue fused kernel - MAX_FUSED_SIZE = 65536 // x.element_size() - BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) - if N > BLOCK_N: - raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") - # Increasing the multiple (e.g. 8) will allow more thread blocks to be launched and hide the - # latency of the gmem reads/writes, but will increase the time of summing up dw / db. - sm_count = torch.cuda.get_device_properties(x.device).multi_processor_count * 8 - _dw = torch.empty((sm_count, N), dtype=torch.float32, device=weight.device) - _db = ( - torch.empty((sm_count, N), dtype=torch.float32, device=bias.device) - if bias is not None - else None - ) - _dw1 = torch.empty_like(_dw) if weight1 is not None else None - _db1 = torch.empty_like(_db) if bias1 is not None else None - rows_per_program = math.ceil(M / sm_count) - grid = (sm_count,) - with torch.cuda.device(x.device.index): - torch.library.wrap_triton(_layer_norm_bwd_kernel)[grid]( - x, - weight, - bias, - y, - dy, - dx, - _dw, - _db, - dresidual, - weight1, - dy1, - dx1, - _dw1, - _db1, - dresidual_in, - rowscale, - seeds, - mean, - rstd, - x.stride(0), - 0 if not recompute_output else y.stride(0), - dy.stride(0), - dx.stride(0), - dresidual.stride(0) if dresidual is not None else 0, - dy1.stride(0) if dy1 is not None else 0, - dx1.stride(0) if dx1 is not None else 0, - dresidual_in.stride(0) if dresidual_in is not None else 0, - M, - N, - eps, - dropout_p, - # Passing bool make torch inductor very unhappy since it then tries to compare to int_max - int(zero_centered_weight), - rows_per_program, - is_rms_norm, - BLOCK_N, - dresidual is not None, - dresidual_in is not None, - bias is not None, - dropout_p > 0.0, - HAS_ROWSCALE=rowscale is not None, - HAS_DY1=dy1 is not None, - HAS_DX1=dx1 is not None, - HAS_B1=bias1 is not None, - RECOMPUTE_OUTPUT=y is not None, - ) - dw = _dw.sum(0).to(weight.dtype) - db = _db.sum(0).to(bias.dtype) if bias is not None else None - dw1 = _dw1.sum(0).to(weight1.dtype) if weight1 is not None else None - db1 = _db1.sum(0).to(bias1.dtype) if bias1 is not None else None - # dresidual_in and dx1 could be None, the wrapper will handle assigning them from dx - return dx, dw, db, dresidual_in, dx1, dw1, db1, y - - -class LayerNormFn(torch.autograd.Function): - - @staticmethod - def forward( - ctx, - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - residual_in_fp32=False, - zero_centered_weight=False, - is_rms_norm=False, - return_dropout_mask=False, - out_dtype=None, - out=None, - residual_out=None - ): - x_shape_og = x.shape - # reshape input data into 2D tensor - x = maybe_contiguous_lastdim(x.reshape(-1, x.shape[-1])) - if residual is not None: - assert residual.shape == x_shape_og - residual = maybe_contiguous_lastdim(residual.reshape(-1, residual.shape[-1])) - if x1 is not None: - assert x1.shape == x_shape_og - assert rowscale is None, "rowscale is not supported with parallel LayerNorm" - x1 = maybe_contiguous_lastdim(x1.reshape(-1, x1.shape[-1])) - weight = weight.contiguous() - bias = maybe_contiguous(bias) - weight1 = maybe_contiguous(weight1) - bias1 = maybe_contiguous(bias1) - if rowscale is not None: - rowscale = rowscale.reshape(-1).contiguous() - residual_dtype = ( - residual.dtype - if residual is not None - else (torch.float32 if residual_in_fp32 else None) - ) - if out is not None: - out = out.reshape(-1, out.shape[-1]) - if residual_out is not None: - residual_out = residual_out.reshape(-1, residual_out.shape[-1]) - y, y1, mean, rstd, residual_out, seeds, dropout_mask, dropout_mask1 = _layer_norm_fwd( - x, - weight, - bias, - eps, - residual, - x1, - weight1, - bias1, - dropout_p=dropout_p, - rowscale=rowscale, - out_dtype=out_dtype, - residual_dtype=residual_dtype, - zero_centered_weight=zero_centered_weight, - is_rms_norm=is_rms_norm, - return_dropout_mask=return_dropout_mask, - out=out, - residual_out=residual_out, - ) - ctx.save_for_backward( - residual_out, weight, bias, weight1, bias1, rowscale, seeds, mean, rstd - ) - ctx.x_shape_og = x_shape_og - ctx.eps = eps - ctx.dropout_p = dropout_p - ctx.is_rms_norm = is_rms_norm - ctx.has_residual = residual is not None - ctx.has_x1 = x1 is not None - ctx.prenorm = prenorm - ctx.x_dtype = x.dtype - ctx.zero_centered_weight = zero_centered_weight - y = y.reshape(x_shape_og) - y1 = y1.reshape(x_shape_og) if y1 is not None else None - residual_out = residual_out.reshape(x_shape_og) if residual_out is not None else None - dropout_mask = dropout_mask.reshape(x_shape_og) if dropout_mask is not None else None - dropout_mask1 = dropout_mask1.reshape(x_shape_og) if dropout_mask1 is not None else None - if not return_dropout_mask: - if weight1 is None: - return y if not prenorm else (y, residual_out) - else: - return (y, y1) if not prenorm else (y, y1, residual_out) - else: - if weight1 is None: - return ( - (y, dropout_mask, dropout_mask1) - if not prenorm - else (y, residual_out, dropout_mask, dropout_mask1) - ) - else: - return ( - (y, y1, dropout_mask, dropout_mask1) - if not prenorm - else (y, y1, residual_out, dropout_mask, dropout_mask1) - ) - - @staticmethod - def backward(ctx, dy, *args): - x, weight, bias, weight1, bias1, rowscale, seeds, mean, rstd = ctx.saved_tensors - dy = dy.reshape(-1, dy.shape[-1]) - if weight1 is not None: - dy1, args = args[0], args[1:] - dy1 = dy1.reshape(-1, dy1.shape[-1]) - assert dy1.shape == x.shape - else: - dy1 = None - if ctx.prenorm: - dresidual = args[0] - dresidual = dresidual.reshape(-1, dresidual.shape[-1]) - assert dresidual.shape == x.shape - else: - dresidual = None - dx, dw, db, dresidual_in, dx1, dw1, db1, _ = _layer_norm_bwd( - dy, - x, - weight, - bias, - ctx.eps, - mean, - rstd, - dresidual, - dy1, - weight1, - bias1, - seeds, - ctx.dropout_p, - rowscale, - ctx.has_residual, - ctx.has_x1, - ctx.zero_centered_weight, - ctx.is_rms_norm, - x_dtype=ctx.x_dtype, - recompute_output=False, - ) - return ( - dx.reshape(ctx.x_shape_og), - dw, - db, - dresidual_in.reshape(ctx.x_shape_og) if ctx.has_residual else None, - dx1.reshape(ctx.x_shape_og) if dx1 is not None else None, - dw1, - db1, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) - - -def layer_norm_fn( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - residual_in_fp32=False, - zero_centered_weight=False, - is_rms_norm=False, - return_dropout_mask=False, - out_dtype=None, - out=None, - residual_out=None -): - return LayerNormFn.apply( - x, - weight, - bias, - residual, - x1, - weight1, - bias1, - eps, - dropout_p, - rowscale, - prenorm, - residual_in_fp32, - zero_centered_weight, - is_rms_norm, - return_dropout_mask, - out_dtype, - out, - residual_out - ) - - -def rms_norm_fn( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - residual_in_fp32=False, - zero_centered_weight=False, - return_dropout_mask=False, - out_dtype=None, - out=None, - residual_out=None -): - return LayerNormFn.apply( - x, - weight, - bias, - residual, - x1, - weight1, - bias1, - eps, - dropout_p, - rowscale, - prenorm, - residual_in_fp32, - zero_centered_weight, - True, - return_dropout_mask, - out_dtype, - out, - residual_out - ) - - -class RMSNorm(torch.nn.Module): - - def __init__(self, hidden_size, eps=1e-5, dropout_p=0.0, zero_centered_weight=False, - device=None, dtype=None): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.eps = eps - if dropout_p > 0.0: - self.drop = torch.nn.Dropout(dropout_p) - else: - self.drop = None - self.zero_centered_weight = zero_centered_weight - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.register_parameter("bias", None) - self.reset_parameters() - - def reset_parameters(self): - if not self.zero_centered_weight: - torch.nn.init.ones_(self.weight) - else: - torch.nn.init.zeros_(self.weight) - - def forward(self, x, residual=None, prenorm=False, residual_in_fp32=False): - return rms_norm_fn( - x, - self.weight, - self.bias, - residual=residual, - eps=self.eps, - dropout_p=self.drop.p if self.drop is not None and self.training else 0.0, - prenorm=prenorm, - residual_in_fp32=residual_in_fp32, - zero_centered_weight=self.zero_centered_weight, - ) - - -class LayerNormLinearFn(torch.autograd.Function): - - @staticmethod - @custom_fwd - def forward( - ctx, - x, - norm_weight, - norm_bias, - linear_weight, - linear_bias, - residual=None, - eps=1e-6, - prenorm=False, - residual_in_fp32=False, - is_rms_norm=False, - ): - x_shape_og = x.shape - # reshape input data into 2D tensor - x = maybe_contiguous_lastdim(x.reshape(-1, x.shape[-1])) - if residual is not None: - assert residual.shape == x_shape_og - residual = maybe_contiguous_lastdim(residual.reshape(-1, residual.shape[-1])) - norm_weight = norm_weight.contiguous() - norm_bias = maybe_contiguous(norm_bias) - residual_dtype = ( - residual.dtype - if residual is not None - else (torch.float32 if residual_in_fp32 else None) - ) - y, _, mean, rstd, residual_out, *rest = _layer_norm_fwd( - x, - norm_weight, - norm_bias, - eps, - residual, - out_dtype=None if not torch.is_autocast_enabled() else torch.get_autocast_dtype("cuda"), - residual_dtype=residual_dtype, - is_rms_norm=is_rms_norm, - ) - y = y.reshape(x_shape_og) - dtype = torch.get_autocast_dtype("cuda") if torch.is_autocast_enabled() else y.dtype - linear_weight = linear_weight.to(dtype) - linear_bias = linear_bias.to(dtype) if linear_bias is not None else None - out = F.linear(y.to(linear_weight.dtype), linear_weight, linear_bias) - # We don't store y, will be recomputed in the backward pass to save memory - ctx.save_for_backward(residual_out, norm_weight, norm_bias, linear_weight, mean, rstd) - ctx.x_shape_og = x_shape_og - ctx.eps = eps - ctx.is_rms_norm = is_rms_norm - ctx.has_residual = residual is not None - ctx.prenorm = prenorm - ctx.x_dtype = x.dtype - ctx.linear_bias_is_none = linear_bias is None - return out if not prenorm else (out, residual_out.reshape(x_shape_og)) - - @staticmethod - @custom_bwd - def backward(ctx, dout, *args): - x, norm_weight, norm_bias, linear_weight, mean, rstd = ctx.saved_tensors - dout = dout.reshape(-1, dout.shape[-1]) - dy = F.linear(dout, linear_weight.t()) - dlinear_bias = None if ctx.linear_bias_is_none else dout.sum(0) - dy = maybe_contiguous_lastdim(dy) - assert dy.shape == x.shape - if ctx.prenorm: - dresidual = args[0] - dresidual = maybe_contiguous_lastdim(dresidual.reshape(-1, dresidual.shape[-1])) - assert dresidual.shape == x.shape - else: - dresidual = None - dx, dnorm_weight, dnorm_bias, dresidual_in, _, _, _, y = _layer_norm_bwd( - dy, - x, - norm_weight, - norm_bias, - ctx.eps, - mean, - rstd, - dresidual=dresidual, - has_residual=ctx.has_residual, - is_rms_norm=ctx.is_rms_norm, - x_dtype=ctx.x_dtype, - recompute_output=True, - ) - dlinear_weight = torch.einsum("bo,bi->oi", dout, y) - return ( - dx.reshape(ctx.x_shape_og), - dnorm_weight, - dnorm_bias, - dlinear_weight, - dlinear_bias, - dresidual_in.reshape(ctx.x_shape_og) if ctx.has_residual else None, - None, - None, - None, - None, - ) - - -def layer_norm_linear_fn( - x, - norm_weight, - norm_bias, - linear_weight, - linear_bias, - residual=None, - eps=1e-6, - prenorm=False, - residual_in_fp32=False, - is_rms_norm=False, -): - return LayerNormLinearFn.apply( - x, - norm_weight, - norm_bias, - linear_weight, - linear_bias, - residual, - eps, - prenorm, - residual_in_fp32, - is_rms_norm, - ) diff --git a/build/torch211-cxx11-cu128-aarch64-linux/__init__.py b/build/torch211-cxx11-cu128-aarch64-linux/__init__.py deleted file mode 100644 index ecc2f9d896b6c93f90b0a1499856dc0612177422..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu128-aarch64-linux/__init__.py +++ /dev/null @@ -1,393 +0,0 @@ -from typing import Optional, List -import torch -from ._ops import ops as flash_attn_ops -from .flash_attn_interface import ( - flash_attn_func, - flash_attn_kvpacked_func, - flash_attn_qkvpacked_func, - flash_attn_varlen_func, - flash_attn_varlen_kvpacked_func, - flash_attn_varlen_qkvpacked_func, - flash_attn_with_kvcache, -) - - -def fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - return_softmax: bool = False, - gen: Optional[torch.Generator] = None, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention. - - Args: - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Optional output tensor, same shape as q - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - return_softmax: Whether to return softmax weights - gen: Optional random number generator - - Returns: - List of tensors: [output, softmax_lse, (softmax if return_softmax)] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.fwd( - q, - k, - v, - out, - alibi_slopes, - p_dropout, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - gen, - ) - - -def varlen_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - out: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - block_table: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - max_seqlen_q: int = 0, - max_seqlen_k: int = 0, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - zero_tensors: bool = False, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - return_softmax: bool = False, - gen: Optional[torch.Generator] = None, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention with variable sequence lengths. - - Args: - q: Query tensor of shape [total_q, num_heads, head_size] - k: Key tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - v: Value tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1] - cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1] - out: Optional output tensor of shape [total_q, num_heads, head_size] - seqused_k: Optional tensor specifying how many keys to use per batch element [batch_size] - leftpad_k: Optional left padding for keys of shape [batch_size] - block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq] - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - max_seqlen_q: Maximum sequence length for queries - max_seqlen_k: Maximum sequence length for keys - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - zero_tensors: Whether to zero tensors before computation - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - return_softmax: Whether to return softmax weights - gen: Optional random number generator - - Returns: - List of tensors: [output, softmax_lse, (softmax if return_softmax)] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.varlen_fwd( - q, - k, - v, - out, - cu_seqlens_q, - cu_seqlens_k, - seqused_k, - leftpad_k, - block_table, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - p_dropout, - softmax_scale, - zero_tensors, - is_causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - gen, - ) - - -def bwd( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - p_dropout: float = 0.0, - 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, - gen: Optional[torch.Generator] = None, - rng_state: Optional[torch.Tensor] = None, -) -> List[torch.Tensor]: - """ - Backward pass for multi-head attention. - - Args: - dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size] - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size] - softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q] - dq: Optional gradient tensor for queries, same shape as q - dk: Optional gradient tensor for keys, same shape as k - dv: Optional gradient tensor for values, same shape as v - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - deterministic: Whether to use deterministic algorithms - gen: Optional random number generator - rng_state: Optional RNG state from forward pass - - Returns: - List of tensors: [dq, dk, dv] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - alibi_slopes, - p_dropout, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - deterministic, - gen, - rng_state, - ) - - -def varlen_bwd( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - max_seqlen_q: int = 0, - max_seqlen_k: int = 0, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - zero_tensors: bool = False, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - deterministic: bool = False, - gen: Optional[torch.Generator] = None, - rng_state: Optional[torch.Tensor] = None, -) -> List[torch.Tensor]: - """ - Backward pass for multi-head attention with variable sequence lengths. - - Args: - dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size] - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size] - softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q] - cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1] - cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1] - dq: Optional gradient tensor for queries, same shape as q - dk: Optional gradient tensor for keys, same shape as k - dv: Optional gradient tensor for values, same shape as v - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - max_seqlen_q: Maximum sequence length for queries - max_seqlen_k: Maximum sequence length for keys - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - zero_tensors: Whether to zero tensors before computation - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - deterministic: Whether to use deterministic algorithms - gen: Optional random number generator - rng_state: Optional RNG state from forward pass - - Returns: - List of tensors: [dq, dk, dv] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.varlen_bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - p_dropout, - softmax_scale, - zero_tensors, - is_causal, - window_size_left, - window_size_right, - softcap, - deterministic, - gen, - rng_state, - ) - - -def fwd_kvcache( - q: torch.Tensor, - kcache: torch.Tensor, - vcache: torch.Tensor, - k: Optional[torch.Tensor] = None, - v: Optional[torch.Tensor] = None, - seqlens_k: Optional[torch.Tensor] = None, - rotary_cos: Optional[torch.Tensor] = None, - rotary_sin: Optional[torch.Tensor] = None, - cache_batch_idx: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - block_table: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - out: 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, - is_rotary_interleaved: bool = False, - num_splits: int = 1, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention with KV cache. - - Args: - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - kcache: Key cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - vcache: Value cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - k: Optional new keys tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size] - v: Optional new values tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size] - seqlens_k: Optional sequence lengths for keys of shape [batch_size] - rotary_cos: Optional rotary cosine tensor of shape [seqlen_ro, rotary_dim/2] - rotary_sin: Optional rotary sine tensor of shape [seqlen_ro, rotary_dim/2] - cache_batch_idx: Optional indices to index into the KV cache - leftpad_k: Optional left padding for keys of shape [batch_size] - block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq] - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - out: Optional output tensor, same shape as q - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - is_rotary_interleaved: Whether rotary embeddings are interleaved - num_splits: Number of splits for computation - - Returns: - List of tensors: [output, softmax_lse] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.fwd_kvcache( - q, - kcache, - vcache, - k, - v, - seqlens_k, - rotary_cos, - rotary_sin, - cache_batch_idx, - leftpad_k, - block_table, - alibi_slopes, - out, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - is_rotary_interleaved, - num_splits, - ) diff --git a/build/torch211-cxx11-cu128-aarch64-linux/_flash_attn2_cuda_f12afc9.abi3.so b/build/torch211-cxx11-cu128-aarch64-linux/_flash_attn2_cuda_f12afc9.abi3.so deleted file mode 100644 index 525fe3c41b97d982b58b371102e8d9c46e4ba32d..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu128-aarch64-linux/_flash_attn2_cuda_f12afc9.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:5f6ebce5c99d8d2cd4a77d5be6f460b0bb03db6fcbe62c9a1a0cafde62e0d94b -size 1035975264 diff --git a/build/torch211-cxx11-cu128-aarch64-linux/_ops.py b/build/torch211-cxx11-cu128-aarch64-linux/_ops.py deleted file mode 100644 index af5b731cb04b5e6ce54e862db0501f3683c3163f..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu128-aarch64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _flash_attn2_cuda_f12afc9 -ops = torch.ops._flash_attn2_cuda_f12afc9 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_flash_attn2_cuda_f12afc9::{op_name}" diff --git a/build/torch211-cxx11-cu128-aarch64-linux/bert_padding.py b/build/torch211-cxx11-cu128-aarch64-linux/bert_padding.py deleted file mode 100644 index 3c2d35159a014a9d03aabead9e52e009168696ea..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu128-aarch64-linux/bert_padding.py +++ /dev/null @@ -1,218 +0,0 @@ -# Adapted from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/padding.py - -import torch -import torch.nn.functional as F -from einops import rearrange, repeat - - -class IndexFirstAxis(torch.autograd.Function): - @staticmethod - def forward(ctx, input, indices): - ctx.save_for_backward(indices) - assert input.ndim >= 2 - ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] - second_dim = other_shape.numel() - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - # return input[indices] - return torch.gather( - rearrange(input, "b ... -> b (...)"), 0, repeat(indices, "z -> z d", d=second_dim) - ).reshape(-1, *other_shape) - - @staticmethod - def backward(ctx, grad_output): - (indices,) = ctx.saved_tensors - assert grad_output.ndim >= 2 - other_shape = grad_output.shape[1:] - grad_output = rearrange(grad_output, "b ... -> b (...)") - grad_input = torch.zeros( - [ctx.first_axis_dim, grad_output.shape[1]], - device=grad_output.device, - dtype=grad_output.dtype, - ) - # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. - # grad_input[indices] = grad_output - grad_input.scatter_(0, repeat(indices, "z -> z d", d=grad_output.shape[1]), grad_output) - return grad_input.reshape(ctx.first_axis_dim, *other_shape), None - - -index_first_axis = IndexFirstAxis.apply - - -class IndexPutFirstAxis(torch.autograd.Function): - @staticmethod - def forward(ctx, values, indices, first_axis_dim): - ctx.save_for_backward(indices) - assert indices.ndim == 1 - assert values.ndim >= 2 - output = torch.zeros( - first_axis_dim, *values.shape[1:], device=values.device, dtype=values.dtype - ) - # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. - output[indices] = values - # output.scatter_(0, repeat(indices, 'z -> z d', d=values.shape[1]), values) - return output - - @staticmethod - def backward(ctx, grad_output): - (indices,) = ctx.saved_tensors - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - grad_values = grad_output[indices] - # grad_values = torch.gather(grad_output, 0, repeat(indices, 'z -> z d', d=grad_output.shape[1])) - return grad_values, None, None - - -index_put_first_axis = IndexPutFirstAxis.apply - - -class IndexFirstAxisResidual(torch.autograd.Function): - @staticmethod - def forward(ctx, input, indices): - ctx.save_for_backward(indices) - assert input.ndim >= 2 - ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] - second_dim = other_shape.numel() - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - output = input[indices] - # We don't want to reshape input (b ... -> b (...)) since it could change the channel_last - # memory format to channel_first. In other words, input might not be contiguous. - # If we don't detach, Pytorch complains about output being a view and is being modified inplace - return output, input.detach() - - @staticmethod - def backward(ctx, grad_output, grad_residual): - (indices,) = ctx.saved_tensors - assert grad_output.ndim >= 2 - other_shape = grad_output.shape[1:] - assert grad_residual.shape[1:] == other_shape - grad_input = grad_residual - # grad_input[indices] += grad_output - indices = indices.reshape(indices.shape[0], *((1,) * (grad_output.ndim - 1))) - indices = indices.expand_as(grad_output) - grad_input.scatter_add_(0, indices, grad_output) - return grad_input.reshape(ctx.first_axis_dim, *other_shape), None - - -index_first_axis_residual = IndexFirstAxisResidual.apply - - -def unpad_input(hidden_states, attention_mask, unused_mask=None): - """ - Arguments: - hidden_states: (batch, seqlen, ...) - attention_mask: (batch, seqlen), bool / int, 1 means valid and 0 means not valid. - unused_mask: (batch, seqlen), bool / int, 1 means the element is allocated but unused. - Return: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens selected in attention_mask + unused_mask. - indices: (total_nnz), the indices of masked tokens from the flattened input sequence. - cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. - max_seqlen_in_batch: int - seqused: (batch), returns the number of tokens selected in attention_mask + unused_mask. - """ - all_masks = (attention_mask + unused_mask) if unused_mask is not None else attention_mask - seqlens_in_batch = all_masks.sum(dim=-1, dtype=torch.int32) - used_seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) - indices = torch.nonzero(all_masks.flatten(), as_tuple=False).flatten() - max_seqlen_in_batch = seqlens_in_batch.max().item() - cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) - # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the - # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim - # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to - # index with integer indices. Moreover, torch's index is a bit slower than it needs to be, - # so we write custom forward and backward to make it a bit faster. - return ( - index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices), - indices, - cu_seqlens, - max_seqlen_in_batch, - used_seqlens_in_batch, - ) - - -def unpad_input_for_concatenated_sequences(hidden_states, attention_mask_in_length): - """ - Supports concatenating short samples in one sequence. The attention_mask_in_length is utilized to mask other short samples. It helps efficient training of variant lengths-based samples (e.g., the supervised fine-tuning task in large language model). - The motivation for this function is explained [here](https://github.com/Dao-AILab/flash-attention/issues/432#issuecomment-1668822286). - - For example, if batch = 3 and seqlen = 6, the attention_mask_in_length is: - ``` - [ - [2, 3, 0, 0, 0, 0], - [3, 2, 0, 0, 0, 0], - [6, 0, 0, 0, 0, 0] - ] - ``` - , which refers to the 3D-attention mask: - ``` - [ - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [0, 0, 1, 0, 0, 0], - [0, 0, 1, 1, 0, 0], - [0, 0, 1, 1, 1, 0], - [0, 0, 0, 0, 0, 1] - ], - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [0, 0, 0, 1, 0, 0], - [0, 0, 0, 1, 1, 0], - [0, 0, 0, 0, 0, 1] - ], - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [1, 1, 1, 1, 0, 0], - [1, 1, 1, 1, 1, 0], - [1, 1, 1, 1, 1, 1] - ] - ] - ```. - - Arguments: - hidden_states: (batch, seqlen, ...) - attention_mask_in_length: (batch, seqlen), int, a nonzero number (e.g., 1, 2, 3, etc.) means length of concatenated sequence in b-th batch, and 0 means none. - Return: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. - indices: (total_nnz), the indices of non-masked tokens from the flattened input sequence. - cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. - max_seqlen_in_batch: int - """ - length = attention_mask_in_length.sum(dim=-1) - seqlen = attention_mask_in_length.size(-1) - attention_mask_2d = torch.arange(seqlen, device=length.device, dtype=length.dtype).expand(len(length), seqlen) < length.unsqueeze(1) - real_indices_idx = torch.nonzero(attention_mask_in_length.flatten(), as_tuple=False).flatten() - seqlens_in_batch = attention_mask_in_length.flatten()[real_indices_idx] - indices = torch.nonzero(attention_mask_2d.flatten(), as_tuple=False).flatten() - max_seqlen_in_batch = seqlens_in_batch.max().item() - cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) - # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the - # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim - # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to - # index with integer indices. Moreover, torch's index is a bit slower than it needs to be, - # so we write custom forward and backward to make it a bit faster. - return ( - index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices), - indices, - cu_seqlens, - max_seqlen_in_batch, - ) - - -def pad_input(hidden_states, indices, batch, seqlen): - """ - Arguments: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. - indices: (total_nnz), the indices that represent the non-masked tokens of the original padded input sequence. - batch: int, batch size for the padded sequence. - seqlen: int, maximum sequence length for the padded sequence. - Return: - hidden_states: (batch, seqlen, ...) - """ - dim = hidden_states.shape[-1] - # output = torch.zeros((batch * seqlen), dim, device=hidden_states.device, dtype=hidden_states.dtype) - # output[indices] = hidden_states - output = index_put_first_axis(hidden_states, indices, batch * seqlen) - return rearrange(output, "(b s) ... -> b s ...", b=batch) diff --git a/build/torch211-cxx11-cu128-aarch64-linux/flash_attn2/__init__.py b/build/torch211-cxx11-cu128-aarch64-linux/flash_attn2/__init__.py deleted file mode 100644 index a9b2672c1cd85b74c1b3ded0fc0b2100e1aeac23..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu128-aarch64-linux/flash_attn2/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -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-aarch64-linux/flash_attn_interface.py b/build/torch211-cxx11-cu128-aarch64-linux/flash_attn_interface.py deleted file mode 100644 index b014e569dc9812a9c840d1043f2cebe767b4d477..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu128-aarch64-linux/flash_attn_interface.py +++ /dev/null @@ -1,1631 +0,0 @@ -# Copyright (c) 2023, Tri Dao. - -from typing import Optional, Sequence, Tuple, Union - -import torch -import torch.nn as nn -import os - -# # isort: off -# # We need to import the CUDA kernels after importing torch -# USE_TRITON_ROCM = os.getenv("FLASH_ATTENTION_TRITON_AMD_ENABLE", "FALSE") == "TRUE" -# if USE_TRITON_ROCM: -# from .flash_attn_triton_amd import interface_fa as flash_attn -# else: -# import flash_attn_2_cuda as flash_attn - - -from . import _ops -from ._ops import ops as flash_attn -from ._ops import add_op_namespace_prefix - -# # isort: on - -def maybe_contiguous(x): - return x.contiguous() if x is not None and x.stride(-1) != 1 else x - - -def _get_device(): - if torch.xpu.is_available(): - return "xpu" - elif torch.cuda.is_available(): - return "cuda" - else: - return "cpu" - - -def _get_block_size_n(device, head_dim, is_dropout, is_causal): - assert head_dim <= 256 - - if device.type == "xpu": - if head_dim <= 96: - return 64 - elif head_dim <= 128: - return 32 - elif head_dim <= 256: - return 64 - else: - return 32 - - # This should match the block sizes in the CUDA kernel - major, minor = torch.cuda.get_device_capability(device) - is_sm8x = major == 8 and minor > 0 # Only include sm86 and sm89, exclude sm80 (A100) - is_sm80 = major == 8 and minor == 0 - is_sm90 = major == 9 and minor == 0 - if head_dim <= 32: - return 128 - if head_dim <= 64: - return 128 if not is_dropout else 64 - elif head_dim <= 96: - return 64 - elif head_dim <= 128: - if is_sm8x: - return 64 if (not is_dropout and is_causal) else 32 - else: - return 64 if not is_dropout else 32 - elif head_dim <= 192: - return 64 - elif head_dim <= 224: - return 64 - elif head_dim <= 256: - return 64 - - -def round_multiple(x, m): - return (x + m - 1) // m * m - - -# torch.compile() support is only enabled for pytorch >= 2.4 -# The reason for this is that we are using the new custom_op and register_fake -# APIs, which support inplace modification of inputs in the function itself -if torch.__version__ >= "2.4.0": - _torch_custom_op_wrapper = torch.library.custom_op - _torch_register_fake_wrapper = torch.library.register_fake -else: - def noop_custom_op_wrapper(name, fn=None, /, *, mutates_args, device_types=None, schema=None): - def wrap(func): - return func - if fn is None: - return wrap - return fn - def noop_register_fake_wrapper(op, fn=None, /, *, lib=None, _stacklevel=1): - def wrap(func): - return func - if fn is None: - return wrap - return fn - _torch_custom_op_wrapper = noop_custom_op_wrapper - _torch_register_fake_wrapper = noop_register_fake_wrapper - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_forward"), mutates_args=(), device_types=_get_device()) -def _flash_attn_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - return_softmax: bool -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - out, softmax_lse, S_dmask, rng_state = flash_attn.fwd( - q, - k, - v, - None, - alibi_slopes, - dropout_p, - softmax_scale, - causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - None, - ) - return out, softmax_lse, S_dmask, rng_state - - -@_torch_register_fake_wrapper(add_op_namespace_prefix("_flash_attn_forward")) -def _flash_attn_forward_fake( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - return_softmax: bool -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - batch_size, seqlen_q, num_heads, head_size = q.shape - seqlen_k = k.shape[1] - out = torch.empty_like(q) - softmax_lse = torch.empty((batch_size, num_heads, seqlen_q), dtype=torch.float32, device=q.device, layout=q.layout) - p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) - if return_softmax: - p = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128), round_multiple(seqlen_k, 128)), dtype=q.dtype, device=q.device, layout=q.layout) - rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) - - return out, softmax_lse, p, rng_state - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_forward = _ops.ops._flash_attn_forward -else: - _wrapped_flash_attn_forward = _flash_attn_forward - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_varlen_forward"), mutates_args=(), device_types=_get_device()) -def _flash_attn_varlen_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - alibi_slopes: Optional[torch.Tensor] = None, - return_softmax: bool = False, - block_table: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - out, softmax_lse, S_dmask, rng_state = flash_attn.varlen_fwd( - q, - k, - v, - None, - cu_seqlens_q, - cu_seqlens_k, - seqused_k, - leftpad_k, - block_table, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - zero_tensors, - causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - None, - ) - # if out.isnan().any() or softmax_lse.isnan().any(): - # breakpoint() - return out, softmax_lse, S_dmask, rng_state - - -@_torch_register_fake_wrapper(add_op_namespace_prefix("_flash_attn_varlen_forward")) -def _flash_attn_varlen_forward_fake( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - alibi_slopes: Optional[torch.Tensor] = None, - return_softmax: bool = False, - block_table: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - paged_kv = block_table is not None - batch_size = cu_seqlens_q.numel() - 1 - total_q, num_heads, _ = q.shape - - out = torch.empty_like(q) - softmax_lse = torch.empty((num_heads, total_q), dtype=torch.float32, device=q.device, layout=q.layout) - p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) - seqlen_q_rounded = round_multiple(max_seqlen_q, 128) - seqlen_k_rounded = round_multiple(max_seqlen_k, 128) - if return_softmax: - p = torch.empty((batch_size, num_heads, seqlen_q_rounded, seqlen_k_rounded), dtype=q.dtype, device=q.device, layout=q.layout) - rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) - return out, softmax_lse, p, rng_state - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_varlen_forward = _ops.ops._flash_attn_varlen_forward -else: - _wrapped_flash_attn_varlen_forward = _flash_attn_varlen_forward - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_backward"), mutates_args=("dq", "dk", "dv"), device_types=_get_device()) -def _flash_attn_backward( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, -) -> 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)] - ( - dq, - dk, - dv, - softmax_d, - ) = flash_attn.bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - alibi_slopes, - dropout_p, - softmax_scale, - causal, - window_size_left, - window_size_right, - softcap, - deterministic, - None, - rng_state, - ) - return softmax_d - - -@_torch_register_fake_wrapper(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, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, -) -> torch.Tensor: - dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] - if dq is None: - dq = torch.empty_like(q) - if dk is None: - dk = torch.empty_like(k) - if dv is None: - dv = torch.empty_like(v) - batch_size, seqlen_q, num_heads, _ = q.shape - softmax_d = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128)), device=q.device, dtype=torch.float32) - - return softmax_d - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_backward = _ops.ops._flash_attn_backward -else: - _wrapped_flash_attn_backward = _flash_attn_backward - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_varlen_backward"), mutates_args=("dq", "dk", "dv"), device_types=_get_device()) -def _flash_attn_varlen_backward( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> 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)] - ( - dq, - dk, - dv, - softmax_d, - ) = flash_attn.varlen_bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - zero_tensors, - causal, - window_size_left, - window_size_right, - softcap, - deterministic, - None, - rng_state, - ) - # if dk.isnan().any() or dk.isnan().any() or dv.isnan().any() or softmax_d.isnan().any(): - # breakpoint() - return softmax_d - - -@_torch_register_fake_wrapper(add_op_namespace_prefix("_flash_attn_varlen_backward")) -def _flash_attn_varlen_backward_fake( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> torch.Tensor: - dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] - batch_size = cu_seqlens_q.numel() - 1 - total_q, num_heads, _ = q.shape - - if dq is None: - dq = torch.empty_like(q) - if dk is None: - dk = torch.empty_like(k) - if dv is None: - dv = torch.empty_like(v) - softmax_d = torch.empty((num_heads, total_q + 128 * batch_size), device=q.device, dtype=torch.float32) - - return softmax_d - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_varlen_backward = _ops.ops._flash_attn_varlen_backward -else: - _wrapped_flash_attn_varlen_backward = _flash_attn_varlen_backward - - -class FlashAttnQKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and qkv.requires_grad - if softmax_scale is None: - softmax_scale = qkv.shape[-1] ** (-0.5) - q, k, v = qkv[:, :, 0].detach(), qkv[:, :, 1].detach(), qkv[:, :, 2].detach() - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) - dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dqkv[:, :, 0], - dqkv[:, :, 1], - dqkv[:, :, 2], - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dqkv = dqkv[..., : dout.shape[-1]] # We could have padded the head dimension - return dqkv, None, None, None, None, None, None, None, None, None - - -class FlashAttnVarlenQKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - cu_seqlens, - max_seqlen, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and qkv.requires_grad - if softmax_scale is None: - softmax_scale = qkv.shape[-1] ** (-0.5) - q, k, v = qkv[:, 0].detach(), qkv[:, 1].detach(), qkv[:, 2].detach() - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens, - cu_seqlens, - max_seqlen, - max_seqlen, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=None, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens, rng_state) - ctx.dropout_p = dropout_p - ctx.max_seqlen = max_seqlen - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens, rng_state = ctx.saved_tensors - qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) - dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dqkv[:, 0], - dqkv[:, 1], - dqkv[:, 2], - cu_seqlens, - cu_seqlens, - ctx.max_seqlen, - ctx.max_seqlen, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - 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 FlashAttnKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - kv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, kv] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - k, v = kv[:, :, 0].detach(), kv[:, :, 1].detach() - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - dq = torch.empty_like(q) - kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) - dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dkv[:, :, 0], - dkv[:, :, 1], - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dkv = dkv[..., : dout.shape[-1]] - return dq, dkv, None, None, None, None, None, None, None, None, None - - -class FlashAttnVarlenKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, kv] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - k, v = kv[:, 0].detach(), kv[:, 1].detach() - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=None, - ) - if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state - ) - ctx.dropout_p = dropout_p - 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.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors - dq = torch.empty_like(q) - kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) - dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dkv[:, 0], - dkv[:, 1], - cu_seqlens_q, - cu_seqlens_k, - ctx.max_seqlen_q, - ctx.max_seqlen_k, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dkv = dkv[..., : dout.shape[-1]] - return dq, dkv, None, None, None, None, None, None, None, None, None, None, None, None, None - - -class FlashAttnFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - k, - v, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, k, v] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dk = dk[..., : dout.shape[-1]] - dv = dv[..., : dout.shape[-1]] - return dq, dk, dv, 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, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - block_table, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, k, v] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=block_table, - ) - if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state - ) - ctx.dropout_p = dropout_p - 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.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors - dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - ctx.max_seqlen_q, - ctx.max_seqlen_k, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dk = dk[..., : dout.shape[-1]] - dv = dv[..., : dout.shape[-1]] - return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None - - -def flash_attn_qkvpacked_func( - qkv, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # <=0.0 means deactivate - alibi_slopes=None, - deterministic=False, - 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, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_kvpacked_func( - q, - kv, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, -): - """dropout_p should be set to 0.0 during evaluation - If 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 K, V. - 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) - kv: (batch_size, seqlen, 2, 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. - softcap: float. Anything > 0 activates softcapping 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). - 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 FlashAttnKVPackedFunc.apply( - q, - kv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_func( - q, - k, - v, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - 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). - 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 FlashAttnFunc.apply( - q, - k, - v, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_qkvpacked_func( - qkv, - cu_seqlens, - max_seqlen, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - 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_varlen_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_varlen_kvpacked_func and flash_attn_varlen_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: (total, 3, nheads, headdim), where total = total number of tokens in the batch. - cu_seqlens: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into qkv. - max_seqlen: int. Maximum sequence length in the batch. - 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenQKVPackedFunc.apply( - qkv, - cu_seqlens, - max_seqlen, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_kvpacked_func( - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, -): - """dropout_p should be set to 0.0 during evaluation - If 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 K, V. - 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: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch. - kv: (total_k, 2, nheads_k, headdim), where total_k = total number of key tokens in the batch. - cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into q. - cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into kv. - max_seqlen_q: int. Maximum query sequence length in the batch. - max_seqlen_k: int. Maximum key sequence length in the batch. - 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 + 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenKVPackedFunc.apply( - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_func( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, - block_table=None, -): - """dropout_p should be set to 0.0 during evaluation - Supports multi-query and grouped-query attention (MQA/GQA) by passing in K, V 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: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch. - k: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. - v: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. - cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into q. - cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into kv. - max_seqlen_q: int. Maximum query sequence length in the batch. - max_seqlen_k: int. Maximum key sequence length in the batch. - 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 + 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenFunc.apply( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - block_table, - False if q.device.type == "cpu" else torch.is_grad_enabled(), - ) - - -def flash_attn_with_kvcache( - q, - k_cache, - v_cache, - k=None, - v=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, - block_table: Optional[torch.Tensor] = None, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - rotary_interleaved=True, - alibi_slopes=None, - num_splits=0, - 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 block_table, - or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_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) if there's no block_table, - or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_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). Similar to k. - 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. - block_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). - 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. - 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" - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-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) - cache_batch_idx = maybe_contiguous(cache_batch_idx) - block_table = maybe_contiguous(block_table) - out, softmax_lse = flash_attn.fwd_kvcache( - q, - k_cache, - v_cache, - k, - v, - cache_seqlens, - rotary_cos, - rotary_sin, - cache_batch_idx, - cache_leftpad, - block_table, - alibi_slopes, - None, - softmax_scale, - causal, - window_size[0], - window_size[1], - softcap, - rotary_interleaved, - num_splits, - ) - return (out, softmax_lse) if return_softmax_lse else out diff --git a/build/torch211-cxx11-cu128-aarch64-linux/layers/__init__.py b/build/torch211-cxx11-cu128-aarch64-linux/layers/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/build/torch211-cxx11-cu128-aarch64-linux/metadata.json b/build/torch211-cxx11-cu128-aarch64-linux/metadata.json deleted file mode 100644 index 9b2b806cfceb63f53e62f9736c4d45f421e4fcf2..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu128-aarch64-linux/metadata.json +++ /dev/null @@ -1,42 +0,0 @@ -{ - "name": "flash-attn2", - "id": "_flash_attn2_cuda_f12afc9", - "version": 1, - "license": "BSD-3-Clause", - "python-depends": [], - "backend": { - "type": "cuda", - "archs": [ - "10.0", - "12.0", - "8.0", - "9.0" - ] - }, - "digest": { - "algorithm": "sha256", - "files": { - "__init__.py": "Cd9MEdd/WUYnBmfiQpC4SyPqAxrTilTjvsvswNju4zc=", - "_flash_attn2_cuda_f12afc9.abi3.so": "X2685cmdjSzUp31b5vRgsLsD22/L5iyaGgyv3mLg2Us=", - "_ops.py": "GQbmTTeSOU+7Q2nI68Hz+1RvEogyfEFccx1ubUYw+2s=", - "bert_padding.py": "gF1EmsdJ+HpQ86MRQ4VxDw+Sb/RVISdQALdNnoByHlw=", - "flash_attn2/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY=", - "flash_attn_interface.py": "W7a9bHR2ly6QL3aKqsBqxWA7M7ZWYcZPmtsPJu6rAfY=", - "layers/__init__.py": "47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU=", - "layers/patch_embed.py": "H58CgME/qSOPTZLOG08wFgrQS1j34pvNwMPrkTj3Ek4=", - "layers/rotary.py": "rFIdNrZQFNawwLD6d7t5LE2m+M93wwDZ8EuExw+gfD8=", - "ops/__init__.py": "47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU=", - "ops/activations.py": "t5lzNg1In8LP6bKeTnyeMizwqjv27JGbJ6ylPdGvZYg=", - "ops/fused_dense.py": "BqT8VTSCGdLJAjh6WBMkCRFYrdVodPeoElWBvH6AIzA=", - "ops/layer_norm.py": "zr7NXIm+2mtEynTp1CS0fbFGI2Mqdp41dY4AfDWF6EQ=", - "ops/rms_norm.py": "XEnihcj0a4aSz4LO55m5iKGVn4HKTeKN8TIyHjuDgxI=", - "ops/triton/__init__.py": "AbpHGcgLb+kRsJGnwFEktk7uzpZOCcBY74+YBdrKVGs=", - "ops/triton/cross_entropy.py": "y56nmuOtpSNxM7sDZcUSgXKIsIEayfuvrU8bAVog/nw=", - "ops/triton/k_activations.py": "+Z3vIyO4JkqBMipKsPvhzmxljtBdIhJCsl/M+/ESqBo=", - "ops/triton/layer_norm.py": "BBQdJiqopbIS12viUqzx3SvOQCNbDSAYLJXhpyMGyC8=", - "ops/triton/linear.py": "OtRvKz8xdpl+7v3q/ZTaS9fdBt9XrzMyapgRr50uBbM=", - "ops/triton/mlp.py": "uWA+D7HyuRUsOT+4nhi3vvCDNZnNEXwYcHMQ02UREfQ=", - "ops/triton/rotary.py": "0YSrlbep0lJmzLpTiANs0slTFeGcQQqR+el6uzbSLvI=" - } - } -} \ No newline at end of file diff --git a/build/torch211-cxx11-cu128-aarch64-linux/metadata.json.sigstore b/build/torch211-cxx11-cu128-aarch64-linux/metadata.json.sigstore deleted file mode 100644 index c1e7b0ce847c23bfe1cac904b24c2e7bc58f56cc..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu128-aarch64-linux/metadata.json.sigstore +++ /dev/null @@ -1 +0,0 @@ -{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHSjCCBtGgAwIBAgIUTXP0YPkRACdQzHCzbBPwyt/CbnYwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwNjIzMjM0MDQyWhcNMjYwNjIzMjM1MDQyWjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEYIBHxpBjz3S87V7Uc7PRsCByiUbIKtsopV0bGdkBNwqKyG85bxIjcgu6yGcOtRUAHHQuVPq5sKVX/4PqafXql6OCBfAwggXsMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUF6435n7YwCVxTpyvhlQsNgxJ198wHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoZjEyYWZjOTJjNDBiYmQ3NmQ5OTBkOTExZWJhYmYyOGU0OTBhN2JkZTATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoZjEyYWZjOTJjNDBiYmQ3NmQ5OTBkOTExZWJhYmYyOGU0OTBhN2JkZTAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoZjEyYWZjOTJjNDBiYmQ3NmQ5OTBkOTExZWJhYmYyOGU0OTBhN2JkZTAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKGYxMmFmYzkyYzQwYmJkNzZkOTkwZDkxMWViYWJmMjhlNDkwYTdiZGUwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMjgwMzYyNjg5MzUvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBigYKKwYBBAHWeQIEAgR8BHoAeAB2AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABnvbbh48AAAQDAEcwRQIgXi188Z8nwooYW3zeED1fuiSgJfD9lvFWrVEj7+r/nEQCIQDC59dR/OZC4X0CXJbZ20qYL9Btf9HeiiYtzIsVhU3MTjAKBggqhkjOPQQDAwNnADBkAjAOLeOHqZYCiD3126Aq3F05/cHqrAciq+8oL8ckuzScwNHO9OV+43G7C3gfr5Axna4CMAuG7cjlIgvpErHf4P6qH1dscnnpZTPYN8uqfb4oTHxooNz5NtqjxtPJUsgspNhD0A=="}, "tlogEntries":[{"logIndex":"1933039424", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1782258043", "inclusionPromise":{"signedEntryTimestamp":"MEUCIQCpbjnlZkzI7PEtkwmQkf1BMkvwek/fH2/5haYYz60jPgIgb/vB9zQ6LD6RrEwseFzRQd6sEezpwMDesKk9spYrl00="}, "inclusionProof":{"logIndex":"1811135162", "rootHash":"q4B2gPr9VUhrNFY6Uz/Ae5Ghy0zlkWLJAlkOXhPe2Kk=", "treeSize":"1811135191", "hashes":["i+TLVttAHYRwCx9xovwZX12c7/3oqZBPin3iwMj/3zo=", "rLhpe7gZhRHUf12DZnrYyu9dsB9RIv+Ab8a5isdR/6U=", "t2YHlmI2WtJYjfnWqPzmv6OiYi8iUleQ1lYZMpuGYu8=", "Ckz51RT5HznFC3ZntJwdZyrXK429Mrp7BHpk3h/BG9E=", "faQSpui6FHHS9m07T5HobnTyppxus7k76qRzwixjkU8=", "AxV3+dqmuxdm0ZwVhcbjvrtDTXJGpXkSLxnPTbp58OM=", "t9ehbaR7xaxG+urfn3oX9xO6He7vHFlzCUHj1Cpe7kc=", "FQILKGmmcKTxlCRQxCi4Glqd2YlZ5AMG0TkLiX9Y80U=", "QVVI+8UorpK5fVvAVaM/9A8oGHpbEDzrgIYi5jMqbqw=", "WD8KwmQq5R8K69koYj5k4K7EKo2VF3LqJcx2xoXcd7o=", "99JTcUz0aJeZZocenAppVOLl8v91qMn0v7CmEjq9lws=", "OSDrt6jTKV6LYAkaG63p2BFz4SmbNmGb3uerhF5pOpQ=", "yehxDZkBWOGIzNDb93qxmQa6VBVo7kacVSoHyXMm3ZI=", "oUh+N71k3jWcTmlDDB6Esa5uqlCnbmzxTmDdut2A16M=", "2BsqkdK2omLvu+CpsFjpwqGi6m22H/7zcwCVkPwslM8=", "S8Oe9ZicoJqbVuL1I0dIdTgVpLc82Cm2CMOcV83sKVs=", "E1O2ilWlpFRLWGdbmJUMDZ2qbGu/eDZ12m6YjQ9/JPA=", "nyv2TDlfkz0XfRwjY98zsbUbOkBZ+Tgz7PML0c2uN14=", "OMfDh4k9YD0E+t4sjqJsh+vg8D1ueo1cC07exyxmr6s=", "+vHKm4wNwGJ46ceyZ8mePD7Vn1SediStNHCStRoTSWQ=", "5HwWjuGNDisetbiWM0cWIXM65TIM3EJlqk13ntCQlcU=", "NDIRlqxO0IaGr94q86MIU7VCAN3yEEgOHS4loya5lu4=", "daxmZaajRpZV+JxHiOYZhJBiSKN5ucqjh2WnGbHhirw=", "DOCeoSMovIvLExkhIvisow9AuNXgeWs4ECkyR6EcqYU="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n1811135191\nq4B2gPr9VUhrNFY6Uz/Ae5Ghy0zlkWLJAlkOXhPe2Kk=\n\n— rekor.sigstore.dev wNI9ajBFAiA7EE7n1BPw8pt1hVcSSAQwQxpozL5mbug7wgQBNlf+RAIhAP21MzLaH0NEMgweVjkvhD71tD9m/sSLZDBScawNb6Lk\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiIwZjA0NWI4ZWQyNjY2YTczZGJiMmU1YWM3NTZmNzg0YWNmMGM3Yzk2NDcwNDMwNTlmYTNlZTE1MTcxZjM3NzI5In19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FUUNJSDJLVHZ0VFFFTnFqbzF1eEVETVorNWFXalg0MUZyVnJZVnMwZURpdnI0U0FpQmpmWVlvTk1qVUhQdmFvalo0OVFNRm5keHpTVnhTR1huVktOdTN5SThsNmc9PSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRha05EUW5SSFowRjNTVUpCWjBsVlZGaFFNRmxRYTFKQlEyUlJla2hEZW1KQ1VIZDVkQzlEWW01WmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDVxU1hwTmFrMHdUVVJSZVZkb1kwNU5hbGwzVG1wSmVrMXFUVEZOUkZGNVYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZaU1VKSWVIQkNhbm96VXpnM1ZqZFZZemRRVW5ORFFubHBWV0pKUzNSemIzQldNR0lLUjJSclFrNTNjVXQ1UnpnMVluaEphbU5uZFRaNVIyTlBkRkpWUVVoSVVYVldVSEUxYzB0V1dDODBVSEZoWmxoeGJEWlBRMEptUVhkbloxaHpUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlZHTmpRekNqVnVOMWwzUTFaNFZIQjVkbWhzVVhOT1ozaEtNVGs0ZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOWFha1Y1V1ZkYWFrOVVTbXBPUkVKcFdXMVJNMDV0VVRWUFZFSnJDazlVUlhoYVYwcG9XVzFaZVU5SFZUQlBWRUpvVGpKS2ExcFVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMXBxUlhsWlYxcHFUMVJLYWs1RVFtbFpiVkV6VG0xUk5VOVVRbXRQVkVWNFdsZEthRmx0V1hrS1QwZFZNRTlVUW1oT01rcHJXbFJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMXBxUlhsWlYxcHFUMVJLYWs1RVFtbFpiVkV6VG0xUk5VOVVRbXNLVDFSRmVGcFhTbWhaYlZsNVQwZFZNRTlVUW1oT01rcHJXbFJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwZFplRTF0Um0wS1dYcHJlVmw2VVhkWmJVcHJUbnBhYTA5VWEzZGFSR3Q0VFZkV2FWbFhTbTFOYW1oc1RrUnJkMWxVWkdsYVIxVjNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOYW1kM1RYcFplVTVxWnpWTmVsVjJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwWjFsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTT0VKSWIwRUtaVUZDTWtGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtNTJZbUpvTkRoQlFVRlJSQXBCUldOM1VsRkpaMWhwTVRnNFdqaHVkMjl2V1ZjemVtVkZSREZtZFdsVFowcG1SRGxzZGtaWGNsWkZhamNyY2k5dVJWRkRTVkZFUXpVNVpGSXZUMXBEQ2pSWU1FTllTbUphTWpCeFdVdzVRblJtT1VobGFXbFpkSHBKYzFab1ZUTk5WR3BCUzBKblozRm9hMnBQVUZGUlJFRjNUbTVCUkVKclFXcEJUMHhsVDBnS2NWcFpRMmxFTXpFeU5rRnhNMFl3TlM5alNIRnlRV05wY1NzNGIwdzRZMnQxZWxOamQwNUlUemxQVmlzME0wYzNRek5uWm5JMVFYaHVZVFJEVFVGMVJ3bzNZMnBzU1dkMmNFVnlTR1kwVURaeFNERmtjMk51Ym5CYVZGQlpUamgxY1daaU5HOVVTSGh2YjA1Nk5VNTBjV3A0ZEZCS1ZYTm5jM0JPYUVRd1FUMDlDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyjADAgEAMIICwQYJKoZIhvcNAQcCoIICsjCCAq4CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQglO01uHSFcZaECDEKPQ7DCcqB1g9Q7D3TkCbaAE+9JTsCFQCSYddF/EKGJi3DKJ5+Js5NaHVUYRgPMjAyNjA2MjMyMzQwNDJaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHbMIIB1wIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MDYyMzIzNDA0MlowLwYJKoZIhvcNAQkEMSIEIFiW6TcbMhwIBGzUomanxoi9pdOpRPdQda4rIdi4jEaOMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRnMGUCMQDUYLCBVAp1lZ0gIH3KQkIhbrLYWl2+FeM23Dht8BK1C82gAyFfHi+XtP4bqvuMYBcCMFbo85kg4r1wdYaIckVhWf76B1YBofVsf/+HM2gfZwdBXbFEi/M3GzgGX9fX/ZYHDg=="}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"DwRbjtJmanPbsuWsdW94Ss8MfJZHBDBZ+j7hUXHzdyk="}, "signature":"MEQCIH2KTvtTQENqjo1uxEDMZ+5aWjX41FrVrYVs0eDivr4SAiBjfYYoNMjUHPvaojZ49QMFndxzSVxSGXnVKNu3yI8l6g=="}} \ No newline at end of file diff --git a/build/torch211-cxx11-cu128-aarch64-linux/ops/__init__.py b/build/torch211-cxx11-cu128-aarch64-linux/ops/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/build/torch211-cxx11-cu128-aarch64-linux/ops/triton/layer_norm.py b/build/torch211-cxx11-cu128-aarch64-linux/ops/triton/layer_norm.py deleted file mode 100644 index 5e6f0b3697b612e6a19031bbc55ac76ad233dc85..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu128-aarch64-linux/ops/triton/layer_norm.py +++ /dev/null @@ -1,1254 +0,0 @@ -# Copyright (c) 2024, Tri Dao. -# Implement dropout + residual + layer_norm / rms_norm. - -# Based on the Triton LayerNorm tutorial: https://triton-lang.org/main/getting-started/tutorials/05-layer-norm.html -# For the backward pass, we keep weight_grad and bias_grad in registers and accumulate. -# This is faster for dimensions up to 8k, but after that it's much slower due to register spilling. -# The models we train have hidden dim up to 8k anyway (e.g. Llama 70B), so this is fine. - -import math -from typing import Optional, List - -import torch -import torch.nn.functional as F -from torch import Tensor - -import triton -import triton.language as tl - -from flash_attn.utils.torch import custom_fwd, custom_bwd -from flash_attn.utils.library import triton_op - -from ..._ops import add_op_namespace_prefix - - -def maybe_contiguous_lastdim(x): - return x.contiguous() if x is not None and x.stride(-1) != 1 else x - - -def maybe_contiguous(x): - return x.contiguous() if x is not None else None - - -def triton_autotune_configs(): - # Return configs with a valid warp count for the current device - configs = [] - # Maximum threads per block is architecture-dependent in theory, but in reality all are 1024 - max_threads_per_block = 1024 - # Default to warp size 32 if not defined by device - warp_size = getattr(torch.cuda.get_device_properties(torch.cuda.current_device()), "warp_size", 32) - # Autotune for warp counts which are powers of 2 and do not exceed thread per block limit - return [triton.Config({}, num_warps=warp_count) for warp_count in [1, 2, 4, 8, 16, 32] - if warp_count * warp_size <= max_threads_per_block] - # return [triton.Config({}, num_warps=8)] - - -def layer_norm_ref( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - zero_centered_weight=False, - dropout_mask=None, - dropout_mask1=None, - upcast=False, -): - dtype = x.dtype - if upcast: - x = x.float() - weight = weight.float() - bias = bias.float() if bias is not None else None - residual = residual.float() if residual is not None else residual - x1 = x1.float() if x1 is not None else None - weight1 = weight1.float() if weight1 is not None else None - bias1 = bias1.float() if bias1 is not None else None - if zero_centered_weight: - weight = weight + 1.0 - if weight1 is not None: - weight1 = weight1 + 1.0 - if x1 is not None: - assert rowscale is None, "rowscale is not supported with parallel LayerNorm" - if rowscale is not None: - x = x * rowscale[..., None] - if dropout_p > 0.0: - if dropout_mask is not None: - x = x.masked_fill(~dropout_mask, 0.0) / (1.0 - dropout_p) - else: - x = F.dropout(x, p=dropout_p) - if x1 is not None: - if dropout_mask1 is not None: - x1 = x1.masked_fill(~dropout_mask1, 0.0) / (1.0 - dropout_p) - else: - x1 = F.dropout(x1, p=dropout_p) - if x1 is not None: - x = x + x1 - if residual is not None: - x = (x + residual).to(x.dtype) - out = F.layer_norm(x.to(weight.dtype), x.shape[-1:], weight=weight, bias=bias, eps=eps).to( - dtype - ) - if weight1 is None: - return out if not prenorm else (out, x) - else: - out1 = F.layer_norm( - x.to(weight1.dtype), x.shape[-1:], weight=weight1, bias=bias1, eps=eps - ).to(dtype) - return (out, out1) if not prenorm else (out, out1, x) - - -def rms_norm_ref( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - zero_centered_weight=False, - dropout_mask=None, - dropout_mask1=None, - upcast=False, -): - dtype = x.dtype - if upcast: - x = x.float() - weight = weight.float() - bias = bias.float() if bias is not None else None - residual = residual.float() if residual is not None else residual - x1 = x1.float() if x1 is not None else None - weight1 = weight1.float() if weight1 is not None else None - bias1 = bias1.float() if bias1 is not None else None - if zero_centered_weight: - weight = weight + 1.0 - if weight1 is not None: - weight1 = weight1 + 1.0 - if x1 is not None: - assert rowscale is None, "rowscale is not supported with parallel LayerNorm" - if rowscale is not None: - x = x * rowscale[..., None] - if dropout_p > 0.0: - if dropout_mask is not None: - x = x.masked_fill(~dropout_mask, 0.0) / (1.0 - dropout_p) - else: - x = F.dropout(x, p=dropout_p) - if x1 is not None: - if dropout_mask1 is not None: - x1 = x1.masked_fill(~dropout_mask1, 0.0) / (1.0 - dropout_p) - else: - x1 = F.dropout(x1, p=dropout_p) - if x1 is not None: - x = x + x1 - if residual is not None: - x = (x + residual).to(x.dtype) - rstd = 1 / torch.sqrt((x.square()).mean(dim=-1, keepdim=True) + eps) - out = ((x * rstd * weight) + bias if bias is not None else (x * rstd * weight)).to(dtype) - if weight1 is None: - return out if not prenorm else (out, x) - else: - out1 = ((x * rstd * weight1) + bias1 if bias1 is not None else (x * rstd * weight1)).to( - dtype - ) - return (out, out1) if not prenorm else (out, out1, x) - - -@triton.autotune( - configs=triton_autotune_configs(), - key=["N", "HAS_RESIDUAL", "STORE_RESIDUAL_OUT", "IS_RMS_NORM", "HAS_BIAS", "HAS_X1", "HAS_W1", "HAS_B1"], -) -# torch compile doesn't like triton.heuristics, so we set these manually when calling the kernel -# @triton.heuristics({"HAS_BIAS": lambda args: args["B"] is not None}) -# @triton.heuristics({"HAS_RESIDUAL": lambda args: args["RESIDUAL"] is not None}) -# @triton.heuristics({"HAS_X1": lambda args: args["X1"] is not None}) -# @triton.heuristics({"HAS_W1": lambda args: args["W1"] is not None}) -# @triton.heuristics({"HAS_B1": lambda args: args["B1"] is not None}) -@triton.jit -def _layer_norm_fwd_1pass_kernel( - X, # pointer to the input - Y, # pointer to the output - W, # pointer to the weights - B, # pointer to the biases - RESIDUAL, # pointer to the residual - X1, - W1, - B1, - Y1, - RESIDUAL_OUT, # pointer to the residual - ROWSCALE, - SEEDS, # Dropout seeds for each row - DROPOUT_MASK, - DROPOUT_MASK1, - Mean, # pointer to the mean - Rstd, # pointer to the 1/std - stride_x_row, # how much to increase the pointer when moving by 1 row - stride_y_row, - stride_res_row, - stride_res_out_row, - stride_x1_row, - stride_y1_row, - M, # number of rows in X - N, # number of columns in X - eps, # epsilon to avoid division by zero - dropout_p, # Dropout probability - zero_centered_weight, # If true, add 1.0 to the weight - IS_RMS_NORM: tl.constexpr, - BLOCK_N: tl.constexpr, - HAS_RESIDUAL: tl.constexpr, - STORE_RESIDUAL_OUT: tl.constexpr, - HAS_BIAS: tl.constexpr, - HAS_DROPOUT: tl.constexpr, - STORE_DROPOUT_MASK: tl.constexpr, - HAS_ROWSCALE: tl.constexpr, - HAS_X1: tl.constexpr, - HAS_W1: tl.constexpr, - HAS_B1: tl.constexpr, -): - # Map the program id to the row of X and Y it should compute. - row = tl.program_id(0) - X += row * stride_x_row - Y += row * stride_y_row - if HAS_RESIDUAL: - RESIDUAL += row * stride_res_row - if STORE_RESIDUAL_OUT: - RESIDUAL_OUT += row * stride_res_out_row - if HAS_X1: - X1 += row * stride_x1_row - if HAS_W1: - Y1 += row * stride_y1_row - # Compute mean and variance - cols = tl.arange(0, BLOCK_N) - x = tl.load(X + cols, mask=cols < N, other=0.0).to(tl.float32) - if HAS_ROWSCALE: - rowscale = tl.load(ROWSCALE + row).to(tl.float32) - x *= rowscale - if HAS_DROPOUT: - # Compute dropout mask - # 7 rounds is good enough, and reduces register pressure - keep_mask = tl.rand(tl.load(SEEDS + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - x = tl.where(keep_mask, x / (1.0 - dropout_p), 0.0) - if STORE_DROPOUT_MASK: - tl.store(DROPOUT_MASK + row * N + cols, keep_mask, mask=cols < N) - if HAS_X1: - x1 = tl.load(X1 + cols, mask=cols < N, other=0.0).to(tl.float32) - if HAS_ROWSCALE: - rowscale = tl.load(ROWSCALE + M + row).to(tl.float32) - x1 *= rowscale - if HAS_DROPOUT: - # Compute dropout mask - # 7 rounds is good enough, and reduces register pressure - keep_mask = ( - tl.rand(tl.load(SEEDS + M + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - ) - x1 = tl.where(keep_mask, x1 / (1.0 - dropout_p), 0.0) - if STORE_DROPOUT_MASK: - tl.store(DROPOUT_MASK1 + row * N + cols, keep_mask, mask=cols < N) - x += x1 - if HAS_RESIDUAL: - residual = tl.load(RESIDUAL + cols, mask=cols < N, other=0.0).to(tl.float32) - x += residual - if STORE_RESIDUAL_OUT: - tl.store(RESIDUAL_OUT + cols, x, mask=cols < N) - if not IS_RMS_NORM: - mean = tl.sum(x, axis=0) / N - tl.store(Mean + row, mean) - xbar = tl.where(cols < N, x - mean, 0.0) - var = tl.sum(xbar * xbar, axis=0) / N - else: - xbar = tl.where(cols < N, x, 0.0) - var = tl.sum(xbar * xbar, axis=0) / N - rstd = 1 / tl.sqrt(var + eps) - tl.store(Rstd + row, rstd) - # Normalize and apply linear transformation - mask = cols < N - w = tl.load(W + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w += 1.0 - if HAS_BIAS: - b = tl.load(B + cols, mask=mask).to(tl.float32) - x_hat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd - y = x_hat * w + b if HAS_BIAS else x_hat * w - # Write output - tl.store(Y + cols, y, mask=mask) - if HAS_W1: - w1 = tl.load(W1 + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w1 += 1.0 - if HAS_B1: - b1 = tl.load(B1 + cols, mask=mask).to(tl.float32) - y1 = x_hat * w1 + b1 if HAS_B1 else x_hat * w1 - tl.store(Y1 + cols, y1, mask=mask) - - -def _layer_norm_fwd( - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - residual: Optional[Tensor] = None, - x1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - out_dtype: Optional[torch.dtype] = None, - residual_dtype: Optional[torch.dtype] = None, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - return_dropout_mask: bool = False, - out: Optional[Tensor] = None, - residual_out: Optional[Tensor] = None -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - # Need to wrap to handle the case where residual_out is a alias of x, which makes torch.library - # and torch.compile unhappy. Also allocate memory for out and residual_out if they are None - # so that _layer_norm_fwd_impl doesn't have to return them. - if out is None: - out = torch.empty_like(x, dtype=x.dtype if out_dtype is None else out_dtype) - if residual is not None: - residual_dtype = residual.dtype - if residual_out is None and ( - residual is not None - or (residual_dtype is not None and residual_dtype != x.dtype) - or dropout_p > 0.0 - or rowscale is not None - or x1 is not None - ): - residual_out = torch.empty_like( - x, dtype=residual_dtype if residual_dtype is not None else x.dtype - ) - else: - residual_out = None - y1, mean, rstd, seeds, dropout_mask, dropout_mask1 = _layer_norm_fwd_impl( - x, - weight, - bias, - eps, - out, - residual=residual, - x1=x1, - weight1=weight1, - bias1=bias1, - dropout_p=dropout_p, - rowscale=rowscale, - zero_centered_weight=zero_centered_weight, - is_rms_norm=is_rms_norm, - return_dropout_mask=return_dropout_mask, - residual_out=residual_out, - ) - # residual_out is None if residual is None and residual_dtype == input_dtype and dropout_p == 0.0 - if residual_out is None: - residual_out = x - return out, y1, mean, rstd, residual_out, seeds, dropout_mask, dropout_mask1 - - -# [2025-04-28] torch.library.triton_op ignores the schema argument, but here we need the schema -# since we're returning a tuple of tensors -@triton_op(add_op_namespace_prefix("layer_norm_fwd_impl"), mutates_args={"out", "residual_out"}, - schema="(Tensor x, Tensor weight, Tensor bias, float eps, Tensor(a!) out, Tensor? residual, Tensor? x1, Tensor? weight1, Tensor? bias1, float dropout_p, Tensor? rowscale, bool zero_centered_weight, bool is_rms_norm, bool return_dropout_mask, Tensor(a!)? residual_out) -> (Tensor y1, Tensor mean, Tensor rstd, Tensor seeds, Tensor dropout_mask, Tensor dropout_mask1)") -def _layer_norm_fwd_impl( - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - out: Tensor, - residual: Optional[Tensor] = None, - x1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - return_dropout_mask: bool = False, - residual_out: Optional[Tensor] = None -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - M, N = x.shape - assert x.stride(-1) == 1 - if residual is not None: - assert residual.stride(-1) == 1 - assert residual.shape == (M, N) - assert weight.shape == (N,) - assert weight.stride(-1) == 1 - if bias is not None: - assert bias.stride(-1) == 1 - assert bias.shape == (N,) - if x1 is not None: - assert x1.shape == x.shape - assert rowscale is None - assert x1.stride(-1) == 1 - if weight1 is not None: - assert weight1.shape == (N,) - assert weight1.stride(-1) == 1 - if bias1 is not None: - assert bias1.shape == (N,) - assert bias1.stride(-1) == 1 - if rowscale is not None: - assert rowscale.is_contiguous() - assert rowscale.shape == (M,) - assert out.shape == x.shape - assert out.stride(-1) == 1 - if residual_out is not None: - assert residual_out.shape == x.shape - assert residual_out.stride(-1) == 1 - if weight1 is not None: - y1 = torch.empty_like(out) - assert y1.stride(-1) == 1 - else: - y1 = None - mean = torch.empty((M,), dtype=torch.float32, device=x.device) if not is_rms_norm else None - rstd = torch.empty((M,), dtype=torch.float32, device=x.device) - if dropout_p > 0.0: - seeds = torch.randint( - 2**32, (M if x1 is None else 2 * M,), device=x.device, dtype=torch.int64 - ) - else: - seeds = None - if return_dropout_mask and dropout_p > 0.0: - dropout_mask = torch.empty(M, N, device=x.device, dtype=torch.bool) - if x1 is not None: - dropout_mask1 = torch.empty(M, N, device=x.device, dtype=torch.bool) - else: - dropout_mask1 = None - else: - dropout_mask, dropout_mask1 = None, None - # Less than 64KB per feature: enqueue fused kernel - MAX_FUSED_SIZE = 65536 // x.element_size() - BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) - if N > BLOCK_N: - raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") - with torch.cuda.device(x.device.index): - torch.library.wrap_triton(_layer_norm_fwd_1pass_kernel)[(M,)]( - x, - out, - weight, - bias, - residual, - x1, - weight1, - bias1, - y1, - residual_out, - rowscale, - seeds, - dropout_mask, - dropout_mask1, - mean, - rstd, - x.stride(0), - out.stride(0), - residual.stride(0) if residual is not None else 0, - residual_out.stride(0) if residual_out is not None else 0, - x1.stride(0) if x1 is not None else 0, - y1.stride(0) if y1 is not None else 0, - M, - N, - eps, - dropout_p, - # Passing bool make torch inductor very unhappy since it then tries to compare to int_max - int(zero_centered_weight), - is_rms_norm, - BLOCK_N, - residual is not None, - residual_out is not None, - bias is not None, - dropout_p > 0.0, - dropout_mask is not None, - rowscale is not None, - HAS_X1=x1 is not None, - HAS_W1=weight1 is not None, - HAS_B1=bias1 is not None, - ) - return y1, mean, rstd, seeds, dropout_mask, dropout_mask1 - - -@triton.autotune( - configs=triton_autotune_configs(), - key=["N", "HAS_DRESIDUAL", "STORE_DRESIDUAL", "IS_RMS_NORM", "HAS_BIAS", "HAS_DROPOUT"], -) -# torch compile doesn't like triton.heuristics, so we set these manually when calling the kernel -# @triton.heuristics({"HAS_BIAS": lambda args: args["B"] is not None}) -# @triton.heuristics({"HAS_DRESIDUAL": lambda args: args["DRESIDUAL"] is not None}) -# @triton.heuristics({"STORE_DRESIDUAL": lambda args: args["DRESIDUAL_IN"] is not None}) -# @triton.heuristics({"HAS_ROWSCALE": lambda args: args["ROWSCALE"] is not None}) -# @triton.heuristics({"HAS_DY1": lambda args: args["DY1"] is not None}) -# @triton.heuristics({"HAS_DX1": lambda args: args["DX1"] is not None}) -# @triton.heuristics({"HAS_B1": lambda args: args["DB1"] is not None}) -# @triton.heuristics({"RECOMPUTE_OUTPUT": lambda args: args["Y"] is not None}) -@triton.jit -def _layer_norm_bwd_kernel( - X, # pointer to the input - W, # pointer to the weights - B, # pointer to the biases - Y, # pointer to the output to be recomputed - DY, # pointer to the output gradient - DX, # pointer to the input gradient - DW, # pointer to the partial sum of weights gradient - DB, # pointer to the partial sum of biases gradient - DRESIDUAL, - W1, - DY1, - DX1, - DW1, - DB1, - DRESIDUAL_IN, - ROWSCALE, - SEEDS, - Mean, # pointer to the mean - Rstd, # pointer to the 1/std - stride_x_row, # how much to increase the pointer when moving by 1 row - stride_y_row, - stride_dy_row, - stride_dx_row, - stride_dres_row, - stride_dy1_row, - stride_dx1_row, - stride_dres_in_row, - M, # number of rows in X - N, # number of columns in X - eps, # epsilon to avoid division by zero - dropout_p, - zero_centered_weight, - rows_per_program, - IS_RMS_NORM: tl.constexpr, - BLOCK_N: tl.constexpr, - HAS_DRESIDUAL: tl.constexpr, - STORE_DRESIDUAL: tl.constexpr, - HAS_BIAS: tl.constexpr, - HAS_DROPOUT: tl.constexpr, - HAS_ROWSCALE: tl.constexpr, - HAS_DY1: tl.constexpr, - HAS_DX1: tl.constexpr, - HAS_B1: tl.constexpr, - RECOMPUTE_OUTPUT: tl.constexpr, -): - # Map the program id to the elements of X, DX, and DY it should compute. - row_block_id = tl.program_id(0) - row_start = row_block_id * rows_per_program - # Do not early exit if row_start >= M, because we need to write DW and DB - cols = tl.arange(0, BLOCK_N) - mask = cols < N - X += row_start * stride_x_row - if HAS_DRESIDUAL: - DRESIDUAL += row_start * stride_dres_row - if STORE_DRESIDUAL: - DRESIDUAL_IN += row_start * stride_dres_in_row - DY += row_start * stride_dy_row - DX += row_start * stride_dx_row - if HAS_DY1: - DY1 += row_start * stride_dy1_row - if HAS_DX1: - DX1 += row_start * stride_dx1_row - if RECOMPUTE_OUTPUT: - Y += row_start * stride_y_row - w = tl.load(W + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w += 1.0 - if RECOMPUTE_OUTPUT and HAS_BIAS: - b = tl.load(B + cols, mask=mask, other=0.0).to(tl.float32) - if HAS_DY1: - w1 = tl.load(W1 + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w1 += 1.0 - dw = tl.zeros((BLOCK_N,), dtype=tl.float32) - if HAS_BIAS: - db = tl.zeros((BLOCK_N,), dtype=tl.float32) - if HAS_DY1: - dw1 = tl.zeros((BLOCK_N,), dtype=tl.float32) - if HAS_B1: - db1 = tl.zeros((BLOCK_N,), dtype=tl.float32) - row_end = min((row_block_id + 1) * rows_per_program, M) - for row in range(row_start, row_end): - # Load data to SRAM - x = tl.load(X + cols, mask=mask, other=0).to(tl.float32) - dy = tl.load(DY + cols, mask=mask, other=0).to(tl.float32) - if HAS_DY1: - dy1 = tl.load(DY1 + cols, mask=mask, other=0).to(tl.float32) - if not IS_RMS_NORM: - mean = tl.load(Mean + row) - rstd = tl.load(Rstd + row) - # Compute dx - xhat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd - xhat = tl.where(mask, xhat, 0.0) - if RECOMPUTE_OUTPUT: - y = xhat * w + b if HAS_BIAS else xhat * w - tl.store(Y + cols, y, mask=mask) - wdy = w * dy - dw += dy * xhat - if HAS_BIAS: - db += dy - if HAS_DY1: - wdy += w1 * dy1 - dw1 += dy1 * xhat - if HAS_B1: - db1 += dy1 - if not IS_RMS_NORM: - c1 = tl.sum(xhat * wdy, axis=0) / N - c2 = tl.sum(wdy, axis=0) / N - dx = (wdy - (xhat * c1 + c2)) * rstd - else: - c1 = tl.sum(xhat * wdy, axis=0) / N - dx = (wdy - xhat * c1) * rstd - if HAS_DRESIDUAL: - dres = tl.load(DRESIDUAL + cols, mask=mask, other=0).to(tl.float32) - dx += dres - # Write dx - if STORE_DRESIDUAL: - tl.store(DRESIDUAL_IN + cols, dx, mask=mask) - if HAS_DX1: - if HAS_DROPOUT: - keep_mask = ( - tl.rand(tl.load(SEEDS + M + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - ) - dx1 = tl.where(keep_mask, dx / (1.0 - dropout_p), 0.0) - else: - dx1 = dx - tl.store(DX1 + cols, dx1, mask=mask) - if HAS_DROPOUT: - keep_mask = tl.rand(tl.load(SEEDS + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - dx = tl.where(keep_mask, dx / (1.0 - dropout_p), 0.0) - if HAS_ROWSCALE: - rowscale = tl.load(ROWSCALE + row).to(tl.float32) - dx *= rowscale - tl.store(DX + cols, dx, mask=mask) - - X += stride_x_row - if HAS_DRESIDUAL: - DRESIDUAL += stride_dres_row - if STORE_DRESIDUAL: - DRESIDUAL_IN += stride_dres_in_row - if RECOMPUTE_OUTPUT: - Y += stride_y_row - DY += stride_dy_row - DX += stride_dx_row - if HAS_DY1: - DY1 += stride_dy1_row - if HAS_DX1: - DX1 += stride_dx1_row - tl.store(DW + row_block_id * N + cols, dw, mask=mask) - if HAS_BIAS: - tl.store(DB + row_block_id * N + cols, db, mask=mask) - if HAS_DY1: - tl.store(DW1 + row_block_id * N + cols, dw1, mask=mask) - if HAS_B1: - tl.store(DB1 + row_block_id * N + cols, db1, mask=mask) - - -def _layer_norm_bwd( - dy: Tensor, - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - mean: Tensor, - rstd: Tensor, - dresidual: Optional[Tensor] = None, - dy1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - seeds: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - has_residual: bool = False, - has_x1: bool = False, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - x_dtype: Optional[torch.dtype] = None, - recompute_output: bool = False, -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - # Need to wrap to handle the case where dresidual_in or dx1 are aliases of x, - # which makes torch.library unhappy - dx, dw, db, dresidual_in, dx1, dw1, db1, y = _layer_norm_bwd_impl( - dy, - x, - weight, - bias, - eps, - mean, - rstd, - dresidual, - dy1, - weight1, - bias1, - seeds, - dropout_p, - rowscale, - has_residual, - has_x1, - zero_centered_weight, - is_rms_norm, - x_dtype=x_dtype, - recompute_output=recompute_output, - ) - # Don't need to compute dresidual_in separately in this case - if has_residual and dx.dtype == x.dtype and dropout_p == 0.0 and rowscale is None: - dresidual_in = dx - if has_x1 and dropout_p == 0.0: - dx1 = dx - return dx, dw, db, dresidual_in, dx1, dw1, db1, y - - - -@triton_op(add_op_namespace_prefix("layer_norm_bwd_impl"), mutates_args={}, - schema="(Tensor dy, Tensor x, Tensor weight, Tensor bias, float eps, Tensor mean, Tensor rstd, Tensor? dresidual, Tensor? dy1, Tensor? weight1, Tensor? bias1, Tensor? seeds, float dropout_p, Tensor? rowscale, bool has_residual, bool has_x1, bool zero_centered_weight, bool is_rms_norm, ScalarType? x_dtype, bool recompute_output) -> (Tensor dx, Tensor dw, Tensor db, Tensor dresidual_in, Tensor dx1, Tensor dw1, Tensor db1, Tensor y)", - allow_decomposition=False, # Don't let torch.compile trace inside - ) -def _layer_norm_bwd_impl( - dy: Tensor, - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - mean: Tensor, - rstd: Tensor, - dresidual: Optional[Tensor] = None, - dy1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - seeds: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - has_residual: bool = False, - has_x1: bool = False, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - x_dtype: Optional[torch.dtype] = None, - recompute_output: bool = False, -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - M, N = x.shape - assert x.stride(-1) == 1 - dy = maybe_contiguous_lastdim(dy) - assert dy.stride(-1) == 1 - assert dy.shape == (M, N) - if dresidual is not None: - dresidual = maybe_contiguous_lastdim(dresidual) - assert dresidual.stride(-1) == 1 - assert dresidual.shape == (M, N) - assert weight.shape == (N,) - assert weight.stride(-1) == 1 - if bias is not None: - assert bias.stride(-1) == 1 - assert bias.shape == (N,) - if dy1 is not None: - dy1 = maybe_contiguous_lastdim(dy1) - assert weight1 is not None - assert dy1.shape == dy.shape - assert dy1.stride(-1) == 1 - if weight1 is not None: - assert weight1.shape == (N,) - assert weight1.stride(-1) == 1 - if bias1 is not None: - assert bias1.shape == (N,) - assert bias1.stride(-1) == 1 - if seeds is not None: - assert seeds.is_contiguous() - assert seeds.shape == (M if not has_x1 else M * 2,) - if rowscale is not None: - assert rowscale.is_contiguous() - assert rowscale.shape == (M,) - # allocate output - dx = ( - torch.empty_like(x) - if x_dtype is None - else torch.empty(M, N, dtype=x_dtype, device=x.device) - ) - dresidual_in = ( - torch.empty_like(x) - if has_residual - and (dx.dtype != x.dtype or dropout_p > 0.0 or rowscale is not None or has_x1) - else None - ) - dx1 = torch.empty_like(dx) if (has_x1 and dropout_p > 0.0) else None - y = torch.empty(M, N, dtype=dy.dtype, device=dy.device) if recompute_output else None - if recompute_output: - assert weight1 is None, "recompute_output is not supported with parallel LayerNorm" - - # Less than 64KB per feature: enqueue fused kernel - MAX_FUSED_SIZE = 65536 // x.element_size() - BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) - if N > BLOCK_N: - raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") - # Increasing the multiple (e.g. 8) will allow more thread blocks to be launched and hide the - # latency of the gmem reads/writes, but will increase the time of summing up dw / db. - sm_count = torch.cuda.get_device_properties(x.device).multi_processor_count * 8 - _dw = torch.empty((sm_count, N), dtype=torch.float32, device=weight.device) - _db = ( - torch.empty((sm_count, N), dtype=torch.float32, device=bias.device) - if bias is not None - else None - ) - _dw1 = torch.empty_like(_dw) if weight1 is not None else None - _db1 = torch.empty_like(_db) if bias1 is not None else None - rows_per_program = math.ceil(M / sm_count) - grid = (sm_count,) - with torch.cuda.device(x.device.index): - torch.library.wrap_triton(_layer_norm_bwd_kernel)[grid]( - x, - weight, - bias, - y, - dy, - dx, - _dw, - _db, - dresidual, - weight1, - dy1, - dx1, - _dw1, - _db1, - dresidual_in, - rowscale, - seeds, - mean, - rstd, - x.stride(0), - 0 if not recompute_output else y.stride(0), - dy.stride(0), - dx.stride(0), - dresidual.stride(0) if dresidual is not None else 0, - dy1.stride(0) if dy1 is not None else 0, - dx1.stride(0) if dx1 is not None else 0, - dresidual_in.stride(0) if dresidual_in is not None else 0, - M, - N, - eps, - dropout_p, - # Passing bool make torch inductor very unhappy since it then tries to compare to int_max - int(zero_centered_weight), - rows_per_program, - is_rms_norm, - BLOCK_N, - dresidual is not None, - dresidual_in is not None, - bias is not None, - dropout_p > 0.0, - HAS_ROWSCALE=rowscale is not None, - HAS_DY1=dy1 is not None, - HAS_DX1=dx1 is not None, - HAS_B1=bias1 is not None, - RECOMPUTE_OUTPUT=y is not None, - ) - dw = _dw.sum(0).to(weight.dtype) - db = _db.sum(0).to(bias.dtype) if bias is not None else None - dw1 = _dw1.sum(0).to(weight1.dtype) if weight1 is not None else None - db1 = _db1.sum(0).to(bias1.dtype) if bias1 is not None else None - # dresidual_in and dx1 could be None, the wrapper will handle assigning them from dx - return dx, dw, db, dresidual_in, dx1, dw1, db1, y - - -class LayerNormFn(torch.autograd.Function): - - @staticmethod - def forward( - ctx, - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - residual_in_fp32=False, - zero_centered_weight=False, - is_rms_norm=False, - return_dropout_mask=False, - out_dtype=None, - out=None, - residual_out=None - ): - x_shape_og = x.shape - # reshape input data into 2D tensor - x = maybe_contiguous_lastdim(x.reshape(-1, x.shape[-1])) - if residual is not None: - assert residual.shape == x_shape_og - residual = maybe_contiguous_lastdim(residual.reshape(-1, residual.shape[-1])) - if x1 is not None: - assert x1.shape == x_shape_og - assert rowscale is None, "rowscale is not supported with parallel LayerNorm" - x1 = maybe_contiguous_lastdim(x1.reshape(-1, x1.shape[-1])) - weight = weight.contiguous() - bias = maybe_contiguous(bias) - weight1 = maybe_contiguous(weight1) - bias1 = maybe_contiguous(bias1) - if rowscale is not None: - rowscale = rowscale.reshape(-1).contiguous() - residual_dtype = ( - residual.dtype - if residual is not None - else (torch.float32 if residual_in_fp32 else None) - ) - if out is not None: - out = out.reshape(-1, out.shape[-1]) - if residual_out is not None: - residual_out = residual_out.reshape(-1, residual_out.shape[-1]) - y, y1, mean, rstd, residual_out, seeds, dropout_mask, dropout_mask1 = _layer_norm_fwd( - x, - weight, - bias, - eps, - residual, - x1, - weight1, - bias1, - dropout_p=dropout_p, - rowscale=rowscale, - out_dtype=out_dtype, - residual_dtype=residual_dtype, - zero_centered_weight=zero_centered_weight, - is_rms_norm=is_rms_norm, - return_dropout_mask=return_dropout_mask, - out=out, - residual_out=residual_out, - ) - ctx.save_for_backward( - residual_out, weight, bias, weight1, bias1, rowscale, seeds, mean, rstd - ) - ctx.x_shape_og = x_shape_og - ctx.eps = eps - ctx.dropout_p = dropout_p - ctx.is_rms_norm = is_rms_norm - ctx.has_residual = residual is not None - ctx.has_x1 = x1 is not None - ctx.prenorm = prenorm - ctx.x_dtype = x.dtype - ctx.zero_centered_weight = zero_centered_weight - y = y.reshape(x_shape_og) - y1 = y1.reshape(x_shape_og) if y1 is not None else None - residual_out = residual_out.reshape(x_shape_og) if residual_out is not None else None - dropout_mask = dropout_mask.reshape(x_shape_og) if dropout_mask is not None else None - dropout_mask1 = dropout_mask1.reshape(x_shape_og) if dropout_mask1 is not None else None - if not return_dropout_mask: - if weight1 is None: - return y if not prenorm else (y, residual_out) - else: - return (y, y1) if not prenorm else (y, y1, residual_out) - else: - if weight1 is None: - return ( - (y, dropout_mask, dropout_mask1) - if not prenorm - else (y, residual_out, dropout_mask, dropout_mask1) - ) - else: - return ( - (y, y1, dropout_mask, dropout_mask1) - if not prenorm - else (y, y1, residual_out, dropout_mask, dropout_mask1) - ) - - @staticmethod - def backward(ctx, dy, *args): - x, weight, bias, weight1, bias1, rowscale, seeds, mean, rstd = ctx.saved_tensors - dy = dy.reshape(-1, dy.shape[-1]) - if weight1 is not None: - dy1, args = args[0], args[1:] - dy1 = dy1.reshape(-1, dy1.shape[-1]) - assert dy1.shape == x.shape - else: - dy1 = None - if ctx.prenorm: - dresidual = args[0] - dresidual = dresidual.reshape(-1, dresidual.shape[-1]) - assert dresidual.shape == x.shape - else: - dresidual = None - dx, dw, db, dresidual_in, dx1, dw1, db1, _ = _layer_norm_bwd( - dy, - x, - weight, - bias, - ctx.eps, - mean, - rstd, - dresidual, - dy1, - weight1, - bias1, - seeds, - ctx.dropout_p, - rowscale, - ctx.has_residual, - ctx.has_x1, - ctx.zero_centered_weight, - ctx.is_rms_norm, - x_dtype=ctx.x_dtype, - recompute_output=False, - ) - return ( - dx.reshape(ctx.x_shape_og), - dw, - db, - dresidual_in.reshape(ctx.x_shape_og) if ctx.has_residual else None, - dx1.reshape(ctx.x_shape_og) if dx1 is not None else None, - dw1, - db1, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) - - -def layer_norm_fn( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - residual_in_fp32=False, - zero_centered_weight=False, - is_rms_norm=False, - return_dropout_mask=False, - out_dtype=None, - out=None, - residual_out=None -): - return LayerNormFn.apply( - x, - weight, - bias, - residual, - x1, - weight1, - bias1, - eps, - dropout_p, - rowscale, - prenorm, - residual_in_fp32, - zero_centered_weight, - is_rms_norm, - return_dropout_mask, - out_dtype, - out, - residual_out - ) - - -def rms_norm_fn( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - residual_in_fp32=False, - zero_centered_weight=False, - return_dropout_mask=False, - out_dtype=None, - out=None, - residual_out=None -): - return LayerNormFn.apply( - x, - weight, - bias, - residual, - x1, - weight1, - bias1, - eps, - dropout_p, - rowscale, - prenorm, - residual_in_fp32, - zero_centered_weight, - True, - return_dropout_mask, - out_dtype, - out, - residual_out - ) - - -class RMSNorm(torch.nn.Module): - - def __init__(self, hidden_size, eps=1e-5, dropout_p=0.0, zero_centered_weight=False, - device=None, dtype=None): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.eps = eps - if dropout_p > 0.0: - self.drop = torch.nn.Dropout(dropout_p) - else: - self.drop = None - self.zero_centered_weight = zero_centered_weight - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.register_parameter("bias", None) - self.reset_parameters() - - def reset_parameters(self): - if not self.zero_centered_weight: - torch.nn.init.ones_(self.weight) - else: - torch.nn.init.zeros_(self.weight) - - def forward(self, x, residual=None, prenorm=False, residual_in_fp32=False): - return rms_norm_fn( - x, - self.weight, - self.bias, - residual=residual, - eps=self.eps, - dropout_p=self.drop.p if self.drop is not None and self.training else 0.0, - prenorm=prenorm, - residual_in_fp32=residual_in_fp32, - zero_centered_weight=self.zero_centered_weight, - ) - - -class LayerNormLinearFn(torch.autograd.Function): - - @staticmethod - @custom_fwd - def forward( - ctx, - x, - norm_weight, - norm_bias, - linear_weight, - linear_bias, - residual=None, - eps=1e-6, - prenorm=False, - residual_in_fp32=False, - is_rms_norm=False, - ): - x_shape_og = x.shape - # reshape input data into 2D tensor - x = maybe_contiguous_lastdim(x.reshape(-1, x.shape[-1])) - if residual is not None: - assert residual.shape == x_shape_og - residual = maybe_contiguous_lastdim(residual.reshape(-1, residual.shape[-1])) - norm_weight = norm_weight.contiguous() - norm_bias = maybe_contiguous(norm_bias) - residual_dtype = ( - residual.dtype - if residual is not None - else (torch.float32 if residual_in_fp32 else None) - ) - y, _, mean, rstd, residual_out, *rest = _layer_norm_fwd( - x, - norm_weight, - norm_bias, - eps, - residual, - out_dtype=None if not torch.is_autocast_enabled() else torch.get_autocast_dtype("cuda"), - residual_dtype=residual_dtype, - is_rms_norm=is_rms_norm, - ) - y = y.reshape(x_shape_og) - dtype = torch.get_autocast_dtype("cuda") if torch.is_autocast_enabled() else y.dtype - linear_weight = linear_weight.to(dtype) - linear_bias = linear_bias.to(dtype) if linear_bias is not None else None - out = F.linear(y.to(linear_weight.dtype), linear_weight, linear_bias) - # We don't store y, will be recomputed in the backward pass to save memory - ctx.save_for_backward(residual_out, norm_weight, norm_bias, linear_weight, mean, rstd) - ctx.x_shape_og = x_shape_og - ctx.eps = eps - ctx.is_rms_norm = is_rms_norm - ctx.has_residual = residual is not None - ctx.prenorm = prenorm - ctx.x_dtype = x.dtype - ctx.linear_bias_is_none = linear_bias is None - return out if not prenorm else (out, residual_out.reshape(x_shape_og)) - - @staticmethod - @custom_bwd - def backward(ctx, dout, *args): - x, norm_weight, norm_bias, linear_weight, mean, rstd = ctx.saved_tensors - dout = dout.reshape(-1, dout.shape[-1]) - dy = F.linear(dout, linear_weight.t()) - dlinear_bias = None if ctx.linear_bias_is_none else dout.sum(0) - dy = maybe_contiguous_lastdim(dy) - assert dy.shape == x.shape - if ctx.prenorm: - dresidual = args[0] - dresidual = maybe_contiguous_lastdim(dresidual.reshape(-1, dresidual.shape[-1])) - assert dresidual.shape == x.shape - else: - dresidual = None - dx, dnorm_weight, dnorm_bias, dresidual_in, _, _, _, y = _layer_norm_bwd( - dy, - x, - norm_weight, - norm_bias, - ctx.eps, - mean, - rstd, - dresidual=dresidual, - has_residual=ctx.has_residual, - is_rms_norm=ctx.is_rms_norm, - x_dtype=ctx.x_dtype, - recompute_output=True, - ) - dlinear_weight = torch.einsum("bo,bi->oi", dout, y) - return ( - dx.reshape(ctx.x_shape_og), - dnorm_weight, - dnorm_bias, - dlinear_weight, - dlinear_bias, - dresidual_in.reshape(ctx.x_shape_og) if ctx.has_residual else None, - None, - None, - None, - None, - ) - - -def layer_norm_linear_fn( - x, - norm_weight, - norm_bias, - linear_weight, - linear_bias, - residual=None, - eps=1e-6, - prenorm=False, - residual_in_fp32=False, - is_rms_norm=False, -): - return LayerNormLinearFn.apply( - x, - norm_weight, - norm_bias, - linear_weight, - linear_bias, - residual, - eps, - prenorm, - residual_in_fp32, - is_rms_norm, - ) diff --git a/build/torch211-cxx11-cu128-x86_64-linux/__init__.py b/build/torch211-cxx11-cu128-x86_64-linux/__init__.py deleted file mode 100644 index ecc2f9d896b6c93f90b0a1499856dc0612177422..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu128-x86_64-linux/__init__.py +++ /dev/null @@ -1,393 +0,0 @@ -from typing import Optional, List -import torch -from ._ops import ops as flash_attn_ops -from .flash_attn_interface import ( - flash_attn_func, - flash_attn_kvpacked_func, - flash_attn_qkvpacked_func, - flash_attn_varlen_func, - flash_attn_varlen_kvpacked_func, - flash_attn_varlen_qkvpacked_func, - flash_attn_with_kvcache, -) - - -def fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - return_softmax: bool = False, - gen: Optional[torch.Generator] = None, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention. - - Args: - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Optional output tensor, same shape as q - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - return_softmax: Whether to return softmax weights - gen: Optional random number generator - - Returns: - List of tensors: [output, softmax_lse, (softmax if return_softmax)] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.fwd( - q, - k, - v, - out, - alibi_slopes, - p_dropout, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - gen, - ) - - -def varlen_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - out: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - block_table: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - max_seqlen_q: int = 0, - max_seqlen_k: int = 0, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - zero_tensors: bool = False, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - return_softmax: bool = False, - gen: Optional[torch.Generator] = None, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention with variable sequence lengths. - - Args: - q: Query tensor of shape [total_q, num_heads, head_size] - k: Key tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - v: Value tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1] - cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1] - out: Optional output tensor of shape [total_q, num_heads, head_size] - seqused_k: Optional tensor specifying how many keys to use per batch element [batch_size] - leftpad_k: Optional left padding for keys of shape [batch_size] - block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq] - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - max_seqlen_q: Maximum sequence length for queries - max_seqlen_k: Maximum sequence length for keys - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - zero_tensors: Whether to zero tensors before computation - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - return_softmax: Whether to return softmax weights - gen: Optional random number generator - - Returns: - List of tensors: [output, softmax_lse, (softmax if return_softmax)] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.varlen_fwd( - q, - k, - v, - out, - cu_seqlens_q, - cu_seqlens_k, - seqused_k, - leftpad_k, - block_table, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - p_dropout, - softmax_scale, - zero_tensors, - is_causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - gen, - ) - - -def bwd( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - p_dropout: float = 0.0, - 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, - gen: Optional[torch.Generator] = None, - rng_state: Optional[torch.Tensor] = None, -) -> List[torch.Tensor]: - """ - Backward pass for multi-head attention. - - Args: - dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size] - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size] - softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q] - dq: Optional gradient tensor for queries, same shape as q - dk: Optional gradient tensor for keys, same shape as k - dv: Optional gradient tensor for values, same shape as v - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - deterministic: Whether to use deterministic algorithms - gen: Optional random number generator - rng_state: Optional RNG state from forward pass - - Returns: - List of tensors: [dq, dk, dv] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - alibi_slopes, - p_dropout, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - deterministic, - gen, - rng_state, - ) - - -def varlen_bwd( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - max_seqlen_q: int = 0, - max_seqlen_k: int = 0, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - zero_tensors: bool = False, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - deterministic: bool = False, - gen: Optional[torch.Generator] = None, - rng_state: Optional[torch.Tensor] = None, -) -> List[torch.Tensor]: - """ - Backward pass for multi-head attention with variable sequence lengths. - - Args: - dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size] - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size] - softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q] - cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1] - cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1] - dq: Optional gradient tensor for queries, same shape as q - dk: Optional gradient tensor for keys, same shape as k - dv: Optional gradient tensor for values, same shape as v - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - max_seqlen_q: Maximum sequence length for queries - max_seqlen_k: Maximum sequence length for keys - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - zero_tensors: Whether to zero tensors before computation - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - deterministic: Whether to use deterministic algorithms - gen: Optional random number generator - rng_state: Optional RNG state from forward pass - - Returns: - List of tensors: [dq, dk, dv] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.varlen_bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - p_dropout, - softmax_scale, - zero_tensors, - is_causal, - window_size_left, - window_size_right, - softcap, - deterministic, - gen, - rng_state, - ) - - -def fwd_kvcache( - q: torch.Tensor, - kcache: torch.Tensor, - vcache: torch.Tensor, - k: Optional[torch.Tensor] = None, - v: Optional[torch.Tensor] = None, - seqlens_k: Optional[torch.Tensor] = None, - rotary_cos: Optional[torch.Tensor] = None, - rotary_sin: Optional[torch.Tensor] = None, - cache_batch_idx: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - block_table: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - out: 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, - is_rotary_interleaved: bool = False, - num_splits: int = 1, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention with KV cache. - - Args: - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - kcache: Key cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - vcache: Value cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - k: Optional new keys tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size] - v: Optional new values tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size] - seqlens_k: Optional sequence lengths for keys of shape [batch_size] - rotary_cos: Optional rotary cosine tensor of shape [seqlen_ro, rotary_dim/2] - rotary_sin: Optional rotary sine tensor of shape [seqlen_ro, rotary_dim/2] - cache_batch_idx: Optional indices to index into the KV cache - leftpad_k: Optional left padding for keys of shape [batch_size] - block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq] - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - out: Optional output tensor, same shape as q - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - is_rotary_interleaved: Whether rotary embeddings are interleaved - num_splits: Number of splits for computation - - Returns: - List of tensors: [output, softmax_lse] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.fwd_kvcache( - q, - kcache, - vcache, - k, - v, - seqlens_k, - rotary_cos, - rotary_sin, - cache_batch_idx, - leftpad_k, - block_table, - alibi_slopes, - out, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - is_rotary_interleaved, - num_splits, - ) diff --git a/build/torch211-cxx11-cu128-x86_64-linux/_flash_attn2_cuda_f12afc9.abi3.so b/build/torch211-cxx11-cu128-x86_64-linux/_flash_attn2_cuda_f12afc9.abi3.so deleted file mode 100644 index d700bc7c927f3538d4a3f896d032886fbace5668..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu128-x86_64-linux/_flash_attn2_cuda_f12afc9.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:acc97ffd2eca4d874d28691eab82e18fe76112a2079adabfccc4a6d90f08045a -size 1035876408 diff --git a/build/torch211-cxx11-cu128-x86_64-linux/_ops.py b/build/torch211-cxx11-cu128-x86_64-linux/_ops.py deleted file mode 100644 index af5b731cb04b5e6ce54e862db0501f3683c3163f..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu128-x86_64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _flash_attn2_cuda_f12afc9 -ops = torch.ops._flash_attn2_cuda_f12afc9 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_flash_attn2_cuda_f12afc9::{op_name}" diff --git a/build/torch211-cxx11-cu128-x86_64-linux/bert_padding.py b/build/torch211-cxx11-cu128-x86_64-linux/bert_padding.py deleted file mode 100644 index 3c2d35159a014a9d03aabead9e52e009168696ea..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu128-x86_64-linux/bert_padding.py +++ /dev/null @@ -1,218 +0,0 @@ -# Adapted from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/padding.py - -import torch -import torch.nn.functional as F -from einops import rearrange, repeat - - -class IndexFirstAxis(torch.autograd.Function): - @staticmethod - def forward(ctx, input, indices): - ctx.save_for_backward(indices) - assert input.ndim >= 2 - ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] - second_dim = other_shape.numel() - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - # return input[indices] - return torch.gather( - rearrange(input, "b ... -> b (...)"), 0, repeat(indices, "z -> z d", d=second_dim) - ).reshape(-1, *other_shape) - - @staticmethod - def backward(ctx, grad_output): - (indices,) = ctx.saved_tensors - assert grad_output.ndim >= 2 - other_shape = grad_output.shape[1:] - grad_output = rearrange(grad_output, "b ... -> b (...)") - grad_input = torch.zeros( - [ctx.first_axis_dim, grad_output.shape[1]], - device=grad_output.device, - dtype=grad_output.dtype, - ) - # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. - # grad_input[indices] = grad_output - grad_input.scatter_(0, repeat(indices, "z -> z d", d=grad_output.shape[1]), grad_output) - return grad_input.reshape(ctx.first_axis_dim, *other_shape), None - - -index_first_axis = IndexFirstAxis.apply - - -class IndexPutFirstAxis(torch.autograd.Function): - @staticmethod - def forward(ctx, values, indices, first_axis_dim): - ctx.save_for_backward(indices) - assert indices.ndim == 1 - assert values.ndim >= 2 - output = torch.zeros( - first_axis_dim, *values.shape[1:], device=values.device, dtype=values.dtype - ) - # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. - output[indices] = values - # output.scatter_(0, repeat(indices, 'z -> z d', d=values.shape[1]), values) - return output - - @staticmethod - def backward(ctx, grad_output): - (indices,) = ctx.saved_tensors - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - grad_values = grad_output[indices] - # grad_values = torch.gather(grad_output, 0, repeat(indices, 'z -> z d', d=grad_output.shape[1])) - return grad_values, None, None - - -index_put_first_axis = IndexPutFirstAxis.apply - - -class IndexFirstAxisResidual(torch.autograd.Function): - @staticmethod - def forward(ctx, input, indices): - ctx.save_for_backward(indices) - assert input.ndim >= 2 - ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] - second_dim = other_shape.numel() - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - output = input[indices] - # We don't want to reshape input (b ... -> b (...)) since it could change the channel_last - # memory format to channel_first. In other words, input might not be contiguous. - # If we don't detach, Pytorch complains about output being a view and is being modified inplace - return output, input.detach() - - @staticmethod - def backward(ctx, grad_output, grad_residual): - (indices,) = ctx.saved_tensors - assert grad_output.ndim >= 2 - other_shape = grad_output.shape[1:] - assert grad_residual.shape[1:] == other_shape - grad_input = grad_residual - # grad_input[indices] += grad_output - indices = indices.reshape(indices.shape[0], *((1,) * (grad_output.ndim - 1))) - indices = indices.expand_as(grad_output) - grad_input.scatter_add_(0, indices, grad_output) - return grad_input.reshape(ctx.first_axis_dim, *other_shape), None - - -index_first_axis_residual = IndexFirstAxisResidual.apply - - -def unpad_input(hidden_states, attention_mask, unused_mask=None): - """ - Arguments: - hidden_states: (batch, seqlen, ...) - attention_mask: (batch, seqlen), bool / int, 1 means valid and 0 means not valid. - unused_mask: (batch, seqlen), bool / int, 1 means the element is allocated but unused. - Return: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens selected in attention_mask + unused_mask. - indices: (total_nnz), the indices of masked tokens from the flattened input sequence. - cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. - max_seqlen_in_batch: int - seqused: (batch), returns the number of tokens selected in attention_mask + unused_mask. - """ - all_masks = (attention_mask + unused_mask) if unused_mask is not None else attention_mask - seqlens_in_batch = all_masks.sum(dim=-1, dtype=torch.int32) - used_seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) - indices = torch.nonzero(all_masks.flatten(), as_tuple=False).flatten() - max_seqlen_in_batch = seqlens_in_batch.max().item() - cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) - # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the - # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim - # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to - # index with integer indices. Moreover, torch's index is a bit slower than it needs to be, - # so we write custom forward and backward to make it a bit faster. - return ( - index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices), - indices, - cu_seqlens, - max_seqlen_in_batch, - used_seqlens_in_batch, - ) - - -def unpad_input_for_concatenated_sequences(hidden_states, attention_mask_in_length): - """ - Supports concatenating short samples in one sequence. The attention_mask_in_length is utilized to mask other short samples. It helps efficient training of variant lengths-based samples (e.g., the supervised fine-tuning task in large language model). - The motivation for this function is explained [here](https://github.com/Dao-AILab/flash-attention/issues/432#issuecomment-1668822286). - - For example, if batch = 3 and seqlen = 6, the attention_mask_in_length is: - ``` - [ - [2, 3, 0, 0, 0, 0], - [3, 2, 0, 0, 0, 0], - [6, 0, 0, 0, 0, 0] - ] - ``` - , which refers to the 3D-attention mask: - ``` - [ - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [0, 0, 1, 0, 0, 0], - [0, 0, 1, 1, 0, 0], - [0, 0, 1, 1, 1, 0], - [0, 0, 0, 0, 0, 1] - ], - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [0, 0, 0, 1, 0, 0], - [0, 0, 0, 1, 1, 0], - [0, 0, 0, 0, 0, 1] - ], - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [1, 1, 1, 1, 0, 0], - [1, 1, 1, 1, 1, 0], - [1, 1, 1, 1, 1, 1] - ] - ] - ```. - - Arguments: - hidden_states: (batch, seqlen, ...) - attention_mask_in_length: (batch, seqlen), int, a nonzero number (e.g., 1, 2, 3, etc.) means length of concatenated sequence in b-th batch, and 0 means none. - Return: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. - indices: (total_nnz), the indices of non-masked tokens from the flattened input sequence. - cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. - max_seqlen_in_batch: int - """ - length = attention_mask_in_length.sum(dim=-1) - seqlen = attention_mask_in_length.size(-1) - attention_mask_2d = torch.arange(seqlen, device=length.device, dtype=length.dtype).expand(len(length), seqlen) < length.unsqueeze(1) - real_indices_idx = torch.nonzero(attention_mask_in_length.flatten(), as_tuple=False).flatten() - seqlens_in_batch = attention_mask_in_length.flatten()[real_indices_idx] - indices = torch.nonzero(attention_mask_2d.flatten(), as_tuple=False).flatten() - max_seqlen_in_batch = seqlens_in_batch.max().item() - cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) - # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the - # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim - # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to - # index with integer indices. Moreover, torch's index is a bit slower than it needs to be, - # so we write custom forward and backward to make it a bit faster. - return ( - index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices), - indices, - cu_seqlens, - max_seqlen_in_batch, - ) - - -def pad_input(hidden_states, indices, batch, seqlen): - """ - Arguments: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. - indices: (total_nnz), the indices that represent the non-masked tokens of the original padded input sequence. - batch: int, batch size for the padded sequence. - seqlen: int, maximum sequence length for the padded sequence. - Return: - hidden_states: (batch, seqlen, ...) - """ - dim = hidden_states.shape[-1] - # output = torch.zeros((batch * seqlen), dim, device=hidden_states.device, dtype=hidden_states.dtype) - # output[indices] = hidden_states - output = index_put_first_axis(hidden_states, indices, batch * seqlen) - return rearrange(output, "(b s) ... -> b s ...", b=batch) diff --git a/build/torch211-cxx11-cu128-x86_64-linux/flash_attn2/__init__.py b/build/torch211-cxx11-cu128-x86_64-linux/flash_attn2/__init__.py deleted file mode 100644 index a9b2672c1cd85b74c1b3ded0fc0b2100e1aeac23..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu128-x86_64-linux/flash_attn2/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -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_interface.py b/build/torch211-cxx11-cu128-x86_64-linux/flash_attn_interface.py deleted file mode 100644 index b014e569dc9812a9c840d1043f2cebe767b4d477..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu128-x86_64-linux/flash_attn_interface.py +++ /dev/null @@ -1,1631 +0,0 @@ -# Copyright (c) 2023, Tri Dao. - -from typing import Optional, Sequence, Tuple, Union - -import torch -import torch.nn as nn -import os - -# # isort: off -# # We need to import the CUDA kernels after importing torch -# USE_TRITON_ROCM = os.getenv("FLASH_ATTENTION_TRITON_AMD_ENABLE", "FALSE") == "TRUE" -# if USE_TRITON_ROCM: -# from .flash_attn_triton_amd import interface_fa as flash_attn -# else: -# import flash_attn_2_cuda as flash_attn - - -from . import _ops -from ._ops import ops as flash_attn -from ._ops import add_op_namespace_prefix - -# # isort: on - -def maybe_contiguous(x): - return x.contiguous() if x is not None and x.stride(-1) != 1 else x - - -def _get_device(): - if torch.xpu.is_available(): - return "xpu" - elif torch.cuda.is_available(): - return "cuda" - else: - return "cpu" - - -def _get_block_size_n(device, head_dim, is_dropout, is_causal): - assert head_dim <= 256 - - if device.type == "xpu": - if head_dim <= 96: - return 64 - elif head_dim <= 128: - return 32 - elif head_dim <= 256: - return 64 - else: - return 32 - - # This should match the block sizes in the CUDA kernel - major, minor = torch.cuda.get_device_capability(device) - is_sm8x = major == 8 and minor > 0 # Only include sm86 and sm89, exclude sm80 (A100) - is_sm80 = major == 8 and minor == 0 - is_sm90 = major == 9 and minor == 0 - if head_dim <= 32: - return 128 - if head_dim <= 64: - return 128 if not is_dropout else 64 - elif head_dim <= 96: - return 64 - elif head_dim <= 128: - if is_sm8x: - return 64 if (not is_dropout and is_causal) else 32 - else: - return 64 if not is_dropout else 32 - elif head_dim <= 192: - return 64 - elif head_dim <= 224: - return 64 - elif head_dim <= 256: - return 64 - - -def round_multiple(x, m): - return (x + m - 1) // m * m - - -# torch.compile() support is only enabled for pytorch >= 2.4 -# The reason for this is that we are using the new custom_op and register_fake -# APIs, which support inplace modification of inputs in the function itself -if torch.__version__ >= "2.4.0": - _torch_custom_op_wrapper = torch.library.custom_op - _torch_register_fake_wrapper = torch.library.register_fake -else: - def noop_custom_op_wrapper(name, fn=None, /, *, mutates_args, device_types=None, schema=None): - def wrap(func): - return func - if fn is None: - return wrap - return fn - def noop_register_fake_wrapper(op, fn=None, /, *, lib=None, _stacklevel=1): - def wrap(func): - return func - if fn is None: - return wrap - return fn - _torch_custom_op_wrapper = noop_custom_op_wrapper - _torch_register_fake_wrapper = noop_register_fake_wrapper - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_forward"), mutates_args=(), device_types=_get_device()) -def _flash_attn_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - return_softmax: bool -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - out, softmax_lse, S_dmask, rng_state = flash_attn.fwd( - q, - k, - v, - None, - alibi_slopes, - dropout_p, - softmax_scale, - causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - None, - ) - return out, softmax_lse, S_dmask, rng_state - - -@_torch_register_fake_wrapper(add_op_namespace_prefix("_flash_attn_forward")) -def _flash_attn_forward_fake( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - return_softmax: bool -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - batch_size, seqlen_q, num_heads, head_size = q.shape - seqlen_k = k.shape[1] - out = torch.empty_like(q) - softmax_lse = torch.empty((batch_size, num_heads, seqlen_q), dtype=torch.float32, device=q.device, layout=q.layout) - p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) - if return_softmax: - p = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128), round_multiple(seqlen_k, 128)), dtype=q.dtype, device=q.device, layout=q.layout) - rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) - - return out, softmax_lse, p, rng_state - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_forward = _ops.ops._flash_attn_forward -else: - _wrapped_flash_attn_forward = _flash_attn_forward - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_varlen_forward"), mutates_args=(), device_types=_get_device()) -def _flash_attn_varlen_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - alibi_slopes: Optional[torch.Tensor] = None, - return_softmax: bool = False, - block_table: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - out, softmax_lse, S_dmask, rng_state = flash_attn.varlen_fwd( - q, - k, - v, - None, - cu_seqlens_q, - cu_seqlens_k, - seqused_k, - leftpad_k, - block_table, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - zero_tensors, - causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - None, - ) - # if out.isnan().any() or softmax_lse.isnan().any(): - # breakpoint() - return out, softmax_lse, S_dmask, rng_state - - -@_torch_register_fake_wrapper(add_op_namespace_prefix("_flash_attn_varlen_forward")) -def _flash_attn_varlen_forward_fake( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - alibi_slopes: Optional[torch.Tensor] = None, - return_softmax: bool = False, - block_table: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - paged_kv = block_table is not None - batch_size = cu_seqlens_q.numel() - 1 - total_q, num_heads, _ = q.shape - - out = torch.empty_like(q) - softmax_lse = torch.empty((num_heads, total_q), dtype=torch.float32, device=q.device, layout=q.layout) - p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) - seqlen_q_rounded = round_multiple(max_seqlen_q, 128) - seqlen_k_rounded = round_multiple(max_seqlen_k, 128) - if return_softmax: - p = torch.empty((batch_size, num_heads, seqlen_q_rounded, seqlen_k_rounded), dtype=q.dtype, device=q.device, layout=q.layout) - rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) - return out, softmax_lse, p, rng_state - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_varlen_forward = _ops.ops._flash_attn_varlen_forward -else: - _wrapped_flash_attn_varlen_forward = _flash_attn_varlen_forward - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_backward"), mutates_args=("dq", "dk", "dv"), device_types=_get_device()) -def _flash_attn_backward( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, -) -> 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)] - ( - dq, - dk, - dv, - softmax_d, - ) = flash_attn.bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - alibi_slopes, - dropout_p, - softmax_scale, - causal, - window_size_left, - window_size_right, - softcap, - deterministic, - None, - rng_state, - ) - return softmax_d - - -@_torch_register_fake_wrapper(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, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, -) -> torch.Tensor: - dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] - if dq is None: - dq = torch.empty_like(q) - if dk is None: - dk = torch.empty_like(k) - if dv is None: - dv = torch.empty_like(v) - batch_size, seqlen_q, num_heads, _ = q.shape - softmax_d = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128)), device=q.device, dtype=torch.float32) - - return softmax_d - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_backward = _ops.ops._flash_attn_backward -else: - _wrapped_flash_attn_backward = _flash_attn_backward - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_varlen_backward"), mutates_args=("dq", "dk", "dv"), device_types=_get_device()) -def _flash_attn_varlen_backward( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> 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)] - ( - dq, - dk, - dv, - softmax_d, - ) = flash_attn.varlen_bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - zero_tensors, - causal, - window_size_left, - window_size_right, - softcap, - deterministic, - None, - rng_state, - ) - # if dk.isnan().any() or dk.isnan().any() or dv.isnan().any() or softmax_d.isnan().any(): - # breakpoint() - return softmax_d - - -@_torch_register_fake_wrapper(add_op_namespace_prefix("_flash_attn_varlen_backward")) -def _flash_attn_varlen_backward_fake( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> torch.Tensor: - dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] - batch_size = cu_seqlens_q.numel() - 1 - total_q, num_heads, _ = q.shape - - if dq is None: - dq = torch.empty_like(q) - if dk is None: - dk = torch.empty_like(k) - if dv is None: - dv = torch.empty_like(v) - softmax_d = torch.empty((num_heads, total_q + 128 * batch_size), device=q.device, dtype=torch.float32) - - return softmax_d - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_varlen_backward = _ops.ops._flash_attn_varlen_backward -else: - _wrapped_flash_attn_varlen_backward = _flash_attn_varlen_backward - - -class FlashAttnQKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and qkv.requires_grad - if softmax_scale is None: - softmax_scale = qkv.shape[-1] ** (-0.5) - q, k, v = qkv[:, :, 0].detach(), qkv[:, :, 1].detach(), qkv[:, :, 2].detach() - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) - dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dqkv[:, :, 0], - dqkv[:, :, 1], - dqkv[:, :, 2], - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dqkv = dqkv[..., : dout.shape[-1]] # We could have padded the head dimension - return dqkv, None, None, None, None, None, None, None, None, None - - -class FlashAttnVarlenQKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - cu_seqlens, - max_seqlen, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and qkv.requires_grad - if softmax_scale is None: - softmax_scale = qkv.shape[-1] ** (-0.5) - q, k, v = qkv[:, 0].detach(), qkv[:, 1].detach(), qkv[:, 2].detach() - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens, - cu_seqlens, - max_seqlen, - max_seqlen, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=None, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens, rng_state) - ctx.dropout_p = dropout_p - ctx.max_seqlen = max_seqlen - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens, rng_state = ctx.saved_tensors - qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) - dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dqkv[:, 0], - dqkv[:, 1], - dqkv[:, 2], - cu_seqlens, - cu_seqlens, - ctx.max_seqlen, - ctx.max_seqlen, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - 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 FlashAttnKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - kv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, kv] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - k, v = kv[:, :, 0].detach(), kv[:, :, 1].detach() - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - dq = torch.empty_like(q) - kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) - dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dkv[:, :, 0], - dkv[:, :, 1], - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dkv = dkv[..., : dout.shape[-1]] - return dq, dkv, None, None, None, None, None, None, None, None, None - - -class FlashAttnVarlenKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, kv] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - k, v = kv[:, 0].detach(), kv[:, 1].detach() - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=None, - ) - if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state - ) - ctx.dropout_p = dropout_p - 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.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors - dq = torch.empty_like(q) - kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) - dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dkv[:, 0], - dkv[:, 1], - cu_seqlens_q, - cu_seqlens_k, - ctx.max_seqlen_q, - ctx.max_seqlen_k, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dkv = dkv[..., : dout.shape[-1]] - return dq, dkv, None, None, None, None, None, None, None, None, None, None, None, None, None - - -class FlashAttnFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - k, - v, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, k, v] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dk = dk[..., : dout.shape[-1]] - dv = dv[..., : dout.shape[-1]] - return dq, dk, dv, 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, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - block_table, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, k, v] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=block_table, - ) - if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state - ) - ctx.dropout_p = dropout_p - 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.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors - dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - ctx.max_seqlen_q, - ctx.max_seqlen_k, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dk = dk[..., : dout.shape[-1]] - dv = dv[..., : dout.shape[-1]] - return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None - - -def flash_attn_qkvpacked_func( - qkv, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # <=0.0 means deactivate - alibi_slopes=None, - deterministic=False, - 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, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_kvpacked_func( - q, - kv, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, -): - """dropout_p should be set to 0.0 during evaluation - If 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 K, V. - 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) - kv: (batch_size, seqlen, 2, 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. - softcap: float. Anything > 0 activates softcapping 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). - 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 FlashAttnKVPackedFunc.apply( - q, - kv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_func( - q, - k, - v, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - 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). - 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 FlashAttnFunc.apply( - q, - k, - v, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_qkvpacked_func( - qkv, - cu_seqlens, - max_seqlen, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - 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_varlen_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_varlen_kvpacked_func and flash_attn_varlen_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: (total, 3, nheads, headdim), where total = total number of tokens in the batch. - cu_seqlens: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into qkv. - max_seqlen: int. Maximum sequence length in the batch. - 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenQKVPackedFunc.apply( - qkv, - cu_seqlens, - max_seqlen, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_kvpacked_func( - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, -): - """dropout_p should be set to 0.0 during evaluation - If 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 K, V. - 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: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch. - kv: (total_k, 2, nheads_k, headdim), where total_k = total number of key tokens in the batch. - cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into q. - cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into kv. - max_seqlen_q: int. Maximum query sequence length in the batch. - max_seqlen_k: int. Maximum key sequence length in the batch. - 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 + 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenKVPackedFunc.apply( - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_func( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, - block_table=None, -): - """dropout_p should be set to 0.0 during evaluation - Supports multi-query and grouped-query attention (MQA/GQA) by passing in K, V 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: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch. - k: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. - v: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. - cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into q. - cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into kv. - max_seqlen_q: int. Maximum query sequence length in the batch. - max_seqlen_k: int. Maximum key sequence length in the batch. - 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 + 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenFunc.apply( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - block_table, - False if q.device.type == "cpu" else torch.is_grad_enabled(), - ) - - -def flash_attn_with_kvcache( - q, - k_cache, - v_cache, - k=None, - v=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, - block_table: Optional[torch.Tensor] = None, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - rotary_interleaved=True, - alibi_slopes=None, - num_splits=0, - 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 block_table, - or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_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) if there's no block_table, - or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_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). Similar to k. - 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. - block_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). - 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. - 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" - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-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) - cache_batch_idx = maybe_contiguous(cache_batch_idx) - block_table = maybe_contiguous(block_table) - out, softmax_lse = flash_attn.fwd_kvcache( - q, - k_cache, - v_cache, - k, - v, - cache_seqlens, - rotary_cos, - rotary_sin, - cache_batch_idx, - cache_leftpad, - block_table, - alibi_slopes, - None, - softmax_scale, - causal, - window_size[0], - window_size[1], - softcap, - rotary_interleaved, - num_splits, - ) - return (out, softmax_lse) if return_softmax_lse else out diff --git a/build/torch211-cxx11-cu128-x86_64-linux/layers/__init__.py b/build/torch211-cxx11-cu128-x86_64-linux/layers/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/build/torch211-cxx11-cu128-x86_64-linux/metadata.json b/build/torch211-cxx11-cu128-x86_64-linux/metadata.json deleted file mode 100644 index 29e721c65440fa2676e369fc2ba5d146d2037d63..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu128-x86_64-linux/metadata.json +++ /dev/null @@ -1,42 +0,0 @@ -{ - "name": "flash-attn2", - "id": "_flash_attn2_cuda_f12afc9", - "version": 1, - "license": "BSD-3-Clause", - "python-depends": [], - "backend": { - "type": "cuda", - "archs": [ - "10.0", - "12.0", - "8.0", - "9.0" - ] - }, - "digest": { - "algorithm": "sha256", - "files": { - "__init__.py": "Cd9MEdd/WUYnBmfiQpC4SyPqAxrTilTjvsvswNju4zc=", - "_flash_attn2_cuda_f12afc9.abi3.so": "rMl//S7KTYdNKGkeq4Lhj+dhEqIHmtq/zMSm2Q8IBFo=", - "_ops.py": "GQbmTTeSOU+7Q2nI68Hz+1RvEogyfEFccx1ubUYw+2s=", - "bert_padding.py": "gF1EmsdJ+HpQ86MRQ4VxDw+Sb/RVISdQALdNnoByHlw=", - "flash_attn2/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY=", - "flash_attn_interface.py": "W7a9bHR2ly6QL3aKqsBqxWA7M7ZWYcZPmtsPJu6rAfY=", - "layers/__init__.py": "47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU=", - "layers/patch_embed.py": "H58CgME/qSOPTZLOG08wFgrQS1j34pvNwMPrkTj3Ek4=", - "layers/rotary.py": "rFIdNrZQFNawwLD6d7t5LE2m+M93wwDZ8EuExw+gfD8=", - "ops/__init__.py": "47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU=", - "ops/activations.py": "t5lzNg1In8LP6bKeTnyeMizwqjv27JGbJ6ylPdGvZYg=", - "ops/fused_dense.py": "BqT8VTSCGdLJAjh6WBMkCRFYrdVodPeoElWBvH6AIzA=", - "ops/layer_norm.py": "zr7NXIm+2mtEynTp1CS0fbFGI2Mqdp41dY4AfDWF6EQ=", - "ops/rms_norm.py": "XEnihcj0a4aSz4LO55m5iKGVn4HKTeKN8TIyHjuDgxI=", - "ops/triton/__init__.py": "AbpHGcgLb+kRsJGnwFEktk7uzpZOCcBY74+YBdrKVGs=", - "ops/triton/cross_entropy.py": "y56nmuOtpSNxM7sDZcUSgXKIsIEayfuvrU8bAVog/nw=", - "ops/triton/k_activations.py": "+Z3vIyO4JkqBMipKsPvhzmxljtBdIhJCsl/M+/ESqBo=", - "ops/triton/layer_norm.py": "BBQdJiqopbIS12viUqzx3SvOQCNbDSAYLJXhpyMGyC8=", - "ops/triton/linear.py": "OtRvKz8xdpl+7v3q/ZTaS9fdBt9XrzMyapgRr50uBbM=", - "ops/triton/mlp.py": "uWA+D7HyuRUsOT+4nhi3vvCDNZnNEXwYcHMQ02UREfQ=", - "ops/triton/rotary.py": "0YSrlbep0lJmzLpTiANs0slTFeGcQQqR+el6uzbSLvI=" - } - } -} \ No newline at end of file diff --git a/build/torch211-cxx11-cu128-x86_64-linux/metadata.json.sigstore b/build/torch211-cxx11-cu128-x86_64-linux/metadata.json.sigstore deleted file mode 100644 index 499da95b2d48ac649c2998bfab63b9ad258c0865..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu128-x86_64-linux/metadata.json.sigstore +++ /dev/null @@ -1 +0,0 @@ -{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json","verificationMaterial":{"certificate":{"rawBytes":"MIIHSzCCBtGgAwIBAgIUbQAqyKNBnYgGPDW5m11uaaQZz6IwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwNjI0MDEwMTIxWhcNMjYwNjI0MDExMTIxWjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAE5k6IsGoStiUZPEPAZucNNy+GeNs+9+S7NRof0v8o1OzPiC0+EnA/L3w2nj83TS059BohE6QKKkRGj/By+gDY/KOCBfAwggXsMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQU0utt0xWYeUfCJ98ExFOwc3Xi00kwHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoZjEyYWZjOTJjNDBiYmQ3NmQ5OTBkOTExZWJhYmYyOGU0OTBhN2JkZTATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoZjEyYWZjOTJjNDBiYmQ3NmQ5OTBkOTExZWJhYmYyOGU0OTBhN2JkZTAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoZjEyYWZjOTJjNDBiYmQ3NmQ5OTBkOTExZWJhYmYyOGU0OTBhN2JkZTAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKGYxMmFmYzkyYzQwYmJkNzZkOTkwZDkxMWViYWJmMjhlNDkwYTdiZGUwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMjgwMzYyNjg5MzUvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBigYKKwYBBAHWeQIEAgR8BHoAeAB2AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABnvclW8QAAAQDAEcwRQIgRogId+lbcg3RzeL6cfjGedn19sP/0saYxYBxN81HmJECIQD7tps9zoD4wMYm77w7/Mi5uAthXGlq3kyO86vzURv7hjAKBggqhkjOPQQDAwNoADBlAjB+GrvkIxgzRHZyHP9SSAOpDswMvT0htJJwfIsZskDwCTFIZtHWaYgTniqcuh7hy00CMQDdEXXWtx4ILLmXZB19TdwMdPsOsNo/Cg6SnT0OhHymi0el6+kR0XB9FoPRyefGHc8="},"tlogEntries":[{"logIndex":"1933772022","logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="},"kindVersion":{"kind":"hashedrekord","version":"0.0.1"},"integratedTime":"1782262881","inclusionPromise":{"signedEntryTimestamp":"MEUCIQC1gg7bySZdqrMEZcaZGoiZ8V04ThIU0o2C1Uid1jB6swIgQqkaGgtNFmz2kpXg9EVi/9CtOQo6T0GdWDGd3pb/KSw="},"inclusionProof":{"logIndex":"1811867760","rootHash":"0JD4J5GKtrvA1RkiRhRLC0HMCnfuvWQiVxBtuVOAC40=","treeSize":"1811867794","hashes":["0REPdOWPzWWG65AvM/K4jifbe9qbo85692ppPQ7XMV0=","FbLNoC5nTMI5ln61oi5UyCt3yacKt2F/pvO0LSTRMPA=","J7fHBVuplmz4rcyR7Vg/o+gANNXiyGpThnFH8VpdfI4=","xNruLQ9I87ltpHWTiDmtlGyVRzBO1vMyl8Waio3Ndis=","Witp3aRYBCw9nnhXbn+/sDP+AnO9cbPnNetUheGD0mQ=","h+YH6XZKZNTNP1PVlQsuW+iRoRbKfjIP1/BE49EI3xo=","db6+ARgXQBGlPZrV6moa1JK5ApoKSVHy0+5txnatNi8=","u0bSKGZwMSREs3Tup7jGbDu7i3jy734yNzJ6qWCIYAk=","fqtytk2if0R+M3kWRpiPUrBEhf2z8y5Li6MeOuWmiks=","mzp9m+2At7GZdgfx0m7RwtO9YzltCZ6UhZBy4p+lUdM=","ipWWLDkqay5JAYe6P4Nrqsck/1C0C5uwqT9QbFsiY1E=","COKvK4TO+dkpqUDGbTbSKewajXzAonx1B1YaRb/ymLw=","dQlR+DXGDKXm5V1fupxaMSOwNZ9F1MybXyyxB96yj80=","JKJdxSZwKPO6QFzLTxHCrUyI5sDU/CwY57zDlA9ni6Q=","OgfnQCh7xr8xUTErIw//YAKaKMDwRJWiThvYExQvLJM=","S8Oe9ZicoJqbVuL1I0dIdTgVpLc82Cm2CMOcV83sKVs=","E1O2ilWlpFRLWGdbmJUMDZ2qbGu/eDZ12m6YjQ9/JPA=","nyv2TDlfkz0XfRwjY98zsbUbOkBZ+Tgz7PML0c2uN14=","OMfDh4k9YD0E+t4sjqJsh+vg8D1ueo1cC07exyxmr6s=","+vHKm4wNwGJ46ceyZ8mePD7Vn1SediStNHCStRoTSWQ=","5HwWjuGNDisetbiWM0cWIXM65TIM3EJlqk13ntCQlcU=","NDIRlqxO0IaGr94q86MIU7VCAN3yEEgOHS4loya5lu4=","daxmZaajRpZV+JxHiOYZhJBiSKN5ucqjh2WnGbHhirw=","DOCeoSMovIvLExkhIvisow9AuNXgeWs4ECkyR6EcqYU="],"checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n1811867794\n0JD4J5GKtrvA1RkiRhRLC0HMCnfuvWQiVxBtuVOAC40=\n\n— rekor.sigstore.dev wNI9ajBFAiBK/8RkmSAV+V2RecRZhkA8T4MFr9tlTCjeuOBPOWwt2QIhANiyEFydPHftjyNj8sw9sVGy4NWXlWgjRcunIUFIdu3P\n"}},"canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiI5Yzc0YTE2N2NmNmM3MjFmOGQ3YTZkOWE2OWY2YjcwMGZiNGZhZTIxODU3MjI1YjUzNGVlZWNhYmVhODg0NzhmIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FWUNJUUNyd1lsNXc2MllKdGRuZEpaNE5idEJJVU9FZHRQbkxvN3Q3L3J5ZHBqZzZRSWhBTnNnOGRKWHpyWHJZTVVXalo5dDY5RTBhVzlRUWlOU3JSa2VvK290dzF3QyIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRla05EUW5SSFowRjNTVUpCWjBsVllsRkJjWGxMVGtKdVdXZEhVRVJYTlcweE1YVmhZVkZhZWpaSmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDVxU1RCTlJFVjNUVlJKZUZkb1kwNU5hbGwzVG1wSk1FMUVSWGhOVkVsNFYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVUxYXpaSmMwZHZVM1JwVlZwUVJWQkJXblZqVGs1NUswZGxUbk1yT1N0VE4wNVNiMllLTUhZNGJ6RlBlbEJwUXpBclJXNUJMMHd6ZHpKdWFqZ3pWRk13TlRsQ2IyaEZObEZMUzJ0U1Iyb3ZRbmtyWjBSWkwwdFBRMEptUVhkbloxaHpUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlV3ZFhSMENqQjRWMWxsVldaRFNqazRSWGhHVDNkak0xaHBNREJyZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOWFha1Y1V1ZkYWFrOVVTbXBPUkVKcFdXMVJNMDV0VVRWUFZFSnJDazlVUlhoYVYwcG9XVzFaZVU5SFZUQlBWRUpvVGpKS2ExcFVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMXBxUlhsWlYxcHFUMVJLYWs1RVFtbFpiVkV6VG0xUk5VOVVRbXRQVkVWNFdsZEthRmx0V1hrS1QwZFZNRTlVUW1oT01rcHJXbFJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMXBxUlhsWlYxcHFUMVJLYWs1RVFtbFpiVkV6VG0xUk5VOVVRbXNLVDFSRmVGcFhTbWhaYlZsNVQwZFZNRTlVUW1oT01rcHJXbFJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwZFplRTF0Um0wS1dYcHJlVmw2VVhkWmJVcHJUbnBhYTA5VWEzZGFSR3Q0VFZkV2FWbFhTbTFOYW1oc1RrUnJkMWxVWkdsYVIxVjNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOYW1kM1RYcFplVTVxWnpWTmVsVjJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwWjFsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTT0VKSWIwRUtaVUZDTWtGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtNTJZMnhYT0ZGQlFVRlJSQXBCUldOM1VsRkpaMUp2WjBsa0syeGlZMmN6VW5wbFREWmpabXBIWldSdU1UbHpVQzh3YzJGWmVGbENlRTQ0TVVodFNrVkRTVkZFTjNSd2N6bDZiMFEwQ25kTldXMDNOM2MzTDAxcE5YVkJkR2hZUjJ4eE0ydDVUemcyZG5wVlVuWTNhR3BCUzBKblozRm9hMnBQVUZGUlJFRjNUbTlCUkVKc1FXcENLMGR5ZG1zS1NYaG5lbEpJV25sSVVEbFRVMEZQY0VSemQwMTJWREJvZEVwS2QyWkpjMXB6YTBSM1ExUkdTVnAwU0ZkaFdXZFVibWx4WTNWb04yaDVNREJEVFZGRVpBcEZXRmhYZEhnMFNVeE1iVmhhUWpFNVZHUjNUV1JRYzA5elRtOHZRMmMyVTI1VU1FOW9TSGx0YVRCbGJEWXJhMUl3V0VJNVJtOVFVbmxsWmtkSVl6ZzlDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}],"timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyTADAgEAMIICwAYJKoZIhvcNAQcCoIICsTCCAq0CAQMxDTALBglghkgBZQMEAgEwgbcGCyqGSIb3DQEJEAEEoIGnBIGkMIGhAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQg0NajHSyXT0F2gt8DXblqwbntJafdBBY96IWS4uS4BUoCFAJIl1wcvJiXDE+cUS2dY6MSUDrZGA8yMDI2MDYyNDAxMDEyMVowAwIBAaAypDAwLjEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MRUwEwYDVQQDEwxzaWdzdG9yZS10c2GgADGCAdswggHXAgEBMFEwOTEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MSAwHgYDVQQDExdzaWdzdG9yZS10c2Etc2VsZnNpZ25lZAIUOhNULwyQYe68wUMvy4qOiyojiwwwCwYJYIZIAWUDBAIBoIH8MBoGCSqGSIb3DQEJAzENBgsqhkiG9w0BCRABBDAcBgkqhkiG9w0BCQUxDxcNMjYwNjI0MDEwMTIxWjAvBgkqhkiG9w0BCQQxIgQguiqfd/SXae+20EO49HmCXjUZD+pXCxHHYJOwRkL9Jz4wgY4GCyqGSIb3DQEJEAIvMX8wfTB7MHkEIIX5J7wHq2LKw7RDVsEO/IGyxog/2nq55thw2dE6zQW3MFUwPaQ7MDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAoGCCqGSM49BAMCBGcwZQIxAIzUdQtRvZt2/6wicGx+t5Gx0lo+YSy8NuEdCV+hiMDgtNiu8gBy1FUUFXi8vqKzDAIwVP8Pb++38GZw6pkvKANWT2qiqa44Np/qID7XVWWY2hXXWOUpxjR8P901Iau33J4T"}]}},"messageSignature":{"messageDigest":{"algorithm":"SHA2_256","digest":"nHShZ89sch+Nem2aafa3APtPriGFciW1NO7sq+qIR48="},"signature":"MEYCIQCrwYl5w62YJtdndJZ4NbtBIUOEdtPnLo7t7/rydpjg6QIhANsg8dJXzrXrYMUWjZ9t69E0aW9QQiNSrRkeo+otw1wC"}} \ No newline at end of file diff --git a/build/torch211-cxx11-cu128-x86_64-linux/ops/__init__.py b/build/torch211-cxx11-cu128-x86_64-linux/ops/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/build/torch211-cxx11-cu128-x86_64-linux/ops/triton/layer_norm.py b/build/torch211-cxx11-cu128-x86_64-linux/ops/triton/layer_norm.py deleted file mode 100644 index 5e6f0b3697b612e6a19031bbc55ac76ad233dc85..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu128-x86_64-linux/ops/triton/layer_norm.py +++ /dev/null @@ -1,1254 +0,0 @@ -# Copyright (c) 2024, Tri Dao. -# Implement dropout + residual + layer_norm / rms_norm. - -# Based on the Triton LayerNorm tutorial: https://triton-lang.org/main/getting-started/tutorials/05-layer-norm.html -# For the backward pass, we keep weight_grad and bias_grad in registers and accumulate. -# This is faster for dimensions up to 8k, but after that it's much slower due to register spilling. -# The models we train have hidden dim up to 8k anyway (e.g. Llama 70B), so this is fine. - -import math -from typing import Optional, List - -import torch -import torch.nn.functional as F -from torch import Tensor - -import triton -import triton.language as tl - -from flash_attn.utils.torch import custom_fwd, custom_bwd -from flash_attn.utils.library import triton_op - -from ..._ops import add_op_namespace_prefix - - -def maybe_contiguous_lastdim(x): - return x.contiguous() if x is not None and x.stride(-1) != 1 else x - - -def maybe_contiguous(x): - return x.contiguous() if x is not None else None - - -def triton_autotune_configs(): - # Return configs with a valid warp count for the current device - configs = [] - # Maximum threads per block is architecture-dependent in theory, but in reality all are 1024 - max_threads_per_block = 1024 - # Default to warp size 32 if not defined by device - warp_size = getattr(torch.cuda.get_device_properties(torch.cuda.current_device()), "warp_size", 32) - # Autotune for warp counts which are powers of 2 and do not exceed thread per block limit - return [triton.Config({}, num_warps=warp_count) for warp_count in [1, 2, 4, 8, 16, 32] - if warp_count * warp_size <= max_threads_per_block] - # return [triton.Config({}, num_warps=8)] - - -def layer_norm_ref( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - zero_centered_weight=False, - dropout_mask=None, - dropout_mask1=None, - upcast=False, -): - dtype = x.dtype - if upcast: - x = x.float() - weight = weight.float() - bias = bias.float() if bias is not None else None - residual = residual.float() if residual is not None else residual - x1 = x1.float() if x1 is not None else None - weight1 = weight1.float() if weight1 is not None else None - bias1 = bias1.float() if bias1 is not None else None - if zero_centered_weight: - weight = weight + 1.0 - if weight1 is not None: - weight1 = weight1 + 1.0 - if x1 is not None: - assert rowscale is None, "rowscale is not supported with parallel LayerNorm" - if rowscale is not None: - x = x * rowscale[..., None] - if dropout_p > 0.0: - if dropout_mask is not None: - x = x.masked_fill(~dropout_mask, 0.0) / (1.0 - dropout_p) - else: - x = F.dropout(x, p=dropout_p) - if x1 is not None: - if dropout_mask1 is not None: - x1 = x1.masked_fill(~dropout_mask1, 0.0) / (1.0 - dropout_p) - else: - x1 = F.dropout(x1, p=dropout_p) - if x1 is not None: - x = x + x1 - if residual is not None: - x = (x + residual).to(x.dtype) - out = F.layer_norm(x.to(weight.dtype), x.shape[-1:], weight=weight, bias=bias, eps=eps).to( - dtype - ) - if weight1 is None: - return out if not prenorm else (out, x) - else: - out1 = F.layer_norm( - x.to(weight1.dtype), x.shape[-1:], weight=weight1, bias=bias1, eps=eps - ).to(dtype) - return (out, out1) if not prenorm else (out, out1, x) - - -def rms_norm_ref( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - zero_centered_weight=False, - dropout_mask=None, - dropout_mask1=None, - upcast=False, -): - dtype = x.dtype - if upcast: - x = x.float() - weight = weight.float() - bias = bias.float() if bias is not None else None - residual = residual.float() if residual is not None else residual - x1 = x1.float() if x1 is not None else None - weight1 = weight1.float() if weight1 is not None else None - bias1 = bias1.float() if bias1 is not None else None - if zero_centered_weight: - weight = weight + 1.0 - if weight1 is not None: - weight1 = weight1 + 1.0 - if x1 is not None: - assert rowscale is None, "rowscale is not supported with parallel LayerNorm" - if rowscale is not None: - x = x * rowscale[..., None] - if dropout_p > 0.0: - if dropout_mask is not None: - x = x.masked_fill(~dropout_mask, 0.0) / (1.0 - dropout_p) - else: - x = F.dropout(x, p=dropout_p) - if x1 is not None: - if dropout_mask1 is not None: - x1 = x1.masked_fill(~dropout_mask1, 0.0) / (1.0 - dropout_p) - else: - x1 = F.dropout(x1, p=dropout_p) - if x1 is not None: - x = x + x1 - if residual is not None: - x = (x + residual).to(x.dtype) - rstd = 1 / torch.sqrt((x.square()).mean(dim=-1, keepdim=True) + eps) - out = ((x * rstd * weight) + bias if bias is not None else (x * rstd * weight)).to(dtype) - if weight1 is None: - return out if not prenorm else (out, x) - else: - out1 = ((x * rstd * weight1) + bias1 if bias1 is not None else (x * rstd * weight1)).to( - dtype - ) - return (out, out1) if not prenorm else (out, out1, x) - - -@triton.autotune( - configs=triton_autotune_configs(), - key=["N", "HAS_RESIDUAL", "STORE_RESIDUAL_OUT", "IS_RMS_NORM", "HAS_BIAS", "HAS_X1", "HAS_W1", "HAS_B1"], -) -# torch compile doesn't like triton.heuristics, so we set these manually when calling the kernel -# @triton.heuristics({"HAS_BIAS": lambda args: args["B"] is not None}) -# @triton.heuristics({"HAS_RESIDUAL": lambda args: args["RESIDUAL"] is not None}) -# @triton.heuristics({"HAS_X1": lambda args: args["X1"] is not None}) -# @triton.heuristics({"HAS_W1": lambda args: args["W1"] is not None}) -# @triton.heuristics({"HAS_B1": lambda args: args["B1"] is not None}) -@triton.jit -def _layer_norm_fwd_1pass_kernel( - X, # pointer to the input - Y, # pointer to the output - W, # pointer to the weights - B, # pointer to the biases - RESIDUAL, # pointer to the residual - X1, - W1, - B1, - Y1, - RESIDUAL_OUT, # pointer to the residual - ROWSCALE, - SEEDS, # Dropout seeds for each row - DROPOUT_MASK, - DROPOUT_MASK1, - Mean, # pointer to the mean - Rstd, # pointer to the 1/std - stride_x_row, # how much to increase the pointer when moving by 1 row - stride_y_row, - stride_res_row, - stride_res_out_row, - stride_x1_row, - stride_y1_row, - M, # number of rows in X - N, # number of columns in X - eps, # epsilon to avoid division by zero - dropout_p, # Dropout probability - zero_centered_weight, # If true, add 1.0 to the weight - IS_RMS_NORM: tl.constexpr, - BLOCK_N: tl.constexpr, - HAS_RESIDUAL: tl.constexpr, - STORE_RESIDUAL_OUT: tl.constexpr, - HAS_BIAS: tl.constexpr, - HAS_DROPOUT: tl.constexpr, - STORE_DROPOUT_MASK: tl.constexpr, - HAS_ROWSCALE: tl.constexpr, - HAS_X1: tl.constexpr, - HAS_W1: tl.constexpr, - HAS_B1: tl.constexpr, -): - # Map the program id to the row of X and Y it should compute. - row = tl.program_id(0) - X += row * stride_x_row - Y += row * stride_y_row - if HAS_RESIDUAL: - RESIDUAL += row * stride_res_row - if STORE_RESIDUAL_OUT: - RESIDUAL_OUT += row * stride_res_out_row - if HAS_X1: - X1 += row * stride_x1_row - if HAS_W1: - Y1 += row * stride_y1_row - # Compute mean and variance - cols = tl.arange(0, BLOCK_N) - x = tl.load(X + cols, mask=cols < N, other=0.0).to(tl.float32) - if HAS_ROWSCALE: - rowscale = tl.load(ROWSCALE + row).to(tl.float32) - x *= rowscale - if HAS_DROPOUT: - # Compute dropout mask - # 7 rounds is good enough, and reduces register pressure - keep_mask = tl.rand(tl.load(SEEDS + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - x = tl.where(keep_mask, x / (1.0 - dropout_p), 0.0) - if STORE_DROPOUT_MASK: - tl.store(DROPOUT_MASK + row * N + cols, keep_mask, mask=cols < N) - if HAS_X1: - x1 = tl.load(X1 + cols, mask=cols < N, other=0.0).to(tl.float32) - if HAS_ROWSCALE: - rowscale = tl.load(ROWSCALE + M + row).to(tl.float32) - x1 *= rowscale - if HAS_DROPOUT: - # Compute dropout mask - # 7 rounds is good enough, and reduces register pressure - keep_mask = ( - tl.rand(tl.load(SEEDS + M + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - ) - x1 = tl.where(keep_mask, x1 / (1.0 - dropout_p), 0.0) - if STORE_DROPOUT_MASK: - tl.store(DROPOUT_MASK1 + row * N + cols, keep_mask, mask=cols < N) - x += x1 - if HAS_RESIDUAL: - residual = tl.load(RESIDUAL + cols, mask=cols < N, other=0.0).to(tl.float32) - x += residual - if STORE_RESIDUAL_OUT: - tl.store(RESIDUAL_OUT + cols, x, mask=cols < N) - if not IS_RMS_NORM: - mean = tl.sum(x, axis=0) / N - tl.store(Mean + row, mean) - xbar = tl.where(cols < N, x - mean, 0.0) - var = tl.sum(xbar * xbar, axis=0) / N - else: - xbar = tl.where(cols < N, x, 0.0) - var = tl.sum(xbar * xbar, axis=0) / N - rstd = 1 / tl.sqrt(var + eps) - tl.store(Rstd + row, rstd) - # Normalize and apply linear transformation - mask = cols < N - w = tl.load(W + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w += 1.0 - if HAS_BIAS: - b = tl.load(B + cols, mask=mask).to(tl.float32) - x_hat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd - y = x_hat * w + b if HAS_BIAS else x_hat * w - # Write output - tl.store(Y + cols, y, mask=mask) - if HAS_W1: - w1 = tl.load(W1 + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w1 += 1.0 - if HAS_B1: - b1 = tl.load(B1 + cols, mask=mask).to(tl.float32) - y1 = x_hat * w1 + b1 if HAS_B1 else x_hat * w1 - tl.store(Y1 + cols, y1, mask=mask) - - -def _layer_norm_fwd( - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - residual: Optional[Tensor] = None, - x1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - out_dtype: Optional[torch.dtype] = None, - residual_dtype: Optional[torch.dtype] = None, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - return_dropout_mask: bool = False, - out: Optional[Tensor] = None, - residual_out: Optional[Tensor] = None -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - # Need to wrap to handle the case where residual_out is a alias of x, which makes torch.library - # and torch.compile unhappy. Also allocate memory for out and residual_out if they are None - # so that _layer_norm_fwd_impl doesn't have to return them. - if out is None: - out = torch.empty_like(x, dtype=x.dtype if out_dtype is None else out_dtype) - if residual is not None: - residual_dtype = residual.dtype - if residual_out is None and ( - residual is not None - or (residual_dtype is not None and residual_dtype != x.dtype) - or dropout_p > 0.0 - or rowscale is not None - or x1 is not None - ): - residual_out = torch.empty_like( - x, dtype=residual_dtype if residual_dtype is not None else x.dtype - ) - else: - residual_out = None - y1, mean, rstd, seeds, dropout_mask, dropout_mask1 = _layer_norm_fwd_impl( - x, - weight, - bias, - eps, - out, - residual=residual, - x1=x1, - weight1=weight1, - bias1=bias1, - dropout_p=dropout_p, - rowscale=rowscale, - zero_centered_weight=zero_centered_weight, - is_rms_norm=is_rms_norm, - return_dropout_mask=return_dropout_mask, - residual_out=residual_out, - ) - # residual_out is None if residual is None and residual_dtype == input_dtype and dropout_p == 0.0 - if residual_out is None: - residual_out = x - return out, y1, mean, rstd, residual_out, seeds, dropout_mask, dropout_mask1 - - -# [2025-04-28] torch.library.triton_op ignores the schema argument, but here we need the schema -# since we're returning a tuple of tensors -@triton_op(add_op_namespace_prefix("layer_norm_fwd_impl"), mutates_args={"out", "residual_out"}, - schema="(Tensor x, Tensor weight, Tensor bias, float eps, Tensor(a!) out, Tensor? residual, Tensor? x1, Tensor? weight1, Tensor? bias1, float dropout_p, Tensor? rowscale, bool zero_centered_weight, bool is_rms_norm, bool return_dropout_mask, Tensor(a!)? residual_out) -> (Tensor y1, Tensor mean, Tensor rstd, Tensor seeds, Tensor dropout_mask, Tensor dropout_mask1)") -def _layer_norm_fwd_impl( - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - out: Tensor, - residual: Optional[Tensor] = None, - x1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - return_dropout_mask: bool = False, - residual_out: Optional[Tensor] = None -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - M, N = x.shape - assert x.stride(-1) == 1 - if residual is not None: - assert residual.stride(-1) == 1 - assert residual.shape == (M, N) - assert weight.shape == (N,) - assert weight.stride(-1) == 1 - if bias is not None: - assert bias.stride(-1) == 1 - assert bias.shape == (N,) - if x1 is not None: - assert x1.shape == x.shape - assert rowscale is None - assert x1.stride(-1) == 1 - if weight1 is not None: - assert weight1.shape == (N,) - assert weight1.stride(-1) == 1 - if bias1 is not None: - assert bias1.shape == (N,) - assert bias1.stride(-1) == 1 - if rowscale is not None: - assert rowscale.is_contiguous() - assert rowscale.shape == (M,) - assert out.shape == x.shape - assert out.stride(-1) == 1 - if residual_out is not None: - assert residual_out.shape == x.shape - assert residual_out.stride(-1) == 1 - if weight1 is not None: - y1 = torch.empty_like(out) - assert y1.stride(-1) == 1 - else: - y1 = None - mean = torch.empty((M,), dtype=torch.float32, device=x.device) if not is_rms_norm else None - rstd = torch.empty((M,), dtype=torch.float32, device=x.device) - if dropout_p > 0.0: - seeds = torch.randint( - 2**32, (M if x1 is None else 2 * M,), device=x.device, dtype=torch.int64 - ) - else: - seeds = None - if return_dropout_mask and dropout_p > 0.0: - dropout_mask = torch.empty(M, N, device=x.device, dtype=torch.bool) - if x1 is not None: - dropout_mask1 = torch.empty(M, N, device=x.device, dtype=torch.bool) - else: - dropout_mask1 = None - else: - dropout_mask, dropout_mask1 = None, None - # Less than 64KB per feature: enqueue fused kernel - MAX_FUSED_SIZE = 65536 // x.element_size() - BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) - if N > BLOCK_N: - raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") - with torch.cuda.device(x.device.index): - torch.library.wrap_triton(_layer_norm_fwd_1pass_kernel)[(M,)]( - x, - out, - weight, - bias, - residual, - x1, - weight1, - bias1, - y1, - residual_out, - rowscale, - seeds, - dropout_mask, - dropout_mask1, - mean, - rstd, - x.stride(0), - out.stride(0), - residual.stride(0) if residual is not None else 0, - residual_out.stride(0) if residual_out is not None else 0, - x1.stride(0) if x1 is not None else 0, - y1.stride(0) if y1 is not None else 0, - M, - N, - eps, - dropout_p, - # Passing bool make torch inductor very unhappy since it then tries to compare to int_max - int(zero_centered_weight), - is_rms_norm, - BLOCK_N, - residual is not None, - residual_out is not None, - bias is not None, - dropout_p > 0.0, - dropout_mask is not None, - rowscale is not None, - HAS_X1=x1 is not None, - HAS_W1=weight1 is not None, - HAS_B1=bias1 is not None, - ) - return y1, mean, rstd, seeds, dropout_mask, dropout_mask1 - - -@triton.autotune( - configs=triton_autotune_configs(), - key=["N", "HAS_DRESIDUAL", "STORE_DRESIDUAL", "IS_RMS_NORM", "HAS_BIAS", "HAS_DROPOUT"], -) -# torch compile doesn't like triton.heuristics, so we set these manually when calling the kernel -# @triton.heuristics({"HAS_BIAS": lambda args: args["B"] is not None}) -# @triton.heuristics({"HAS_DRESIDUAL": lambda args: args["DRESIDUAL"] is not None}) -# @triton.heuristics({"STORE_DRESIDUAL": lambda args: args["DRESIDUAL_IN"] is not None}) -# @triton.heuristics({"HAS_ROWSCALE": lambda args: args["ROWSCALE"] is not None}) -# @triton.heuristics({"HAS_DY1": lambda args: args["DY1"] is not None}) -# @triton.heuristics({"HAS_DX1": lambda args: args["DX1"] is not None}) -# @triton.heuristics({"HAS_B1": lambda args: args["DB1"] is not None}) -# @triton.heuristics({"RECOMPUTE_OUTPUT": lambda args: args["Y"] is not None}) -@triton.jit -def _layer_norm_bwd_kernel( - X, # pointer to the input - W, # pointer to the weights - B, # pointer to the biases - Y, # pointer to the output to be recomputed - DY, # pointer to the output gradient - DX, # pointer to the input gradient - DW, # pointer to the partial sum of weights gradient - DB, # pointer to the partial sum of biases gradient - DRESIDUAL, - W1, - DY1, - DX1, - DW1, - DB1, - DRESIDUAL_IN, - ROWSCALE, - SEEDS, - Mean, # pointer to the mean - Rstd, # pointer to the 1/std - stride_x_row, # how much to increase the pointer when moving by 1 row - stride_y_row, - stride_dy_row, - stride_dx_row, - stride_dres_row, - stride_dy1_row, - stride_dx1_row, - stride_dres_in_row, - M, # number of rows in X - N, # number of columns in X - eps, # epsilon to avoid division by zero - dropout_p, - zero_centered_weight, - rows_per_program, - IS_RMS_NORM: tl.constexpr, - BLOCK_N: tl.constexpr, - HAS_DRESIDUAL: tl.constexpr, - STORE_DRESIDUAL: tl.constexpr, - HAS_BIAS: tl.constexpr, - HAS_DROPOUT: tl.constexpr, - HAS_ROWSCALE: tl.constexpr, - HAS_DY1: tl.constexpr, - HAS_DX1: tl.constexpr, - HAS_B1: tl.constexpr, - RECOMPUTE_OUTPUT: tl.constexpr, -): - # Map the program id to the elements of X, DX, and DY it should compute. - row_block_id = tl.program_id(0) - row_start = row_block_id * rows_per_program - # Do not early exit if row_start >= M, because we need to write DW and DB - cols = tl.arange(0, BLOCK_N) - mask = cols < N - X += row_start * stride_x_row - if HAS_DRESIDUAL: - DRESIDUAL += row_start * stride_dres_row - if STORE_DRESIDUAL: - DRESIDUAL_IN += row_start * stride_dres_in_row - DY += row_start * stride_dy_row - DX += row_start * stride_dx_row - if HAS_DY1: - DY1 += row_start * stride_dy1_row - if HAS_DX1: - DX1 += row_start * stride_dx1_row - if RECOMPUTE_OUTPUT: - Y += row_start * stride_y_row - w = tl.load(W + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w += 1.0 - if RECOMPUTE_OUTPUT and HAS_BIAS: - b = tl.load(B + cols, mask=mask, other=0.0).to(tl.float32) - if HAS_DY1: - w1 = tl.load(W1 + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w1 += 1.0 - dw = tl.zeros((BLOCK_N,), dtype=tl.float32) - if HAS_BIAS: - db = tl.zeros((BLOCK_N,), dtype=tl.float32) - if HAS_DY1: - dw1 = tl.zeros((BLOCK_N,), dtype=tl.float32) - if HAS_B1: - db1 = tl.zeros((BLOCK_N,), dtype=tl.float32) - row_end = min((row_block_id + 1) * rows_per_program, M) - for row in range(row_start, row_end): - # Load data to SRAM - x = tl.load(X + cols, mask=mask, other=0).to(tl.float32) - dy = tl.load(DY + cols, mask=mask, other=0).to(tl.float32) - if HAS_DY1: - dy1 = tl.load(DY1 + cols, mask=mask, other=0).to(tl.float32) - if not IS_RMS_NORM: - mean = tl.load(Mean + row) - rstd = tl.load(Rstd + row) - # Compute dx - xhat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd - xhat = tl.where(mask, xhat, 0.0) - if RECOMPUTE_OUTPUT: - y = xhat * w + b if HAS_BIAS else xhat * w - tl.store(Y + cols, y, mask=mask) - wdy = w * dy - dw += dy * xhat - if HAS_BIAS: - db += dy - if HAS_DY1: - wdy += w1 * dy1 - dw1 += dy1 * xhat - if HAS_B1: - db1 += dy1 - if not IS_RMS_NORM: - c1 = tl.sum(xhat * wdy, axis=0) / N - c2 = tl.sum(wdy, axis=0) / N - dx = (wdy - (xhat * c1 + c2)) * rstd - else: - c1 = tl.sum(xhat * wdy, axis=0) / N - dx = (wdy - xhat * c1) * rstd - if HAS_DRESIDUAL: - dres = tl.load(DRESIDUAL + cols, mask=mask, other=0).to(tl.float32) - dx += dres - # Write dx - if STORE_DRESIDUAL: - tl.store(DRESIDUAL_IN + cols, dx, mask=mask) - if HAS_DX1: - if HAS_DROPOUT: - keep_mask = ( - tl.rand(tl.load(SEEDS + M + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - ) - dx1 = tl.where(keep_mask, dx / (1.0 - dropout_p), 0.0) - else: - dx1 = dx - tl.store(DX1 + cols, dx1, mask=mask) - if HAS_DROPOUT: - keep_mask = tl.rand(tl.load(SEEDS + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - dx = tl.where(keep_mask, dx / (1.0 - dropout_p), 0.0) - if HAS_ROWSCALE: - rowscale = tl.load(ROWSCALE + row).to(tl.float32) - dx *= rowscale - tl.store(DX + cols, dx, mask=mask) - - X += stride_x_row - if HAS_DRESIDUAL: - DRESIDUAL += stride_dres_row - if STORE_DRESIDUAL: - DRESIDUAL_IN += stride_dres_in_row - if RECOMPUTE_OUTPUT: - Y += stride_y_row - DY += stride_dy_row - DX += stride_dx_row - if HAS_DY1: - DY1 += stride_dy1_row - if HAS_DX1: - DX1 += stride_dx1_row - tl.store(DW + row_block_id * N + cols, dw, mask=mask) - if HAS_BIAS: - tl.store(DB + row_block_id * N + cols, db, mask=mask) - if HAS_DY1: - tl.store(DW1 + row_block_id * N + cols, dw1, mask=mask) - if HAS_B1: - tl.store(DB1 + row_block_id * N + cols, db1, mask=mask) - - -def _layer_norm_bwd( - dy: Tensor, - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - mean: Tensor, - rstd: Tensor, - dresidual: Optional[Tensor] = None, - dy1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - seeds: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - has_residual: bool = False, - has_x1: bool = False, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - x_dtype: Optional[torch.dtype] = None, - recompute_output: bool = False, -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - # Need to wrap to handle the case where dresidual_in or dx1 are aliases of x, - # which makes torch.library unhappy - dx, dw, db, dresidual_in, dx1, dw1, db1, y = _layer_norm_bwd_impl( - dy, - x, - weight, - bias, - eps, - mean, - rstd, - dresidual, - dy1, - weight1, - bias1, - seeds, - dropout_p, - rowscale, - has_residual, - has_x1, - zero_centered_weight, - is_rms_norm, - x_dtype=x_dtype, - recompute_output=recompute_output, - ) - # Don't need to compute dresidual_in separately in this case - if has_residual and dx.dtype == x.dtype and dropout_p == 0.0 and rowscale is None: - dresidual_in = dx - if has_x1 and dropout_p == 0.0: - dx1 = dx - return dx, dw, db, dresidual_in, dx1, dw1, db1, y - - - -@triton_op(add_op_namespace_prefix("layer_norm_bwd_impl"), mutates_args={}, - schema="(Tensor dy, Tensor x, Tensor weight, Tensor bias, float eps, Tensor mean, Tensor rstd, Tensor? dresidual, Tensor? dy1, Tensor? weight1, Tensor? bias1, Tensor? seeds, float dropout_p, Tensor? rowscale, bool has_residual, bool has_x1, bool zero_centered_weight, bool is_rms_norm, ScalarType? x_dtype, bool recompute_output) -> (Tensor dx, Tensor dw, Tensor db, Tensor dresidual_in, Tensor dx1, Tensor dw1, Tensor db1, Tensor y)", - allow_decomposition=False, # Don't let torch.compile trace inside - ) -def _layer_norm_bwd_impl( - dy: Tensor, - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - mean: Tensor, - rstd: Tensor, - dresidual: Optional[Tensor] = None, - dy1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - seeds: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - has_residual: bool = False, - has_x1: bool = False, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - x_dtype: Optional[torch.dtype] = None, - recompute_output: bool = False, -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - M, N = x.shape - assert x.stride(-1) == 1 - dy = maybe_contiguous_lastdim(dy) - assert dy.stride(-1) == 1 - assert dy.shape == (M, N) - if dresidual is not None: - dresidual = maybe_contiguous_lastdim(dresidual) - assert dresidual.stride(-1) == 1 - assert dresidual.shape == (M, N) - assert weight.shape == (N,) - assert weight.stride(-1) == 1 - if bias is not None: - assert bias.stride(-1) == 1 - assert bias.shape == (N,) - if dy1 is not None: - dy1 = maybe_contiguous_lastdim(dy1) - assert weight1 is not None - assert dy1.shape == dy.shape - assert dy1.stride(-1) == 1 - if weight1 is not None: - assert weight1.shape == (N,) - assert weight1.stride(-1) == 1 - if bias1 is not None: - assert bias1.shape == (N,) - assert bias1.stride(-1) == 1 - if seeds is not None: - assert seeds.is_contiguous() - assert seeds.shape == (M if not has_x1 else M * 2,) - if rowscale is not None: - assert rowscale.is_contiguous() - assert rowscale.shape == (M,) - # allocate output - dx = ( - torch.empty_like(x) - if x_dtype is None - else torch.empty(M, N, dtype=x_dtype, device=x.device) - ) - dresidual_in = ( - torch.empty_like(x) - if has_residual - and (dx.dtype != x.dtype or dropout_p > 0.0 or rowscale is not None or has_x1) - else None - ) - dx1 = torch.empty_like(dx) if (has_x1 and dropout_p > 0.0) else None - y = torch.empty(M, N, dtype=dy.dtype, device=dy.device) if recompute_output else None - if recompute_output: - assert weight1 is None, "recompute_output is not supported with parallel LayerNorm" - - # Less than 64KB per feature: enqueue fused kernel - MAX_FUSED_SIZE = 65536 // x.element_size() - BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) - if N > BLOCK_N: - raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") - # Increasing the multiple (e.g. 8) will allow more thread blocks to be launched and hide the - # latency of the gmem reads/writes, but will increase the time of summing up dw / db. - sm_count = torch.cuda.get_device_properties(x.device).multi_processor_count * 8 - _dw = torch.empty((sm_count, N), dtype=torch.float32, device=weight.device) - _db = ( - torch.empty((sm_count, N), dtype=torch.float32, device=bias.device) - if bias is not None - else None - ) - _dw1 = torch.empty_like(_dw) if weight1 is not None else None - _db1 = torch.empty_like(_db) if bias1 is not None else None - rows_per_program = math.ceil(M / sm_count) - grid = (sm_count,) - with torch.cuda.device(x.device.index): - torch.library.wrap_triton(_layer_norm_bwd_kernel)[grid]( - x, - weight, - bias, - y, - dy, - dx, - _dw, - _db, - dresidual, - weight1, - dy1, - dx1, - _dw1, - _db1, - dresidual_in, - rowscale, - seeds, - mean, - rstd, - x.stride(0), - 0 if not recompute_output else y.stride(0), - dy.stride(0), - dx.stride(0), - dresidual.stride(0) if dresidual is not None else 0, - dy1.stride(0) if dy1 is not None else 0, - dx1.stride(0) if dx1 is not None else 0, - dresidual_in.stride(0) if dresidual_in is not None else 0, - M, - N, - eps, - dropout_p, - # Passing bool make torch inductor very unhappy since it then tries to compare to int_max - int(zero_centered_weight), - rows_per_program, - is_rms_norm, - BLOCK_N, - dresidual is not None, - dresidual_in is not None, - bias is not None, - dropout_p > 0.0, - HAS_ROWSCALE=rowscale is not None, - HAS_DY1=dy1 is not None, - HAS_DX1=dx1 is not None, - HAS_B1=bias1 is not None, - RECOMPUTE_OUTPUT=y is not None, - ) - dw = _dw.sum(0).to(weight.dtype) - db = _db.sum(0).to(bias.dtype) if bias is not None else None - dw1 = _dw1.sum(0).to(weight1.dtype) if weight1 is not None else None - db1 = _db1.sum(0).to(bias1.dtype) if bias1 is not None else None - # dresidual_in and dx1 could be None, the wrapper will handle assigning them from dx - return dx, dw, db, dresidual_in, dx1, dw1, db1, y - - -class LayerNormFn(torch.autograd.Function): - - @staticmethod - def forward( - ctx, - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - residual_in_fp32=False, - zero_centered_weight=False, - is_rms_norm=False, - return_dropout_mask=False, - out_dtype=None, - out=None, - residual_out=None - ): - x_shape_og = x.shape - # reshape input data into 2D tensor - x = maybe_contiguous_lastdim(x.reshape(-1, x.shape[-1])) - if residual is not None: - assert residual.shape == x_shape_og - residual = maybe_contiguous_lastdim(residual.reshape(-1, residual.shape[-1])) - if x1 is not None: - assert x1.shape == x_shape_og - assert rowscale is None, "rowscale is not supported with parallel LayerNorm" - x1 = maybe_contiguous_lastdim(x1.reshape(-1, x1.shape[-1])) - weight = weight.contiguous() - bias = maybe_contiguous(bias) - weight1 = maybe_contiguous(weight1) - bias1 = maybe_contiguous(bias1) - if rowscale is not None: - rowscale = rowscale.reshape(-1).contiguous() - residual_dtype = ( - residual.dtype - if residual is not None - else (torch.float32 if residual_in_fp32 else None) - ) - if out is not None: - out = out.reshape(-1, out.shape[-1]) - if residual_out is not None: - residual_out = residual_out.reshape(-1, residual_out.shape[-1]) - y, y1, mean, rstd, residual_out, seeds, dropout_mask, dropout_mask1 = _layer_norm_fwd( - x, - weight, - bias, - eps, - residual, - x1, - weight1, - bias1, - dropout_p=dropout_p, - rowscale=rowscale, - out_dtype=out_dtype, - residual_dtype=residual_dtype, - zero_centered_weight=zero_centered_weight, - is_rms_norm=is_rms_norm, - return_dropout_mask=return_dropout_mask, - out=out, - residual_out=residual_out, - ) - ctx.save_for_backward( - residual_out, weight, bias, weight1, bias1, rowscale, seeds, mean, rstd - ) - ctx.x_shape_og = x_shape_og - ctx.eps = eps - ctx.dropout_p = dropout_p - ctx.is_rms_norm = is_rms_norm - ctx.has_residual = residual is not None - ctx.has_x1 = x1 is not None - ctx.prenorm = prenorm - ctx.x_dtype = x.dtype - ctx.zero_centered_weight = zero_centered_weight - y = y.reshape(x_shape_og) - y1 = y1.reshape(x_shape_og) if y1 is not None else None - residual_out = residual_out.reshape(x_shape_og) if residual_out is not None else None - dropout_mask = dropout_mask.reshape(x_shape_og) if dropout_mask is not None else None - dropout_mask1 = dropout_mask1.reshape(x_shape_og) if dropout_mask1 is not None else None - if not return_dropout_mask: - if weight1 is None: - return y if not prenorm else (y, residual_out) - else: - return (y, y1) if not prenorm else (y, y1, residual_out) - else: - if weight1 is None: - return ( - (y, dropout_mask, dropout_mask1) - if not prenorm - else (y, residual_out, dropout_mask, dropout_mask1) - ) - else: - return ( - (y, y1, dropout_mask, dropout_mask1) - if not prenorm - else (y, y1, residual_out, dropout_mask, dropout_mask1) - ) - - @staticmethod - def backward(ctx, dy, *args): - x, weight, bias, weight1, bias1, rowscale, seeds, mean, rstd = ctx.saved_tensors - dy = dy.reshape(-1, dy.shape[-1]) - if weight1 is not None: - dy1, args = args[0], args[1:] - dy1 = dy1.reshape(-1, dy1.shape[-1]) - assert dy1.shape == x.shape - else: - dy1 = None - if ctx.prenorm: - dresidual = args[0] - dresidual = dresidual.reshape(-1, dresidual.shape[-1]) - assert dresidual.shape == x.shape - else: - dresidual = None - dx, dw, db, dresidual_in, dx1, dw1, db1, _ = _layer_norm_bwd( - dy, - x, - weight, - bias, - ctx.eps, - mean, - rstd, - dresidual, - dy1, - weight1, - bias1, - seeds, - ctx.dropout_p, - rowscale, - ctx.has_residual, - ctx.has_x1, - ctx.zero_centered_weight, - ctx.is_rms_norm, - x_dtype=ctx.x_dtype, - recompute_output=False, - ) - return ( - dx.reshape(ctx.x_shape_og), - dw, - db, - dresidual_in.reshape(ctx.x_shape_og) if ctx.has_residual else None, - dx1.reshape(ctx.x_shape_og) if dx1 is not None else None, - dw1, - db1, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) - - -def layer_norm_fn( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - residual_in_fp32=False, - zero_centered_weight=False, - is_rms_norm=False, - return_dropout_mask=False, - out_dtype=None, - out=None, - residual_out=None -): - return LayerNormFn.apply( - x, - weight, - bias, - residual, - x1, - weight1, - bias1, - eps, - dropout_p, - rowscale, - prenorm, - residual_in_fp32, - zero_centered_weight, - is_rms_norm, - return_dropout_mask, - out_dtype, - out, - residual_out - ) - - -def rms_norm_fn( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - residual_in_fp32=False, - zero_centered_weight=False, - return_dropout_mask=False, - out_dtype=None, - out=None, - residual_out=None -): - return LayerNormFn.apply( - x, - weight, - bias, - residual, - x1, - weight1, - bias1, - eps, - dropout_p, - rowscale, - prenorm, - residual_in_fp32, - zero_centered_weight, - True, - return_dropout_mask, - out_dtype, - out, - residual_out - ) - - -class RMSNorm(torch.nn.Module): - - def __init__(self, hidden_size, eps=1e-5, dropout_p=0.0, zero_centered_weight=False, - device=None, dtype=None): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.eps = eps - if dropout_p > 0.0: - self.drop = torch.nn.Dropout(dropout_p) - else: - self.drop = None - self.zero_centered_weight = zero_centered_weight - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.register_parameter("bias", None) - self.reset_parameters() - - def reset_parameters(self): - if not self.zero_centered_weight: - torch.nn.init.ones_(self.weight) - else: - torch.nn.init.zeros_(self.weight) - - def forward(self, x, residual=None, prenorm=False, residual_in_fp32=False): - return rms_norm_fn( - x, - self.weight, - self.bias, - residual=residual, - eps=self.eps, - dropout_p=self.drop.p if self.drop is not None and self.training else 0.0, - prenorm=prenorm, - residual_in_fp32=residual_in_fp32, - zero_centered_weight=self.zero_centered_weight, - ) - - -class LayerNormLinearFn(torch.autograd.Function): - - @staticmethod - @custom_fwd - def forward( - ctx, - x, - norm_weight, - norm_bias, - linear_weight, - linear_bias, - residual=None, - eps=1e-6, - prenorm=False, - residual_in_fp32=False, - is_rms_norm=False, - ): - x_shape_og = x.shape - # reshape input data into 2D tensor - x = maybe_contiguous_lastdim(x.reshape(-1, x.shape[-1])) - if residual is not None: - assert residual.shape == x_shape_og - residual = maybe_contiguous_lastdim(residual.reshape(-1, residual.shape[-1])) - norm_weight = norm_weight.contiguous() - norm_bias = maybe_contiguous(norm_bias) - residual_dtype = ( - residual.dtype - if residual is not None - else (torch.float32 if residual_in_fp32 else None) - ) - y, _, mean, rstd, residual_out, *rest = _layer_norm_fwd( - x, - norm_weight, - norm_bias, - eps, - residual, - out_dtype=None if not torch.is_autocast_enabled() else torch.get_autocast_dtype("cuda"), - residual_dtype=residual_dtype, - is_rms_norm=is_rms_norm, - ) - y = y.reshape(x_shape_og) - dtype = torch.get_autocast_dtype("cuda") if torch.is_autocast_enabled() else y.dtype - linear_weight = linear_weight.to(dtype) - linear_bias = linear_bias.to(dtype) if linear_bias is not None else None - out = F.linear(y.to(linear_weight.dtype), linear_weight, linear_bias) - # We don't store y, will be recomputed in the backward pass to save memory - ctx.save_for_backward(residual_out, norm_weight, norm_bias, linear_weight, mean, rstd) - ctx.x_shape_og = x_shape_og - ctx.eps = eps - ctx.is_rms_norm = is_rms_norm - ctx.has_residual = residual is not None - ctx.prenorm = prenorm - ctx.x_dtype = x.dtype - ctx.linear_bias_is_none = linear_bias is None - return out if not prenorm else (out, residual_out.reshape(x_shape_og)) - - @staticmethod - @custom_bwd - def backward(ctx, dout, *args): - x, norm_weight, norm_bias, linear_weight, mean, rstd = ctx.saved_tensors - dout = dout.reshape(-1, dout.shape[-1]) - dy = F.linear(dout, linear_weight.t()) - dlinear_bias = None if ctx.linear_bias_is_none else dout.sum(0) - dy = maybe_contiguous_lastdim(dy) - assert dy.shape == x.shape - if ctx.prenorm: - dresidual = args[0] - dresidual = maybe_contiguous_lastdim(dresidual.reshape(-1, dresidual.shape[-1])) - assert dresidual.shape == x.shape - else: - dresidual = None - dx, dnorm_weight, dnorm_bias, dresidual_in, _, _, _, y = _layer_norm_bwd( - dy, - x, - norm_weight, - norm_bias, - ctx.eps, - mean, - rstd, - dresidual=dresidual, - has_residual=ctx.has_residual, - is_rms_norm=ctx.is_rms_norm, - x_dtype=ctx.x_dtype, - recompute_output=True, - ) - dlinear_weight = torch.einsum("bo,bi->oi", dout, y) - return ( - dx.reshape(ctx.x_shape_og), - dnorm_weight, - dnorm_bias, - dlinear_weight, - dlinear_bias, - dresidual_in.reshape(ctx.x_shape_og) if ctx.has_residual else None, - None, - None, - None, - None, - ) - - -def layer_norm_linear_fn( - x, - norm_weight, - norm_bias, - linear_weight, - linear_bias, - residual=None, - eps=1e-6, - prenorm=False, - residual_in_fp32=False, - is_rms_norm=False, -): - return LayerNormLinearFn.apply( - x, - norm_weight, - norm_bias, - linear_weight, - linear_bias, - residual, - eps, - prenorm, - residual_in_fp32, - is_rms_norm, - ) diff --git a/build/torch211-cxx11-cu130-aarch64-linux/__init__.py b/build/torch211-cxx11-cu130-aarch64-linux/__init__.py deleted file mode 100644 index ecc2f9d896b6c93f90b0a1499856dc0612177422..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu130-aarch64-linux/__init__.py +++ /dev/null @@ -1,393 +0,0 @@ -from typing import Optional, List -import torch -from ._ops import ops as flash_attn_ops -from .flash_attn_interface import ( - flash_attn_func, - flash_attn_kvpacked_func, - flash_attn_qkvpacked_func, - flash_attn_varlen_func, - flash_attn_varlen_kvpacked_func, - flash_attn_varlen_qkvpacked_func, - flash_attn_with_kvcache, -) - - -def fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - return_softmax: bool = False, - gen: Optional[torch.Generator] = None, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention. - - Args: - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Optional output tensor, same shape as q - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - return_softmax: Whether to return softmax weights - gen: Optional random number generator - - Returns: - List of tensors: [output, softmax_lse, (softmax if return_softmax)] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.fwd( - q, - k, - v, - out, - alibi_slopes, - p_dropout, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - gen, - ) - - -def varlen_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - out: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - block_table: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - max_seqlen_q: int = 0, - max_seqlen_k: int = 0, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - zero_tensors: bool = False, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - return_softmax: bool = False, - gen: Optional[torch.Generator] = None, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention with variable sequence lengths. - - Args: - q: Query tensor of shape [total_q, num_heads, head_size] - k: Key tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - v: Value tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1] - cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1] - out: Optional output tensor of shape [total_q, num_heads, head_size] - seqused_k: Optional tensor specifying how many keys to use per batch element [batch_size] - leftpad_k: Optional left padding for keys of shape [batch_size] - block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq] - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - max_seqlen_q: Maximum sequence length for queries - max_seqlen_k: Maximum sequence length for keys - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - zero_tensors: Whether to zero tensors before computation - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - return_softmax: Whether to return softmax weights - gen: Optional random number generator - - Returns: - List of tensors: [output, softmax_lse, (softmax if return_softmax)] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.varlen_fwd( - q, - k, - v, - out, - cu_seqlens_q, - cu_seqlens_k, - seqused_k, - leftpad_k, - block_table, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - p_dropout, - softmax_scale, - zero_tensors, - is_causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - gen, - ) - - -def bwd( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - p_dropout: float = 0.0, - 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, - gen: Optional[torch.Generator] = None, - rng_state: Optional[torch.Tensor] = None, -) -> List[torch.Tensor]: - """ - Backward pass for multi-head attention. - - Args: - dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size] - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size] - softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q] - dq: Optional gradient tensor for queries, same shape as q - dk: Optional gradient tensor for keys, same shape as k - dv: Optional gradient tensor for values, same shape as v - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - deterministic: Whether to use deterministic algorithms - gen: Optional random number generator - rng_state: Optional RNG state from forward pass - - Returns: - List of tensors: [dq, dk, dv] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - alibi_slopes, - p_dropout, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - deterministic, - gen, - rng_state, - ) - - -def varlen_bwd( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - max_seqlen_q: int = 0, - max_seqlen_k: int = 0, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - zero_tensors: bool = False, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - deterministic: bool = False, - gen: Optional[torch.Generator] = None, - rng_state: Optional[torch.Tensor] = None, -) -> List[torch.Tensor]: - """ - Backward pass for multi-head attention with variable sequence lengths. - - Args: - dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size] - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size] - softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q] - cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1] - cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1] - dq: Optional gradient tensor for queries, same shape as q - dk: Optional gradient tensor for keys, same shape as k - dv: Optional gradient tensor for values, same shape as v - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - max_seqlen_q: Maximum sequence length for queries - max_seqlen_k: Maximum sequence length for keys - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - zero_tensors: Whether to zero tensors before computation - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - deterministic: Whether to use deterministic algorithms - gen: Optional random number generator - rng_state: Optional RNG state from forward pass - - Returns: - List of tensors: [dq, dk, dv] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.varlen_bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - p_dropout, - softmax_scale, - zero_tensors, - is_causal, - window_size_left, - window_size_right, - softcap, - deterministic, - gen, - rng_state, - ) - - -def fwd_kvcache( - q: torch.Tensor, - kcache: torch.Tensor, - vcache: torch.Tensor, - k: Optional[torch.Tensor] = None, - v: Optional[torch.Tensor] = None, - seqlens_k: Optional[torch.Tensor] = None, - rotary_cos: Optional[torch.Tensor] = None, - rotary_sin: Optional[torch.Tensor] = None, - cache_batch_idx: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - block_table: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - out: 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, - is_rotary_interleaved: bool = False, - num_splits: int = 1, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention with KV cache. - - Args: - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - kcache: Key cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - vcache: Value cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - k: Optional new keys tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size] - v: Optional new values tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size] - seqlens_k: Optional sequence lengths for keys of shape [batch_size] - rotary_cos: Optional rotary cosine tensor of shape [seqlen_ro, rotary_dim/2] - rotary_sin: Optional rotary sine tensor of shape [seqlen_ro, rotary_dim/2] - cache_batch_idx: Optional indices to index into the KV cache - leftpad_k: Optional left padding for keys of shape [batch_size] - block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq] - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - out: Optional output tensor, same shape as q - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - is_rotary_interleaved: Whether rotary embeddings are interleaved - num_splits: Number of splits for computation - - Returns: - List of tensors: [output, softmax_lse] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.fwd_kvcache( - q, - kcache, - vcache, - k, - v, - seqlens_k, - rotary_cos, - rotary_sin, - cache_batch_idx, - leftpad_k, - block_table, - alibi_slopes, - out, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - is_rotary_interleaved, - num_splits, - ) diff --git a/build/torch211-cxx11-cu130-aarch64-linux/_flash_attn2_cuda_f12afc9.abi3.so b/build/torch211-cxx11-cu130-aarch64-linux/_flash_attn2_cuda_f12afc9.abi3.so deleted file mode 100644 index 7d1f1429f5ccef549382724c85c559b8a3712b9a..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu130-aarch64-linux/_flash_attn2_cuda_f12afc9.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:a2e2580829f9caf390c941f16b60e4ca57b43838c4edebb254aa099371ecabdf -size 1006694184 diff --git a/build/torch211-cxx11-cu130-aarch64-linux/_ops.py b/build/torch211-cxx11-cu130-aarch64-linux/_ops.py deleted file mode 100644 index af5b731cb04b5e6ce54e862db0501f3683c3163f..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu130-aarch64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _flash_attn2_cuda_f12afc9 -ops = torch.ops._flash_attn2_cuda_f12afc9 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_flash_attn2_cuda_f12afc9::{op_name}" diff --git a/build/torch211-cxx11-cu130-aarch64-linux/bert_padding.py b/build/torch211-cxx11-cu130-aarch64-linux/bert_padding.py deleted file mode 100644 index 3c2d35159a014a9d03aabead9e52e009168696ea..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu130-aarch64-linux/bert_padding.py +++ /dev/null @@ -1,218 +0,0 @@ -# Adapted from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/padding.py - -import torch -import torch.nn.functional as F -from einops import rearrange, repeat - - -class IndexFirstAxis(torch.autograd.Function): - @staticmethod - def forward(ctx, input, indices): - ctx.save_for_backward(indices) - assert input.ndim >= 2 - ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] - second_dim = other_shape.numel() - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - # return input[indices] - return torch.gather( - rearrange(input, "b ... -> b (...)"), 0, repeat(indices, "z -> z d", d=second_dim) - ).reshape(-1, *other_shape) - - @staticmethod - def backward(ctx, grad_output): - (indices,) = ctx.saved_tensors - assert grad_output.ndim >= 2 - other_shape = grad_output.shape[1:] - grad_output = rearrange(grad_output, "b ... -> b (...)") - grad_input = torch.zeros( - [ctx.first_axis_dim, grad_output.shape[1]], - device=grad_output.device, - dtype=grad_output.dtype, - ) - # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. - # grad_input[indices] = grad_output - grad_input.scatter_(0, repeat(indices, "z -> z d", d=grad_output.shape[1]), grad_output) - return grad_input.reshape(ctx.first_axis_dim, *other_shape), None - - -index_first_axis = IndexFirstAxis.apply - - -class IndexPutFirstAxis(torch.autograd.Function): - @staticmethod - def forward(ctx, values, indices, first_axis_dim): - ctx.save_for_backward(indices) - assert indices.ndim == 1 - assert values.ndim >= 2 - output = torch.zeros( - first_axis_dim, *values.shape[1:], device=values.device, dtype=values.dtype - ) - # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. - output[indices] = values - # output.scatter_(0, repeat(indices, 'z -> z d', d=values.shape[1]), values) - return output - - @staticmethod - def backward(ctx, grad_output): - (indices,) = ctx.saved_tensors - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - grad_values = grad_output[indices] - # grad_values = torch.gather(grad_output, 0, repeat(indices, 'z -> z d', d=grad_output.shape[1])) - return grad_values, None, None - - -index_put_first_axis = IndexPutFirstAxis.apply - - -class IndexFirstAxisResidual(torch.autograd.Function): - @staticmethod - def forward(ctx, input, indices): - ctx.save_for_backward(indices) - assert input.ndim >= 2 - ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] - second_dim = other_shape.numel() - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - output = input[indices] - # We don't want to reshape input (b ... -> b (...)) since it could change the channel_last - # memory format to channel_first. In other words, input might not be contiguous. - # If we don't detach, Pytorch complains about output being a view and is being modified inplace - return output, input.detach() - - @staticmethod - def backward(ctx, grad_output, grad_residual): - (indices,) = ctx.saved_tensors - assert grad_output.ndim >= 2 - other_shape = grad_output.shape[1:] - assert grad_residual.shape[1:] == other_shape - grad_input = grad_residual - # grad_input[indices] += grad_output - indices = indices.reshape(indices.shape[0], *((1,) * (grad_output.ndim - 1))) - indices = indices.expand_as(grad_output) - grad_input.scatter_add_(0, indices, grad_output) - return grad_input.reshape(ctx.first_axis_dim, *other_shape), None - - -index_first_axis_residual = IndexFirstAxisResidual.apply - - -def unpad_input(hidden_states, attention_mask, unused_mask=None): - """ - Arguments: - hidden_states: (batch, seqlen, ...) - attention_mask: (batch, seqlen), bool / int, 1 means valid and 0 means not valid. - unused_mask: (batch, seqlen), bool / int, 1 means the element is allocated but unused. - Return: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens selected in attention_mask + unused_mask. - indices: (total_nnz), the indices of masked tokens from the flattened input sequence. - cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. - max_seqlen_in_batch: int - seqused: (batch), returns the number of tokens selected in attention_mask + unused_mask. - """ - all_masks = (attention_mask + unused_mask) if unused_mask is not None else attention_mask - seqlens_in_batch = all_masks.sum(dim=-1, dtype=torch.int32) - used_seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) - indices = torch.nonzero(all_masks.flatten(), as_tuple=False).flatten() - max_seqlen_in_batch = seqlens_in_batch.max().item() - cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) - # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the - # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim - # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to - # index with integer indices. Moreover, torch's index is a bit slower than it needs to be, - # so we write custom forward and backward to make it a bit faster. - return ( - index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices), - indices, - cu_seqlens, - max_seqlen_in_batch, - used_seqlens_in_batch, - ) - - -def unpad_input_for_concatenated_sequences(hidden_states, attention_mask_in_length): - """ - Supports concatenating short samples in one sequence. The attention_mask_in_length is utilized to mask other short samples. It helps efficient training of variant lengths-based samples (e.g., the supervised fine-tuning task in large language model). - The motivation for this function is explained [here](https://github.com/Dao-AILab/flash-attention/issues/432#issuecomment-1668822286). - - For example, if batch = 3 and seqlen = 6, the attention_mask_in_length is: - ``` - [ - [2, 3, 0, 0, 0, 0], - [3, 2, 0, 0, 0, 0], - [6, 0, 0, 0, 0, 0] - ] - ``` - , which refers to the 3D-attention mask: - ``` - [ - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [0, 0, 1, 0, 0, 0], - [0, 0, 1, 1, 0, 0], - [0, 0, 1, 1, 1, 0], - [0, 0, 0, 0, 0, 1] - ], - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [0, 0, 0, 1, 0, 0], - [0, 0, 0, 1, 1, 0], - [0, 0, 0, 0, 0, 1] - ], - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [1, 1, 1, 1, 0, 0], - [1, 1, 1, 1, 1, 0], - [1, 1, 1, 1, 1, 1] - ] - ] - ```. - - Arguments: - hidden_states: (batch, seqlen, ...) - attention_mask_in_length: (batch, seqlen), int, a nonzero number (e.g., 1, 2, 3, etc.) means length of concatenated sequence in b-th batch, and 0 means none. - Return: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. - indices: (total_nnz), the indices of non-masked tokens from the flattened input sequence. - cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. - max_seqlen_in_batch: int - """ - length = attention_mask_in_length.sum(dim=-1) - seqlen = attention_mask_in_length.size(-1) - attention_mask_2d = torch.arange(seqlen, device=length.device, dtype=length.dtype).expand(len(length), seqlen) < length.unsqueeze(1) - real_indices_idx = torch.nonzero(attention_mask_in_length.flatten(), as_tuple=False).flatten() - seqlens_in_batch = attention_mask_in_length.flatten()[real_indices_idx] - indices = torch.nonzero(attention_mask_2d.flatten(), as_tuple=False).flatten() - max_seqlen_in_batch = seqlens_in_batch.max().item() - cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) - # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the - # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim - # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to - # index with integer indices. Moreover, torch's index is a bit slower than it needs to be, - # so we write custom forward and backward to make it a bit faster. - return ( - index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices), - indices, - cu_seqlens, - max_seqlen_in_batch, - ) - - -def pad_input(hidden_states, indices, batch, seqlen): - """ - Arguments: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. - indices: (total_nnz), the indices that represent the non-masked tokens of the original padded input sequence. - batch: int, batch size for the padded sequence. - seqlen: int, maximum sequence length for the padded sequence. - Return: - hidden_states: (batch, seqlen, ...) - """ - dim = hidden_states.shape[-1] - # output = torch.zeros((batch * seqlen), dim, device=hidden_states.device, dtype=hidden_states.dtype) - # output[indices] = hidden_states - output = index_put_first_axis(hidden_states, indices, batch * seqlen) - return rearrange(output, "(b s) ... -> b s ...", b=batch) diff --git a/build/torch211-cxx11-cu130-aarch64-linux/flash_attn2/__init__.py b/build/torch211-cxx11-cu130-aarch64-linux/flash_attn2/__init__.py deleted file mode 100644 index a9b2672c1cd85b74c1b3ded0fc0b2100e1aeac23..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu130-aarch64-linux/flash_attn2/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -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-aarch64-linux/flash_attn_interface.py b/build/torch211-cxx11-cu130-aarch64-linux/flash_attn_interface.py deleted file mode 100644 index b014e569dc9812a9c840d1043f2cebe767b4d477..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu130-aarch64-linux/flash_attn_interface.py +++ /dev/null @@ -1,1631 +0,0 @@ -# Copyright (c) 2023, Tri Dao. - -from typing import Optional, Sequence, Tuple, Union - -import torch -import torch.nn as nn -import os - -# # isort: off -# # We need to import the CUDA kernels after importing torch -# USE_TRITON_ROCM = os.getenv("FLASH_ATTENTION_TRITON_AMD_ENABLE", "FALSE") == "TRUE" -# if USE_TRITON_ROCM: -# from .flash_attn_triton_amd import interface_fa as flash_attn -# else: -# import flash_attn_2_cuda as flash_attn - - -from . import _ops -from ._ops import ops as flash_attn -from ._ops import add_op_namespace_prefix - -# # isort: on - -def maybe_contiguous(x): - return x.contiguous() if x is not None and x.stride(-1) != 1 else x - - -def _get_device(): - if torch.xpu.is_available(): - return "xpu" - elif torch.cuda.is_available(): - return "cuda" - else: - return "cpu" - - -def _get_block_size_n(device, head_dim, is_dropout, is_causal): - assert head_dim <= 256 - - if device.type == "xpu": - if head_dim <= 96: - return 64 - elif head_dim <= 128: - return 32 - elif head_dim <= 256: - return 64 - else: - return 32 - - # This should match the block sizes in the CUDA kernel - major, minor = torch.cuda.get_device_capability(device) - is_sm8x = major == 8 and minor > 0 # Only include sm86 and sm89, exclude sm80 (A100) - is_sm80 = major == 8 and minor == 0 - is_sm90 = major == 9 and minor == 0 - if head_dim <= 32: - return 128 - if head_dim <= 64: - return 128 if not is_dropout else 64 - elif head_dim <= 96: - return 64 - elif head_dim <= 128: - if is_sm8x: - return 64 if (not is_dropout and is_causal) else 32 - else: - return 64 if not is_dropout else 32 - elif head_dim <= 192: - return 64 - elif head_dim <= 224: - return 64 - elif head_dim <= 256: - return 64 - - -def round_multiple(x, m): - return (x + m - 1) // m * m - - -# torch.compile() support is only enabled for pytorch >= 2.4 -# The reason for this is that we are using the new custom_op and register_fake -# APIs, which support inplace modification of inputs in the function itself -if torch.__version__ >= "2.4.0": - _torch_custom_op_wrapper = torch.library.custom_op - _torch_register_fake_wrapper = torch.library.register_fake -else: - def noop_custom_op_wrapper(name, fn=None, /, *, mutates_args, device_types=None, schema=None): - def wrap(func): - return func - if fn is None: - return wrap - return fn - def noop_register_fake_wrapper(op, fn=None, /, *, lib=None, _stacklevel=1): - def wrap(func): - return func - if fn is None: - return wrap - return fn - _torch_custom_op_wrapper = noop_custom_op_wrapper - _torch_register_fake_wrapper = noop_register_fake_wrapper - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_forward"), mutates_args=(), device_types=_get_device()) -def _flash_attn_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - return_softmax: bool -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - out, softmax_lse, S_dmask, rng_state = flash_attn.fwd( - q, - k, - v, - None, - alibi_slopes, - dropout_p, - softmax_scale, - causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - None, - ) - return out, softmax_lse, S_dmask, rng_state - - -@_torch_register_fake_wrapper(add_op_namespace_prefix("_flash_attn_forward")) -def _flash_attn_forward_fake( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - return_softmax: bool -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - batch_size, seqlen_q, num_heads, head_size = q.shape - seqlen_k = k.shape[1] - out = torch.empty_like(q) - softmax_lse = torch.empty((batch_size, num_heads, seqlen_q), dtype=torch.float32, device=q.device, layout=q.layout) - p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) - if return_softmax: - p = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128), round_multiple(seqlen_k, 128)), dtype=q.dtype, device=q.device, layout=q.layout) - rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) - - return out, softmax_lse, p, rng_state - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_forward = _ops.ops._flash_attn_forward -else: - _wrapped_flash_attn_forward = _flash_attn_forward - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_varlen_forward"), mutates_args=(), device_types=_get_device()) -def _flash_attn_varlen_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - alibi_slopes: Optional[torch.Tensor] = None, - return_softmax: bool = False, - block_table: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - out, softmax_lse, S_dmask, rng_state = flash_attn.varlen_fwd( - q, - k, - v, - None, - cu_seqlens_q, - cu_seqlens_k, - seqused_k, - leftpad_k, - block_table, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - zero_tensors, - causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - None, - ) - # if out.isnan().any() or softmax_lse.isnan().any(): - # breakpoint() - return out, softmax_lse, S_dmask, rng_state - - -@_torch_register_fake_wrapper(add_op_namespace_prefix("_flash_attn_varlen_forward")) -def _flash_attn_varlen_forward_fake( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - alibi_slopes: Optional[torch.Tensor] = None, - return_softmax: bool = False, - block_table: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - paged_kv = block_table is not None - batch_size = cu_seqlens_q.numel() - 1 - total_q, num_heads, _ = q.shape - - out = torch.empty_like(q) - softmax_lse = torch.empty((num_heads, total_q), dtype=torch.float32, device=q.device, layout=q.layout) - p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) - seqlen_q_rounded = round_multiple(max_seqlen_q, 128) - seqlen_k_rounded = round_multiple(max_seqlen_k, 128) - if return_softmax: - p = torch.empty((batch_size, num_heads, seqlen_q_rounded, seqlen_k_rounded), dtype=q.dtype, device=q.device, layout=q.layout) - rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) - return out, softmax_lse, p, rng_state - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_varlen_forward = _ops.ops._flash_attn_varlen_forward -else: - _wrapped_flash_attn_varlen_forward = _flash_attn_varlen_forward - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_backward"), mutates_args=("dq", "dk", "dv"), device_types=_get_device()) -def _flash_attn_backward( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, -) -> 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)] - ( - dq, - dk, - dv, - softmax_d, - ) = flash_attn.bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - alibi_slopes, - dropout_p, - softmax_scale, - causal, - window_size_left, - window_size_right, - softcap, - deterministic, - None, - rng_state, - ) - return softmax_d - - -@_torch_register_fake_wrapper(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, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, -) -> torch.Tensor: - dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] - if dq is None: - dq = torch.empty_like(q) - if dk is None: - dk = torch.empty_like(k) - if dv is None: - dv = torch.empty_like(v) - batch_size, seqlen_q, num_heads, _ = q.shape - softmax_d = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128)), device=q.device, dtype=torch.float32) - - return softmax_d - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_backward = _ops.ops._flash_attn_backward -else: - _wrapped_flash_attn_backward = _flash_attn_backward - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_varlen_backward"), mutates_args=("dq", "dk", "dv"), device_types=_get_device()) -def _flash_attn_varlen_backward( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> 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)] - ( - dq, - dk, - dv, - softmax_d, - ) = flash_attn.varlen_bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - zero_tensors, - causal, - window_size_left, - window_size_right, - softcap, - deterministic, - None, - rng_state, - ) - # if dk.isnan().any() or dk.isnan().any() or dv.isnan().any() or softmax_d.isnan().any(): - # breakpoint() - return softmax_d - - -@_torch_register_fake_wrapper(add_op_namespace_prefix("_flash_attn_varlen_backward")) -def _flash_attn_varlen_backward_fake( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> torch.Tensor: - dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] - batch_size = cu_seqlens_q.numel() - 1 - total_q, num_heads, _ = q.shape - - if dq is None: - dq = torch.empty_like(q) - if dk is None: - dk = torch.empty_like(k) - if dv is None: - dv = torch.empty_like(v) - softmax_d = torch.empty((num_heads, total_q + 128 * batch_size), device=q.device, dtype=torch.float32) - - return softmax_d - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_varlen_backward = _ops.ops._flash_attn_varlen_backward -else: - _wrapped_flash_attn_varlen_backward = _flash_attn_varlen_backward - - -class FlashAttnQKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and qkv.requires_grad - if softmax_scale is None: - softmax_scale = qkv.shape[-1] ** (-0.5) - q, k, v = qkv[:, :, 0].detach(), qkv[:, :, 1].detach(), qkv[:, :, 2].detach() - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) - dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dqkv[:, :, 0], - dqkv[:, :, 1], - dqkv[:, :, 2], - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dqkv = dqkv[..., : dout.shape[-1]] # We could have padded the head dimension - return dqkv, None, None, None, None, None, None, None, None, None - - -class FlashAttnVarlenQKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - cu_seqlens, - max_seqlen, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and qkv.requires_grad - if softmax_scale is None: - softmax_scale = qkv.shape[-1] ** (-0.5) - q, k, v = qkv[:, 0].detach(), qkv[:, 1].detach(), qkv[:, 2].detach() - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens, - cu_seqlens, - max_seqlen, - max_seqlen, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=None, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens, rng_state) - ctx.dropout_p = dropout_p - ctx.max_seqlen = max_seqlen - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens, rng_state = ctx.saved_tensors - qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) - dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dqkv[:, 0], - dqkv[:, 1], - dqkv[:, 2], - cu_seqlens, - cu_seqlens, - ctx.max_seqlen, - ctx.max_seqlen, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - 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 FlashAttnKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - kv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, kv] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - k, v = kv[:, :, 0].detach(), kv[:, :, 1].detach() - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - dq = torch.empty_like(q) - kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) - dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dkv[:, :, 0], - dkv[:, :, 1], - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dkv = dkv[..., : dout.shape[-1]] - return dq, dkv, None, None, None, None, None, None, None, None, None - - -class FlashAttnVarlenKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, kv] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - k, v = kv[:, 0].detach(), kv[:, 1].detach() - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=None, - ) - if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state - ) - ctx.dropout_p = dropout_p - 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.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors - dq = torch.empty_like(q) - kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) - dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dkv[:, 0], - dkv[:, 1], - cu_seqlens_q, - cu_seqlens_k, - ctx.max_seqlen_q, - ctx.max_seqlen_k, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dkv = dkv[..., : dout.shape[-1]] - return dq, dkv, None, None, None, None, None, None, None, None, None, None, None, None, None - - -class FlashAttnFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - k, - v, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, k, v] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dk = dk[..., : dout.shape[-1]] - dv = dv[..., : dout.shape[-1]] - return dq, dk, dv, 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, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - block_table, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, k, v] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=block_table, - ) - if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state - ) - ctx.dropout_p = dropout_p - 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.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors - dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - ctx.max_seqlen_q, - ctx.max_seqlen_k, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dk = dk[..., : dout.shape[-1]] - dv = dv[..., : dout.shape[-1]] - return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None - - -def flash_attn_qkvpacked_func( - qkv, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # <=0.0 means deactivate - alibi_slopes=None, - deterministic=False, - 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, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_kvpacked_func( - q, - kv, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, -): - """dropout_p should be set to 0.0 during evaluation - If 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 K, V. - 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) - kv: (batch_size, seqlen, 2, 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. - softcap: float. Anything > 0 activates softcapping 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). - 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 FlashAttnKVPackedFunc.apply( - q, - kv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_func( - q, - k, - v, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - 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). - 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 FlashAttnFunc.apply( - q, - k, - v, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_qkvpacked_func( - qkv, - cu_seqlens, - max_seqlen, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - 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_varlen_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_varlen_kvpacked_func and flash_attn_varlen_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: (total, 3, nheads, headdim), where total = total number of tokens in the batch. - cu_seqlens: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into qkv. - max_seqlen: int. Maximum sequence length in the batch. - 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenQKVPackedFunc.apply( - qkv, - cu_seqlens, - max_seqlen, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_kvpacked_func( - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, -): - """dropout_p should be set to 0.0 during evaluation - If 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 K, V. - 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: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch. - kv: (total_k, 2, nheads_k, headdim), where total_k = total number of key tokens in the batch. - cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into q. - cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into kv. - max_seqlen_q: int. Maximum query sequence length in the batch. - max_seqlen_k: int. Maximum key sequence length in the batch. - 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 + 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenKVPackedFunc.apply( - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_func( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, - block_table=None, -): - """dropout_p should be set to 0.0 during evaluation - Supports multi-query and grouped-query attention (MQA/GQA) by passing in K, V 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: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch. - k: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. - v: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. - cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into q. - cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into kv. - max_seqlen_q: int. Maximum query sequence length in the batch. - max_seqlen_k: int. Maximum key sequence length in the batch. - 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 + 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenFunc.apply( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - block_table, - False if q.device.type == "cpu" else torch.is_grad_enabled(), - ) - - -def flash_attn_with_kvcache( - q, - k_cache, - v_cache, - k=None, - v=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, - block_table: Optional[torch.Tensor] = None, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - rotary_interleaved=True, - alibi_slopes=None, - num_splits=0, - 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 block_table, - or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_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) if there's no block_table, - or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_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). Similar to k. - 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. - block_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). - 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. - 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" - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-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) - cache_batch_idx = maybe_contiguous(cache_batch_idx) - block_table = maybe_contiguous(block_table) - out, softmax_lse = flash_attn.fwd_kvcache( - q, - k_cache, - v_cache, - k, - v, - cache_seqlens, - rotary_cos, - rotary_sin, - cache_batch_idx, - cache_leftpad, - block_table, - alibi_slopes, - None, - softmax_scale, - causal, - window_size[0], - window_size[1], - softcap, - rotary_interleaved, - num_splits, - ) - return (out, softmax_lse) if return_softmax_lse else out diff --git a/build/torch211-cxx11-cu130-aarch64-linux/layers/__init__.py b/build/torch211-cxx11-cu130-aarch64-linux/layers/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/build/torch211-cxx11-cu130-aarch64-linux/layers/patch_embed.py b/build/torch211-cxx11-cu130-aarch64-linux/layers/patch_embed.py deleted file mode 100644 index 05562f8e8bcdb58e947c6f402a49eacd2d031871..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu130-aarch64-linux/layers/patch_embed.py +++ /dev/null @@ -1,67 +0,0 @@ -# We use the same API as https://github.com/rwightman/pytorch-image-models/blob/v0.6.11/timm/models/layers/patch_embed.py -# But we use nn.Linear instead of Conv2d and it's about 8x faster. - -from functools import partial - -import torch.nn as nn -from einops import rearrange -from torch import _assert -from torch.nn.modules.utils import _pair - -try: - from flash_attn.ops.fused_dense import FusedDense -except ImportError: - FusedDense = None - - -class PatchEmbed(nn.Module): - """2D Image to Patch Embedding""" - - def __init__( - self, - img_size=224, - patch_size=16, - in_chans=3, - embed_dim=768, - norm_layer=None, - flatten=True, - bias=True, - fused_bias_fc=False, - ): - super().__init__() - img_size = _pair(img_size) - patch_size = _pair(patch_size) - self.img_size = img_size - self.patch_size = patch_size - self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1]) - self.num_patches = self.grid_size[0] * self.grid_size[1] - self.flatten = flatten - if fused_bias_fc and FusedDense is None: - raise ImportError("fused_dense is not installed") - - linear_cls = nn.Linear if not fused_bias_fc or not bias else FusedDense - self.proj = linear_cls(in_chans * patch_size[0] * patch_size[1], embed_dim, bias=bias) - self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() - - def forward(self, x): - _, _, H, W = x.shape - _assert( - H == self.img_size[0], - f"Input image height ({H}) doesn't match model ({self.img_size[0]}).", - ) - _assert( - W == self.img_size[1], - f"Input image width ({W}) doesn't match model ({self.img_size[1]}).", - ) - x = self.proj( - rearrange( - x, - "b c (h p1) (w p2) -> b h w (c p1 p2)", - p1=self.patch_size[0], - p2=self.patch_size[1], - ) - ) - if self.flatten: - x = rearrange(x, "b h w c -> b (h w) c") - x = self.norm(x) - return x diff --git a/build/torch211-cxx11-cu130-aarch64-linux/layers/rotary.py b/build/torch211-cxx11-cu130-aarch64-linux/layers/rotary.py deleted file mode 100644 index d1bfc21fc7de1dd287e8f382847b194a48075981..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu130-aarch64-linux/layers/rotary.py +++ /dev/null @@ -1,483 +0,0 @@ -# Copyright (c) 2025, Tri Dao - -import math -from functools import partial -from typing import Optional, Tuple, Union - -import torch -from torch import Tensor - -from einops import rearrange, repeat -# from flash_attn.ops.triton.rotary import apply_rotary -from ..ops.triton.rotary import apply_rotary - - -def rotate_half(x, interleaved=False): - if not interleaved: - x1, x2 = x.chunk(2, dim=-1) - return torch.cat((-x2, x1), dim=-1) - else: - x1, x2 = x[..., ::2], x[..., 1::2] - return rearrange(torch.stack((-x2, x1), dim=-1), "... d two -> ... (d two)", two=2) - - -def apply_rotary_emb_torch(x, cos, sin, interleaved=False): - """ - x: (batch_size, seqlen, nheads, headdim) - cos, sin: (seqlen, rotary_dim / 2) or (batch_size, seqlen, rotary_dim / 2) - """ - ro_dim = cos.shape[-1] * 2 - assert ro_dim <= x.shape[-1] - cos = repeat(cos, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)") - sin = repeat(sin, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)") - return torch.cat( - [x[..., :ro_dim] * cos + rotate_half(x[..., :ro_dim], interleaved) * sin, x[..., ro_dim:]], - dim=-1, - ) - - -class ApplyRotaryEmb(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x, - cos, - sin, - interleaved=False, - inplace=False, - seqlen_offsets: Union[int, Tensor] = 0, - cu_seqlens: Optional[Tensor] = None, - max_seqlen: Optional[int] = None, - ): - out = apply_rotary( - x, - cos, - sin, - seqlen_offsets=seqlen_offsets, - cu_seqlens=cu_seqlens, - max_seqlen=max_seqlen, - interleaved=interleaved, - inplace=inplace, - ) - if isinstance(seqlen_offsets, int): - ctx.save_for_backward(cos, sin, cu_seqlens) # Can't save int with save_for_backward - ctx.seqlen_offsets = seqlen_offsets - else: - ctx.save_for_backward(cos, sin, cu_seqlens, seqlen_offsets) - ctx.seqlen_offsets = None - ctx.interleaved = interleaved - ctx.inplace = inplace - ctx.max_seqlen = max_seqlen - return out if not inplace else x - - @staticmethod - def backward(ctx, do): - seqlen_offsets = ctx.seqlen_offsets - if seqlen_offsets is None: - cos, sin, cu_seqlens, seqlen_offsets = ctx.saved_tensors - else: - cos, sin, cu_seqlens = ctx.saved_tensors - dx = apply_rotary( - do, - cos, - sin, - seqlen_offsets=seqlen_offsets, - cu_seqlens=cu_seqlens, - max_seqlen=ctx.max_seqlen, - interleaved=ctx.interleaved, - inplace=ctx.inplace, - conjugate=True, - ) - return dx, None, None, None, None, None, None, None - - -def apply_rotary_emb( - x, - cos, - sin, - interleaved=False, - inplace=False, - seqlen_offsets: Union[int, Tensor] = 0, - cu_seqlens: Optional[Tensor] = None, - max_seqlen: Optional[int] = None, -): - """ - Arguments: - x: (batch_size, seqlen, nheads, headdim) if cu_seqlens is None - else (total_seqlen, nheads, headdim) - cos, sin: (seqlen_rotary, rotary_dim / 2) - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead - of 1st half and 2nd half (GPT-NeoX style). - inplace: if True, apply rotary embedding in-place. - seqlen_offsets: (batch_size,) or int. Each sequence in x is shifted by this amount. - Most commonly used in inference when we have KV cache. - cu_seqlens: (batch + 1,) or None - max_seqlen: int - Return: - out: (batch_size, seqlen, nheads, headdim) if cu_seqlens is None - else (total_seqlen, nheads, headdim) - rotary_dim must be <= headdim - Apply rotary embedding to the first rotary_dim of x. - """ - return ApplyRotaryEmb.apply( - x, cos, sin, interleaved, inplace, seqlen_offsets, cu_seqlens, max_seqlen - ) - - -# For backward compatibility -apply_rotary_emb_func = apply_rotary_emb - - -def _apply_rotary_emb_qkv( - qkv, - cos, - sin, - cos_k=None, - sin_k=None, - interleaved=False, - inplace=False, - conjugate=False, - seqlen_offsets: Union[int, Tensor] = 0, - num_heads_q: Optional[int] = None, -): - apply_rotary_fn = partial( - apply_rotary, - interleaved=interleaved, - inplace=inplace, - conjugate=conjugate, - seqlen_offsets=seqlen_offsets - ) - if cos_k is None and sin_k is None and qkv.is_contiguous(): - # Call 1 kernel instead of 2 kernels - # We need qkv to be contiguous so that when we reshape to combine (3, nheads) - # dimensions, we get the same tensor - if qkv.dim() == 5: - batch, seqlen, three, nheads, headdim = qkv.shape - assert three == 3 - # qk = rearrange(qkv[:, :, :2], "b s t h d -> b s (t h) d") - qk = qkv[:, :, :2].reshape(batch, seqlen, -1, headdim) - qk = apply_rotary_fn(qk, cos, sin) - else: - assert qkv.dim() == 4 - assert num_heads_q is not None - num_heads_k = (qkv.shape[2] - num_heads_q) // 2 - assert qkv.shape[2] == num_heads_q + 2 * num_heads_k - qk = qkv[:, :, :num_heads_q + num_heads_k] - qk = apply_rotary_fn(qk, cos, sin) - if not inplace: - if qkv.dim() == 5: - qkv = torch.cat([rearrange(qk, "b s (t h) d -> b s t h d", t=2), qkv[:, :, 2:]], dim=2) - else: - qkv = torch.cat([qk, qkv[:, :, num_heads_q + num_heads_k :]], dim=2) - else: - cos_k = cos if cos_k is None else cos_k - sin_k = sin if sin_k is None else sin_k - if qkv.dim() == 5: - batch, seqlen, three, nheads, headdim = qkv.shape - assert three == 3 - q, k = qkv[:, :, 0], qkv[:, :, 1] - else: - assert qkv.dim() == 4 - assert num_heads_q is not None - num_heads_k = (qkv.shape[2] - num_heads_q) // 2 - assert qkv.shape[2] == num_heads_q + 2 * num_heads_k - q, k = qkv[:, :, :num_heads_q], qkv[:, :, num_heads_q : num_heads_q + num_heads_k] - q = apply_rotary_fn(q, cos, sin) - k = apply_rotary_fn(k, cos_k, sin_k) - if not inplace: - if qkv.dim() == 5: - qkv = torch.stack([q, k, qkv[:, :, 2]], dim=2) - else: - qkv = torch.cat([q, k, qkv[:, :, num_heads_q + num_heads_k:]], dim=2) - return qkv - - -class ApplyRotaryEmbQKV_(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - cos, - sin, - cos_k=None, - sin_k=None, - interleaved=False, - seqlen_offsets: Union[int, torch.Tensor] = 0, - num_heads_q: Optional[int] = None, - ): - # apply_rotary_emb_qkv_inplace( - qkv = _apply_rotary_emb_qkv( - qkv, cos, sin, cos_k, sin_k, interleaved=interleaved, inplace=True, - seqlen_offsets=seqlen_offsets, num_heads_q=num_heads_q, - ) - if isinstance(seqlen_offsets, int): - ctx.save_for_backward(cos, sin, cos_k, sin_k) - ctx.seqlen_offsets = seqlen_offsets - else: - ctx.save_for_backward(cos, sin, cos_k, sin_k, seqlen_offsets) - ctx.seqlen_offsets = None - ctx.interleaved = interleaved - ctx.num_heads_q = num_heads_q - return qkv - - @staticmethod - def backward(ctx, dqkv): - seqlen_offsets = ctx.seqlen_offsets - if seqlen_offsets is None: - cos, sin, cos_k, sin_k, seqlen_offsets = ctx.saved_tensors - else: - cos, sin, cos_k, sin_k = ctx.saved_tensors - dqkv = _apply_rotary_emb_qkv( - dqkv, cos, sin, cos_k, sin_k, interleaved=ctx.interleaved, inplace=True, - seqlen_offsets=seqlen_offsets, num_heads_q=ctx.num_heads_q, conjugate=True, - ) - return dqkv, None, None, None, None, None, None, None - - -def apply_rotary_emb_qkv_( - qkv, - cos, - sin, - cos_k=None, - sin_k=None, - interleaved=False, - seqlen_offsets: Union[int, torch.Tensor] = 0, - num_heads_q: Optional[int] = None, -): - """ - Arguments: - qkv: (batch_size, seqlen, 3, nheads, headdim) or (batch_size, seqlen, num_heads_q + 2 * num_heads_k, headdim). - If qkv has shape (batch_size, seqlen, num_heads_q + 2 * num_heads_k, headdim) (e.g. MQA / GQA), - then num_heads_q must be provided. - cos, sin: (seqlen, rotary_dim / 2) - cos_k, sin_k: (seqlen, rotary_dim / 2), optional - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead of - 1st half and 2nd half (GPT-NeoX style). - seqlen_offsets: (batch_size,) or int. Each sequence in Q and K is shifted by this amount. - Most commonly used in inference when we have KV cache. - Return: - qkv: (batch_size, seqlen, 3, nheads, headdim) or (batch_size, seqlen, num_heads_q + 2 * num_heads_k, headdim) - rotary_dim must be <= headdim - Apply rotary embedding *inplace* to the first rotary_dim of Q and K. - """ - return ApplyRotaryEmbQKV_.apply( - qkv, cos, sin, cos_k, sin_k, interleaved, seqlen_offsets, num_heads_q - ) - - -class ApplyRotaryEmbKV_(torch.autograd.Function): - - @staticmethod - def forward(ctx, kv, cos, sin, interleaved=False, seqlen_offsets: Union[int, torch.Tensor] = 0): - batch, seqlen, two, nheads, headdim = kv.shape - assert two == 2 - k = kv[:, :, 0] - apply_rotary( - k, cos, sin, seqlen_offsets=seqlen_offsets, interleaved=interleaved, inplace=True - ) - if isinstance(seqlen_offsets, int): - ctx.save_for_backward(cos, sin) # Can't save int with save_for_backward - ctx.seqlen_offsets = seqlen_offsets - else: - ctx.save_for_backward(cos, sin, seqlen_offsets) - ctx.seqlen_offsets = None - ctx.interleaved = interleaved - return kv - - @staticmethod - def backward(ctx, dkv): - seqlen_offsets = ctx.seqlen_offsets - if seqlen_offsets is None: - cos, sin, seqlen_offsets = ctx.saved_tensors - else: - cos, sin = ctx.saved_tensors - apply_rotary( - dkv[:, :, 0], - cos, - sin, - seqlen_offsets=seqlen_offsets, - interleaved=ctx.interleaved, - inplace=True, - conjugate=True, - ) - return dkv, None, None, None, None - - -apply_rotary_emb_kv_ = ApplyRotaryEmbKV_.apply - - -def apply_rotary_emb_kv_( - kv, - cos, - sin, - interleaved=False, - seqlen_offsets: Union[int, torch.Tensor] = 0, -): - """ - Arguments: - kv: (batch_size, seqlen, 2, nheads, headdim) - cos, sin: (seqlen, rotary_dim / 2) - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead of - 1st half and 2nd half (GPT-NeoX style). - seqlen_offsets: (batch_size,) or int. Each sequence in Q and K is shifted by this amount. - Most commonly used in inference when we have KV cache. - Return: - kv: (batch_size, seqlen, 2, nheads, headdim) - rotary_dim must be <= headdim - Apply rotary embedding *inplace* to the first rotary_dim of K. - """ - return ApplyRotaryEmbKV_.apply(kv, cos, sin, interleaved, seqlen_offsets) - - -class RotaryEmbedding(torch.nn.Module): - """ - The rotary position embeddings from RoFormer_ (Su et. al). - A crucial insight from the method is that the query and keys are - transformed by rotation matrices which depend on the relative positions. - - Other implementations are available in the Rotary Transformer repo_ and in - GPT-NeoX_, GPT-NeoX was an inspiration - - .. _RoFormer: https://arxiv.org/abs/2104.09864 - .. _repo: https://github.com/ZhuiyiTechnology/roformer - .. _GPT-NeoX: https://github.com/EleutherAI/gpt-neox - - If scale_base is not None, this implements XPos (Sun et al., https://arxiv.org/abs/2212.10554). - A recommended value for scale_base is 512: https://github.com/HazyResearch/flash-attention/issues/96 - Reference: https://github.com/sunyt32/torchscale/blob/main/torchscale/component/xpos_relative_position.py - """ - - def __init__( - self, - dim: int, - base=10000.0, - interleaved=False, - scale_base=None, - device=None, - ): - """ - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead - of 1st half and 2nd half (GPT-NeoX style). - """ - super().__init__() - self.dim = dim - self.base = float(base) - # Generate and save the inverse frequency buffer (non trainable) - inv_freq = self._compute_inv_freq(device) - self.register_buffer("inv_freq", inv_freq, persistent=False) - self.interleaved = interleaved - self.scale_base = scale_base - scale = ( - (torch.arange(0, dim, 2, device=device, dtype=torch.float32) + 0.4 * dim) / (1.4 * dim) - if scale_base is not None - else None - ) - self.register_buffer("scale", scale, persistent=False) - - self._seq_len_cached = 0 - self._cos_cached = None - self._sin_cached = None - self._cos_k_cached = None - self._sin_k_cached = None - - def _compute_inv_freq(self, device=None): - return 1.0 / ( - self.base - ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim) - ) - - def _update_cos_sin_cache(self, seqlen, device=None, dtype=None): - # Reset the tables if the sequence length has changed, - # if we're on a new device (possibly due to tracing for instance), - # or if we're switching from inference mode to training - if ( - seqlen > self._seq_len_cached - or self._cos_cached is None - or self._cos_cached.device != device - or self._cos_cached.dtype != dtype - or (self.training and self._cos_cached.is_inference()) - ): - self._seq_len_cached = seqlen - # We want fp32 here, not self.inv_freq.dtype, since the model could be loaded in bf16 - # And the output of arange can be quite large, so bf16 would lose a lot of precision. - t = torch.arange(seqlen, device=device, dtype=torch.float32) - # We want fp32 here as well since inv_freq will be multiplied with t, and the output - # will be large. Having it in bf16 will lose a lot of precision and cause the - # cos & sin output to change significantly. - # We want to recompute self.inv_freq if it was not loaded in fp32 - if self.inv_freq.dtype != torch.float32: - inv_freq = self._compute_inv_freq(device=device) - else: - inv_freq = self.inv_freq - # Don't do einsum, it converts fp32 to bf16 under AMP - # freqs = torch.einsum("i,j->ij", t, self.inv_freq) - freqs = torch.outer(t, inv_freq) - if self.scale is None: - self._cos_cached = torch.cos(freqs).to(dtype) - self._sin_cached = torch.sin(freqs).to(dtype) - else: - power = ( - torch.arange(seqlen, dtype=self.scale.dtype, device=self.scale.device) - - seqlen // 2 - ) / self.scale_base - scale = self.scale.to(device=power.device) ** rearrange(power, "s -> s 1") - # We want the multiplication by scale to happen in fp32 - self._cos_cached = (torch.cos(freqs) * scale).to(dtype) - self._sin_cached = (torch.sin(freqs) * scale).to(dtype) - self._cos_k_cached = (torch.cos(freqs) / scale).to(dtype) - self._sin_k_cached = (torch.sin(freqs) / scale).to(dtype) - - def forward( - self, - qkv: torch.Tensor, - kv: Optional[torch.Tensor] = None, - seqlen_offset: Union[int, torch.Tensor] = 0, - max_seqlen: Optional[int] = None, - num_heads_q: Optional[int] = None, - ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: - """ - qkv: (batch, seqlen, 3, nheads, headdim) or (batch, seqlen, num_heads_q + 2 * num_heads_k, headdim) - if kv is none, else it's just q of shape (batch, seqlen, nheads, headdim). - If qkv has shape (batch, seqlen, num_heads_q + 2 * num_heads_k, headdim) (e.g. MQA / GQA), - then num_heads_q must be provided. - kv: (batch, seqlen, 2, nheads, headdim) - seqlen_offset: (batch_size,) or int. Each sequence in x is shifted by this amount. - Most commonly used in inference when we have KV cache. - If it's a tensor of shape (batch_size,), then to update the cos / sin cache, one - should pass in max_seqlen, which will update the cos / sin cache up to that length. - Apply rotary embedding *inplace* to qkv and / or kv. - """ - seqlen = qkv.shape[1] - if max_seqlen is not None: - self._update_cos_sin_cache(max_seqlen, device=qkv.device, dtype=qkv.dtype) - elif isinstance(seqlen_offset, int): - self._update_cos_sin_cache(seqlen + seqlen_offset, device=qkv.device, dtype=qkv.dtype) - if kv is None: - return apply_rotary_emb_qkv_( - qkv, - self._cos_cached, - self._sin_cached, - self._cos_k_cached if self.scale is not None else None, - self._sin_k_cached if self.scale is not None else None, - interleaved=self.interleaved, - seqlen_offsets=seqlen_offset, - num_heads_q=num_heads_q, - ) - else: - q = qkv - q = apply_rotary_emb_func( - q, - self._cos_cached, - self._sin_cached, - interleaved=self.interleaved, - inplace=True, - seqlen_offsets=seqlen_offset, - ) - kv = apply_rotary_emb_kv_( - kv, - self._cos_cached if self.scale is None else self._cos_k_cached, - self._sin_cached if self.scale is None else self._sin_k_cached, - interleaved=self.interleaved, - seqlen_offsets=seqlen_offset, - ) - return q, kv diff --git a/build/torch211-cxx11-cu130-aarch64-linux/metadata.json b/build/torch211-cxx11-cu130-aarch64-linux/metadata.json deleted file mode 100644 index d9fb2ed8da88a9f96c6195300019cfcfb8df882f..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu130-aarch64-linux/metadata.json +++ /dev/null @@ -1,42 +0,0 @@ -{ - "name": "flash-attn2", - "id": "_flash_attn2_cuda_f12afc9", - "version": 1, - "license": "BSD-3-Clause", - "python-depends": [], - "backend": { - "type": "cuda", - "archs": [ - "10.0", - "12.0", - "8.0", - "9.0" - ] - }, - "digest": { - "algorithm": "sha256", - "files": { - "__init__.py": "Cd9MEdd/WUYnBmfiQpC4SyPqAxrTilTjvsvswNju4zc=", - "_flash_attn2_cuda_f12afc9.abi3.so": "ouJYCCn5yvOQyUHxa2Dkyle0ODjE7euyVKoJk3Hsq98=", - "_ops.py": "GQbmTTeSOU+7Q2nI68Hz+1RvEogyfEFccx1ubUYw+2s=", - "bert_padding.py": "gF1EmsdJ+HpQ86MRQ4VxDw+Sb/RVISdQALdNnoByHlw=", - "flash_attn2/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY=", - "flash_attn_interface.py": "W7a9bHR2ly6QL3aKqsBqxWA7M7ZWYcZPmtsPJu6rAfY=", - "layers/__init__.py": "47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU=", - "layers/patch_embed.py": "H58CgME/qSOPTZLOG08wFgrQS1j34pvNwMPrkTj3Ek4=", - "layers/rotary.py": "rFIdNrZQFNawwLD6d7t5LE2m+M93wwDZ8EuExw+gfD8=", - "ops/__init__.py": "47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU=", - "ops/activations.py": "t5lzNg1In8LP6bKeTnyeMizwqjv27JGbJ6ylPdGvZYg=", - "ops/fused_dense.py": "BqT8VTSCGdLJAjh6WBMkCRFYrdVodPeoElWBvH6AIzA=", - "ops/layer_norm.py": "zr7NXIm+2mtEynTp1CS0fbFGI2Mqdp41dY4AfDWF6EQ=", - "ops/rms_norm.py": "XEnihcj0a4aSz4LO55m5iKGVn4HKTeKN8TIyHjuDgxI=", - "ops/triton/__init__.py": "AbpHGcgLb+kRsJGnwFEktk7uzpZOCcBY74+YBdrKVGs=", - "ops/triton/cross_entropy.py": "y56nmuOtpSNxM7sDZcUSgXKIsIEayfuvrU8bAVog/nw=", - "ops/triton/k_activations.py": "+Z3vIyO4JkqBMipKsPvhzmxljtBdIhJCsl/M+/ESqBo=", - "ops/triton/layer_norm.py": "BBQdJiqopbIS12viUqzx3SvOQCNbDSAYLJXhpyMGyC8=", - "ops/triton/linear.py": "OtRvKz8xdpl+7v3q/ZTaS9fdBt9XrzMyapgRr50uBbM=", - "ops/triton/mlp.py": "uWA+D7HyuRUsOT+4nhi3vvCDNZnNEXwYcHMQ02UREfQ=", - "ops/triton/rotary.py": "0YSrlbep0lJmzLpTiANs0slTFeGcQQqR+el6uzbSLvI=" - } - } -} \ No newline at end of file diff --git a/build/torch211-cxx11-cu130-aarch64-linux/metadata.json.sigstore b/build/torch211-cxx11-cu130-aarch64-linux/metadata.json.sigstore deleted file mode 100644 index 34188f3c45609f0b1c780891e1a7b6af4e8da0f4..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu130-aarch64-linux/metadata.json.sigstore +++ /dev/null @@ -1 +0,0 @@ -{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHSTCCBtCgAwIBAgIUBt+cPpQMGmVjdUU+aRbepBgEU20wCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwNjIzMjM0MDQzWhcNMjYwNjIzMjM1MDQzWjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEILCE0vAyXVJMqxYciiTIDekNzwxddo3bGaEYl8hIPD1MTOOO3COd6rHUXiFH7v0C1mmFvFJivDy4rSTHPU3zn6OCBe8wggXrMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUJMBZA8FLkQ8wMfkYxJ/pYYjc00EwHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoZjEyYWZjOTJjNDBiYmQ3NmQ5OTBkOTExZWJhYmYyOGU0OTBhN2JkZTATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoZjEyYWZjOTJjNDBiYmQ3NmQ5OTBkOTExZWJhYmYyOGU0OTBhN2JkZTAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoZjEyYWZjOTJjNDBiYmQ3NmQ5OTBkOTExZWJhYmYyOGU0OTBhN2JkZTAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKGYxMmFmYzkyYzQwYmJkNzZkOTkwZDkxMWViYWJmMjhlNDkwYTdiZGUwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMjgwMzYyNjg5MzUvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBiQYKKwYBBAHWeQIEAgR7BHkAdwB1AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABnvbbi78AAAQDAEYwRAIgehu39h4wKNr53LgTfrgJiJziVXS7CH3S2tgexfStIg0CICH5AW9jlS7FCdqBfnsXRReBxTs/Ah7Wurzg/LUNz6BJMAoGCCqGSM49BAMDA2cAMGQCMF/0yKtI+Hqlec5XFWy8YvmhSK5qfANqG16dXUMBL6rltjRwsCcZUkdxdUBzf4qMVQIwfVoID66JqIGQiq3bsS7G5N9yhe/2GnQH5Lc5QMCx3Wu8ajGqZ+2opcqDQnw2XYlQ"}, "tlogEntries":[{"logIndex":"1933039631", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1782258044", "inclusionPromise":{"signedEntryTimestamp":"MEYCIQCSqJFaclTWEywFkuH32fFDVgKG5fLSRfn2/1SSzBAvfAIhAIs5ibRVfSotQ7FOP9yzDGOm2ngxx1e7ufswfBBoPBih"}, "inclusionProof":{"logIndex":"1811135369", "rootHash":"30tL16t+k88p9VBS3Y62WEnWK/tiPCyK+Gn55o6p8Og=", "treeSize":"1811135376", "hashes":["dklmBCH6VlNl8PVXzXiowdxxAw9zdCsmkE/whnHyY2w=", "NhzNTpD0QwBWBEj34npH7QysnYSTYVhQgmhKOU9S0f4=", "XRy6LLO8KxfgQhBegL9TG500TrwGpdXN6CKH/djLzTk=", "fBzysZXOfeWKNIjJAUJKSWe23EzkNqWh0jsPHHV4+Ws=", "nXobHpXWfSQK0Xc/1LYwO037xKTIYgRh9DWPXinj280=", "Aakktr4S+GSdq6/pfGIo0B3cIQgfAs66LkJAq1RdLpA=", "QVVI+8UorpK5fVvAVaM/9A8oGHpbEDzrgIYi5jMqbqw=", "WD8KwmQq5R8K69koYj5k4K7EKo2VF3LqJcx2xoXcd7o=", "99JTcUz0aJeZZocenAppVOLl8v91qMn0v7CmEjq9lws=", "OSDrt6jTKV6LYAkaG63p2BFz4SmbNmGb3uerhF5pOpQ=", "yehxDZkBWOGIzNDb93qxmQa6VBVo7kacVSoHyXMm3ZI=", "oUh+N71k3jWcTmlDDB6Esa5uqlCnbmzxTmDdut2A16M=", "2BsqkdK2omLvu+CpsFjpwqGi6m22H/7zcwCVkPwslM8=", "S8Oe9ZicoJqbVuL1I0dIdTgVpLc82Cm2CMOcV83sKVs=", "E1O2ilWlpFRLWGdbmJUMDZ2qbGu/eDZ12m6YjQ9/JPA=", "nyv2TDlfkz0XfRwjY98zsbUbOkBZ+Tgz7PML0c2uN14=", "OMfDh4k9YD0E+t4sjqJsh+vg8D1ueo1cC07exyxmr6s=", "+vHKm4wNwGJ46ceyZ8mePD7Vn1SediStNHCStRoTSWQ=", "5HwWjuGNDisetbiWM0cWIXM65TIM3EJlqk13ntCQlcU=", "NDIRlqxO0IaGr94q86MIU7VCAN3yEEgOHS4loya5lu4=", "daxmZaajRpZV+JxHiOYZhJBiSKN5ucqjh2WnGbHhirw=", "DOCeoSMovIvLExkhIvisow9AuNXgeWs4ECkyR6EcqYU="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n1811135376\n30tL16t+k88p9VBS3Y62WEnWK/tiPCyK+Gn55o6p8Og=\n\n— rekor.sigstore.dev wNI9ajBFAiEAmgPY10ETeqfqN8PS/147/J/q79HLTmtl2JzEoVWXyRYCIEkmjePNGSsGhWLBxsmac2tU+3CittHUOliBZo64i0N0\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiJjZjUxZTgyMWY3NmU2MzU5ZGYwMzFiMjAwZjQyZTE2NjUwYzQ4NzNiMDQxNmFmNDQxODkxOTFlYzZlMTUwN2IwIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJSDJ6enc4Nnp2Vnhlb2szOUttTTl4U3RYYXFVcDNyZjg3d1B3UEZFR2hQVEFpRUF2TVBWVk5MWG5JSHVaa2VKVmhwYklvOTBWZU5keWlweStUWUFrcWtiaVowPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRWRU5EUW5SRFowRjNTVUpCWjBsVlFuUXJZMUJ3VVUxSGJWWnFaRlZWSzJGU1ltVndRbWRGVlRJd2QwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDVxU1hwTmFrMHdUVVJSZWxkb1kwNU5hbGwzVG1wSmVrMXFUVEZOUkZGNlYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZKVEVORk1IWkJlVmhXU2sxeGVGbGphV2xVU1VSbGEwNTZkM2hrWkc4ellrZGhSVmtLYkRob1NWQkVNVTFVVDA5UE0wTlBaRFp5U0ZWWWFVWklOM1l3UXpGdGJVWjJSa3BwZGtSNU5ISlRWRWhRVlRONmJqWlBRMEpsT0hkbloxaHlUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlZLVFVKYUNrRTRSa3hyVVRoM1RXWnJXWGhLTDNCWldXcGpNREJGZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOWFha1Y1V1ZkYWFrOVVTbXBPUkVKcFdXMVJNMDV0VVRWUFZFSnJDazlVUlhoYVYwcG9XVzFaZVU5SFZUQlBWRUpvVGpKS2ExcFVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMXBxUlhsWlYxcHFUMVJLYWs1RVFtbFpiVkV6VG0xUk5VOVVRbXRQVkVWNFdsZEthRmx0V1hrS1QwZFZNRTlVUW1oT01rcHJXbFJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMXBxUlhsWlYxcHFUMVJLYWs1RVFtbFpiVkV6VG0xUk5VOVVRbXNLVDFSRmVGcFhTbWhaYlZsNVQwZFZNRTlVUW1oT01rcHJXbFJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwZFplRTF0Um0wS1dYcHJlVmw2VVhkWmJVcHJUbnBhYTA5VWEzZGFSR3Q0VFZkV2FWbFhTbTFOYW1oc1RrUnJkMWxVWkdsYVIxVjNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOYW1kM1RYcFplVTVxWnpWTmVsVjJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwVVZsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTTjBKSWEwRUtaSGRDTVVGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtNTJZbUpwTnpoQlFVRlJSQXBCUlZsM1VrRkpaMlZvZFRNNWFEUjNTMDV5TlROTVoxUm1jbWRLYVVwNmFWWllVemREU0ROVE1uUm5aWGhtVTNSSlp6QkRTVU5JTlVGWE9XcHNVemRHQ2tOa2NVSm1ibk5ZVWxKbFFuaFVjeTlCYURkWGRYSjZaeTlNVlU1Nk5rSktUVUZ2UjBORGNVZFRUVFE1UWtGTlJFRXlZMEZOUjFGRFRVWXZNSGxMZEVrS0swaHhiR1ZqTlZoR1YzazRXWFp0YUZOTE5YRm1RVTV4UnpFMlpGaFZUVUpNTm5Kc2RHcFNkM05EWTFwVmEyUjRaRlZDZW1ZMGNVMVdVVWwzWmxadlNRcEVOalpLY1VsSFVXbHhNMkp6VXpkSE5VNDVlV2hsTHpKSGJsRklOVXhqTlZGTlEzZ3pWM1U0WVdwSGNWb3JNbTl3WTNGRVVXNTNNbGhaYkZFS0xTMHRMUzFGVGtRZ1EwVlNWRWxHU1VOQlZFVXRMUzB0TFFvPSJ9fX19"}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyjADAgEAMIICwQYJKoZIhvcNAQcCoIICsjCCAq4CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgIhaaPWkEg27bfk58I1PMf5vZeiUT8uM/bcEoHdBy6QACFQDhvZ1sqed8sGKz0KhRydsN6F+1UhgPMjAyNjA2MjMyMzQwNDNaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHbMIIB1wIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MDYyMzIzNDA0M1owLwYJKoZIhvcNAQkEMSIEIBsUyDZ/TjjxsOzSJvPqhnngqL1a6wfNn6Spqi79ZrCpMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRnMGUCMQDbl0CMFNuNUWJDDlTsjknOCzrABjSV60j0pujMDXwqFR86Cl4MZHKOVP8pXn6J6c8CMCLl0qUI8vckzMuIuenjbikh8qz3CcbSNl9SKcfUJk7IaU1UoOps5zOZb7eeKMwHnw=="}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"z1HoIfduY1nfAxsgD0LhZlDEhzsEFq9EGJGR7G4VB7A="}, "signature":"MEUCIH2zzw86zvVxeok39KmM9xStXaqUp3rf87wPwPFEGhPTAiEAvMPVVNLXnIHuZkeJVhpbIo90VeNdyipy+TYAkqkbiZ0="}} \ No newline at end of file diff --git a/build/torch211-cxx11-cu130-aarch64-linux/ops/__init__.py b/build/torch211-cxx11-cu130-aarch64-linux/ops/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/build/torch211-cxx11-cu130-aarch64-linux/ops/activations.py b/build/torch211-cxx11-cu130-aarch64-linux/ops/activations.py deleted file mode 100644 index 7c09649fc41e12d5a360c5672825d8380bc7ec80..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu130-aarch64-linux/ops/activations.py +++ /dev/null @@ -1,135 +0,0 @@ -# Copied from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/model/layers/activations.py -import math - -import torch -import torch.nn as nn -import torch.nn.functional as F - -# 1/sqrt(2*pi)-> 0.3989423 -# 1/sqrt(2) -> 0.70710678 -# sqrt(2/pi) -> 0.79788456 - -# this function is tanh approximation of gelu -# actual gelu is: -# x * 0.5 * (1.0 + torch.erf(x * 0.70710678)) -@torch.jit.script -def bias_gelu(y, bias): - x = bias + y - return (x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))).to(dtype=y.dtype) - - -# gradient of tanh approximation of gelu -# gradient of actual gelu is: -# 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x) -@torch.jit.script -def bias_gelu_back(g, y, bias): - """Assume that y has shape (B, D) and bias has shape (D)""" - x = bias + y - tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)) - # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243 - ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( - 1 + tanh_out - ) - grad_y = ff * g - return grad_y.to(dtype=y.dtype), grad_y.sum(dim=(0), dtype=bias.dtype) - - -class GeLUFunction(torch.autograd.Function): - @staticmethod - # bias is an optional argument - def forward(ctx, input, bias): - ctx.save_for_backward(input, bias) - return bias_gelu(input, bias) - - @staticmethod - def backward(ctx, grad_output): - input, bias = ctx.saved_tensors - tmp = bias_gelu_back(grad_output, input, bias) - return tmp, tmp - - -bias_gelu_impl = GeLUFunction.apply - -# this function is tanh approximation of gelu -# actual gelu is: -# x * 0.5 * (1.0 + torch.erf(x * 0.70710678)) -@torch.jit.script -def gelu_fwd(x): - return (x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))).to(dtype=x.dtype) - - -# gradient of tanh approximation of gelu -# gradient of actual gelu is: -# 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x) -@torch.jit.script -def gelu_bwd(g, x): - tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)) - # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243 - ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( - 1 + tanh_out - ) - return (ff * g).to(dtype=x.dtype) - - -class FastGeLUFunction(torch.autograd.Function): - @staticmethod - # bias is an optional argument - def forward(ctx, input): - ctx.save_for_backward(input) - return gelu_fwd(input) - - @staticmethod - def backward(ctx, grad_output): - (input,) = ctx.saved_tensors - tmp = gelu_bwd(grad_output, input) - return tmp - - -fast_gelu_impl = FastGeLUFunction.apply - - -@torch.jit.script -def relu_bwd(g, x): - return torch.where(x >= 0, g, 0.0).to(dtype=x.dtype) - - -@torch.jit.script -def sqrelu_fwd(x): - r = F.relu(x) - return (r * r).to(dtype=x.dtype) - - -@torch.jit.script -def sqrelu_bwd(g, x): - return (2.0 * g * F.relu(x)).to(dtype=x.dtype) - - -swiglu_fwd_codestring = """ -template T swiglu_fwd(T x, T y) { - return float(x) * float(y) / (1.0f + ::exp(-float(x))); -} -""" -swiglu_bwd_codestring = """ -template void swiglu_bwd(T x, T y, T g, T& dx, T& dy) { - float x_sigmoid = 1.0f / (1.0f + ::exp(-float(x))); - dx = x_sigmoid * (1 + float(x) * (1.0f - x_sigmoid)) * float(g) * float(y); - dy = float(x) * x_sigmoid * float(g); -} -""" -swiglu_fwd = torch.cuda.jiterator._create_jit_fn(swiglu_fwd_codestring) -swiglu_bwd = torch.cuda.jiterator._create_multi_output_jit_fn(swiglu_bwd_codestring, num_outputs=2) - - -class SwiGLUFunction(torch.autograd.Function): - - @staticmethod - def forward(ctx, x, y): - ctx.save_for_backward(x, y) - return swiglu_fwd(x, y) - - @staticmethod - def backward(ctx, dout): - x, y = ctx.saved_tensors - return swiglu_bwd(x, y, dout) - -swiglu = SwiGLUFunction.apply diff --git a/build/torch211-cxx11-cu130-aarch64-linux/ops/fused_dense.py b/build/torch211-cxx11-cu130-aarch64-linux/ops/fused_dense.py deleted file mode 100644 index 6b4033d134e4093fe278f7b3f8c7d3128ce9f36d..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu130-aarch64-linux/ops/fused_dense.py +++ /dev/null @@ -1,688 +0,0 @@ -# Copyright (c) 2023, Tri Dao. -# Inspired by https://github.com/NVIDIA/apex/blob/master/apex/fused_dense/fused_dense.py -# We make it work with pytorch amp and with bfloat16. -# The TensorParallel linear modules are inspired by https://github.com/NVIDIA/apex/blob/master/apex/transformer/tensor_parallel/layers.py -from functools import partial -from typing import Optional - -# import fused_dense_cuda # from apex -import fused_dense_lib as fused_dense_cuda -import torch -import torch.nn as nn -import torch.nn.functional as F -from torch import Tensor -from torch.distributed import ProcessGroup - -from flash_attn.utils.torch import custom_fwd, custom_bwd -from flash_attn.ops.activations import gelu_bwd, relu_bwd, sqrelu_bwd, sqrelu_fwd -from flash_attn.utils.distributed import ( - all_gather_raw, - all_reduce, - all_reduce_raw, - reduce_scatter, - reduce_scatter_raw, -) - - -class FusedDenseFunc(torch.autograd.Function): - @staticmethod - @custom_fwd - def forward( - ctx, x, weight, bias, return_residual=False, process_group=None, sequence_parallel=True - ): - """ - If process_group is not None and sequence_parallel=True, we're doing Tensor Parallel - with sequence parallelism: we do an all_gather_raw of x before doing the matmul. - """ - ctx.compute_weight_gradient = weight.requires_grad - ctx.return_residual = return_residual - ctx.process_group = process_group - ctx.sequence_parallel = sequence_parallel - - if torch.is_autocast_enabled(): - x = x.to(dtype=torch.get_autocast_gpu_dtype()) - x = x.contiguous() - if process_group is not None and sequence_parallel: - # We want to kick off the all_gather early, before weight dtype conversion - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - else: - total_x = x - - if torch.is_autocast_enabled(): - weight = weight.to(dtype=torch.get_autocast_gpu_dtype()) - bias = bias.to(dtype=torch.get_autocast_gpu_dtype()) if bias is not None else None - weight = weight.contiguous() - if process_group is not None and sequence_parallel: - handle_x.wait() - batch_shape, n = total_x.shape[:-1], total_x.shape[-1] - batch_dim = batch_shape.numel() - # https://github.com/pytorch/pytorch/blob/5b51849b48a7dbccd297286cc0110def4706f9e7/aten/src/ATen/native/cuda/Blas.cpp#L174 - if min(batch_dim, n, *weight.shape) > 65535 * 32: - raise RuntimeError("fused_dense only supports matrix dims <= 2M") - output = F.linear(total_x, weight, bias) - if ctx.compute_weight_gradient: - ctx.save_for_backward(x, weight) - else: - ctx.save_for_backward(weight) - return output if not return_residual else (output, x) - - @staticmethod - @custom_bwd - def backward(ctx, grad_output, *args): - grad_output = grad_output.contiguous() - if ctx.return_residual: - (grad_input,) = args - grad_input = grad_input.contiguous() - process_group = ctx.process_group - sequence_parallel = ctx.sequence_parallel - if ctx.compute_weight_gradient: - x, weight = ctx.saved_tensors - if process_group is not None and sequence_parallel: - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - else: - total_x = x - else: - (weight,) = ctx.saved_tensors - total_x = None - batch_shape = grad_output.shape[:-1] - batch_dim = batch_shape.numel() - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - if ctx.needs_input_grad[0]: - if not ctx.return_residual: - grad_input = F.linear(grad_output, weight.t()) - else: - grad_input = torch.addmm( - grad_input.reshape(batch_dim, grad_input.shape[-1]), grad_output, weight - ) - grad_input = grad_input.reshape(*batch_shape, grad_input.shape[-1]) - if process_group is not None: - reduce_fn = reduce_scatter_raw if sequence_parallel else all_reduce_raw - grad_input, handle_grad_input = reduce_fn(grad_input, process_group, async_op=True) - else: - grad_input = None - if ctx.needs_input_grad[1]: - assert ctx.compute_weight_gradient - if process_group is not None and sequence_parallel: - handle_x.wait() - grad_weight, grad_bias = fused_dense_cuda.linear_bias_wgrad( - total_x.reshape(batch_dim, total_x.shape[-1]), grad_output, ctx.needs_input_grad[2] - ) - else: - grad_weight = None - grad_bias = grad_output if ctx.needs_input_grad[2] else None - if process_group is not None and ctx.needs_input_grad[0]: - handle_grad_input.wait() - return grad_input, grad_weight, grad_bias, None, None, None - - -def fused_dense_func( - x: Tensor, - weight: Tensor, - bias: Optional[Tensor] = None, - return_residual: bool = False, - process_group: Optional[ProcessGroup] = None, - sequence_parallel: bool = True, -): - dtype_eligible = x.dtype in [torch.float16, torch.bfloat16] or ( - x.dtype == torch.float32 and torch.is_autocast_enabled() - ) - if x.is_cuda and weight.is_cuda and (bias is None or bias.is_cuda) and dtype_eligible: - return FusedDenseFunc.apply( - x, weight, bias, return_residual, process_group, sequence_parallel - ) - else: - assert process_group is None - out = F.linear(x, weight, bias) - return out if not return_residual else (out, x) - - -class FusedDense(nn.Linear): - def __init__( - self, - in_features: int, - out_features: int, - bias: bool = True, - return_residual: bool = False, - device=None, - dtype=None, - ) -> None: - super().__init__(in_features, out_features, bias=bias, device=device, dtype=dtype) - self.return_residual = return_residual - - def forward(self, x, process_group=None): - """ - If process_group is not None, we're doing Tensor Parallel with sequence parallelism: - we do an all_gather of x before doing the matmul. - """ - return fused_dense_func( - x, - self.weight, - self.bias, - return_residual=self.return_residual, - process_group=process_group, - ) - - -class ColumnParallelLinear(nn.Linear): - def __init__( - self, - in_features: int, - out_features: int, - process_group: ProcessGroup, - bias: bool = True, - sequence_parallel=True, - multiple_of=1, - device=None, - dtype=None, - ) -> None: - world_size = torch.distributed.get_world_size(process_group) - if out_features % multiple_of: - raise ValueError(f"out_features ({out_features}) must be a multiple of {multiple_of}") - multiple = out_features // multiple_of - # We want to split @multiple across world_size, but it could be an uneven split - div = multiple // world_size - mod = multiple % world_size - # The first @mod ranks get @div + 1 copies, the rest get @div copies - local_multiple = div + int(torch.distributed.get_rank(process_group) < mod) - super().__init__( - in_features, local_multiple * multiple_of, bias=bias, device=device, dtype=dtype - ) - self.process_group = process_group - self.sequence_parallel = sequence_parallel - - def forward(self, x): - # If self.sequence_parallel is True, we're doing Tensor Parallel with sequence parallelism: - # we do an all_gather of x before doing the matmul. - # If not, then the input is already gathered. - return fused_dense_func( - x, - self.weight, - self.bias, - process_group=self.process_group, - sequence_parallel=self.sequence_parallel, - ) - - -class RowParallelLinear(nn.Linear): - def __init__( - self, - in_features: int, - out_features: int, - process_group: ProcessGroup, - bias: bool = True, - sequence_parallel=True, - multiple_of=1, - device=None, - dtype=None, - ) -> None: - world_size = torch.distributed.get_world_size(process_group) - rank = torch.distributed.get_rank(process_group) - if in_features % multiple_of: - raise ValueError(f"in_features ({in_features}) must be a multiple of {multiple_of}") - multiple = in_features // multiple_of - # We want to split @multiple across world_size, but it could be an uneven split - div = multiple // world_size - mod = multiple % world_size - # The first @mod ranks get @div + 1 copies, the rest get @div copies - local_multiple = div + int(torch.distributed.get_rank(process_group) < mod) - # Only rank 0 will have bias - super().__init__( - local_multiple * multiple_of, - out_features, - bias=bias and rank == 0, - device=device, - dtype=dtype, - ) - self.process_group = process_group - self.sequence_parallel = sequence_parallel - - def forward(self, x): - """ - We're doing Tensor Parallel with sequence parallelism: we do the matmul and then - a reduce_scatter of the result. - """ - out = fused_dense_func(x, self.weight, self.bias) - reduce_fn = reduce_scatter if self.sequence_parallel else all_reduce - return reduce_fn(out, self.process_group) - - -class FusedMLPFunc(torch.autograd.Function): - @staticmethod - @custom_fwd - def forward( - ctx, - x, - weight1, - bias1, - weight2, - bias2, - activation="gelu_approx", - save_pre_act=True, - return_residual=False, - checkpoint_lvl=0, - heuristic=0, - process_group=None, - sequence_parallel=True, - ): - """ - If process_group is not None and sequence_parallel=True, we're doing Tensor Parallel - with sequence parallelism: we do an all_gather of x before doing the matmul. - If sequence_parallel=False, then the input is already gathered. - - checkpoint_lvl: - 0: no recomputation in the bwd - 1: recompute gelu_out / relu_out in the bwd - 2: recompute pre_act and gelu_out / relu_out in the bwd - """ - assert -1 <= heuristic <= 4 - assert activation in ["gelu_approx", "relu", "sqrelu"] - if activation == "sqrelu": - assert heuristic == -1 - if not save_pre_act: - checkpoint_lvl = 2 - assert checkpoint_lvl in [0, 1, 2] - ctx.return_residual = return_residual - ctx.process_group = process_group - ctx.sequence_parallel = sequence_parallel - ctx.checkpoint_lvl = checkpoint_lvl - ctx.activation = activation - ctx.heuristic = heuristic - - if torch.is_autocast_enabled(): - x = x.to(dtype=torch.get_autocast_gpu_dtype()) - x = x.contiguous() - if process_group is not None and sequence_parallel: - # We want to kick off the all_gather early, before weight dtype conversion - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - else: - total_x = x - - if torch.is_autocast_enabled(): - dtype = torch.get_autocast_gpu_dtype() - weight1, weight2 = [a.to(dtype=dtype) for a in [weight1, weight2]] - bias1 = bias1.to(dtype=dtype) if bias1 is not None else None - bias2 = bias2.to(dtype=dtype) if bias2 is not None else None - weight1 = weight1.contiguous() - bias1 = bias1.contiguous() if bias1 is not None else None - weight2 = weight2.contiguous() - bias2 = bias2.contiguous() if bias2 is not None else None - if process_group is not None and sequence_parallel: - handle_x.wait() - batch_shape, n = total_x.shape[:-1], total_x.shape[-1] - batch_dim = batch_shape.numel() - # https://github.com/pytorch/pytorch/blob/5b51849b48a7dbccd297286cc0110def4706f9e7/aten/src/ATen/native/cuda/Blas.cpp#L174 - if min(batch_dim, n, *weight1.shape, *weight2.shape) > 65535 * 32: - raise RuntimeError("fused_dense only supports matrix dims <= 2M") - if heuristic == -1: - pre_act = F.linear(total_x, weight1, bias1) - activation_fn = ( - partial(F.gelu, approximate="tanh") - if activation == "gelu_approx" - else (sqrelu_fwd if activation == "sqrelu" else F.relu) - ) - with torch.jit.fuser("fuser2"): - output1 = activation_fn(pre_act) - # This is before adding bias1 - # pre_act = F.linear(total_x.reshape(batch_dim, n), weight1) - # with torch.jit.fuser('fuser2'): - # output1 = bias_gelu(pre_act, bias1) - else: - is_gelu = activation == "gelu_approx" - output1, *rest = fused_dense_cuda.linear_act_forward( - total_x.reshape(batch_dim, n), weight1, bias1, is_gelu, save_pre_act, heuristic - ) - if save_pre_act: - pre_act = rest[0] - output2 = F.linear(output1, weight2, bias2) - if checkpoint_lvl == 0 or (checkpoint_lvl == 1 and activation == "relu"): - # For RELU the pre_act is very small (just a bit-mask) so we just save it - ctx.save_for_backward(x, weight1, weight2, pre_act, output1) - elif checkpoint_lvl == 1: - ctx.save_for_backward(x, weight1, weight2, pre_act) - elif checkpoint_lvl == 2: - ctx.save_for_backward(x, weight1, weight2, bias1) - output2 = output2.reshape(*batch_shape, output2.shape[-1]) - return output2 if not return_residual else (output2, x) - - @staticmethod - @custom_bwd - def backward(ctx, grad_output, *args): - grad_output = grad_output.contiguous() - checkpoint_lvl = ctx.checkpoint_lvl - activation = ctx.activation - activation_fn = ( - partial(F.gelu, approximate="tanh") - if activation == "gelu_approx" - else (sqrelu_fwd if activation == "sqrelu" else F.relu) - ) - if ctx.return_residual: - (grad_input,) = args - grad_input = grad_input.contiguous() - process_group = ctx.process_group - sequence_parallel = ctx.sequence_parallel - x, weight1, weight2, *rest = ctx.saved_tensors - if process_group is None or not sequence_parallel: - total_x = x - batch_shape = grad_output.shape[:-1] - batch_dim = batch_shape.numel() - if checkpoint_lvl in [0, 1]: - if process_group is not None and sequence_parallel: - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - if checkpoint_lvl == 0 or (checkpoint_lvl == 1 and activation == "relu"): - pre_act, output1 = rest - elif checkpoint_lvl == 1: - (pre_act,) = rest - with torch.jit.fuser("fuser2"): - output1 = activation_fn(pre_act) - elif checkpoint_lvl == 2: - (bias1,) = rest - if process_group is not None and sequence_parallel: - total_x, _ = all_gather_raw(x, process_group) - if ctx.heuristic == -1: - pre_act = F.linear(total_x, weight1, bias1) - with torch.jit.fuser("fuser2"): - output1 = activation_fn(pre_act) - else: - output1, pre_act = fused_dense_cuda.linear_act_forward( - total_x.reshape(batch_dim, total_x.shape[-1]), - weight1, - bias1, - activation == "gelu_approx", - True, - ctx.heuristic, - ) - - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - output1 = output1.reshape(batch_dim, output1.shape[-1]) - pre_act = pre_act.reshape(batch_dim, pre_act.shape[-1]) - if ctx.needs_input_grad[3]: - grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad( - output1, grad_output, ctx.needs_input_grad[4] - ) - else: - grad_weight2 = None - grad_bias2 = grad_output if ctx.needs_input_grad[4] else None - if ctx.heuristic == -1: - # grad_pre_act = matmul_dgelu(grad_output, weight2, pre_act) - grad_output1 = F.linear(grad_output, weight2.t()) - activation_grad_fn = ( - gelu_bwd - if activation == "gelu_approx" - else (sqrelu_bwd if activation == "sqrelu" else relu_bwd) - ) - with torch.jit.fuser("fuser2"): - grad_pre_act = activation_grad_fn(grad_output1, pre_act) - else: - # The cublasLt epilogue has to compute both gelu/relu grad and bias grad, we can't - # just compute gelu/relu grad - grad_pre_act, grad_bias1 = fused_dense_cuda.bias_act_linear_dgrad_bgrad( - weight2, grad_output, pre_act, activation == "gelu_approx", ctx.heuristic - ) - if not ctx.needs_input_grad[2]: - grad_bias1 = None - if ctx.needs_input_grad[0]: - if not ctx.return_residual: - grad_input = F.linear(grad_pre_act, weight1.t()) - else: - grad_input = torch.addmm( - grad_input.reshape(batch_dim, grad_input.shape[-1]), grad_pre_act, weight1 - ) - grad_input = grad_input.reshape(*batch_shape, grad_input.shape[-1]) - if process_group is not None: - reduce_fn = reduce_scatter_raw if sequence_parallel else all_reduce_raw - grad_input, handle_grad_input = reduce_fn(grad_input, process_group, async_op=True) - else: - grad_input = None - if ctx.heuristic == -1: - if ctx.needs_input_grad[1]: - if process_group is not None and sequence_parallel and checkpoint_lvl != 2: - handle_x.wait() - grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_wgrad( - total_x.reshape(batch_dim, total_x.shape[-1]), - grad_pre_act, - ctx.needs_input_grad[2], - ) - else: - grad_weight1 = None - grad_bias1 = grad_pre_act if ctx.needs_input_grad[2] else None - else: - if ctx.needs_input_grad[1]: - if process_group is not None and sequence_parallel and checkpoint_lvl != 2: - handle_x.wait() - grad_weight1 = F.linear( - grad_pre_act.t(), total_x.reshape(batch_dim, total_x.shape[-1]).t() - ) - else: - grad_weight1 = None - if process_group is not None and ctx.needs_input_grad[0]: - handle_grad_input.wait() - return ( - grad_input, - grad_weight1, - grad_bias1, - grad_weight2, - grad_bias2, - None, - None, - None, - None, - None, - None, - None, - ) - - -def fused_mlp_func( - x: Tensor, - weight1: Tensor, - weight2: Tensor, - bias1: Optional[Tensor] = None, - bias2: Optional[Tensor] = None, - activation: str = "gelu_approx", - save_pre_act: bool = True, - return_residual: bool = False, - checkpoint_lvl: int = 0, - heuristic: int = 0, - process_group: Optional[ProcessGroup] = None, - sequence_parallel: bool = True, -): - assert activation in ["gelu_approx", "relu", "sqrelu"] - dtype_eligible = x.dtype in [torch.float16, torch.bfloat16] or ( - x.dtype == torch.float32 and torch.is_autocast_enabled() - ) - # If we save pre-activation, dimension must be divisible by 128 (relu) or 8 (gelu) - dim_eligible = not save_pre_act or (x.shape[-1] % (128 if activation == "relu" else 8) == 0) - if ( - x.is_cuda - and weight1.is_cuda - and weight2.is_cuda - and (bias1 is None or bias1.is_cuda) - and (bias2 is None or bias2.is_cuda) - and dtype_eligible - and dim_eligible - ): - return FusedMLPFunc.apply( - x, - weight1, - bias1, - weight2, - bias2, - activation, - save_pre_act, - return_residual, - checkpoint_lvl, - heuristic, - process_group, - sequence_parallel, - ) - else: - assert process_group is None - pre_act = F.linear(x, weight1, bias1) - activation_fn = ( - partial(F.gelu, approximate="tanh") - if activation == "gelu_approx" - else partial(F.relu, inplace=True) - ) - output1 = activation_fn(pre_act) - output2 = F.linear(output1, weight2, bias2) - return output2 if not return_residual else (output2, x) - - -class FusedMLP(nn.Module): - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - bias1=True, - bias2=True, - activation="gelu_approx", - return_residual=False, - checkpoint_lvl=0, - heuristic="auto", - device=None, - dtype=None, - ): - """ - If process_group is not None, we're doing Tensor Parallel with sequence parallelism: - we do an all_gather of x before doing the matmul, gelu, then matmul. - Finally we do a reduce_scatter of the output. - - checkpoint_lvl (increasing lvl means slower but more memory saving): - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute pre_act and gelu_out in the bwd - heuristic: - -1: don't fuse gemm + gelu (separate kernel) - 0..4: use this heuristic for the algo section in the fused gemm + gelu - 'auto': heuristic will be picked automatically: - For CUDA >= 11.8, we set heuristic=0 for both fp16 and bf16 for best perf. - For CUDA <= 11.7, we set heuristic=1 for fp16 and heuristic=-1 for bf16. - For H100, we set heuristic=-1 for both fp16 and bf16 as the fused cuBlasLt implementation - is slower than the unfused version. - return_residual: whether to return the input x along with the output. This is for - performance reason: for post-norm architecture, returning the input allows us - to fuse the backward of nn.Linear with the residual connection. - """ - assert checkpoint_lvl in [0, 1, 2] - assert activation in ["gelu_approx", "relu", "sqrelu"] - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features * 4 - self.activation = activation - self.return_residual = return_residual - self.checkpoint_lvl = checkpoint_lvl - self.heuristic = heuristic if activation != "sqrelu" else -1 - self.fc1 = nn.Linear(in_features, hidden_features, bias=bias1, **factory_kwargs) - self.fc2 = nn.Linear(hidden_features, out_features, bias=bias2, **factory_kwargs) - - def forward(self, x, process_group=None): - dtype = x.dtype if not torch.is_autocast_enabled() else torch.get_autocast_gpu_dtype() - if self.heuristic == "auto": - if self.activation == "gelu_approx": - if torch.cuda.get_device_capability("cuda") == (9, 0): - heuristic = -1 - else: - cuda_ver = tuple(map(int, torch.version.cuda.split("."))) - heuristic = 0 if cuda_ver >= (11, 8) else (1 if dtype == torch.float16 else -1) - else: - heuristic = 0 - else: - heuristic = self.heuristic - out = fused_mlp_func( - x, - self.fc1.weight, - self.fc2.weight, - self.fc1.bias, - self.fc2.bias, - activation=self.activation, - save_pre_act=self.training, - return_residual=self.return_residual, - checkpoint_lvl=self.checkpoint_lvl, - heuristic=heuristic, - process_group=process_group, - ) - if self.return_residual: - out, x = out - if process_group is not None: - out = reduce_scatter(out, process_group) - return out if not self.return_residual else (out, x) - - -class ParallelFusedMLP(nn.Module): - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - activation="gelu_approx", - process_group: ProcessGroup = None, - bias1=True, - bias2=True, - sequence_parallel=True, - checkpoint_lvl=0, - heuristic="auto", - device=None, - dtype=None, - ): - """ - process_group is required. We're doing Tensor Parallel with sequence parallelism: - we do an all_gather of x before doing the matmul, gelu, then matmul. - Finally we do a reduce_scatter of the output. - - checkpoint_lvl (increasing lvl means slower but more memory saving): - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute pre_act and gelu_out in the bwd - heuristic: - -1: don't fuse gemm + gelu (separate kernel) - 0..4: use this heuristic for the algo section in the fused gemm + gelu - 'auto': heuristic will be picked automatically: - For CUDA >= 11.8, we set heuristic=0 for both fp16 and bf16 for best perf. - For CUDA <= 11.7, we set heuristic=1 for fp16 and heuristic=-1 for bf16. - """ - assert checkpoint_lvl in [0, 1, 2] - assert activation in ["gelu_approx", "relu", "sqrelu"] - assert process_group is not None - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features * 4 - self.activation = activation - self.process_group = process_group - self.sequence_parallel = sequence_parallel - self.checkpoint_lvl = checkpoint_lvl - self.heuristic = heuristic if activation != "sqrelu" else -1 - self.fc1 = ColumnParallelLinear( - in_features, hidden_features, process_group, bias=bias1, **factory_kwargs - ) - self.fc2 = RowParallelLinear( - hidden_features, out_features, process_group, bias=bias2, **factory_kwargs - ) - - def forward(self, x): - dtype = x.dtype if not torch.is_autocast_enabled() else torch.get_autocast_gpu_dtype() - if self.heuristic == "auto": - if self.activation == "gelu_approx": - cuda_ver = tuple(map(int, torch.version.cuda.split("."))) - heuristic = 0 if cuda_ver >= (11, 8) else (1 if dtype == torch.float16 else -1) - else: - heuristic = 0 - else: - heuristic = self.heuristic - out = fused_mlp_func( - x, - self.fc1.weight, - self.fc2.weight, - self.fc1.bias, - self.fc2.bias, - activation=self.activation, - save_pre_act=self.training, - checkpoint_lvl=self.checkpoint_lvl, - heuristic=heuristic, - process_group=self.process_group, - sequence_parallel=self.sequence_parallel, - ) - reduce_fn = reduce_scatter if self.sequence_parallel else all_reduce - return reduce_fn(out, self.process_group) diff --git a/build/torch211-cxx11-cu130-aarch64-linux/ops/layer_norm.py b/build/torch211-cxx11-cu130-aarch64-linux/ops/layer_norm.py deleted file mode 100644 index 4b6cd798fd02844ef9cd3897f8ab95e490e638bf..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu130-aarch64-linux/ops/layer_norm.py +++ /dev/null @@ -1,800 +0,0 @@ -# Copyright (c) 2022, Tri Dao. -# Adapted from https://github.com/NVIDIA/apex/blob/master/apex/contrib/layer_norm/layer_norm.py - -import dropout_layer_norm -import torch -from torch.nn import init - - -def maybe_align(x, alignment_in_bytes=16): - """Assume that x already has last dim divisible by alignment_in_bytes""" - # TD [2023-07-04] I'm not 100% sure that clone will align the memory - # https://discuss.pytorch.org/t/how-to-ensure-that-tensor-data-ptr-is-aligned-to-16-bytes/183440 - return x if x.data_ptr() % alignment_in_bytes == 0 else x.clone() - - -def _dropout_add_layer_norm_forward( - x0, - residual, - gamma, - beta, - rowscale, - colscale, - dropout_p, - epsilon, - residual_in_fp32=False, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes""" - hidden_size = gamma.numel() - x0mat = x0.view((-1, hidden_size)) - residualmat = residual.view((-1, hidden_size)) if residual is not None else None - rowscale = rowscale.view(-1) if rowscale is not None else None - zmat, xmat, dmask, mu, rsigma = dropout_layer_norm.dropout_add_ln_fwd( - x0mat, - residualmat, - gamma, - beta, - rowscale, - colscale, - None, - None, - dropout_p, - epsilon, - 1.0, - 0, - None, - residual_in_fp32, - is_rms_norm, - ) - # dmask is None if dropout_p == 0.0 - # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype - return zmat, xmat if xmat is not None else x0mat, dmask, mu, rsigma - - -def _dropout_add_layer_norm_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - rowscale, - colscale, - dropout_p, - has_residual, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes - dx == None means that it was a post-norm architecture - (x = drop(x0) + residual was not returned in the fwd). - x0 must not be None if we have colscale. - """ - hidden_size = gamma.numel() - xmat = x.view((-1, hidden_size)) - dzmat = dz.view(xmat.shape) - dxmat = dx.view(xmat.shape) if dx is not None else None - x0mat = x0.view((-1, hidden_size)) if x0 is not None else None - rowscale = rowscale.view(-1) if rowscale is not None else None - if colscale is not None: - assert x0 is not None, "x0 is required to compute the gradient of colscale" - dx0mat, dresidualmat, dgamma, dbeta, _, _, *rest = dropout_layer_norm.dropout_add_ln_bwd( - dzmat, - dxmat, - xmat, - x0mat, - dmask, - mu, - rsigma, - gamma, - rowscale, - colscale, - None, - None, - dropout_p, - 1.0, - 0, - has_residual, - is_rms_norm, - ) - # dresidualmat is None if not has_residual - if colscale is None: - return dx0mat, dresidualmat, dgamma, dbeta - else: - dcolscale = rest[0] - return dx0mat, dresidualmat, dgamma, dbeta, dcolscale - - -def _dropout_add_layer_norm_subset_forward( - x0, - residual, - gamma, - beta, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32=False, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes""" - hidden_size = gamma.numel() - x0mat = x0.view((-1, hidden_size)) - residualmat = residual.view((-1, hidden_size)) if residual is not None else None - x0_subset = x0_subset.view(-1) if x0_subset is not None else None - out_subset = out_subset.view(-1) if out_subset is not None else None - zmat, xmat, dmask, mu, rsigma = dropout_layer_norm.dropout_add_ln_fwd( - x0mat, - residualmat, - gamma, - beta, - None, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - None, - residual_in_fp32, - is_rms_norm, - ) - # dmask is None if dropout_p == 0.0 - # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype - return zmat, xmat if xmat is not None else x0mat, dmask, mu, rsigma - - -def _dropout_add_layer_norm_subset_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - colscale, - x0_subset, - out_subset, - dropout_p, - rowscale_const, - x0_numrows, - has_residual, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes - dx == None means that it was a post-norm architecture - (x = drop(x0) + residual was not returned in the fwd). - x0 must not be None if we have colscale. - """ - hidden_size = gamma.numel() - xmat = x.view((-1, hidden_size)) - dzmat = dz.view(-1, hidden_size) - dxmat = dx.view(xmat.shape) if dx is not None else None - x0mat = x0.view((-1, hidden_size)) if x0 is not None else None - x0_subset = x0_subset.view(-1) if x0_subset is not None else None - out_subset = out_subset.view(-1) if out_subset is not None else None - if colscale is not None: - assert x0 is not None, "x0 is required to compute the gradient of colscale" - dx0mat, dresidualmat, dgamma, dbeta, _, _, *rest = dropout_layer_norm.dropout_add_ln_bwd( - dzmat, - dxmat, - xmat, - x0mat, - dmask, - mu, - rsigma, - gamma, - None, - colscale, - x0_subset, - out_subset, - dropout_p, - rowscale_const, - x0_numrows, - has_residual, - is_rms_norm, - ) - # dresidualmat is None if not has_residual - if colscale is None: - return dx0mat, dresidualmat, dgamma, dbeta - else: - dcolscale = rest[0] - return dx0mat, dresidualmat, dgamma, dbeta, dcolscale - - -def _dropout_add_layer_norm_parallel_residual_forward( - x0, - x1, - residual, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - residual_in_fp32=False, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes""" - hidden_size = gamma0.numel() - x0mat = x0.view((-1, hidden_size)) - x1mat = x1.view((-1, hidden_size)) if x1 is not None else None - residualmat = residual.view((-1, hidden_size)) if residual is not None else None - ( - z0mat, - z1mat, - xmat, - dmask0, - dmask1, - mu, - rsigma, - ) = dropout_layer_norm.dropout_add_ln_parallel_residual_fwd( - x0mat, - x1mat, - residualmat, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - None, - residual_in_fp32, - is_rms_norm, - ) - # dmask0 and dmask1 are None if dropout_p == 0.0 - # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype - return z0mat, z1mat, xmat if xmat is not None else x0mat, dmask0, dmask1, mu, rsigma - - -def _dropout_add_layer_norm_parallel_residual_backward( - dz0, - dz1, - dx, - x, - dmask0, - dmask1, - mu, - rsigma, - gamma0, - gamma1, - dropout_p, - has_x1, - has_residual, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes - dx == None means that it was a post-norm architecture - (x = drop(x0) + residual was not returned in the fwd). - """ - hidden_size = gamma0.numel() - xmat = x.view((-1, hidden_size)) - dz0mat = dz0.view(xmat.shape) - dz1mat = dz1.view(xmat.shape) if dz1 is not None else None - dxmat = dx.view(xmat.shape) if dx is not None else None - ( - dx0mat, - dx1mat, - dresidualmat, - dgamma0, - dbeta0, - dgamma1, - dbeta1, - *rest, - ) = dropout_layer_norm.dropout_add_ln_parallel_residual_bwd( - dz0mat, - dz1mat, - dxmat, - xmat, - dmask0, - dmask1, - mu, - rsigma, - gamma0, - gamma1, - dropout_p, - has_x1, - has_residual, - is_rms_norm, - ) - # dresidualmat is None if not has_residual - return dx0mat, dx1mat, dresidualmat, dgamma0, dbeta0, dgamma1, dbeta1 - - -class DropoutAddLayerNormFn(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x0, - residual, - gamma, - beta, - rowscale, - colscale, - dropout_p, - epsilon, - residual_in_fp32=False, - prenorm=False, - is_rms_norm=False, - return_dmask=False, - ): - x0 = maybe_align(x0.contiguous(), 16) - residual = maybe_align(residual.contiguous(), 16) if residual is not None else None - gamma = maybe_align(gamma.contiguous(), 16) - beta = maybe_align(beta.contiguous(), 16) if beta is not None else None - rowscale = maybe_align(rowscale.contiguous(), 16) if rowscale is not None else None - colscale = maybe_align(colscale.contiguous(), 16) if colscale is not None else None - zmat, xmat, dmask, mu, rsigma = _dropout_add_layer_norm_forward( - x0, - residual, - gamma, - beta, - rowscale, - colscale, - dropout_p, - epsilon, - residual_in_fp32, - is_rms_norm, - ) - # Only need to save x0 if we need to compute gradient wrt colscale - x0_saved = x0 if colscale is not None else None - ctx.save_for_backward( - xmat.view(x0.shape), x0_saved, dmask, gamma, mu, rsigma, rowscale, colscale - ) - ctx.prenorm = prenorm - ctx.dropout_p = dropout_p - ctx.has_residual = residual is not None - ctx.is_rms_norm = is_rms_norm - ctx.has_beta = beta is not None - if not return_dmask: - return ( - zmat.view(x0.shape) if not prenorm else (zmat.view(x0.shape), xmat.view(x0.shape)) - ) - else: - dmask = ( - dmask.view(x0.shape) - if dropout_p > 0.0 - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - ctx.mark_non_differentiable(dmask) - return ( - (zmat.view(x0.shape), dmask) - if not prenorm - else (zmat.view(x0.shape), xmat.view(x0.shape), dmask) - ) - - @staticmethod - def backward(ctx, dz, *args): - # assert dz.is_contiguous() - dz = maybe_align(dz.contiguous(), 16) # this happens! - dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None - x, x0, dmask, gamma, mu, rsigma, rowscale, colscale = ctx.saved_tensors - # x0 is None if colscale is None - dropout_p = ctx.dropout_p - has_residual = ctx.has_residual - dx0mat, dresidualmat, dgamma, dbeta, *rest = _dropout_add_layer_norm_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - rowscale, - colscale, - dropout_p, - has_residual, - ctx.is_rms_norm, - ) - dx0 = dx0mat.view(x.shape) - dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None - dcolscale = rest[0] if colscale is not None else None - return ( - dx0, - dresidual, - dgamma, - dbeta if ctx.has_beta else None, - None, - dcolscale, - None, - None, - None, - None, - None, - None, - ) - - -class DropoutAddLayerNormSubsetFn(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x0, - residual, - gamma, - beta, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32=False, - prenorm=False, - is_rms_norm=False, - return_dmask=False, - ): - x0 = maybe_align(x0.contiguous(), 16) - residual = maybe_align(residual.contiguous(), 16) if residual is not None else None - gamma = maybe_align(gamma.contiguous(), 16) - beta = maybe_align(beta.contiguous(), 16) if beta is not None else None - colscale = maybe_align(colscale.contiguous(), 16) if colscale is not None else None - zmat, xmat, dmask, mu, rsigma = _dropout_add_layer_norm_subset_forward( - x0, - residual, - gamma, - beta, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32, - is_rms_norm, - ) - # Only need to save x0 if we need to compute gradient wrt colscale - x0_saved = x0 if colscale is not None else None - x_shape = (-1, *x0.shape[1:]) - ctx.save_for_backward( - xmat.view(x_shape), x0_saved, dmask, gamma, mu, rsigma, colscale, x0_subset, out_subset - ) - ctx.prenorm = prenorm - ctx.dropout_p = dropout_p - ctx.rowscale_const = rowscale_const - ctx.x0_numrows = x0.shape[:-1].numel() - ctx.has_residual = residual is not None - ctx.is_rms_norm = is_rms_norm - ctx.has_beta = beta is not None - z_shape = (-1, *x0.shape[1:]) - if not return_dmask: - return zmat.view(z_shape) if not prenorm else (zmat.view(z_shape), xmat.view(x0.shape)) - else: - z = zmat.view(z_shape) - dmask = ( - dmask.view(x0.shape) - if dropout_p > 0.0 - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - ctx.mark_non_differentiable(dmask) - return (z, dmask) if not prenorm else (z, xmat.view(x_shape), dmask) - - @staticmethod - def backward(ctx, dz, *args): - # assert dz.is_contiguous() - dz = maybe_align(dz.contiguous(), 16) # this happens! - dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None - x, x0, dmask, gamma, mu, rsigma, colscale, x0_subset, out_subset = ctx.saved_tensors - # x0 is None if colscale is None - dropout_p = ctx.dropout_p - has_residual = ctx.has_residual - dx0mat, dresidualmat, dgamma, dbeta, *rest = _dropout_add_layer_norm_subset_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - colscale, - x0_subset, - out_subset, - dropout_p, - ctx.rowscale_const, - ctx.x0_numrows, - has_residual, - ctx.is_rms_norm, - ) - dx0 = dx0mat.view(-1, *x.shape[1:]) - dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None - dcolscale = rest[0] if colscale is not None else None - return ( - dx0, - dresidual, - dgamma, - dbeta if ctx.has_beta else None, - dcolscale, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) - - -class DropoutAddLayerNormParallelResidualFn(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x0, - x1, - residual, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - residual_in_fp32=False, - prenorm=False, - is_rms_norm=False, - return_dmask=False, - ): - x0 = maybe_align(x0.contiguous(), 16) - x1 = maybe_align(x1.contiguous(), 16) if x1 is not None else None - residual = maybe_align(residual.contiguous(), 16) if residual is not None else None - gamma0 = maybe_align(gamma0.contiguous(), 16) - beta0 = maybe_align(beta0.contiguous(), 16) if beta0 is not None else None - gamma1 = maybe_align(gamma1.contiguous(), 16) if gamma1 is not None else None - beta1 = maybe_align(beta1.contiguous(), 16) if beta1 is not None else None - ( - z0mat, - z1mat, - xmat, - dmask0, - dmask1, - mu, - rsigma, - ) = _dropout_add_layer_norm_parallel_residual_forward( - x0, - x1, - residual, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - residual_in_fp32, - is_rms_norm, - ) - ctx.save_for_backward(xmat.view(x0.shape), dmask0, dmask1, gamma0, gamma1, mu, rsigma) - ctx.prenorm = prenorm - ctx.dropout_p = dropout_p - ctx.has_x1 = x1 is not None - ctx.has_residual = residual is not None - ctx.is_rms_norm = is_rms_norm - ctx.has_beta = beta0 is not None - z = (z0mat.view(x0.shape), z1mat.view(x0.shape) if z1mat is not None else None) - if not return_dmask: - return z if not prenorm else (*z, xmat.view(x0.shape)) - else: - dmask0 = ( - dmask0.view(x0.shape) - if dropout_p > 0.0 - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - dmask1 = ( - dmask1.view(x0.shape) - if dropout_p > 0.0 and x1 is not None - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - ctx.mark_non_differentiable(dmask0) - ctx.mark_non_differentiable(dmask1) - return ( - (*z, dmask0, dmask1) if not prenorm else (*z, xmat.view(x0.shape), dmask0, dmask1) - ) - - @staticmethod - def backward(ctx, dz0, dz1, *args): - dz0 = maybe_align(dz0.contiguous(), 16) # this happens! - dz1 = maybe_align(dz1.contiguous(), 16) if dz1 is not None else None - dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None - x, dmask0, dmask1, gamma0, gamma1, mu, rsigma = ctx.saved_tensors - dropout_p = ctx.dropout_p - has_x1 = ctx.has_x1 - has_residual = ctx.has_residual - ( - dx0mat, - dx1mat, - dresidualmat, - dgamma0, - dbeta0, - dgamma1, - dbeta1, - ) = _dropout_add_layer_norm_parallel_residual_backward( - dz0, - dz1, - dx, - x, - dmask0, - dmask1, - mu, - rsigma, - gamma0, - gamma1, - dropout_p, - has_x1, - has_residual, - ctx.is_rms_norm, - ) - dx0 = dx0mat.view(x.shape) - dx1 = dx1mat.view(x.shape) if dx1mat is not None else None - dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None - return ( - dx0, - dx1, - dresidual, - dgamma0, - dbeta0 if ctx.has_beta else None, - dgamma1, - dbeta1 if ctx.has_beta else None, - None, - None, - None, - None, - None, - None, - ) - - -def layer_norm(x, weight, bias, epsilon): - return DropoutAddLayerNormFn.apply(x, None, weight, bias, None, None, 0.0, epsilon, False) - - -def dropout_add_layer_norm( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - rowscale=None, - layerscale=None, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormFn.apply( - x0, - residual, - weight, - bias, - rowscale, - layerscale, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - False, - return_dropout_mask, - ) - - -def dropout_add_layer_norm_subset( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - layerscale=None, - x0_subset=None, - out_subset=None, - rowscale_const=1.0, - out_numrows=0, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormSubsetFn.apply( - x0, - residual, - weight, - bias, - layerscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32, - prenorm, - False, - return_dropout_mask, - ) - - -def dropout_add_layer_norm_parallel_residual( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormParallelResidualFn.apply( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - False, - return_dropout_mask, - ) - - -class DropoutAddLayerNorm(torch.nn.Module): - def __init__( - self, - hidden_size, - prenorm=False, - p=0.0, - eps=1e-5, - residual_in_fp32=False, - device=None, - dtype=None, - ): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.prenorm = prenorm - self.p = p - self.eps = eps - self.residual_in_fp32 = residual_in_fp32 - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.bias = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.reset_parameters() - - def reset_parameters(self): - init.ones_(self.weight) - init.zeros_(self.bias) - - def forward(self, x0, residual=None): - return dropout_add_layer_norm( - x0, - residual, - self.weight, - self.bias, - self.p if self.training else 0.0, - self.eps, - prenorm=self.prenorm, - residual_in_fp32=self.residual_in_fp32, - ) diff --git a/build/torch211-cxx11-cu130-aarch64-linux/ops/rms_norm.py b/build/torch211-cxx11-cu130-aarch64-linux/ops/rms_norm.py deleted file mode 100644 index 068348d61290e3839dd082b540d898578ba1e8e2..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu130-aarch64-linux/ops/rms_norm.py +++ /dev/null @@ -1,174 +0,0 @@ -# Copyright (c) 2022, Tri Dao. -# Adapted from https://github.com/NVIDIA/apex/blob/master/apex/contrib/layer_norm/layer_norm.py - -import torch -from torch.nn import init - -from flash_attn.ops.layer_norm import ( - DropoutAddLayerNormFn, - DropoutAddLayerNormParallelResidualFn, - DropoutAddLayerNormSubsetFn, -) - - -def rms_norm(x, weight, epsilon): - return DropoutAddLayerNormFn.apply( - x, None, weight, None, None, None, 0.0, epsilon, False, False, True - ) - - -def dropout_add_rms_norm( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - rowscale=None, - layerscale=None, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormFn.apply( - x0, - residual, - weight, - bias, - rowscale, - layerscale, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - True, - return_dropout_mask, - ) - - -def dropout_add_rms_norm_subset( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - layerscale=None, - x0_subset=None, - out_subset=None, - rowscale_const=1.0, - out_numrows=0, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormSubsetFn.apply( - x0, - residual, - weight, - bias, - layerscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32, - prenorm, - True, - return_dropout_mask, - ) - - -def dropout_add_rms_norm_parallel_residual( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormParallelResidualFn.apply( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - True, - return_dropout_mask, - ) - - -class RMSNorm(torch.nn.Module): - def __init__(self, hidden_size, eps=1e-5, device=None, dtype=None): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.eps = eps - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.register_parameter("bias", None) - self.reset_parameters() - - def reset_parameters(self): - init.ones_(self.weight) - - def forward(self, x): - return rms_norm(x, self.weight, self.eps) - - -class DropoutAddRMSNorm(torch.nn.Module): - def __init__( - self, - hidden_size, - prenorm=False, - p=0.0, - eps=1e-5, - residual_in_fp32=False, - device=None, - dtype=None, - ): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.prenorm = prenorm - self.p = p - self.eps = eps - self.residual_in_fp32 = residual_in_fp32 - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.register_parameter("bias", None) - self.reset_parameters() - - def reset_parameters(self): - init.ones_(self.weight) - - def forward(self, x0, residual=None): - return dropout_add_rms_norm( - x0, - residual, - self.weight, - None, - self.p if self.training else 0.0, - self.eps, - prenorm=self.prenorm, - residual_in_fp32=self.residual_in_fp32, - ) diff --git a/build/torch211-cxx11-cu130-aarch64-linux/ops/triton/__init__.py b/build/torch211-cxx11-cu130-aarch64-linux/ops/triton/__init__.py deleted file mode 100644 index 8b137891791fe96927ad78e64b0aad7bded08bdc..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu130-aarch64-linux/ops/triton/__init__.py +++ /dev/null @@ -1 +0,0 @@ - diff --git a/build/torch211-cxx11-cu130-aarch64-linux/ops/triton/cross_entropy.py b/build/torch211-cxx11-cu130-aarch64-linux/ops/triton/cross_entropy.py deleted file mode 100644 index 1b5a415b73f236f3e05fb14b9141959559e18526..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu130-aarch64-linux/ops/triton/cross_entropy.py +++ /dev/null @@ -1,330 +0,0 @@ -# Copyright (c) 2023, Tri Dao. - -from typing import Tuple, Optional, Union - -import torch -import torch.nn.functional as F - -import triton -import triton.language as tl - -# `all_gather_into_tensor` and `reduce_scatter_tensor` are new placeholders for -# `_all_gather_base` and `_reduce_scatter_base`. They require the most recent -# version of PyTorch. The following 2 lines are for backward compatibility with -# older PyTorch. -if "all_gather_into_tensor" not in dir(torch.distributed): - torch.distributed.all_gather_into_tensor = torch.distributed._all_gather_base - - -@triton.heuristics( - { - "HAS_SMOOTHING": lambda args: args["smoothing"] > 0.0, - } -) -@triton.jit -def cross_entropy_fwd_kernel( - loss_ptr, # data ptrs - lse_ptr, - z_loss_ptr, - logits_ptr, - labels_ptr, - smoothing, - logit_scale, - lse_square_scale, - ignore_index, - total_classes, - class_start_idx, # Useful for tensor parallel when each rank only has a subset of classes - n_cols, # shapes - logits_row_stride, # strides - BLOCK_SIZE: tl.constexpr, - HAS_SMOOTHING: tl.constexpr, - # if SPLIT (e.g. tensor parallel), don't include the LSE in the loss since it's not the final LSE - SPLIT: tl.constexpr, - PRECOMPUTED_LSE: tl.constexpr, # If LSE is already computed (also no smoothing and logit_scale == 1.0) -): - row_idx = tl.program_id(0) - logits_ptr = logits_ptr + row_idx * logits_row_stride.to(tl.int64) - sum_logits = 0.0 # For smoothing - if not PRECOMPUTED_LSE: - # Statistics for online softmax - m_i = -float("inf") - l_i = 0.0 - for col_offset in range(0, n_cols, BLOCK_SIZE): - cols = col_offset + tl.arange(0, BLOCK_SIZE) - logits = tl.load(logits_ptr + cols, mask=cols < n_cols, other=-float("inf")).to( - tl.float32 - ) * logit_scale - if HAS_SMOOTHING: - sum_logits += tl.sum(tl.where(cols < n_cols, logits, 0.0)) - m_i_new = tl.maximum(m_i, tl.max(logits)) - l_i = tl.exp(m_i - m_i_new) * l_i + tl.sum(tl.exp(logits - m_i_new)) - m_i = m_i_new - lse = tl.log(l_i) + m_i - tl.store(lse_ptr + row_idx, lse) - else: - lse = tl.load(lse_ptr + row_idx) - label_idx = tl.load(labels_ptr + row_idx) - if label_idx == ignore_index: - loss = 0.0 - z_loss = 0.0 - else: - label_idx -= class_start_idx - if label_idx >= 0 and label_idx < n_cols: - logits_label = tl.load(logits_ptr + label_idx) * logit_scale - if HAS_SMOOTHING: - loss = ( - (lse if not SPLIT else 0.0) - - smoothing * sum_logits / total_classes - - (1 - smoothing) * logits_label - ) - else: - loss = (lse if not SPLIT else 0.0) - logits_label - else: - # If label is out of bounds, we set the CE loss to 0.0. But we still want the smoothing loss - if HAS_SMOOTHING: - loss = smoothing * ((lse if not SPLIT else 0.0) - sum_logits / total_classes) - else: - loss = 0.0 - if not SPLIT: - z_loss = lse_square_scale * lse * lse - loss += z_loss - else: - z_loss = 0.0 - tl.store(loss_ptr + row_idx, loss) - if not SPLIT: - tl.store(z_loss_ptr + row_idx, z_loss) - - -@triton.heuristics( - { - "HAS_SMOOTHING": lambda args: args["smoothing"] > 0.0, - } -) -@triton.jit -def cross_entropy_bwd_kernel( - dlogits_ptr, # data ptrs - dloss_ptr, - logits_ptr, - lse_ptr, - labels_ptr, - smoothing, - logit_scale, - lse_square_scale, - ignore_index, - total_classes, - class_start_idx, # Useful for tensor parallel when each rank only has a subset of classes - n_cols, # shapes - logits_row_stride, # strides - dlogits_row_stride, - dloss_row_stride, - BLOCK_SIZE: tl.constexpr, - HAS_SMOOTHING: tl.constexpr, -): - row_idx = tl.program_id(0) - col_block_idx = tl.program_id(1) - logits_ptr = logits_ptr + row_idx * logits_row_stride.to(tl.int64) - dlogits_ptr = dlogits_ptr + row_idx * dlogits_row_stride.to(tl.int64) - col_offsets = col_block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - label_idx = tl.load(labels_ptr + row_idx) - if label_idx != ignore_index: - dloss = tl.load(dloss_ptr + row_idx * dloss_row_stride) - else: - dloss = 0.0 - logits = tl.load(logits_ptr + col_offsets, mask=col_offsets < n_cols, other=-float("inf")).to( - tl.float32 - ) * logit_scale - lse = tl.load(lse_ptr + row_idx) - probs = tl.exp(logits - lse) - probs += 2.0 * lse_square_scale * lse * probs - label_idx -= class_start_idx - if HAS_SMOOTHING: - smooth_positive = 1.0 - smoothing - smooth_negative = smoothing / total_classes - probs = tl.where(col_offsets == label_idx, probs - smooth_positive, probs) - smooth_negative - else: - probs = tl.where(col_offsets == label_idx, probs - 1.0, probs) - tl.store(dlogits_ptr + col_offsets, (dloss * logit_scale) * probs, mask=col_offsets < n_cols) - - -class CrossEntropyLoss(torch.autograd.Function): - - @staticmethod - def forward( - ctx, - logits, - labels, - precomputed_lse=None, - smoothing=0.0, - logit_scale=1.0, - lse_square_scale=0.0, - ignore_index=-100, - inplace_backward=False, - process_group=None, - ): - # For some reason Triton generates wrong code when labels has dtype long and its address - # is not aligned to 16 bytes. The ld.global.b64 seems to load the wrong label index. - if labels.dtype == torch.long and labels.data_ptr() % 16 != 0: - labels = F.pad(labels, (0, 1))[..., :-1] - assert labels.data_ptr() % 16 == 0 - assert logit_scale > 0.0 - n_rows, n_cols = logits.shape - assert labels.shape == (n_rows,) - world_size = 1 if process_group is None else torch.distributed.get_world_size(process_group) - total_classes = world_size * n_cols - rank = 0 if process_group is None else torch.distributed.get_rank(process_group) - class_start_idx = rank * n_cols - use_precomputed_lse = precomputed_lse is not None and logit_scale == 1.0 and smoothing == 0.0 - - if logits.stride(-1) != 1: - logits = logits.contiguous() - MAX_BLOCK_SIZE = 16 * 1024 - BLOCK_SIZE = min(triton.next_power_of_2(n_cols), MAX_BLOCK_SIZE) - num_warps = ( - 4 - if BLOCK_SIZE < 2048 - else (8 if BLOCK_SIZE < 8192 else (16 if BLOCK_SIZE < 128 * 1024 else 32)) - ) - losses = torch.empty(n_rows, dtype=torch.float, device=logits.device) - if use_precomputed_lse: - assert precomputed_lse.shape == (n_rows,) - lse = precomputed_lse.contiguous() - else: - lse = torch.empty(n_rows, dtype=torch.float, device=logits.device) - z_losses = torch.empty(n_rows, dtype=torch.float, device=logits.device) - # Need this, otherwise Triton tries to launch from cuda:0 and we get - # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?) - with torch.cuda.device(logits.device.index): - cross_entropy_fwd_kernel[(n_rows,)]( - losses, # data ptrs - lse, - z_losses, - logits, - labels, - smoothing, - logit_scale, - lse_square_scale, - ignore_index, - total_classes, - class_start_idx, - n_cols, # shapes - logits.stride(0), # strides - BLOCK_SIZE=BLOCK_SIZE, # constants - SPLIT=world_size > 1, - PRECOMPUTED_LSE=use_precomputed_lse, - num_warps=num_warps, - ) - - if world_size > 1: - # If there's no smoothing, if labels are in the vocab of this partition, losses contains - # - predicted logit, and 0 otherwise. - # If there's smoothing=0.1, for labels in the vocab of this partition, losses contains - # -0.9 * predicted logit - 0.1 * sum logit / total_classes. - # For labels not in the vocab of this partition, losses contains - # -0.1 * sum logit / total_classes. - if world_size > 1: - lse_allgather = torch.empty(world_size, n_rows, dtype=lse.dtype, device=lse.device) - torch.distributed.all_gather_into_tensor(lse_allgather, lse, group=process_group) - handle_losses = torch.distributed.all_reduce( - losses, op=torch.distributed.ReduceOp.SUM, group=process_group, async_op=True - ) - lse = torch.logsumexp(lse_allgather, dim=0) - handle_losses.wait() - # After the allreduce, if there's no smoothing, the total losses are - predicted_logit, - # we just have to add the (global) lse. - # If there's smoothing=0.1, the total losses are - # -0.9 * predicted_logit - 0.1 * sum logit / total_classes. - # Again, we just have to add the (global) lse. - losses += lse - if lse_square_scale != 0.0: - z_losses = lse_square_scale * lse.square() - z_losses.masked_fill_(labels == ignore_index, 0.0) - losses += z_losses - else: - z_losses = torch.zeros_like(losses) - losses.masked_fill_(labels == ignore_index, 0.0) - - ctx.save_for_backward(logits, lse, labels) - ctx.mark_non_differentiable(z_losses) - ctx.smoothing = smoothing - ctx.logit_scale = logit_scale - ctx.lse_square_scale = lse_square_scale - ctx.ignore_index = ignore_index - ctx.total_classes = total_classes - ctx.class_start_idx = class_start_idx - ctx.inplace_backward = inplace_backward - return losses, z_losses - - @staticmethod - def backward(ctx, grad_losses, grad_z_losses): - del grad_z_losses # z_losses are only for logging. - - logits, lse, labels = ctx.saved_tensors - dlogits = logits if ctx.inplace_backward else torch.empty_like(logits) - n_rows, n_cols = logits.shape - BLOCK_SIZE = min(triton.next_power_of_2(n_cols), 4 * 1024) - num_warps = 4 if BLOCK_SIZE < 2048 else (8 if BLOCK_SIZE < 8192 else 16) - grid = lambda META: (n_rows, triton.cdiv(n_cols, META["BLOCK_SIZE"])) # noqa - # Need this, otherwise Triton tries to launch from cuda:0 and we get - # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?) - with torch.cuda.device(logits.device.index): - cross_entropy_bwd_kernel[grid]( - dlogits, # data ptrs - grad_losses, - logits, - lse, - labels, - ctx.smoothing, - ctx.logit_scale, - ctx.lse_square_scale, - ctx.ignore_index, - ctx.total_classes, - ctx.class_start_idx, - n_cols, # shapes - logits.stride(0), # strides - dlogits.stride(0), - grad_losses.stride(0), - BLOCK_SIZE=BLOCK_SIZE, # constants - num_warps=num_warps, - ) - return dlogits, None, None, None, None, None, None, None, None, None - - -def cross_entropy_loss( - logits: torch.Tensor, - labels: torch.Tensor, - precomputed_lse: Optional[torch.Tensor] = None, - label_smoothing: float = 0.0, - logit_scale: float = 1.0, - lse_square_scale: float = 0.0, - ignore_index=-100, - inplace_backward: bool = False, - process_group=None, -) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Arguments: - logits: (batch, vocab_size) - labels: (batch,) - label_smoothing: float - logit_scale: float. Multiply logits by this scale before calculating the loss. - lse_square_scale: float. If > 0, we add lse_square_scale * lse(logits) ^ 2 to the loss. - This is also referred to as "z-loss". - ignore_index: int. If labels == ignore_index, the loss is set to 0.0. - inplace_backward: bool. If True, we do the backward pass in-place by modifying the logits. - This saves memory. - process_group: if not None, we're doing Tensor Parallel: each process is responsible for - one part of the vocab. The loss will be aggregated across processes. - Returns: - losses: (batch,), float - z_losses: (batch,), float - """ - return CrossEntropyLoss.apply( - logits, - labels, - precomputed_lse, - label_smoothing, - logit_scale, - lse_square_scale, - ignore_index, - inplace_backward, - process_group, - ) diff --git a/build/torch211-cxx11-cu130-aarch64-linux/ops/triton/k_activations.py b/build/torch211-cxx11-cu130-aarch64-linux/ops/triton/k_activations.py deleted file mode 100644 index efb83c358eb4a85d069ee340a3c83f418f9a805b..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu130-aarch64-linux/ops/triton/k_activations.py +++ /dev/null @@ -1,162 +0,0 @@ -# Adapted from https://github.com/facebookresearch/xformers/blob/main/xformers/triton/k_activations.py -# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved. -# -# This source code is licensed under the BSD license found in the -# LICENSE file in the root directory of this source tree. - -import math -from enum import Enum -from typing import Optional - -import triton -import triton.language as tl - -_sqrt2pi = math.sqrt(2.0 / math.pi) -_sqrt1_2 = math.sqrt(1.0 / 2) -_gaussian_pdf_normalization = 1.0 / math.sqrt(2 * math.pi) - - -class Activation(str, Enum): - SquaredReLU = "squared_relu" - GeLU = "gelu" - GeLUApprox = "gelu_approx" - LeakyReLU = "leaky_relu" - ReLU = "relu" - - -def get_triton_activation_kernel(activation: Optional[Activation]): - return ( - { - Activation.ReLU: relu, - Activation.LeakyReLU: leaky_relu, - Activation.GeLU: gelu, - Activation.GeLUApprox: gelu_approx, - Activation.SquaredReLU: squared_relu, - }[activation] - if activation - else None - ) - - -def get_triton_activation_bwd_kernel(activation: Optional[Activation]): - return ( - { - Activation.ReLU: relu_grad, - Activation.LeakyReLU: leaky_relu_grad, - Activation.GeLU: gelu_grad, - Activation.GeLUApprox: gelu_approx_grad, - Activation.SquaredReLU: squared_relu_grad, - }[activation] - if activation - else None - ) - - -@triton.jit -def tanh(x): - # Tanh is just a scaled sigmoid - return 2 * tl.sigmoid(2 * x) - 1 - - -@triton.jit -def cosh(x): - exp_x = tl.exp(x) - return (exp_x + 1.0 / exp_x) * 0.5 - - -# a Triton implementation of the most used activations -# See for instance http://arxiv.org/abs/1606.08415 for an overview - -# ReLU -@triton.jit -def relu(x): - """ - ReLU_ activation function - - .. _ReLU: https://pytorch.org/docs/stable/generated/torch.nn.ReLU.html - """ - zero = 0.0 - return tl.where(x >= 0, x, zero.to(x.dtype)) - - -@triton.jit -def relu_grad(x): - # ReLU is different from other activations - # in that it does not require the input to retrospectively compute its gradient - # here the input is the downstream gradient, and we return the upstream gradient directly - zero = 0.0 - one = 1.0 - return tl.where(x >= 0, one.to(x.dtype), zero.to(x.dtype)) - - -@triton.jit -def squared_relu(x): - """ - Squared ReLU activation, as proposed in the Primer_ paper. - - .. _Primer: https://arxiv.org/abs/2109.08668 - """ - x_ = relu(x) - return (x_ * x_).to(x.dtype) - - -@triton.jit -def squared_relu_grad(x): - return tl.where(x >= 0, 2.0 * x, 0.0) - - -# Leaky ReLU -@triton.jit -def leaky_relu(x): - """ - LeakyReLU_ activation - - .. _LeakyReLU: https://pytorch.org/docs/stable/generated/torch.nn.LeakyReLU.html - """ - scale = 0.01 + 0.0 - scale = scale.to(x.dtype) - return tl.where(x >= 0, x, scale * x) - - -@triton.jit -def leaky_relu_grad(x): - min_grad = 0.01 - max_grad = 1 - - min_grad = min_grad.to(x.dtype) - max_grad = max_grad.to(x.dtype) - - return tl.where(x >= 0, max_grad, min_grad) - - -@triton.jit -def gelu(x): - """Gaussian Error Linear Unit (GELU)""" - return x * 0.5 * (1.0 + tl.libdevice.erf(x * _sqrt1_2)) - - -@triton.jit -def gelu_grad(x): - cdf = 0.5 * (1.0 + tl.libdevice.erf(x * _sqrt1_2)) - pdf = tl.exp(-0.5 * x * x) * _gaussian_pdf_normalization - return cdf + x * pdf - - -@triton.jit -def gelu_approx(x): - """ - GeLU_ activation - Gaussian error linear unit, with tanh approximation - - .. _GeLU: https://arxiv.org/pdf/1606.08415.pdf - """ - return 0.5 * x * (1.0 + tanh(_sqrt2pi * x * (1.0 + 0.044715 * x * x))) - - -@triton.jit -def gelu_approx_grad(x): - # CREDITS: Fast implementation proposed in - # https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/fused_bias_gelu.py#L30 - tanh_out = tanh(0.79788456 * x * (1 + 0.044715 * x * x)) - return 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( - 1 + tanh_out - ) diff --git a/build/torch211-cxx11-cu130-aarch64-linux/ops/triton/layer_norm.py b/build/torch211-cxx11-cu130-aarch64-linux/ops/triton/layer_norm.py deleted file mode 100644 index 5e6f0b3697b612e6a19031bbc55ac76ad233dc85..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu130-aarch64-linux/ops/triton/layer_norm.py +++ /dev/null @@ -1,1254 +0,0 @@ -# Copyright (c) 2024, Tri Dao. -# Implement dropout + residual + layer_norm / rms_norm. - -# Based on the Triton LayerNorm tutorial: https://triton-lang.org/main/getting-started/tutorials/05-layer-norm.html -# For the backward pass, we keep weight_grad and bias_grad in registers and accumulate. -# This is faster for dimensions up to 8k, but after that it's much slower due to register spilling. -# The models we train have hidden dim up to 8k anyway (e.g. Llama 70B), so this is fine. - -import math -from typing import Optional, List - -import torch -import torch.nn.functional as F -from torch import Tensor - -import triton -import triton.language as tl - -from flash_attn.utils.torch import custom_fwd, custom_bwd -from flash_attn.utils.library import triton_op - -from ..._ops import add_op_namespace_prefix - - -def maybe_contiguous_lastdim(x): - return x.contiguous() if x is not None and x.stride(-1) != 1 else x - - -def maybe_contiguous(x): - return x.contiguous() if x is not None else None - - -def triton_autotune_configs(): - # Return configs with a valid warp count for the current device - configs = [] - # Maximum threads per block is architecture-dependent in theory, but in reality all are 1024 - max_threads_per_block = 1024 - # Default to warp size 32 if not defined by device - warp_size = getattr(torch.cuda.get_device_properties(torch.cuda.current_device()), "warp_size", 32) - # Autotune for warp counts which are powers of 2 and do not exceed thread per block limit - return [triton.Config({}, num_warps=warp_count) for warp_count in [1, 2, 4, 8, 16, 32] - if warp_count * warp_size <= max_threads_per_block] - # return [triton.Config({}, num_warps=8)] - - -def layer_norm_ref( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - zero_centered_weight=False, - dropout_mask=None, - dropout_mask1=None, - upcast=False, -): - dtype = x.dtype - if upcast: - x = x.float() - weight = weight.float() - bias = bias.float() if bias is not None else None - residual = residual.float() if residual is not None else residual - x1 = x1.float() if x1 is not None else None - weight1 = weight1.float() if weight1 is not None else None - bias1 = bias1.float() if bias1 is not None else None - if zero_centered_weight: - weight = weight + 1.0 - if weight1 is not None: - weight1 = weight1 + 1.0 - if x1 is not None: - assert rowscale is None, "rowscale is not supported with parallel LayerNorm" - if rowscale is not None: - x = x * rowscale[..., None] - if dropout_p > 0.0: - if dropout_mask is not None: - x = x.masked_fill(~dropout_mask, 0.0) / (1.0 - dropout_p) - else: - x = F.dropout(x, p=dropout_p) - if x1 is not None: - if dropout_mask1 is not None: - x1 = x1.masked_fill(~dropout_mask1, 0.0) / (1.0 - dropout_p) - else: - x1 = F.dropout(x1, p=dropout_p) - if x1 is not None: - x = x + x1 - if residual is not None: - x = (x + residual).to(x.dtype) - out = F.layer_norm(x.to(weight.dtype), x.shape[-1:], weight=weight, bias=bias, eps=eps).to( - dtype - ) - if weight1 is None: - return out if not prenorm else (out, x) - else: - out1 = F.layer_norm( - x.to(weight1.dtype), x.shape[-1:], weight=weight1, bias=bias1, eps=eps - ).to(dtype) - return (out, out1) if not prenorm else (out, out1, x) - - -def rms_norm_ref( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - zero_centered_weight=False, - dropout_mask=None, - dropout_mask1=None, - upcast=False, -): - dtype = x.dtype - if upcast: - x = x.float() - weight = weight.float() - bias = bias.float() if bias is not None else None - residual = residual.float() if residual is not None else residual - x1 = x1.float() if x1 is not None else None - weight1 = weight1.float() if weight1 is not None else None - bias1 = bias1.float() if bias1 is not None else None - if zero_centered_weight: - weight = weight + 1.0 - if weight1 is not None: - weight1 = weight1 + 1.0 - if x1 is not None: - assert rowscale is None, "rowscale is not supported with parallel LayerNorm" - if rowscale is not None: - x = x * rowscale[..., None] - if dropout_p > 0.0: - if dropout_mask is not None: - x = x.masked_fill(~dropout_mask, 0.0) / (1.0 - dropout_p) - else: - x = F.dropout(x, p=dropout_p) - if x1 is not None: - if dropout_mask1 is not None: - x1 = x1.masked_fill(~dropout_mask1, 0.0) / (1.0 - dropout_p) - else: - x1 = F.dropout(x1, p=dropout_p) - if x1 is not None: - x = x + x1 - if residual is not None: - x = (x + residual).to(x.dtype) - rstd = 1 / torch.sqrt((x.square()).mean(dim=-1, keepdim=True) + eps) - out = ((x * rstd * weight) + bias if bias is not None else (x * rstd * weight)).to(dtype) - if weight1 is None: - return out if not prenorm else (out, x) - else: - out1 = ((x * rstd * weight1) + bias1 if bias1 is not None else (x * rstd * weight1)).to( - dtype - ) - return (out, out1) if not prenorm else (out, out1, x) - - -@triton.autotune( - configs=triton_autotune_configs(), - key=["N", "HAS_RESIDUAL", "STORE_RESIDUAL_OUT", "IS_RMS_NORM", "HAS_BIAS", "HAS_X1", "HAS_W1", "HAS_B1"], -) -# torch compile doesn't like triton.heuristics, so we set these manually when calling the kernel -# @triton.heuristics({"HAS_BIAS": lambda args: args["B"] is not None}) -# @triton.heuristics({"HAS_RESIDUAL": lambda args: args["RESIDUAL"] is not None}) -# @triton.heuristics({"HAS_X1": lambda args: args["X1"] is not None}) -# @triton.heuristics({"HAS_W1": lambda args: args["W1"] is not None}) -# @triton.heuristics({"HAS_B1": lambda args: args["B1"] is not None}) -@triton.jit -def _layer_norm_fwd_1pass_kernel( - X, # pointer to the input - Y, # pointer to the output - W, # pointer to the weights - B, # pointer to the biases - RESIDUAL, # pointer to the residual - X1, - W1, - B1, - Y1, - RESIDUAL_OUT, # pointer to the residual - ROWSCALE, - SEEDS, # Dropout seeds for each row - DROPOUT_MASK, - DROPOUT_MASK1, - Mean, # pointer to the mean - Rstd, # pointer to the 1/std - stride_x_row, # how much to increase the pointer when moving by 1 row - stride_y_row, - stride_res_row, - stride_res_out_row, - stride_x1_row, - stride_y1_row, - M, # number of rows in X - N, # number of columns in X - eps, # epsilon to avoid division by zero - dropout_p, # Dropout probability - zero_centered_weight, # If true, add 1.0 to the weight - IS_RMS_NORM: tl.constexpr, - BLOCK_N: tl.constexpr, - HAS_RESIDUAL: tl.constexpr, - STORE_RESIDUAL_OUT: tl.constexpr, - HAS_BIAS: tl.constexpr, - HAS_DROPOUT: tl.constexpr, - STORE_DROPOUT_MASK: tl.constexpr, - HAS_ROWSCALE: tl.constexpr, - HAS_X1: tl.constexpr, - HAS_W1: tl.constexpr, - HAS_B1: tl.constexpr, -): - # Map the program id to the row of X and Y it should compute. - row = tl.program_id(0) - X += row * stride_x_row - Y += row * stride_y_row - if HAS_RESIDUAL: - RESIDUAL += row * stride_res_row - if STORE_RESIDUAL_OUT: - RESIDUAL_OUT += row * stride_res_out_row - if HAS_X1: - X1 += row * stride_x1_row - if HAS_W1: - Y1 += row * stride_y1_row - # Compute mean and variance - cols = tl.arange(0, BLOCK_N) - x = tl.load(X + cols, mask=cols < N, other=0.0).to(tl.float32) - if HAS_ROWSCALE: - rowscale = tl.load(ROWSCALE + row).to(tl.float32) - x *= rowscale - if HAS_DROPOUT: - # Compute dropout mask - # 7 rounds is good enough, and reduces register pressure - keep_mask = tl.rand(tl.load(SEEDS + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - x = tl.where(keep_mask, x / (1.0 - dropout_p), 0.0) - if STORE_DROPOUT_MASK: - tl.store(DROPOUT_MASK + row * N + cols, keep_mask, mask=cols < N) - if HAS_X1: - x1 = tl.load(X1 + cols, mask=cols < N, other=0.0).to(tl.float32) - if HAS_ROWSCALE: - rowscale = tl.load(ROWSCALE + M + row).to(tl.float32) - x1 *= rowscale - if HAS_DROPOUT: - # Compute dropout mask - # 7 rounds is good enough, and reduces register pressure - keep_mask = ( - tl.rand(tl.load(SEEDS + M + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - ) - x1 = tl.where(keep_mask, x1 / (1.0 - dropout_p), 0.0) - if STORE_DROPOUT_MASK: - tl.store(DROPOUT_MASK1 + row * N + cols, keep_mask, mask=cols < N) - x += x1 - if HAS_RESIDUAL: - residual = tl.load(RESIDUAL + cols, mask=cols < N, other=0.0).to(tl.float32) - x += residual - if STORE_RESIDUAL_OUT: - tl.store(RESIDUAL_OUT + cols, x, mask=cols < N) - if not IS_RMS_NORM: - mean = tl.sum(x, axis=0) / N - tl.store(Mean + row, mean) - xbar = tl.where(cols < N, x - mean, 0.0) - var = tl.sum(xbar * xbar, axis=0) / N - else: - xbar = tl.where(cols < N, x, 0.0) - var = tl.sum(xbar * xbar, axis=0) / N - rstd = 1 / tl.sqrt(var + eps) - tl.store(Rstd + row, rstd) - # Normalize and apply linear transformation - mask = cols < N - w = tl.load(W + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w += 1.0 - if HAS_BIAS: - b = tl.load(B + cols, mask=mask).to(tl.float32) - x_hat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd - y = x_hat * w + b if HAS_BIAS else x_hat * w - # Write output - tl.store(Y + cols, y, mask=mask) - if HAS_W1: - w1 = tl.load(W1 + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w1 += 1.0 - if HAS_B1: - b1 = tl.load(B1 + cols, mask=mask).to(tl.float32) - y1 = x_hat * w1 + b1 if HAS_B1 else x_hat * w1 - tl.store(Y1 + cols, y1, mask=mask) - - -def _layer_norm_fwd( - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - residual: Optional[Tensor] = None, - x1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - out_dtype: Optional[torch.dtype] = None, - residual_dtype: Optional[torch.dtype] = None, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - return_dropout_mask: bool = False, - out: Optional[Tensor] = None, - residual_out: Optional[Tensor] = None -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - # Need to wrap to handle the case where residual_out is a alias of x, which makes torch.library - # and torch.compile unhappy. Also allocate memory for out and residual_out if they are None - # so that _layer_norm_fwd_impl doesn't have to return them. - if out is None: - out = torch.empty_like(x, dtype=x.dtype if out_dtype is None else out_dtype) - if residual is not None: - residual_dtype = residual.dtype - if residual_out is None and ( - residual is not None - or (residual_dtype is not None and residual_dtype != x.dtype) - or dropout_p > 0.0 - or rowscale is not None - or x1 is not None - ): - residual_out = torch.empty_like( - x, dtype=residual_dtype if residual_dtype is not None else x.dtype - ) - else: - residual_out = None - y1, mean, rstd, seeds, dropout_mask, dropout_mask1 = _layer_norm_fwd_impl( - x, - weight, - bias, - eps, - out, - residual=residual, - x1=x1, - weight1=weight1, - bias1=bias1, - dropout_p=dropout_p, - rowscale=rowscale, - zero_centered_weight=zero_centered_weight, - is_rms_norm=is_rms_norm, - return_dropout_mask=return_dropout_mask, - residual_out=residual_out, - ) - # residual_out is None if residual is None and residual_dtype == input_dtype and dropout_p == 0.0 - if residual_out is None: - residual_out = x - return out, y1, mean, rstd, residual_out, seeds, dropout_mask, dropout_mask1 - - -# [2025-04-28] torch.library.triton_op ignores the schema argument, but here we need the schema -# since we're returning a tuple of tensors -@triton_op(add_op_namespace_prefix("layer_norm_fwd_impl"), mutates_args={"out", "residual_out"}, - schema="(Tensor x, Tensor weight, Tensor bias, float eps, Tensor(a!) out, Tensor? residual, Tensor? x1, Tensor? weight1, Tensor? bias1, float dropout_p, Tensor? rowscale, bool zero_centered_weight, bool is_rms_norm, bool return_dropout_mask, Tensor(a!)? residual_out) -> (Tensor y1, Tensor mean, Tensor rstd, Tensor seeds, Tensor dropout_mask, Tensor dropout_mask1)") -def _layer_norm_fwd_impl( - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - out: Tensor, - residual: Optional[Tensor] = None, - x1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - return_dropout_mask: bool = False, - residual_out: Optional[Tensor] = None -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - M, N = x.shape - assert x.stride(-1) == 1 - if residual is not None: - assert residual.stride(-1) == 1 - assert residual.shape == (M, N) - assert weight.shape == (N,) - assert weight.stride(-1) == 1 - if bias is not None: - assert bias.stride(-1) == 1 - assert bias.shape == (N,) - if x1 is not None: - assert x1.shape == x.shape - assert rowscale is None - assert x1.stride(-1) == 1 - if weight1 is not None: - assert weight1.shape == (N,) - assert weight1.stride(-1) == 1 - if bias1 is not None: - assert bias1.shape == (N,) - assert bias1.stride(-1) == 1 - if rowscale is not None: - assert rowscale.is_contiguous() - assert rowscale.shape == (M,) - assert out.shape == x.shape - assert out.stride(-1) == 1 - if residual_out is not None: - assert residual_out.shape == x.shape - assert residual_out.stride(-1) == 1 - if weight1 is not None: - y1 = torch.empty_like(out) - assert y1.stride(-1) == 1 - else: - y1 = None - mean = torch.empty((M,), dtype=torch.float32, device=x.device) if not is_rms_norm else None - rstd = torch.empty((M,), dtype=torch.float32, device=x.device) - if dropout_p > 0.0: - seeds = torch.randint( - 2**32, (M if x1 is None else 2 * M,), device=x.device, dtype=torch.int64 - ) - else: - seeds = None - if return_dropout_mask and dropout_p > 0.0: - dropout_mask = torch.empty(M, N, device=x.device, dtype=torch.bool) - if x1 is not None: - dropout_mask1 = torch.empty(M, N, device=x.device, dtype=torch.bool) - else: - dropout_mask1 = None - else: - dropout_mask, dropout_mask1 = None, None - # Less than 64KB per feature: enqueue fused kernel - MAX_FUSED_SIZE = 65536 // x.element_size() - BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) - if N > BLOCK_N: - raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") - with torch.cuda.device(x.device.index): - torch.library.wrap_triton(_layer_norm_fwd_1pass_kernel)[(M,)]( - x, - out, - weight, - bias, - residual, - x1, - weight1, - bias1, - y1, - residual_out, - rowscale, - seeds, - dropout_mask, - dropout_mask1, - mean, - rstd, - x.stride(0), - out.stride(0), - residual.stride(0) if residual is not None else 0, - residual_out.stride(0) if residual_out is not None else 0, - x1.stride(0) if x1 is not None else 0, - y1.stride(0) if y1 is not None else 0, - M, - N, - eps, - dropout_p, - # Passing bool make torch inductor very unhappy since it then tries to compare to int_max - int(zero_centered_weight), - is_rms_norm, - BLOCK_N, - residual is not None, - residual_out is not None, - bias is not None, - dropout_p > 0.0, - dropout_mask is not None, - rowscale is not None, - HAS_X1=x1 is not None, - HAS_W1=weight1 is not None, - HAS_B1=bias1 is not None, - ) - return y1, mean, rstd, seeds, dropout_mask, dropout_mask1 - - -@triton.autotune( - configs=triton_autotune_configs(), - key=["N", "HAS_DRESIDUAL", "STORE_DRESIDUAL", "IS_RMS_NORM", "HAS_BIAS", "HAS_DROPOUT"], -) -# torch compile doesn't like triton.heuristics, so we set these manually when calling the kernel -# @triton.heuristics({"HAS_BIAS": lambda args: args["B"] is not None}) -# @triton.heuristics({"HAS_DRESIDUAL": lambda args: args["DRESIDUAL"] is not None}) -# @triton.heuristics({"STORE_DRESIDUAL": lambda args: args["DRESIDUAL_IN"] is not None}) -# @triton.heuristics({"HAS_ROWSCALE": lambda args: args["ROWSCALE"] is not None}) -# @triton.heuristics({"HAS_DY1": lambda args: args["DY1"] is not None}) -# @triton.heuristics({"HAS_DX1": lambda args: args["DX1"] is not None}) -# @triton.heuristics({"HAS_B1": lambda args: args["DB1"] is not None}) -# @triton.heuristics({"RECOMPUTE_OUTPUT": lambda args: args["Y"] is not None}) -@triton.jit -def _layer_norm_bwd_kernel( - X, # pointer to the input - W, # pointer to the weights - B, # pointer to the biases - Y, # pointer to the output to be recomputed - DY, # pointer to the output gradient - DX, # pointer to the input gradient - DW, # pointer to the partial sum of weights gradient - DB, # pointer to the partial sum of biases gradient - DRESIDUAL, - W1, - DY1, - DX1, - DW1, - DB1, - DRESIDUAL_IN, - ROWSCALE, - SEEDS, - Mean, # pointer to the mean - Rstd, # pointer to the 1/std - stride_x_row, # how much to increase the pointer when moving by 1 row - stride_y_row, - stride_dy_row, - stride_dx_row, - stride_dres_row, - stride_dy1_row, - stride_dx1_row, - stride_dres_in_row, - M, # number of rows in X - N, # number of columns in X - eps, # epsilon to avoid division by zero - dropout_p, - zero_centered_weight, - rows_per_program, - IS_RMS_NORM: tl.constexpr, - BLOCK_N: tl.constexpr, - HAS_DRESIDUAL: tl.constexpr, - STORE_DRESIDUAL: tl.constexpr, - HAS_BIAS: tl.constexpr, - HAS_DROPOUT: tl.constexpr, - HAS_ROWSCALE: tl.constexpr, - HAS_DY1: tl.constexpr, - HAS_DX1: tl.constexpr, - HAS_B1: tl.constexpr, - RECOMPUTE_OUTPUT: tl.constexpr, -): - # Map the program id to the elements of X, DX, and DY it should compute. - row_block_id = tl.program_id(0) - row_start = row_block_id * rows_per_program - # Do not early exit if row_start >= M, because we need to write DW and DB - cols = tl.arange(0, BLOCK_N) - mask = cols < N - X += row_start * stride_x_row - if HAS_DRESIDUAL: - DRESIDUAL += row_start * stride_dres_row - if STORE_DRESIDUAL: - DRESIDUAL_IN += row_start * stride_dres_in_row - DY += row_start * stride_dy_row - DX += row_start * stride_dx_row - if HAS_DY1: - DY1 += row_start * stride_dy1_row - if HAS_DX1: - DX1 += row_start * stride_dx1_row - if RECOMPUTE_OUTPUT: - Y += row_start * stride_y_row - w = tl.load(W + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w += 1.0 - if RECOMPUTE_OUTPUT and HAS_BIAS: - b = tl.load(B + cols, mask=mask, other=0.0).to(tl.float32) - if HAS_DY1: - w1 = tl.load(W1 + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w1 += 1.0 - dw = tl.zeros((BLOCK_N,), dtype=tl.float32) - if HAS_BIAS: - db = tl.zeros((BLOCK_N,), dtype=tl.float32) - if HAS_DY1: - dw1 = tl.zeros((BLOCK_N,), dtype=tl.float32) - if HAS_B1: - db1 = tl.zeros((BLOCK_N,), dtype=tl.float32) - row_end = min((row_block_id + 1) * rows_per_program, M) - for row in range(row_start, row_end): - # Load data to SRAM - x = tl.load(X + cols, mask=mask, other=0).to(tl.float32) - dy = tl.load(DY + cols, mask=mask, other=0).to(tl.float32) - if HAS_DY1: - dy1 = tl.load(DY1 + cols, mask=mask, other=0).to(tl.float32) - if not IS_RMS_NORM: - mean = tl.load(Mean + row) - rstd = tl.load(Rstd + row) - # Compute dx - xhat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd - xhat = tl.where(mask, xhat, 0.0) - if RECOMPUTE_OUTPUT: - y = xhat * w + b if HAS_BIAS else xhat * w - tl.store(Y + cols, y, mask=mask) - wdy = w * dy - dw += dy * xhat - if HAS_BIAS: - db += dy - if HAS_DY1: - wdy += w1 * dy1 - dw1 += dy1 * xhat - if HAS_B1: - db1 += dy1 - if not IS_RMS_NORM: - c1 = tl.sum(xhat * wdy, axis=0) / N - c2 = tl.sum(wdy, axis=0) / N - dx = (wdy - (xhat * c1 + c2)) * rstd - else: - c1 = tl.sum(xhat * wdy, axis=0) / N - dx = (wdy - xhat * c1) * rstd - if HAS_DRESIDUAL: - dres = tl.load(DRESIDUAL + cols, mask=mask, other=0).to(tl.float32) - dx += dres - # Write dx - if STORE_DRESIDUAL: - tl.store(DRESIDUAL_IN + cols, dx, mask=mask) - if HAS_DX1: - if HAS_DROPOUT: - keep_mask = ( - tl.rand(tl.load(SEEDS + M + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - ) - dx1 = tl.where(keep_mask, dx / (1.0 - dropout_p), 0.0) - else: - dx1 = dx - tl.store(DX1 + cols, dx1, mask=mask) - if HAS_DROPOUT: - keep_mask = tl.rand(tl.load(SEEDS + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - dx = tl.where(keep_mask, dx / (1.0 - dropout_p), 0.0) - if HAS_ROWSCALE: - rowscale = tl.load(ROWSCALE + row).to(tl.float32) - dx *= rowscale - tl.store(DX + cols, dx, mask=mask) - - X += stride_x_row - if HAS_DRESIDUAL: - DRESIDUAL += stride_dres_row - if STORE_DRESIDUAL: - DRESIDUAL_IN += stride_dres_in_row - if RECOMPUTE_OUTPUT: - Y += stride_y_row - DY += stride_dy_row - DX += stride_dx_row - if HAS_DY1: - DY1 += stride_dy1_row - if HAS_DX1: - DX1 += stride_dx1_row - tl.store(DW + row_block_id * N + cols, dw, mask=mask) - if HAS_BIAS: - tl.store(DB + row_block_id * N + cols, db, mask=mask) - if HAS_DY1: - tl.store(DW1 + row_block_id * N + cols, dw1, mask=mask) - if HAS_B1: - tl.store(DB1 + row_block_id * N + cols, db1, mask=mask) - - -def _layer_norm_bwd( - dy: Tensor, - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - mean: Tensor, - rstd: Tensor, - dresidual: Optional[Tensor] = None, - dy1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - seeds: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - has_residual: bool = False, - has_x1: bool = False, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - x_dtype: Optional[torch.dtype] = None, - recompute_output: bool = False, -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - # Need to wrap to handle the case where dresidual_in or dx1 are aliases of x, - # which makes torch.library unhappy - dx, dw, db, dresidual_in, dx1, dw1, db1, y = _layer_norm_bwd_impl( - dy, - x, - weight, - bias, - eps, - mean, - rstd, - dresidual, - dy1, - weight1, - bias1, - seeds, - dropout_p, - rowscale, - has_residual, - has_x1, - zero_centered_weight, - is_rms_norm, - x_dtype=x_dtype, - recompute_output=recompute_output, - ) - # Don't need to compute dresidual_in separately in this case - if has_residual and dx.dtype == x.dtype and dropout_p == 0.0 and rowscale is None: - dresidual_in = dx - if has_x1 and dropout_p == 0.0: - dx1 = dx - return dx, dw, db, dresidual_in, dx1, dw1, db1, y - - - -@triton_op(add_op_namespace_prefix("layer_norm_bwd_impl"), mutates_args={}, - schema="(Tensor dy, Tensor x, Tensor weight, Tensor bias, float eps, Tensor mean, Tensor rstd, Tensor? dresidual, Tensor? dy1, Tensor? weight1, Tensor? bias1, Tensor? seeds, float dropout_p, Tensor? rowscale, bool has_residual, bool has_x1, bool zero_centered_weight, bool is_rms_norm, ScalarType? x_dtype, bool recompute_output) -> (Tensor dx, Tensor dw, Tensor db, Tensor dresidual_in, Tensor dx1, Tensor dw1, Tensor db1, Tensor y)", - allow_decomposition=False, # Don't let torch.compile trace inside - ) -def _layer_norm_bwd_impl( - dy: Tensor, - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - mean: Tensor, - rstd: Tensor, - dresidual: Optional[Tensor] = None, - dy1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - seeds: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - has_residual: bool = False, - has_x1: bool = False, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - x_dtype: Optional[torch.dtype] = None, - recompute_output: bool = False, -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - M, N = x.shape - assert x.stride(-1) == 1 - dy = maybe_contiguous_lastdim(dy) - assert dy.stride(-1) == 1 - assert dy.shape == (M, N) - if dresidual is not None: - dresidual = maybe_contiguous_lastdim(dresidual) - assert dresidual.stride(-1) == 1 - assert dresidual.shape == (M, N) - assert weight.shape == (N,) - assert weight.stride(-1) == 1 - if bias is not None: - assert bias.stride(-1) == 1 - assert bias.shape == (N,) - if dy1 is not None: - dy1 = maybe_contiguous_lastdim(dy1) - assert weight1 is not None - assert dy1.shape == dy.shape - assert dy1.stride(-1) == 1 - if weight1 is not None: - assert weight1.shape == (N,) - assert weight1.stride(-1) == 1 - if bias1 is not None: - assert bias1.shape == (N,) - assert bias1.stride(-1) == 1 - if seeds is not None: - assert seeds.is_contiguous() - assert seeds.shape == (M if not has_x1 else M * 2,) - if rowscale is not None: - assert rowscale.is_contiguous() - assert rowscale.shape == (M,) - # allocate output - dx = ( - torch.empty_like(x) - if x_dtype is None - else torch.empty(M, N, dtype=x_dtype, device=x.device) - ) - dresidual_in = ( - torch.empty_like(x) - if has_residual - and (dx.dtype != x.dtype or dropout_p > 0.0 or rowscale is not None or has_x1) - else None - ) - dx1 = torch.empty_like(dx) if (has_x1 and dropout_p > 0.0) else None - y = torch.empty(M, N, dtype=dy.dtype, device=dy.device) if recompute_output else None - if recompute_output: - assert weight1 is None, "recompute_output is not supported with parallel LayerNorm" - - # Less than 64KB per feature: enqueue fused kernel - MAX_FUSED_SIZE = 65536 // x.element_size() - BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) - if N > BLOCK_N: - raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") - # Increasing the multiple (e.g. 8) will allow more thread blocks to be launched and hide the - # latency of the gmem reads/writes, but will increase the time of summing up dw / db. - sm_count = torch.cuda.get_device_properties(x.device).multi_processor_count * 8 - _dw = torch.empty((sm_count, N), dtype=torch.float32, device=weight.device) - _db = ( - torch.empty((sm_count, N), dtype=torch.float32, device=bias.device) - if bias is not None - else None - ) - _dw1 = torch.empty_like(_dw) if weight1 is not None else None - _db1 = torch.empty_like(_db) if bias1 is not None else None - rows_per_program = math.ceil(M / sm_count) - grid = (sm_count,) - with torch.cuda.device(x.device.index): - torch.library.wrap_triton(_layer_norm_bwd_kernel)[grid]( - x, - weight, - bias, - y, - dy, - dx, - _dw, - _db, - dresidual, - weight1, - dy1, - dx1, - _dw1, - _db1, - dresidual_in, - rowscale, - seeds, - mean, - rstd, - x.stride(0), - 0 if not recompute_output else y.stride(0), - dy.stride(0), - dx.stride(0), - dresidual.stride(0) if dresidual is not None else 0, - dy1.stride(0) if dy1 is not None else 0, - dx1.stride(0) if dx1 is not None else 0, - dresidual_in.stride(0) if dresidual_in is not None else 0, - M, - N, - eps, - dropout_p, - # Passing bool make torch inductor very unhappy since it then tries to compare to int_max - int(zero_centered_weight), - rows_per_program, - is_rms_norm, - BLOCK_N, - dresidual is not None, - dresidual_in is not None, - bias is not None, - dropout_p > 0.0, - HAS_ROWSCALE=rowscale is not None, - HAS_DY1=dy1 is not None, - HAS_DX1=dx1 is not None, - HAS_B1=bias1 is not None, - RECOMPUTE_OUTPUT=y is not None, - ) - dw = _dw.sum(0).to(weight.dtype) - db = _db.sum(0).to(bias.dtype) if bias is not None else None - dw1 = _dw1.sum(0).to(weight1.dtype) if weight1 is not None else None - db1 = _db1.sum(0).to(bias1.dtype) if bias1 is not None else None - # dresidual_in and dx1 could be None, the wrapper will handle assigning them from dx - return dx, dw, db, dresidual_in, dx1, dw1, db1, y - - -class LayerNormFn(torch.autograd.Function): - - @staticmethod - def forward( - ctx, - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - residual_in_fp32=False, - zero_centered_weight=False, - is_rms_norm=False, - return_dropout_mask=False, - out_dtype=None, - out=None, - residual_out=None - ): - x_shape_og = x.shape - # reshape input data into 2D tensor - x = maybe_contiguous_lastdim(x.reshape(-1, x.shape[-1])) - if residual is not None: - assert residual.shape == x_shape_og - residual = maybe_contiguous_lastdim(residual.reshape(-1, residual.shape[-1])) - if x1 is not None: - assert x1.shape == x_shape_og - assert rowscale is None, "rowscale is not supported with parallel LayerNorm" - x1 = maybe_contiguous_lastdim(x1.reshape(-1, x1.shape[-1])) - weight = weight.contiguous() - bias = maybe_contiguous(bias) - weight1 = maybe_contiguous(weight1) - bias1 = maybe_contiguous(bias1) - if rowscale is not None: - rowscale = rowscale.reshape(-1).contiguous() - residual_dtype = ( - residual.dtype - if residual is not None - else (torch.float32 if residual_in_fp32 else None) - ) - if out is not None: - out = out.reshape(-1, out.shape[-1]) - if residual_out is not None: - residual_out = residual_out.reshape(-1, residual_out.shape[-1]) - y, y1, mean, rstd, residual_out, seeds, dropout_mask, dropout_mask1 = _layer_norm_fwd( - x, - weight, - bias, - eps, - residual, - x1, - weight1, - bias1, - dropout_p=dropout_p, - rowscale=rowscale, - out_dtype=out_dtype, - residual_dtype=residual_dtype, - zero_centered_weight=zero_centered_weight, - is_rms_norm=is_rms_norm, - return_dropout_mask=return_dropout_mask, - out=out, - residual_out=residual_out, - ) - ctx.save_for_backward( - residual_out, weight, bias, weight1, bias1, rowscale, seeds, mean, rstd - ) - ctx.x_shape_og = x_shape_og - ctx.eps = eps - ctx.dropout_p = dropout_p - ctx.is_rms_norm = is_rms_norm - ctx.has_residual = residual is not None - ctx.has_x1 = x1 is not None - ctx.prenorm = prenorm - ctx.x_dtype = x.dtype - ctx.zero_centered_weight = zero_centered_weight - y = y.reshape(x_shape_og) - y1 = y1.reshape(x_shape_og) if y1 is not None else None - residual_out = residual_out.reshape(x_shape_og) if residual_out is not None else None - dropout_mask = dropout_mask.reshape(x_shape_og) if dropout_mask is not None else None - dropout_mask1 = dropout_mask1.reshape(x_shape_og) if dropout_mask1 is not None else None - if not return_dropout_mask: - if weight1 is None: - return y if not prenorm else (y, residual_out) - else: - return (y, y1) if not prenorm else (y, y1, residual_out) - else: - if weight1 is None: - return ( - (y, dropout_mask, dropout_mask1) - if not prenorm - else (y, residual_out, dropout_mask, dropout_mask1) - ) - else: - return ( - (y, y1, dropout_mask, dropout_mask1) - if not prenorm - else (y, y1, residual_out, dropout_mask, dropout_mask1) - ) - - @staticmethod - def backward(ctx, dy, *args): - x, weight, bias, weight1, bias1, rowscale, seeds, mean, rstd = ctx.saved_tensors - dy = dy.reshape(-1, dy.shape[-1]) - if weight1 is not None: - dy1, args = args[0], args[1:] - dy1 = dy1.reshape(-1, dy1.shape[-1]) - assert dy1.shape == x.shape - else: - dy1 = None - if ctx.prenorm: - dresidual = args[0] - dresidual = dresidual.reshape(-1, dresidual.shape[-1]) - assert dresidual.shape == x.shape - else: - dresidual = None - dx, dw, db, dresidual_in, dx1, dw1, db1, _ = _layer_norm_bwd( - dy, - x, - weight, - bias, - ctx.eps, - mean, - rstd, - dresidual, - dy1, - weight1, - bias1, - seeds, - ctx.dropout_p, - rowscale, - ctx.has_residual, - ctx.has_x1, - ctx.zero_centered_weight, - ctx.is_rms_norm, - x_dtype=ctx.x_dtype, - recompute_output=False, - ) - return ( - dx.reshape(ctx.x_shape_og), - dw, - db, - dresidual_in.reshape(ctx.x_shape_og) if ctx.has_residual else None, - dx1.reshape(ctx.x_shape_og) if dx1 is not None else None, - dw1, - db1, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) - - -def layer_norm_fn( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - residual_in_fp32=False, - zero_centered_weight=False, - is_rms_norm=False, - return_dropout_mask=False, - out_dtype=None, - out=None, - residual_out=None -): - return LayerNormFn.apply( - x, - weight, - bias, - residual, - x1, - weight1, - bias1, - eps, - dropout_p, - rowscale, - prenorm, - residual_in_fp32, - zero_centered_weight, - is_rms_norm, - return_dropout_mask, - out_dtype, - out, - residual_out - ) - - -def rms_norm_fn( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - residual_in_fp32=False, - zero_centered_weight=False, - return_dropout_mask=False, - out_dtype=None, - out=None, - residual_out=None -): - return LayerNormFn.apply( - x, - weight, - bias, - residual, - x1, - weight1, - bias1, - eps, - dropout_p, - rowscale, - prenorm, - residual_in_fp32, - zero_centered_weight, - True, - return_dropout_mask, - out_dtype, - out, - residual_out - ) - - -class RMSNorm(torch.nn.Module): - - def __init__(self, hidden_size, eps=1e-5, dropout_p=0.0, zero_centered_weight=False, - device=None, dtype=None): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.eps = eps - if dropout_p > 0.0: - self.drop = torch.nn.Dropout(dropout_p) - else: - self.drop = None - self.zero_centered_weight = zero_centered_weight - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.register_parameter("bias", None) - self.reset_parameters() - - def reset_parameters(self): - if not self.zero_centered_weight: - torch.nn.init.ones_(self.weight) - else: - torch.nn.init.zeros_(self.weight) - - def forward(self, x, residual=None, prenorm=False, residual_in_fp32=False): - return rms_norm_fn( - x, - self.weight, - self.bias, - residual=residual, - eps=self.eps, - dropout_p=self.drop.p if self.drop is not None and self.training else 0.0, - prenorm=prenorm, - residual_in_fp32=residual_in_fp32, - zero_centered_weight=self.zero_centered_weight, - ) - - -class LayerNormLinearFn(torch.autograd.Function): - - @staticmethod - @custom_fwd - def forward( - ctx, - x, - norm_weight, - norm_bias, - linear_weight, - linear_bias, - residual=None, - eps=1e-6, - prenorm=False, - residual_in_fp32=False, - is_rms_norm=False, - ): - x_shape_og = x.shape - # reshape input data into 2D tensor - x = maybe_contiguous_lastdim(x.reshape(-1, x.shape[-1])) - if residual is not None: - assert residual.shape == x_shape_og - residual = maybe_contiguous_lastdim(residual.reshape(-1, residual.shape[-1])) - norm_weight = norm_weight.contiguous() - norm_bias = maybe_contiguous(norm_bias) - residual_dtype = ( - residual.dtype - if residual is not None - else (torch.float32 if residual_in_fp32 else None) - ) - y, _, mean, rstd, residual_out, *rest = _layer_norm_fwd( - x, - norm_weight, - norm_bias, - eps, - residual, - out_dtype=None if not torch.is_autocast_enabled() else torch.get_autocast_dtype("cuda"), - residual_dtype=residual_dtype, - is_rms_norm=is_rms_norm, - ) - y = y.reshape(x_shape_og) - dtype = torch.get_autocast_dtype("cuda") if torch.is_autocast_enabled() else y.dtype - linear_weight = linear_weight.to(dtype) - linear_bias = linear_bias.to(dtype) if linear_bias is not None else None - out = F.linear(y.to(linear_weight.dtype), linear_weight, linear_bias) - # We don't store y, will be recomputed in the backward pass to save memory - ctx.save_for_backward(residual_out, norm_weight, norm_bias, linear_weight, mean, rstd) - ctx.x_shape_og = x_shape_og - ctx.eps = eps - ctx.is_rms_norm = is_rms_norm - ctx.has_residual = residual is not None - ctx.prenorm = prenorm - ctx.x_dtype = x.dtype - ctx.linear_bias_is_none = linear_bias is None - return out if not prenorm else (out, residual_out.reshape(x_shape_og)) - - @staticmethod - @custom_bwd - def backward(ctx, dout, *args): - x, norm_weight, norm_bias, linear_weight, mean, rstd = ctx.saved_tensors - dout = dout.reshape(-1, dout.shape[-1]) - dy = F.linear(dout, linear_weight.t()) - dlinear_bias = None if ctx.linear_bias_is_none else dout.sum(0) - dy = maybe_contiguous_lastdim(dy) - assert dy.shape == x.shape - if ctx.prenorm: - dresidual = args[0] - dresidual = maybe_contiguous_lastdim(dresidual.reshape(-1, dresidual.shape[-1])) - assert dresidual.shape == x.shape - else: - dresidual = None - dx, dnorm_weight, dnorm_bias, dresidual_in, _, _, _, y = _layer_norm_bwd( - dy, - x, - norm_weight, - norm_bias, - ctx.eps, - mean, - rstd, - dresidual=dresidual, - has_residual=ctx.has_residual, - is_rms_norm=ctx.is_rms_norm, - x_dtype=ctx.x_dtype, - recompute_output=True, - ) - dlinear_weight = torch.einsum("bo,bi->oi", dout, y) - return ( - dx.reshape(ctx.x_shape_og), - dnorm_weight, - dnorm_bias, - dlinear_weight, - dlinear_bias, - dresidual_in.reshape(ctx.x_shape_og) if ctx.has_residual else None, - None, - None, - None, - None, - ) - - -def layer_norm_linear_fn( - x, - norm_weight, - norm_bias, - linear_weight, - linear_bias, - residual=None, - eps=1e-6, - prenorm=False, - residual_in_fp32=False, - is_rms_norm=False, -): - return LayerNormLinearFn.apply( - x, - norm_weight, - norm_bias, - linear_weight, - linear_bias, - residual, - eps, - prenorm, - residual_in_fp32, - is_rms_norm, - ) diff --git a/build/torch211-cxx11-cu130-aarch64-linux/ops/triton/linear.py b/build/torch211-cxx11-cu130-aarch64-linux/ops/triton/linear.py deleted file mode 100644 index a8966dbc345ab0e593df0124451ee7be3dae131a..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu130-aarch64-linux/ops/triton/linear.py +++ /dev/null @@ -1,594 +0,0 @@ -# Adapted from https://github.com/ELS-RD/kernl/blob/main/src/kernl/implementations/linear_layer.py -# and https://github.com/openai/triton/blob/master/python/triton/ops/matmul.py -from typing import Optional - -import torch -import triton -import triton.language as tl -from triton.ops.matmul_perf_model import early_config_prune, estimate_matmul_time - -from flash_attn.ops.triton.k_activations import ( - gelu, - gelu_approx, - gelu_approx_grad, - gelu_grad, - squared_relu, - squared_relu_grad, -) - -# CREDITS: Initially inspired by the Triton tutorial on matrix multiplications - - -def init_to_zero(name): - return lambda nargs: nargs[name].zero_() - - -def get_configs_io_bound(): - configs = [] - for num_stages in [2, 3, 4, 5, 6]: - for block_m in [16, 32]: - for block_k in [32, 64]: - for block_n in [32, 64, 128, 256]: - num_warps = 2 if block_n <= 64 else 4 - configs.append( - triton.Config( - { - "BLOCK_M": block_m, - "BLOCK_N": block_n, - "BLOCK_K": block_k, - "SPLIT_K": 1, - }, - num_stages=num_stages, - num_warps=num_warps, - ) - ) - # split_k not used - # for split_k in [2, 4, 8, 16]: - # configs.append(triton.Config( - # {'BLOCK_M': block_m, 'BLOCK_N': block_n, 'BLOCK_K': block_k, 'SPLIT_K': split_k}, - # num_stages=num_stages, num_warps=num_warps, pre_hook=init_to_zero('C'))) - return configs - - -@triton.autotune( - configs=[ - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - # good for int8 - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=4, - num_warps=4, - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - ] - + get_configs_io_bound(), - key=["CACHE_KEY_M", "CACHE_KEY_N", "CACHE_KEY_K"], - prune_configs_by={ - "early_config_prune": early_config_prune, - "perf_model": estimate_matmul_time, - "top_k": 10, - }, -) -@triton.heuristics( - { - "EVEN_K": lambda args: args["K"] % (args["BLOCK_K"] * args["SPLIT_K"]) == 0, - } -) -@triton.jit -def kernel_fwd( - C, # Pointers to matrices - ACT_INPUT, - A, - B, - bias, - # Matrix dimensions - M, - N, - K, - CACHE_KEY_M, - CACHE_KEY_N, - CACHE_KEY_K, - # The stride variables represent how much to increase the ptr by when moving by 1 - # element in a particular dimension. E.g. stride_am is how much to increase a_ptr - # by to get the element one row down (A has M rows) - stride_cm, - # stride_cn, # Assume that stride_cn == 1 - stride_am, - stride_ak, - stride_bn, - stride_bk, - # Meta-parameters - BLOCK_M: tl.constexpr, - GROUP_M: tl.constexpr, - BLOCK_N: tl.constexpr, - BLOCK_K: tl.constexpr, - # split k not used, not performant with activation, kept because early_config_prune is expecting it - SPLIT_K: tl.constexpr, - EVEN_K: tl.constexpr, - A_ROWMAJOR: tl.constexpr, - B_COLMAJOR: tl.constexpr, - BIAS: tl.constexpr, - SAVE_ACT_INPUT: tl.constexpr, - ACTIVATION: tl.constexpr, -): - - """ - Kernel for computing Out = activation(A x W + C) - - Input has shape (M, K) - - Weight has shape (K, N) - - Bias has shape (N,) - - Output has shape (M, N) - - ActInputs (optional) has shape (M, N) - 'ActInputs' optionally saves the A x W + C intermediate for backward computations - This kernel will consolidate over K - """ - - pid = tl.program_id(axis=0) - - grid_m = (M + BLOCK_M - 1) // BLOCK_M - grid_n = (N + BLOCK_N - 1) // BLOCK_N - # re-order program ID for better L2 performance - width = GROUP_M * grid_n - group_id = pid // width - group_size = min(grid_m - group_id * GROUP_M, GROUP_M) - pid_m = group_id * GROUP_M + (pid % group_size) - pid_n = (pid % width) // (group_size) - - # now compute the block that each program will go through - # rm (resp. rn) denotes a range of indices - # for rows (resp. col) of C - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - # trick to avoid masking on M and N axis - ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M) - rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N) - rk = tl.arange(0, BLOCK_K) - - if A_ROWMAJOR: - A = A + (ram[:, None] * stride_am + rk[None, :]) - else: - A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak) - if B_COLMAJOR: - B = B + (rk[:, None] + rbn[None, :] * stride_bn) - else: - B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn) - - acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) - - for k in range(K, 0, -BLOCK_K): - if EVEN_K: - a = tl.load(A) - b = tl.load(B) - else: - a = tl.load(A, mask=rk[None, :] < k, other=0.0) - b = tl.load(B, mask=rk[:, None] < k, other=0.0) - acc += tl.dot(a, b) - - if A_ROWMAJOR: - A += BLOCK_K - else: - A += BLOCK_K * stride_ak - if B_COLMAJOR: - B += BLOCK_K - else: - B += BLOCK_K * stride_bk - - # Putting bias after the matmul (instead of before) is faster, idk why - if BIAS: - bias = tl.load(bias + rn, mask=rn < N, other=0.0).to(tl.float32) - acc += bias[None, :] - - # optional: save the activation inputs - if SAVE_ACT_INPUT: - # act_in_ptrs = ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :] * stride_cn - act_in_ptrs = ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :] - tl.store(act_in_ptrs, acc) - - # optional: fused activation (while the data is in shared memory) - if ACTIVATION == "gelu": - acc = gelu(acc) - elif ACTIVATION == "gelu_approx": - acc = gelu_approx(acc) - elif ACTIVATION == "squared_relu": - acc = squared_relu(acc) - # rematerialize rm and rn to save registers - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - - # write back result - # C = C + rm[:, None] * stride_cm + rn[None, :] * stride_cn - C = C + rm[:, None] * stride_cm + rn[None, :] - mask = (rm < M)[:, None] & (rn < N)[None, :] - tl.store(C, acc) - - -def triton_linear_act( - x: torch.Tensor, - weight: torch.Tensor, - bias: Optional[torch.Tensor] = None, - activation: str = "id", - save_act_input: bool = False, -) -> torch.Tensor: - """ - Compute e = activation(x @ weight.T + bias). - This wrapper kicks the `kernel_fwd` Triton kernel - :param x: input tensor - :param weight: weight matrix - :param bias: an optional bias tensor - :param activation: Activation name. Needs to be a Triton kernel. - :param act_input: an optional tensor to save the activation inputs (for backward) - :return: result tensor - """ - # if torch.is_autocast_enabled(): - # dtype = torch.get_autocast_gpu_dtype() - # x, weight, bias = [a.to(dtype=dtype) for a in [x, weight, bias]] - - assert activation in ["id", "gelu", "gelu_approx", "squared_relu"] - - batch_shape, n = x.shape[:-1], x.shape[-1] - batch_dim = batch_shape.numel() - x_reshaped = x.reshape(batch_dim, n) - - if x_reshaped.stride(0) > 1 and x_reshaped.stride(1) > 1: - x_reshaped = x_reshaped.contiguous() - if weight.stride(0) > 1 and weight.stride(1) > 1: - weight = weight.contiguous() - bias = bias.contiguous() if bias is not None else None - - assert ( - x.dtype == weight.dtype - ), f"Input and weight must have the same dtype, got {x.dtype} and {weight.dtype}" - if bias is not None: - assert ( - x.dtype == bias.dtype - ), f"Input and bias must have the same dtype, got {x.dtype} and {bias.dtype}" - assert ( - x_reshaped.shape[1] == weight.shape[1] - ), f"Incompatible dimensions: {x_reshaped.shape} - {weight.shape}" - - assert ( - bias is None or bias.shape[0] == weight.shape[0] - ), "Incompatible dimensions in between weight and bias" - - M, K = x_reshaped.shape - N, K = weight.shape - - output = torch.empty((M, N), device=x.device, dtype=x.dtype) - act_input = torch.empty_like(output) if save_act_input else None - - # 1D launch kernel where each block gets its own program. - grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),) # noqa - - kernel_fwd[grid]( - output, - act_input, - x_reshaped, - weight, # data ptrs - bias if bias is not None else x, # auto skip bias if not present - M, # shapes - N, - K, - M // 32, # key for triton cache (limit number of compilations) - N // 32, - K // 32, - stride_cm=output.stride(0), # strides - # stride_cn=output.stride(1), - stride_am=x_reshaped.stride(0), - stride_ak=x_reshaped.stride(1), - stride_bk=weight.stride(1), - stride_bn=weight.stride(0), - BIAS=bias is not None, # optional fused bias - SAVE_ACT_INPUT=save_act_input, # optional save activation inputs - ACTIVATION=activation, # optional fused activation - A_ROWMAJOR=x_reshaped.stride(1) == 1, - B_COLMAJOR=weight.stride(1) == 1, - GROUP_M=8, # speed optimization: group the programs - ) - - if not save_act_input: - return output.reshape(*batch_shape, output.shape[-1]) - else: - return ( - output.reshape(*batch_shape, output.shape[-1]), - act_input.reshape(*batch_shape, act_input.shape[-1]), - ) - - -@triton.autotune( - configs=[ - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - # good for int8 - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=4, - num_warps=4, - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - ] - + get_configs_io_bound(), - key=["CACHE_KEY_M", "CACHE_KEY_N", "CACHE_KEY_K"], - prune_configs_by={ - "early_config_prune": early_config_prune, - "perf_model": estimate_matmul_time, - "top_k": 10, - }, -) -@triton.heuristics( - { - "EVEN_K": lambda args: args["K"] % (args["BLOCK_K"] * args["SPLIT_K"]) == 0, - } -) -@triton.jit -def kernel_bwd( - C, # Pointers to matrices - ACT_INPUT, - A, - B, - # Matrix dimensions - M, - N, - K, - CACHE_KEY_M, - CACHE_KEY_N, - CACHE_KEY_K, - # The stride variables represent how much to increase the ptr by when moving by 1 - # element in a particular dimension. E.g. stride_am is how much to increase a_ptr - # by to get the element one row down (A has M rows) - stride_cm, - # stride_cn, # Assume that stride_cn == 1 - stride_am, - stride_ak, - stride_bk, - stride_bn, - # Meta-parameters - BLOCK_M: tl.constexpr, - GROUP_M: tl.constexpr, - BLOCK_N: tl.constexpr, - BLOCK_K: tl.constexpr, - # split k not used, not performant with activation, kept because early_config_prune is expecting it - SPLIT_K: tl.constexpr, - EVEN_K: tl.constexpr, - ACTIVATION: tl.constexpr, -): - - """ - Kernel for computing Out = activation(A x W + C) - - Input has shape (M, K) - - Weight has shape (K, N) - - Output has shape (M, N) - - ActInputs (optional) has shape (M, N) - 'ActInputs' optionally saves the A x W + C intermediate for backward computations - This kernel will consolidate over K - """ - - pid = tl.program_id(axis=0) - - grid_m = (M + BLOCK_M - 1) // BLOCK_M - grid_n = (N + BLOCK_N - 1) // BLOCK_N - # re-order program ID for better L2 performance - width = GROUP_M * grid_n - group_id = pid // width - group_size = min(grid_m - group_id * GROUP_M, GROUP_M) - pid_m = group_id * GROUP_M + (pid % group_size) - pid_n = (pid % width) // (group_size) - - # now compute the block that each program will go through - # rm (resp. rn) denotes a range of indices - # for rows (resp. col) of C - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - # trick to avoid masking on M and N axis - ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M) - rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N) - rk = tl.arange(0, BLOCK_K) - - A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak) - B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn) - - acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) - - for k in range(K, 0, -BLOCK_K): - if EVEN_K: - a = tl.load(A) - b = tl.load(B) - else: - a = tl.load(A, mask=rk[None, :] < k, other=0.0) - b = tl.load(B, mask=rk[:, None] < k, other=0.0) - acc += tl.dot(a, b) - - A += BLOCK_K * stride_ak - B += BLOCK_K * stride_bk - - # optional: fused activation (while the data is in shared memory) - if ACTIVATION != "id": - act_in_ptrs = ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :] - act_input = tl.load(act_in_ptrs).to(acc.dtype) - if ACTIVATION == "gelu": - acc *= gelu_grad(act_input) - elif ACTIVATION == "gelu_approx": - acc *= gelu_approx_grad(act_input) - elif ACTIVATION == "squared_relu": - acc *= squared_relu_grad(act_input) - - # rematerialize rm and rn to save registers - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - - # write back result - C = C + rm[:, None] * stride_cm + rn[None, :] - mask = (rm < M)[:, None] & (rn < N)[None, :] - tl.store(C, acc, mask=mask) - - -def triton_dgrad_act( - grad_output: torch.Tensor, - weight: torch.Tensor, - activation: str = "id", - act_input: Optional[torch.Tensor] = None, -) -> torch.Tensor: - """ - Compute e = activation(grad_output @ weight + bias). - This wrapper kicks the `kernel_fwd` Triton kernel - :param grad_output: input tensor - :param weight: weight matrix - :param activation: Activation name. Needs to be a Triton kernel. - :param act_input: an optional tensor to save the activation inputs (for backward) - :return: result tensor - """ - assert activation in ["id", "gelu", "gelu_approx", "squared_relu"] - - batch_shape, n = grad_output.shape[:-1], grad_output.shape[-1] - batch_dim = batch_shape.numel() - grad_output_reshaped = grad_output.reshape(batch_dim, n) - - if grad_output_reshaped.stride(0) > 1 and grad_output_reshaped.stride(1) > 1: - grad_output_reshaped = grad_output_reshaped.contiguous() - if weight.stride(0) > 1 and weight.stride(1) > 1: - weight = weight.contiguous() - - assert ( - grad_output.dtype == weight.dtype - ), f"grad_output and weight must have the same dtype, got {grad_output.dtype} and {weight.dtype}" - assert ( - grad_output_reshaped.shape[1] == weight.shape[0] - ), f"Incompatible dimensions: {grad_output_reshaped.shape} - {weight.shape}" - if activation != "id": - assert act_input is not None, f"act_input is required for activation {activation}" - - # M, N, K in bwd are different from M, N, K in fwd - M, K = grad_output_reshaped.shape - K, N = weight.shape - - grad_input = torch.empty((M, N), device=grad_output.device, dtype=grad_output.dtype) - - # 1D launch kernel where each block gets its own program. - grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),) # noqa - - kernel_bwd[grid]( - grad_input, - act_input, - grad_output_reshaped, - weight, # data ptrs - M, # shapes - N, - K, - M // 32, # key for triton cache (limit number of compilations) - N // 32, - K // 32, - stride_cm=grad_input.stride(0), # strides - # stride_cn=grad_input.stride(1), - stride_am=grad_output_reshaped.stride(0), - stride_ak=grad_output_reshaped.stride(1), - stride_bk=weight.stride(0), - stride_bn=weight.stride(1), - ACTIVATION=activation, # optional fused activation - GROUP_M=8, # speed optimization: group the programs - ) - - return grad_input.reshape(*batch_shape, grad_input.shape[-1]) diff --git a/build/torch211-cxx11-cu130-aarch64-linux/ops/triton/mlp.py b/build/torch211-cxx11-cu130-aarch64-linux/ops/triton/mlp.py deleted file mode 100644 index 059f4f8a5e174c1f4824e43d313fca18eaa799b8..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu130-aarch64-linux/ops/triton/mlp.py +++ /dev/null @@ -1,149 +0,0 @@ -# The triton fused matmul + sqrelu is faster for fp16 but slower for bf16, compared -# to naive implementation. -import fused_dense_lib as fused_dense_cuda -import torch -import torch.nn as nn -import torch.nn.functional as F - -from flash_attn.utils.torch import custom_fwd, custom_bwd -from flash_attn.ops.activations import sqrelu_bwd, sqrelu_fwd -from flash_attn.ops.triton.linear import triton_dgrad_act, triton_linear_act - - -class FusedDenseSqreluDenseFunc(torch.autograd.Function): - @staticmethod - @custom_fwd - def forward(ctx, x, weight1, bias1, weight2, bias2, checkpoint_lvl=0): - """checkpoint_lvl: - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute act_input and gelu_out in the bwd - """ - if torch.is_autocast_enabled(): - dtype = torch.get_autocast_gpu_dtype() - x, weight1, bias1, weight2, bias2 = [ - a.to(dtype=dtype) for a in [x, weight1, bias1, weight2, bias2] - ] - is_bf16 = x.dtype == torch.bfloat16 - assert checkpoint_lvl in [0, 1, 2] - x = x.contiguous() - weight1 = weight1.contiguous() - bias1 = bias1.contiguous() - weight2 = weight2.contiguous() - bias2 = bias2.contiguous() - batch_shape, n = x.shape[:-1], x.shape[-1] - batch_dim = batch_shape.numel() - if is_bf16: - act_input = fused_dense_cuda.linear_bias_forward( - x.reshape(batch_dim, n), weight1, bias1 - ) - output1 = sqrelu_fwd(act_input) - else: - save_act_input = checkpoint_lvl != 2 - result = triton_linear_act( - x.reshape(batch_dim, n), - weight1, - bias1, - activation="squared_relu", - save_act_input=save_act_input, - ) - if save_act_input: - output1, act_input = result - else: - output1 = result - output2 = fused_dense_cuda.linear_bias_forward(output1, weight2, bias2) - ctx.checkpoint_lvl = checkpoint_lvl - if checkpoint_lvl == 0: - ctx.save_for_backward(x, weight1, bias1, weight2, act_input, output1) - elif checkpoint_lvl == 1: - ctx.save_for_backward(x, weight1, bias1, weight2, act_input) - elif checkpoint_lvl == 2: - ctx.save_for_backward(x, weight1, bias1, weight2) - return output2.reshape(*batch_shape, output2.shape[-1]) - - @staticmethod - @custom_bwd - def backward(ctx, grad_output): - grad_output = grad_output.contiguous() - checkpoint_lvl = ctx.checkpoint_lvl - x, weight1, bias1, weight2, *rest = ctx.saved_tensors - batch_shape, n = x.shape[:-1], x.shape[-1] - batch_dim = batch_shape.numel() - is_bf16 = x.dtype == torch.bfloat16 - if checkpoint_lvl == 0: - act_input, output1 = rest - elif checkpoint_lvl == 1: - (act_input,) = rest - output1 = sqrelu_fwd(act_input) - elif checkpoint_lvl == 2: - if is_bf16: - act_input = fused_dense_cuda.linear_bias_forward( - x.reshape(batch_dim, n), weight1, bias1 - ) - output1 = sqrelu_fwd(act_input) - else: - output1, act_input = triton_linear_act( - x.reshape(batch_dim, n), - weight1, - bias1, - activation="squared_relu", - save_act_input=True, - ) - - if is_bf16: - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad(output1, grad_output) - grad_output1 = grad_output @ weight2 - grad_act_input = sqrelu_bwd(grad_output1, act_input) - grad_input, grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_backward( - x.reshape(batch_dim, n), weight1, grad_act_input - ) - else: - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad(output1, grad_output) - grad_act_input = triton_dgrad_act( - grad_output, weight2, activation="squared_relu", act_input=act_input - ) - grad_input, grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_backward( - x.reshape(batch_dim, n), weight1, grad_act_input - ) - return grad_input.reshape_as(x), grad_weight1, grad_bias1, grad_weight2, grad_bias2, None - - -fused_dense_sqrelu_dense_function = FusedDenseSqreluDenseFunc.apply - - -class FusedDenseSqreluDense(nn.Module): - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - bias1=True, - bias2=True, - checkpoint_lvl=0, - device=None, - dtype=None, - ): - """ - checkpoint_lvl (increasing lvl means slower but more memory saving): - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute gelu_in and gelu_out in the bwd - """ - assert checkpoint_lvl in [0, 1, 2] - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features * 4 - assert bias1 == True, "DenseSqreluDense module without bias is currently not supported" - assert bias2 == True, "DenseSqreluDense module without bias is currently not supported" - self.checkpoint_lvl = checkpoint_lvl - self.fc1 = nn.Linear(in_features, hidden_features, bias=bias1, **factory_kwargs) - self.fc2 = nn.Linear(hidden_features, out_features, bias=bias2, **factory_kwargs) - - def forward(self, x): - assert x.is_cuda - return fused_dense_sqrelu_dense_function( - x, self.fc1.weight, self.fc1.bias, self.fc2.weight, self.fc2.bias, self.checkpoint_lvl - ) diff --git a/build/torch211-cxx11-cu130-aarch64-linux/ops/triton/rotary.py b/build/torch211-cxx11-cu130-aarch64-linux/ops/triton/rotary.py deleted file mode 100644 index 6ad3ab412a084270120167ad32b67ad1720f2c37..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu130-aarch64-linux/ops/triton/rotary.py +++ /dev/null @@ -1,186 +0,0 @@ -# Copyright (c) 2025, Tri Dao. -# As of 2025-04-23, we require triton >= 3.0 - -from typing import Optional, Union - -import torch - -import triton -import triton.language as tl - - -@triton.jit -def rotary_kernel( - OUT, # Pointers to matrices - X, - COS, - SIN, - CU_SEQLENS, - SEQLEN_OFFSETS, # this could be int or a pointer - # Matrix dimensions - seqlen, - nheads, - seqlen_ro, - # strides - stride_out_batch, - stride_out_seqlen, - stride_out_nheads, - stride_out_headdim, - stride_x_batch, - stride_x_seqlen, - stride_x_nheads, - stride_x_headdim, - # Meta-parameters - # We want ROTARY_DIM to be constexpr, otherwise the triton compiler doesn't know that - # the mask is constant every 8 elements, and it will generate LDG.16 instead of LDG.128 - ROTARY_DIM: tl.constexpr, - IS_SEQLEN_OFFSETS_TENSOR: tl.constexpr, - IS_VARLEN: tl.constexpr, - INTERLEAVED: tl.constexpr, - CONJUGATE: tl.constexpr, - BLOCK_H: tl.constexpr, - BLOCK_M: tl.constexpr, -): - BLOCK_K: tl.constexpr = triton.next_power_of_2(ROTARY_DIM) - ROTARY_DIM_HALF = ROTARY_DIM // 2 - pid_head = tl.program_id(axis=0) - pid_m = tl.program_id(axis=1) - pid_batch = tl.program_id(axis=2) - - if not IS_VARLEN: - X = X + pid_batch * stride_x_batch - OUT = OUT + pid_batch * stride_out_batch - else: - start_idx = tl.load(CU_SEQLENS + pid_batch) - seqlen = tl.load(CU_SEQLENS + pid_batch + 1) - start_idx - X = X + start_idx * stride_x_seqlen - OUT = OUT + start_idx * stride_out_seqlen - - if pid_m * BLOCK_M >= seqlen: - return - - rh = pid_head * BLOCK_H + tl.arange(0, BLOCK_H) - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - if not IS_SEQLEN_OFFSETS_TENSOR: - rm_cs = rm + SEQLEN_OFFSETS - else: - rm_cs = rm + tl.load(SEQLEN_OFFSETS + pid_batch) - - rk_half = tl.arange(0, BLOCK_K // 2) - COS = COS + (rm_cs[:, None] * ROTARY_DIM_HALF + rk_half[None, :]) - SIN = SIN + (rm_cs[:, None] * ROTARY_DIM_HALF + rk_half[None, :]) - mask_cs = (rm_cs[:, None] < seqlen_ro) & (rk_half[None, :] < ROTARY_DIM_HALF) - cos = tl.load(COS, mask=mask_cs, other=1.0).to(tl.float32) - sin = tl.load(SIN, mask=mask_cs, other=0.0).to(tl.float32) - if CONJUGATE: - sin = -sin - - if not INTERLEAVED: - # Load the 1st and 2nd halves of X, do calculation, then store to 1st and 2nd halves of OUT - X = X + (rh[:, None, None] * stride_x_nheads + rm[None, :, None] * stride_x_seqlen + rk_half[None, None, :] * stride_x_headdim) - OUT = OUT + (rh[:, None, None] * stride_out_nheads + rm[None, :, None] * stride_out_seqlen + rk_half[None, None, :] * stride_out_headdim) - mask = (rh[:, None, None] < nheads) & (rm[None, :, None] < seqlen) & (rk_half[None, None, :] < ROTARY_DIM_HALF) - x0 = tl.load(X, mask=mask, other=0.0).to(tl.float32) - x1 = tl.load(X + ROTARY_DIM_HALF * stride_x_headdim, mask=mask, other=0.0,).to(tl.float32) - o0 = x0 * cos - x1 * sin - o1 = x0 * sin + x1 * cos - tl.store(OUT, o0, mask=mask) - tl.store(OUT + ROTARY_DIM_HALF * stride_out_headdim, o1, mask=mask) - else: - rk = tl.arange(0, BLOCK_K) - X = X + (rh[:, None, None] * stride_x_nheads + rm[None, :, None] * stride_x_seqlen + rk[None, None, :] * stride_x_headdim) - OUT = OUT + (rh[:, None, None] * stride_out_nheads + rm[None, :, None] * stride_out_seqlen + rk[None, None, :] * stride_out_headdim) - mask = (rh[:, None, None] < nheads) & (rm[None, :, None] < seqlen) & (rk[None, None, :] < ROTARY_DIM) - x = tl.load(X, mask=mask, other=0.0).to(tl.float32) - x0, x1 = tl.split(tl.reshape(x, [BLOCK_H, BLOCK_M, BLOCK_K // 2, 2])) - o0 = x0 * cos - x1 * sin - o1 = x0 * sin + x1 * cos - o = tl.reshape(tl.join(o0, o1), [BLOCK_H, BLOCK_M, BLOCK_K]) - tl.store(OUT, o, mask=mask) - - -def apply_rotary( - x: torch.Tensor, - cos: torch.Tensor, - sin: torch.Tensor, - seqlen_offsets: Union[int, torch.Tensor] = 0, - cu_seqlens: Optional[torch.Tensor] = None, - max_seqlen: Optional[int] = None, - interleaved=False, - inplace=False, - conjugate=False, -) -> torch.Tensor: - """ - Arguments: - x: (batch, seqlen, nheads, headdim) if cu_seqlens is None - else (total_seqlen, nheads, headdim). - cos: (seqlen_ro, rotary_dim / 2) - sin: (seqlen_ro, rotary_dim / 2) - seqlen_offsets: integer or integer tensor of size (batch,) - cu_seqlens: (batch + 1,) or None - max_seqlen: int - Returns: - y: (batch, seqlen, nheads, headdim) - """ - is_varlen = cu_seqlens is not None - if not is_varlen: - batch, seqlen, nheads, headdim = x.shape - else: - assert max_seqlen is not None, "If cu_seqlens is passed in, then max_seqlen must be passed" - total_seqlen, nheads, headdim = x.shape - batch_p_1 = cu_seqlens.shape[0] - batch = batch_p_1 - 1 - seqlen = max_seqlen - seqlen_ro, rotary_dim = cos.shape - assert sin.shape == cos.shape - rotary_dim *= 2 - assert rotary_dim <= headdim, "rotary_dim must be <= headdim" - assert headdim <= 256, "Only support headdim <= 256" - assert seqlen_ro >= seqlen, "seqlen_ro must be >= seqlen" - - cos, sin = cos.contiguous(), sin.contiguous() - if isinstance(seqlen_offsets, torch.Tensor): - assert seqlen_offsets.shape == (batch,) - assert seqlen_offsets.dtype in [torch.int32, torch.int64] - seqlen_offsets = seqlen_offsets.contiguous() - else: - assert seqlen_offsets + seqlen <= seqlen_ro - - output = torch.empty_like(x) if not inplace else x - if rotary_dim < headdim and not inplace: - output[..., rotary_dim:].copy_(x[..., rotary_dim:]) - - grid = lambda META: (triton.cdiv(nheads, META["BLOCK_H"]), triton.cdiv(seqlen, META["BLOCK_M"]), batch) # noqa - BLOCK_M = 8 if rotary_dim <= 128 else 4 - - # Need this, otherwise Triton tries to launch from cuda:0 and we get - # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?) - device_ctx = torch.cuda.device(x.device.index) if x.device.type == 'cuda' else torch.xpu.device(x.device.index) - with device_ctx: - torch.library.wrap_triton(rotary_kernel)[grid]( - output, # data ptrs - x, - cos, - sin, - cu_seqlens, - seqlen_offsets, - seqlen, # shapes - nheads, - seqlen_ro, - output.stride(0) if not is_varlen else 0, # batch_strides if not varlen else 0 - output.stride(-3), # seqlen_stride or total_seqlen_stride - output.stride(-2), # nheads_stride - output.stride(-1), # headdim_stride - x.stride(0) if not is_varlen else 0, # batch_strides if not varlen else 0 - x.stride(-3), # seqlen stride or total_seqlen_stride - x.stride(-2), # nheads stride - x.stride(-1), # headdim stride - rotary_dim, - isinstance(seqlen_offsets, torch.Tensor), - is_varlen, - interleaved, - conjugate, - BLOCK_M=BLOCK_M, - BLOCK_H=2, - ) - return output diff --git a/build/torch211-cxx11-cu130-x86_64-linux/__init__.py b/build/torch211-cxx11-cu130-x86_64-linux/__init__.py deleted file mode 100644 index ecc2f9d896b6c93f90b0a1499856dc0612177422..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu130-x86_64-linux/__init__.py +++ /dev/null @@ -1,393 +0,0 @@ -from typing import Optional, List -import torch -from ._ops import ops as flash_attn_ops -from .flash_attn_interface import ( - flash_attn_func, - flash_attn_kvpacked_func, - flash_attn_qkvpacked_func, - flash_attn_varlen_func, - flash_attn_varlen_kvpacked_func, - flash_attn_varlen_qkvpacked_func, - flash_attn_with_kvcache, -) - - -def fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - return_softmax: bool = False, - gen: Optional[torch.Generator] = None, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention. - - Args: - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Optional output tensor, same shape as q - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - return_softmax: Whether to return softmax weights - gen: Optional random number generator - - Returns: - List of tensors: [output, softmax_lse, (softmax if return_softmax)] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.fwd( - q, - k, - v, - out, - alibi_slopes, - p_dropout, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - gen, - ) - - -def varlen_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - out: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - block_table: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - max_seqlen_q: int = 0, - max_seqlen_k: int = 0, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - zero_tensors: bool = False, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - return_softmax: bool = False, - gen: Optional[torch.Generator] = None, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention with variable sequence lengths. - - Args: - q: Query tensor of shape [total_q, num_heads, head_size] - k: Key tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - v: Value tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1] - cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1] - out: Optional output tensor of shape [total_q, num_heads, head_size] - seqused_k: Optional tensor specifying how many keys to use per batch element [batch_size] - leftpad_k: Optional left padding for keys of shape [batch_size] - block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq] - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - max_seqlen_q: Maximum sequence length for queries - max_seqlen_k: Maximum sequence length for keys - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - zero_tensors: Whether to zero tensors before computation - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - return_softmax: Whether to return softmax weights - gen: Optional random number generator - - Returns: - List of tensors: [output, softmax_lse, (softmax if return_softmax)] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.varlen_fwd( - q, - k, - v, - out, - cu_seqlens_q, - cu_seqlens_k, - seqused_k, - leftpad_k, - block_table, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - p_dropout, - softmax_scale, - zero_tensors, - is_causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - gen, - ) - - -def bwd( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - p_dropout: float = 0.0, - 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, - gen: Optional[torch.Generator] = None, - rng_state: Optional[torch.Tensor] = None, -) -> List[torch.Tensor]: - """ - Backward pass for multi-head attention. - - Args: - dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size] - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size] - softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q] - dq: Optional gradient tensor for queries, same shape as q - dk: Optional gradient tensor for keys, same shape as k - dv: Optional gradient tensor for values, same shape as v - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - deterministic: Whether to use deterministic algorithms - gen: Optional random number generator - rng_state: Optional RNG state from forward pass - - Returns: - List of tensors: [dq, dk, dv] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - alibi_slopes, - p_dropout, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - deterministic, - gen, - rng_state, - ) - - -def varlen_bwd( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - max_seqlen_q: int = 0, - max_seqlen_k: int = 0, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - zero_tensors: bool = False, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - deterministic: bool = False, - gen: Optional[torch.Generator] = None, - rng_state: Optional[torch.Tensor] = None, -) -> List[torch.Tensor]: - """ - Backward pass for multi-head attention with variable sequence lengths. - - Args: - dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size] - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size] - softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q] - cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1] - cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1] - dq: Optional gradient tensor for queries, same shape as q - dk: Optional gradient tensor for keys, same shape as k - dv: Optional gradient tensor for values, same shape as v - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - max_seqlen_q: Maximum sequence length for queries - max_seqlen_k: Maximum sequence length for keys - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - zero_tensors: Whether to zero tensors before computation - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - deterministic: Whether to use deterministic algorithms - gen: Optional random number generator - rng_state: Optional RNG state from forward pass - - Returns: - List of tensors: [dq, dk, dv] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.varlen_bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - p_dropout, - softmax_scale, - zero_tensors, - is_causal, - window_size_left, - window_size_right, - softcap, - deterministic, - gen, - rng_state, - ) - - -def fwd_kvcache( - q: torch.Tensor, - kcache: torch.Tensor, - vcache: torch.Tensor, - k: Optional[torch.Tensor] = None, - v: Optional[torch.Tensor] = None, - seqlens_k: Optional[torch.Tensor] = None, - rotary_cos: Optional[torch.Tensor] = None, - rotary_sin: Optional[torch.Tensor] = None, - cache_batch_idx: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - block_table: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - out: 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, - is_rotary_interleaved: bool = False, - num_splits: int = 1, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention with KV cache. - - Args: - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - kcache: Key cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - vcache: Value cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - k: Optional new keys tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size] - v: Optional new values tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size] - seqlens_k: Optional sequence lengths for keys of shape [batch_size] - rotary_cos: Optional rotary cosine tensor of shape [seqlen_ro, rotary_dim/2] - rotary_sin: Optional rotary sine tensor of shape [seqlen_ro, rotary_dim/2] - cache_batch_idx: Optional indices to index into the KV cache - leftpad_k: Optional left padding for keys of shape [batch_size] - block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq] - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - out: Optional output tensor, same shape as q - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - is_rotary_interleaved: Whether rotary embeddings are interleaved - num_splits: Number of splits for computation - - Returns: - List of tensors: [output, softmax_lse] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.fwd_kvcache( - q, - kcache, - vcache, - k, - v, - seqlens_k, - rotary_cos, - rotary_sin, - cache_batch_idx, - leftpad_k, - block_table, - alibi_slopes, - out, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - is_rotary_interleaved, - num_splits, - ) diff --git a/build/torch211-cxx11-cu130-x86_64-linux/_flash_attn2_cuda_f12afc9.abi3.so b/build/torch211-cxx11-cu130-x86_64-linux/_flash_attn2_cuda_f12afc9.abi3.so deleted file mode 100644 index 84acd36e2fa4a974c39313122d1a2162a105480b..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu130-x86_64-linux/_flash_attn2_cuda_f12afc9.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:7510e8726dd3b2b5cf46d24133ec10457e6de1f05e7376f9b0aa33837e98575d -size 1007144456 diff --git a/build/torch211-cxx11-cu130-x86_64-linux/_ops.py b/build/torch211-cxx11-cu130-x86_64-linux/_ops.py deleted file mode 100644 index af5b731cb04b5e6ce54e862db0501f3683c3163f..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu130-x86_64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _flash_attn2_cuda_f12afc9 -ops = torch.ops._flash_attn2_cuda_f12afc9 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_flash_attn2_cuda_f12afc9::{op_name}" diff --git a/build/torch211-cxx11-cu130-x86_64-linux/bert_padding.py b/build/torch211-cxx11-cu130-x86_64-linux/bert_padding.py deleted file mode 100644 index 3c2d35159a014a9d03aabead9e52e009168696ea..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu130-x86_64-linux/bert_padding.py +++ /dev/null @@ -1,218 +0,0 @@ -# Adapted from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/padding.py - -import torch -import torch.nn.functional as F -from einops import rearrange, repeat - - -class IndexFirstAxis(torch.autograd.Function): - @staticmethod - def forward(ctx, input, indices): - ctx.save_for_backward(indices) - assert input.ndim >= 2 - ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] - second_dim = other_shape.numel() - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - # return input[indices] - return torch.gather( - rearrange(input, "b ... -> b (...)"), 0, repeat(indices, "z -> z d", d=second_dim) - ).reshape(-1, *other_shape) - - @staticmethod - def backward(ctx, grad_output): - (indices,) = ctx.saved_tensors - assert grad_output.ndim >= 2 - other_shape = grad_output.shape[1:] - grad_output = rearrange(grad_output, "b ... -> b (...)") - grad_input = torch.zeros( - [ctx.first_axis_dim, grad_output.shape[1]], - device=grad_output.device, - dtype=grad_output.dtype, - ) - # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. - # grad_input[indices] = grad_output - grad_input.scatter_(0, repeat(indices, "z -> z d", d=grad_output.shape[1]), grad_output) - return grad_input.reshape(ctx.first_axis_dim, *other_shape), None - - -index_first_axis = IndexFirstAxis.apply - - -class IndexPutFirstAxis(torch.autograd.Function): - @staticmethod - def forward(ctx, values, indices, first_axis_dim): - ctx.save_for_backward(indices) - assert indices.ndim == 1 - assert values.ndim >= 2 - output = torch.zeros( - first_axis_dim, *values.shape[1:], device=values.device, dtype=values.dtype - ) - # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. - output[indices] = values - # output.scatter_(0, repeat(indices, 'z -> z d', d=values.shape[1]), values) - return output - - @staticmethod - def backward(ctx, grad_output): - (indices,) = ctx.saved_tensors - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - grad_values = grad_output[indices] - # grad_values = torch.gather(grad_output, 0, repeat(indices, 'z -> z d', d=grad_output.shape[1])) - return grad_values, None, None - - -index_put_first_axis = IndexPutFirstAxis.apply - - -class IndexFirstAxisResidual(torch.autograd.Function): - @staticmethod - def forward(ctx, input, indices): - ctx.save_for_backward(indices) - assert input.ndim >= 2 - ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] - second_dim = other_shape.numel() - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - output = input[indices] - # We don't want to reshape input (b ... -> b (...)) since it could change the channel_last - # memory format to channel_first. In other words, input might not be contiguous. - # If we don't detach, Pytorch complains about output being a view and is being modified inplace - return output, input.detach() - - @staticmethod - def backward(ctx, grad_output, grad_residual): - (indices,) = ctx.saved_tensors - assert grad_output.ndim >= 2 - other_shape = grad_output.shape[1:] - assert grad_residual.shape[1:] == other_shape - grad_input = grad_residual - # grad_input[indices] += grad_output - indices = indices.reshape(indices.shape[0], *((1,) * (grad_output.ndim - 1))) - indices = indices.expand_as(grad_output) - grad_input.scatter_add_(0, indices, grad_output) - return grad_input.reshape(ctx.first_axis_dim, *other_shape), None - - -index_first_axis_residual = IndexFirstAxisResidual.apply - - -def unpad_input(hidden_states, attention_mask, unused_mask=None): - """ - Arguments: - hidden_states: (batch, seqlen, ...) - attention_mask: (batch, seqlen), bool / int, 1 means valid and 0 means not valid. - unused_mask: (batch, seqlen), bool / int, 1 means the element is allocated but unused. - Return: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens selected in attention_mask + unused_mask. - indices: (total_nnz), the indices of masked tokens from the flattened input sequence. - cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. - max_seqlen_in_batch: int - seqused: (batch), returns the number of tokens selected in attention_mask + unused_mask. - """ - all_masks = (attention_mask + unused_mask) if unused_mask is not None else attention_mask - seqlens_in_batch = all_masks.sum(dim=-1, dtype=torch.int32) - used_seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) - indices = torch.nonzero(all_masks.flatten(), as_tuple=False).flatten() - max_seqlen_in_batch = seqlens_in_batch.max().item() - cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) - # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the - # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim - # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to - # index with integer indices. Moreover, torch's index is a bit slower than it needs to be, - # so we write custom forward and backward to make it a bit faster. - return ( - index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices), - indices, - cu_seqlens, - max_seqlen_in_batch, - used_seqlens_in_batch, - ) - - -def unpad_input_for_concatenated_sequences(hidden_states, attention_mask_in_length): - """ - Supports concatenating short samples in one sequence. The attention_mask_in_length is utilized to mask other short samples. It helps efficient training of variant lengths-based samples (e.g., the supervised fine-tuning task in large language model). - The motivation for this function is explained [here](https://github.com/Dao-AILab/flash-attention/issues/432#issuecomment-1668822286). - - For example, if batch = 3 and seqlen = 6, the attention_mask_in_length is: - ``` - [ - [2, 3, 0, 0, 0, 0], - [3, 2, 0, 0, 0, 0], - [6, 0, 0, 0, 0, 0] - ] - ``` - , which refers to the 3D-attention mask: - ``` - [ - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [0, 0, 1, 0, 0, 0], - [0, 0, 1, 1, 0, 0], - [0, 0, 1, 1, 1, 0], - [0, 0, 0, 0, 0, 1] - ], - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [0, 0, 0, 1, 0, 0], - [0, 0, 0, 1, 1, 0], - [0, 0, 0, 0, 0, 1] - ], - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [1, 1, 1, 1, 0, 0], - [1, 1, 1, 1, 1, 0], - [1, 1, 1, 1, 1, 1] - ] - ] - ```. - - Arguments: - hidden_states: (batch, seqlen, ...) - attention_mask_in_length: (batch, seqlen), int, a nonzero number (e.g., 1, 2, 3, etc.) means length of concatenated sequence in b-th batch, and 0 means none. - Return: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. - indices: (total_nnz), the indices of non-masked tokens from the flattened input sequence. - cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. - max_seqlen_in_batch: int - """ - length = attention_mask_in_length.sum(dim=-1) - seqlen = attention_mask_in_length.size(-1) - attention_mask_2d = torch.arange(seqlen, device=length.device, dtype=length.dtype).expand(len(length), seqlen) < length.unsqueeze(1) - real_indices_idx = torch.nonzero(attention_mask_in_length.flatten(), as_tuple=False).flatten() - seqlens_in_batch = attention_mask_in_length.flatten()[real_indices_idx] - indices = torch.nonzero(attention_mask_2d.flatten(), as_tuple=False).flatten() - max_seqlen_in_batch = seqlens_in_batch.max().item() - cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) - # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the - # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim - # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to - # index with integer indices. Moreover, torch's index is a bit slower than it needs to be, - # so we write custom forward and backward to make it a bit faster. - return ( - index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices), - indices, - cu_seqlens, - max_seqlen_in_batch, - ) - - -def pad_input(hidden_states, indices, batch, seqlen): - """ - Arguments: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. - indices: (total_nnz), the indices that represent the non-masked tokens of the original padded input sequence. - batch: int, batch size for the padded sequence. - seqlen: int, maximum sequence length for the padded sequence. - Return: - hidden_states: (batch, seqlen, ...) - """ - dim = hidden_states.shape[-1] - # output = torch.zeros((batch * seqlen), dim, device=hidden_states.device, dtype=hidden_states.dtype) - # output[indices] = hidden_states - output = index_put_first_axis(hidden_states, indices, batch * seqlen) - return rearrange(output, "(b s) ... -> b s ...", b=batch) diff --git a/build/torch211-cxx11-cu130-x86_64-linux/flash_attn2/__init__.py b/build/torch211-cxx11-cu130-x86_64-linux/flash_attn2/__init__.py deleted file mode 100644 index a9b2672c1cd85b74c1b3ded0fc0b2100e1aeac23..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu130-x86_64-linux/flash_attn2/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -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_interface.py b/build/torch211-cxx11-cu130-x86_64-linux/flash_attn_interface.py deleted file mode 100644 index b014e569dc9812a9c840d1043f2cebe767b4d477..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu130-x86_64-linux/flash_attn_interface.py +++ /dev/null @@ -1,1631 +0,0 @@ -# Copyright (c) 2023, Tri Dao. - -from typing import Optional, Sequence, Tuple, Union - -import torch -import torch.nn as nn -import os - -# # isort: off -# # We need to import the CUDA kernels after importing torch -# USE_TRITON_ROCM = os.getenv("FLASH_ATTENTION_TRITON_AMD_ENABLE", "FALSE") == "TRUE" -# if USE_TRITON_ROCM: -# from .flash_attn_triton_amd import interface_fa as flash_attn -# else: -# import flash_attn_2_cuda as flash_attn - - -from . import _ops -from ._ops import ops as flash_attn -from ._ops import add_op_namespace_prefix - -# # isort: on - -def maybe_contiguous(x): - return x.contiguous() if x is not None and x.stride(-1) != 1 else x - - -def _get_device(): - if torch.xpu.is_available(): - return "xpu" - elif torch.cuda.is_available(): - return "cuda" - else: - return "cpu" - - -def _get_block_size_n(device, head_dim, is_dropout, is_causal): - assert head_dim <= 256 - - if device.type == "xpu": - if head_dim <= 96: - return 64 - elif head_dim <= 128: - return 32 - elif head_dim <= 256: - return 64 - else: - return 32 - - # This should match the block sizes in the CUDA kernel - major, minor = torch.cuda.get_device_capability(device) - is_sm8x = major == 8 and minor > 0 # Only include sm86 and sm89, exclude sm80 (A100) - is_sm80 = major == 8 and minor == 0 - is_sm90 = major == 9 and minor == 0 - if head_dim <= 32: - return 128 - if head_dim <= 64: - return 128 if not is_dropout else 64 - elif head_dim <= 96: - return 64 - elif head_dim <= 128: - if is_sm8x: - return 64 if (not is_dropout and is_causal) else 32 - else: - return 64 if not is_dropout else 32 - elif head_dim <= 192: - return 64 - elif head_dim <= 224: - return 64 - elif head_dim <= 256: - return 64 - - -def round_multiple(x, m): - return (x + m - 1) // m * m - - -# torch.compile() support is only enabled for pytorch >= 2.4 -# The reason for this is that we are using the new custom_op and register_fake -# APIs, which support inplace modification of inputs in the function itself -if torch.__version__ >= "2.4.0": - _torch_custom_op_wrapper = torch.library.custom_op - _torch_register_fake_wrapper = torch.library.register_fake -else: - def noop_custom_op_wrapper(name, fn=None, /, *, mutates_args, device_types=None, schema=None): - def wrap(func): - return func - if fn is None: - return wrap - return fn - def noop_register_fake_wrapper(op, fn=None, /, *, lib=None, _stacklevel=1): - def wrap(func): - return func - if fn is None: - return wrap - return fn - _torch_custom_op_wrapper = noop_custom_op_wrapper - _torch_register_fake_wrapper = noop_register_fake_wrapper - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_forward"), mutates_args=(), device_types=_get_device()) -def _flash_attn_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - return_softmax: bool -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - out, softmax_lse, S_dmask, rng_state = flash_attn.fwd( - q, - k, - v, - None, - alibi_slopes, - dropout_p, - softmax_scale, - causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - None, - ) - return out, softmax_lse, S_dmask, rng_state - - -@_torch_register_fake_wrapper(add_op_namespace_prefix("_flash_attn_forward")) -def _flash_attn_forward_fake( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - return_softmax: bool -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - batch_size, seqlen_q, num_heads, head_size = q.shape - seqlen_k = k.shape[1] - out = torch.empty_like(q) - softmax_lse = torch.empty((batch_size, num_heads, seqlen_q), dtype=torch.float32, device=q.device, layout=q.layout) - p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) - if return_softmax: - p = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128), round_multiple(seqlen_k, 128)), dtype=q.dtype, device=q.device, layout=q.layout) - rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) - - return out, softmax_lse, p, rng_state - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_forward = _ops.ops._flash_attn_forward -else: - _wrapped_flash_attn_forward = _flash_attn_forward - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_varlen_forward"), mutates_args=(), device_types=_get_device()) -def _flash_attn_varlen_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - alibi_slopes: Optional[torch.Tensor] = None, - return_softmax: bool = False, - block_table: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - out, softmax_lse, S_dmask, rng_state = flash_attn.varlen_fwd( - q, - k, - v, - None, - cu_seqlens_q, - cu_seqlens_k, - seqused_k, - leftpad_k, - block_table, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - zero_tensors, - causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - None, - ) - # if out.isnan().any() or softmax_lse.isnan().any(): - # breakpoint() - return out, softmax_lse, S_dmask, rng_state - - -@_torch_register_fake_wrapper(add_op_namespace_prefix("_flash_attn_varlen_forward")) -def _flash_attn_varlen_forward_fake( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - alibi_slopes: Optional[torch.Tensor] = None, - return_softmax: bool = False, - block_table: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - paged_kv = block_table is not None - batch_size = cu_seqlens_q.numel() - 1 - total_q, num_heads, _ = q.shape - - out = torch.empty_like(q) - softmax_lse = torch.empty((num_heads, total_q), dtype=torch.float32, device=q.device, layout=q.layout) - p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) - seqlen_q_rounded = round_multiple(max_seqlen_q, 128) - seqlen_k_rounded = round_multiple(max_seqlen_k, 128) - if return_softmax: - p = torch.empty((batch_size, num_heads, seqlen_q_rounded, seqlen_k_rounded), dtype=q.dtype, device=q.device, layout=q.layout) - rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) - return out, softmax_lse, p, rng_state - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_varlen_forward = _ops.ops._flash_attn_varlen_forward -else: - _wrapped_flash_attn_varlen_forward = _flash_attn_varlen_forward - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_backward"), mutates_args=("dq", "dk", "dv"), device_types=_get_device()) -def _flash_attn_backward( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, -) -> 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)] - ( - dq, - dk, - dv, - softmax_d, - ) = flash_attn.bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - alibi_slopes, - dropout_p, - softmax_scale, - causal, - window_size_left, - window_size_right, - softcap, - deterministic, - None, - rng_state, - ) - return softmax_d - - -@_torch_register_fake_wrapper(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, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, -) -> torch.Tensor: - dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] - if dq is None: - dq = torch.empty_like(q) - if dk is None: - dk = torch.empty_like(k) - if dv is None: - dv = torch.empty_like(v) - batch_size, seqlen_q, num_heads, _ = q.shape - softmax_d = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128)), device=q.device, dtype=torch.float32) - - return softmax_d - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_backward = _ops.ops._flash_attn_backward -else: - _wrapped_flash_attn_backward = _flash_attn_backward - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_varlen_backward"), mutates_args=("dq", "dk", "dv"), device_types=_get_device()) -def _flash_attn_varlen_backward( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> 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)] - ( - dq, - dk, - dv, - softmax_d, - ) = flash_attn.varlen_bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - zero_tensors, - causal, - window_size_left, - window_size_right, - softcap, - deterministic, - None, - rng_state, - ) - # if dk.isnan().any() or dk.isnan().any() or dv.isnan().any() or softmax_d.isnan().any(): - # breakpoint() - return softmax_d - - -@_torch_register_fake_wrapper(add_op_namespace_prefix("_flash_attn_varlen_backward")) -def _flash_attn_varlen_backward_fake( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> torch.Tensor: - dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] - batch_size = cu_seqlens_q.numel() - 1 - total_q, num_heads, _ = q.shape - - if dq is None: - dq = torch.empty_like(q) - if dk is None: - dk = torch.empty_like(k) - if dv is None: - dv = torch.empty_like(v) - softmax_d = torch.empty((num_heads, total_q + 128 * batch_size), device=q.device, dtype=torch.float32) - - return softmax_d - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_varlen_backward = _ops.ops._flash_attn_varlen_backward -else: - _wrapped_flash_attn_varlen_backward = _flash_attn_varlen_backward - - -class FlashAttnQKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and qkv.requires_grad - if softmax_scale is None: - softmax_scale = qkv.shape[-1] ** (-0.5) - q, k, v = qkv[:, :, 0].detach(), qkv[:, :, 1].detach(), qkv[:, :, 2].detach() - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) - dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dqkv[:, :, 0], - dqkv[:, :, 1], - dqkv[:, :, 2], - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dqkv = dqkv[..., : dout.shape[-1]] # We could have padded the head dimension - return dqkv, None, None, None, None, None, None, None, None, None - - -class FlashAttnVarlenQKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - cu_seqlens, - max_seqlen, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and qkv.requires_grad - if softmax_scale is None: - softmax_scale = qkv.shape[-1] ** (-0.5) - q, k, v = qkv[:, 0].detach(), qkv[:, 1].detach(), qkv[:, 2].detach() - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens, - cu_seqlens, - max_seqlen, - max_seqlen, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=None, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens, rng_state) - ctx.dropout_p = dropout_p - ctx.max_seqlen = max_seqlen - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens, rng_state = ctx.saved_tensors - qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) - dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dqkv[:, 0], - dqkv[:, 1], - dqkv[:, 2], - cu_seqlens, - cu_seqlens, - ctx.max_seqlen, - ctx.max_seqlen, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - 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 FlashAttnKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - kv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, kv] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - k, v = kv[:, :, 0].detach(), kv[:, :, 1].detach() - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - dq = torch.empty_like(q) - kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) - dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dkv[:, :, 0], - dkv[:, :, 1], - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dkv = dkv[..., : dout.shape[-1]] - return dq, dkv, None, None, None, None, None, None, None, None, None - - -class FlashAttnVarlenKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, kv] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - k, v = kv[:, 0].detach(), kv[:, 1].detach() - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=None, - ) - if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state - ) - ctx.dropout_p = dropout_p - 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.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors - dq = torch.empty_like(q) - kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) - dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dkv[:, 0], - dkv[:, 1], - cu_seqlens_q, - cu_seqlens_k, - ctx.max_seqlen_q, - ctx.max_seqlen_k, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dkv = dkv[..., : dout.shape[-1]] - return dq, dkv, None, None, None, None, None, None, None, None, None, None, None, None, None - - -class FlashAttnFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - k, - v, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, k, v] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dk = dk[..., : dout.shape[-1]] - dv = dv[..., : dout.shape[-1]] - return dq, dk, dv, 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, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - block_table, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, k, v] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=block_table, - ) - if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state - ) - ctx.dropout_p = dropout_p - 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.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors - dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - ctx.max_seqlen_q, - ctx.max_seqlen_k, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dk = dk[..., : dout.shape[-1]] - dv = dv[..., : dout.shape[-1]] - return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None - - -def flash_attn_qkvpacked_func( - qkv, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # <=0.0 means deactivate - alibi_slopes=None, - deterministic=False, - 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, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_kvpacked_func( - q, - kv, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, -): - """dropout_p should be set to 0.0 during evaluation - If 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 K, V. - 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) - kv: (batch_size, seqlen, 2, 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. - softcap: float. Anything > 0 activates softcapping 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). - 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 FlashAttnKVPackedFunc.apply( - q, - kv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_func( - q, - k, - v, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - 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). - 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 FlashAttnFunc.apply( - q, - k, - v, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_qkvpacked_func( - qkv, - cu_seqlens, - max_seqlen, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - 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_varlen_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_varlen_kvpacked_func and flash_attn_varlen_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: (total, 3, nheads, headdim), where total = total number of tokens in the batch. - cu_seqlens: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into qkv. - max_seqlen: int. Maximum sequence length in the batch. - 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenQKVPackedFunc.apply( - qkv, - cu_seqlens, - max_seqlen, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_kvpacked_func( - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, -): - """dropout_p should be set to 0.0 during evaluation - If 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 K, V. - 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: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch. - kv: (total_k, 2, nheads_k, headdim), where total_k = total number of key tokens in the batch. - cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into q. - cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into kv. - max_seqlen_q: int. Maximum query sequence length in the batch. - max_seqlen_k: int. Maximum key sequence length in the batch. - 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 + 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenKVPackedFunc.apply( - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_func( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, - block_table=None, -): - """dropout_p should be set to 0.0 during evaluation - Supports multi-query and grouped-query attention (MQA/GQA) by passing in K, V 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: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch. - k: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. - v: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. - cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into q. - cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into kv. - max_seqlen_q: int. Maximum query sequence length in the batch. - max_seqlen_k: int. Maximum key sequence length in the batch. - 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 + 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenFunc.apply( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - block_table, - False if q.device.type == "cpu" else torch.is_grad_enabled(), - ) - - -def flash_attn_with_kvcache( - q, - k_cache, - v_cache, - k=None, - v=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, - block_table: Optional[torch.Tensor] = None, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - rotary_interleaved=True, - alibi_slopes=None, - num_splits=0, - 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 block_table, - or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_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) if there's no block_table, - or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_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). Similar to k. - 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. - block_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). - 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. - 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" - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-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) - cache_batch_idx = maybe_contiguous(cache_batch_idx) - block_table = maybe_contiguous(block_table) - out, softmax_lse = flash_attn.fwd_kvcache( - q, - k_cache, - v_cache, - k, - v, - cache_seqlens, - rotary_cos, - rotary_sin, - cache_batch_idx, - cache_leftpad, - block_table, - alibi_slopes, - None, - softmax_scale, - causal, - window_size[0], - window_size[1], - softcap, - rotary_interleaved, - num_splits, - ) - return (out, softmax_lse) if return_softmax_lse else out diff --git a/build/torch211-cxx11-cu130-x86_64-linux/layers/__init__.py b/build/torch211-cxx11-cu130-x86_64-linux/layers/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/build/torch211-cxx11-cu130-x86_64-linux/layers/patch_embed.py b/build/torch211-cxx11-cu130-x86_64-linux/layers/patch_embed.py deleted file mode 100644 index 05562f8e8bcdb58e947c6f402a49eacd2d031871..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu130-x86_64-linux/layers/patch_embed.py +++ /dev/null @@ -1,67 +0,0 @@ -# We use the same API as https://github.com/rwightman/pytorch-image-models/blob/v0.6.11/timm/models/layers/patch_embed.py -# But we use nn.Linear instead of Conv2d and it's about 8x faster. - -from functools import partial - -import torch.nn as nn -from einops import rearrange -from torch import _assert -from torch.nn.modules.utils import _pair - -try: - from flash_attn.ops.fused_dense import FusedDense -except ImportError: - FusedDense = None - - -class PatchEmbed(nn.Module): - """2D Image to Patch Embedding""" - - def __init__( - self, - img_size=224, - patch_size=16, - in_chans=3, - embed_dim=768, - norm_layer=None, - flatten=True, - bias=True, - fused_bias_fc=False, - ): - super().__init__() - img_size = _pair(img_size) - patch_size = _pair(patch_size) - self.img_size = img_size - self.patch_size = patch_size - self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1]) - self.num_patches = self.grid_size[0] * self.grid_size[1] - self.flatten = flatten - if fused_bias_fc and FusedDense is None: - raise ImportError("fused_dense is not installed") - - linear_cls = nn.Linear if not fused_bias_fc or not bias else FusedDense - self.proj = linear_cls(in_chans * patch_size[0] * patch_size[1], embed_dim, bias=bias) - self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() - - def forward(self, x): - _, _, H, W = x.shape - _assert( - H == self.img_size[0], - f"Input image height ({H}) doesn't match model ({self.img_size[0]}).", - ) - _assert( - W == self.img_size[1], - f"Input image width ({W}) doesn't match model ({self.img_size[1]}).", - ) - x = self.proj( - rearrange( - x, - "b c (h p1) (w p2) -> b h w (c p1 p2)", - p1=self.patch_size[0], - p2=self.patch_size[1], - ) - ) - if self.flatten: - x = rearrange(x, "b h w c -> b (h w) c") - x = self.norm(x) - return x diff --git a/build/torch211-cxx11-cu130-x86_64-linux/layers/rotary.py b/build/torch211-cxx11-cu130-x86_64-linux/layers/rotary.py deleted file mode 100644 index d1bfc21fc7de1dd287e8f382847b194a48075981..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu130-x86_64-linux/layers/rotary.py +++ /dev/null @@ -1,483 +0,0 @@ -# Copyright (c) 2025, Tri Dao - -import math -from functools import partial -from typing import Optional, Tuple, Union - -import torch -from torch import Tensor - -from einops import rearrange, repeat -# from flash_attn.ops.triton.rotary import apply_rotary -from ..ops.triton.rotary import apply_rotary - - -def rotate_half(x, interleaved=False): - if not interleaved: - x1, x2 = x.chunk(2, dim=-1) - return torch.cat((-x2, x1), dim=-1) - else: - x1, x2 = x[..., ::2], x[..., 1::2] - return rearrange(torch.stack((-x2, x1), dim=-1), "... d two -> ... (d two)", two=2) - - -def apply_rotary_emb_torch(x, cos, sin, interleaved=False): - """ - x: (batch_size, seqlen, nheads, headdim) - cos, sin: (seqlen, rotary_dim / 2) or (batch_size, seqlen, rotary_dim / 2) - """ - ro_dim = cos.shape[-1] * 2 - assert ro_dim <= x.shape[-1] - cos = repeat(cos, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)") - sin = repeat(sin, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)") - return torch.cat( - [x[..., :ro_dim] * cos + rotate_half(x[..., :ro_dim], interleaved) * sin, x[..., ro_dim:]], - dim=-1, - ) - - -class ApplyRotaryEmb(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x, - cos, - sin, - interleaved=False, - inplace=False, - seqlen_offsets: Union[int, Tensor] = 0, - cu_seqlens: Optional[Tensor] = None, - max_seqlen: Optional[int] = None, - ): - out = apply_rotary( - x, - cos, - sin, - seqlen_offsets=seqlen_offsets, - cu_seqlens=cu_seqlens, - max_seqlen=max_seqlen, - interleaved=interleaved, - inplace=inplace, - ) - if isinstance(seqlen_offsets, int): - ctx.save_for_backward(cos, sin, cu_seqlens) # Can't save int with save_for_backward - ctx.seqlen_offsets = seqlen_offsets - else: - ctx.save_for_backward(cos, sin, cu_seqlens, seqlen_offsets) - ctx.seqlen_offsets = None - ctx.interleaved = interleaved - ctx.inplace = inplace - ctx.max_seqlen = max_seqlen - return out if not inplace else x - - @staticmethod - def backward(ctx, do): - seqlen_offsets = ctx.seqlen_offsets - if seqlen_offsets is None: - cos, sin, cu_seqlens, seqlen_offsets = ctx.saved_tensors - else: - cos, sin, cu_seqlens = ctx.saved_tensors - dx = apply_rotary( - do, - cos, - sin, - seqlen_offsets=seqlen_offsets, - cu_seqlens=cu_seqlens, - max_seqlen=ctx.max_seqlen, - interleaved=ctx.interleaved, - inplace=ctx.inplace, - conjugate=True, - ) - return dx, None, None, None, None, None, None, None - - -def apply_rotary_emb( - x, - cos, - sin, - interleaved=False, - inplace=False, - seqlen_offsets: Union[int, Tensor] = 0, - cu_seqlens: Optional[Tensor] = None, - max_seqlen: Optional[int] = None, -): - """ - Arguments: - x: (batch_size, seqlen, nheads, headdim) if cu_seqlens is None - else (total_seqlen, nheads, headdim) - cos, sin: (seqlen_rotary, rotary_dim / 2) - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead - of 1st half and 2nd half (GPT-NeoX style). - inplace: if True, apply rotary embedding in-place. - seqlen_offsets: (batch_size,) or int. Each sequence in x is shifted by this amount. - Most commonly used in inference when we have KV cache. - cu_seqlens: (batch + 1,) or None - max_seqlen: int - Return: - out: (batch_size, seqlen, nheads, headdim) if cu_seqlens is None - else (total_seqlen, nheads, headdim) - rotary_dim must be <= headdim - Apply rotary embedding to the first rotary_dim of x. - """ - return ApplyRotaryEmb.apply( - x, cos, sin, interleaved, inplace, seqlen_offsets, cu_seqlens, max_seqlen - ) - - -# For backward compatibility -apply_rotary_emb_func = apply_rotary_emb - - -def _apply_rotary_emb_qkv( - qkv, - cos, - sin, - cos_k=None, - sin_k=None, - interleaved=False, - inplace=False, - conjugate=False, - seqlen_offsets: Union[int, Tensor] = 0, - num_heads_q: Optional[int] = None, -): - apply_rotary_fn = partial( - apply_rotary, - interleaved=interleaved, - inplace=inplace, - conjugate=conjugate, - seqlen_offsets=seqlen_offsets - ) - if cos_k is None and sin_k is None and qkv.is_contiguous(): - # Call 1 kernel instead of 2 kernels - # We need qkv to be contiguous so that when we reshape to combine (3, nheads) - # dimensions, we get the same tensor - if qkv.dim() == 5: - batch, seqlen, three, nheads, headdim = qkv.shape - assert three == 3 - # qk = rearrange(qkv[:, :, :2], "b s t h d -> b s (t h) d") - qk = qkv[:, :, :2].reshape(batch, seqlen, -1, headdim) - qk = apply_rotary_fn(qk, cos, sin) - else: - assert qkv.dim() == 4 - assert num_heads_q is not None - num_heads_k = (qkv.shape[2] - num_heads_q) // 2 - assert qkv.shape[2] == num_heads_q + 2 * num_heads_k - qk = qkv[:, :, :num_heads_q + num_heads_k] - qk = apply_rotary_fn(qk, cos, sin) - if not inplace: - if qkv.dim() == 5: - qkv = torch.cat([rearrange(qk, "b s (t h) d -> b s t h d", t=2), qkv[:, :, 2:]], dim=2) - else: - qkv = torch.cat([qk, qkv[:, :, num_heads_q + num_heads_k :]], dim=2) - else: - cos_k = cos if cos_k is None else cos_k - sin_k = sin if sin_k is None else sin_k - if qkv.dim() == 5: - batch, seqlen, three, nheads, headdim = qkv.shape - assert three == 3 - q, k = qkv[:, :, 0], qkv[:, :, 1] - else: - assert qkv.dim() == 4 - assert num_heads_q is not None - num_heads_k = (qkv.shape[2] - num_heads_q) // 2 - assert qkv.shape[2] == num_heads_q + 2 * num_heads_k - q, k = qkv[:, :, :num_heads_q], qkv[:, :, num_heads_q : num_heads_q + num_heads_k] - q = apply_rotary_fn(q, cos, sin) - k = apply_rotary_fn(k, cos_k, sin_k) - if not inplace: - if qkv.dim() == 5: - qkv = torch.stack([q, k, qkv[:, :, 2]], dim=2) - else: - qkv = torch.cat([q, k, qkv[:, :, num_heads_q + num_heads_k:]], dim=2) - return qkv - - -class ApplyRotaryEmbQKV_(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - cos, - sin, - cos_k=None, - sin_k=None, - interleaved=False, - seqlen_offsets: Union[int, torch.Tensor] = 0, - num_heads_q: Optional[int] = None, - ): - # apply_rotary_emb_qkv_inplace( - qkv = _apply_rotary_emb_qkv( - qkv, cos, sin, cos_k, sin_k, interleaved=interleaved, inplace=True, - seqlen_offsets=seqlen_offsets, num_heads_q=num_heads_q, - ) - if isinstance(seqlen_offsets, int): - ctx.save_for_backward(cos, sin, cos_k, sin_k) - ctx.seqlen_offsets = seqlen_offsets - else: - ctx.save_for_backward(cos, sin, cos_k, sin_k, seqlen_offsets) - ctx.seqlen_offsets = None - ctx.interleaved = interleaved - ctx.num_heads_q = num_heads_q - return qkv - - @staticmethod - def backward(ctx, dqkv): - seqlen_offsets = ctx.seqlen_offsets - if seqlen_offsets is None: - cos, sin, cos_k, sin_k, seqlen_offsets = ctx.saved_tensors - else: - cos, sin, cos_k, sin_k = ctx.saved_tensors - dqkv = _apply_rotary_emb_qkv( - dqkv, cos, sin, cos_k, sin_k, interleaved=ctx.interleaved, inplace=True, - seqlen_offsets=seqlen_offsets, num_heads_q=ctx.num_heads_q, conjugate=True, - ) - return dqkv, None, None, None, None, None, None, None - - -def apply_rotary_emb_qkv_( - qkv, - cos, - sin, - cos_k=None, - sin_k=None, - interleaved=False, - seqlen_offsets: Union[int, torch.Tensor] = 0, - num_heads_q: Optional[int] = None, -): - """ - Arguments: - qkv: (batch_size, seqlen, 3, nheads, headdim) or (batch_size, seqlen, num_heads_q + 2 * num_heads_k, headdim). - If qkv has shape (batch_size, seqlen, num_heads_q + 2 * num_heads_k, headdim) (e.g. MQA / GQA), - then num_heads_q must be provided. - cos, sin: (seqlen, rotary_dim / 2) - cos_k, sin_k: (seqlen, rotary_dim / 2), optional - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead of - 1st half and 2nd half (GPT-NeoX style). - seqlen_offsets: (batch_size,) or int. Each sequence in Q and K is shifted by this amount. - Most commonly used in inference when we have KV cache. - Return: - qkv: (batch_size, seqlen, 3, nheads, headdim) or (batch_size, seqlen, num_heads_q + 2 * num_heads_k, headdim) - rotary_dim must be <= headdim - Apply rotary embedding *inplace* to the first rotary_dim of Q and K. - """ - return ApplyRotaryEmbQKV_.apply( - qkv, cos, sin, cos_k, sin_k, interleaved, seqlen_offsets, num_heads_q - ) - - -class ApplyRotaryEmbKV_(torch.autograd.Function): - - @staticmethod - def forward(ctx, kv, cos, sin, interleaved=False, seqlen_offsets: Union[int, torch.Tensor] = 0): - batch, seqlen, two, nheads, headdim = kv.shape - assert two == 2 - k = kv[:, :, 0] - apply_rotary( - k, cos, sin, seqlen_offsets=seqlen_offsets, interleaved=interleaved, inplace=True - ) - if isinstance(seqlen_offsets, int): - ctx.save_for_backward(cos, sin) # Can't save int with save_for_backward - ctx.seqlen_offsets = seqlen_offsets - else: - ctx.save_for_backward(cos, sin, seqlen_offsets) - ctx.seqlen_offsets = None - ctx.interleaved = interleaved - return kv - - @staticmethod - def backward(ctx, dkv): - seqlen_offsets = ctx.seqlen_offsets - if seqlen_offsets is None: - cos, sin, seqlen_offsets = ctx.saved_tensors - else: - cos, sin = ctx.saved_tensors - apply_rotary( - dkv[:, :, 0], - cos, - sin, - seqlen_offsets=seqlen_offsets, - interleaved=ctx.interleaved, - inplace=True, - conjugate=True, - ) - return dkv, None, None, None, None - - -apply_rotary_emb_kv_ = ApplyRotaryEmbKV_.apply - - -def apply_rotary_emb_kv_( - kv, - cos, - sin, - interleaved=False, - seqlen_offsets: Union[int, torch.Tensor] = 0, -): - """ - Arguments: - kv: (batch_size, seqlen, 2, nheads, headdim) - cos, sin: (seqlen, rotary_dim / 2) - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead of - 1st half and 2nd half (GPT-NeoX style). - seqlen_offsets: (batch_size,) or int. Each sequence in Q and K is shifted by this amount. - Most commonly used in inference when we have KV cache. - Return: - kv: (batch_size, seqlen, 2, nheads, headdim) - rotary_dim must be <= headdim - Apply rotary embedding *inplace* to the first rotary_dim of K. - """ - return ApplyRotaryEmbKV_.apply(kv, cos, sin, interleaved, seqlen_offsets) - - -class RotaryEmbedding(torch.nn.Module): - """ - The rotary position embeddings from RoFormer_ (Su et. al). - A crucial insight from the method is that the query and keys are - transformed by rotation matrices which depend on the relative positions. - - Other implementations are available in the Rotary Transformer repo_ and in - GPT-NeoX_, GPT-NeoX was an inspiration - - .. _RoFormer: https://arxiv.org/abs/2104.09864 - .. _repo: https://github.com/ZhuiyiTechnology/roformer - .. _GPT-NeoX: https://github.com/EleutherAI/gpt-neox - - If scale_base is not None, this implements XPos (Sun et al., https://arxiv.org/abs/2212.10554). - A recommended value for scale_base is 512: https://github.com/HazyResearch/flash-attention/issues/96 - Reference: https://github.com/sunyt32/torchscale/blob/main/torchscale/component/xpos_relative_position.py - """ - - def __init__( - self, - dim: int, - base=10000.0, - interleaved=False, - scale_base=None, - device=None, - ): - """ - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead - of 1st half and 2nd half (GPT-NeoX style). - """ - super().__init__() - self.dim = dim - self.base = float(base) - # Generate and save the inverse frequency buffer (non trainable) - inv_freq = self._compute_inv_freq(device) - self.register_buffer("inv_freq", inv_freq, persistent=False) - self.interleaved = interleaved - self.scale_base = scale_base - scale = ( - (torch.arange(0, dim, 2, device=device, dtype=torch.float32) + 0.4 * dim) / (1.4 * dim) - if scale_base is not None - else None - ) - self.register_buffer("scale", scale, persistent=False) - - self._seq_len_cached = 0 - self._cos_cached = None - self._sin_cached = None - self._cos_k_cached = None - self._sin_k_cached = None - - def _compute_inv_freq(self, device=None): - return 1.0 / ( - self.base - ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim) - ) - - def _update_cos_sin_cache(self, seqlen, device=None, dtype=None): - # Reset the tables if the sequence length has changed, - # if we're on a new device (possibly due to tracing for instance), - # or if we're switching from inference mode to training - if ( - seqlen > self._seq_len_cached - or self._cos_cached is None - or self._cos_cached.device != device - or self._cos_cached.dtype != dtype - or (self.training and self._cos_cached.is_inference()) - ): - self._seq_len_cached = seqlen - # We want fp32 here, not self.inv_freq.dtype, since the model could be loaded in bf16 - # And the output of arange can be quite large, so bf16 would lose a lot of precision. - t = torch.arange(seqlen, device=device, dtype=torch.float32) - # We want fp32 here as well since inv_freq will be multiplied with t, and the output - # will be large. Having it in bf16 will lose a lot of precision and cause the - # cos & sin output to change significantly. - # We want to recompute self.inv_freq if it was not loaded in fp32 - if self.inv_freq.dtype != torch.float32: - inv_freq = self._compute_inv_freq(device=device) - else: - inv_freq = self.inv_freq - # Don't do einsum, it converts fp32 to bf16 under AMP - # freqs = torch.einsum("i,j->ij", t, self.inv_freq) - freqs = torch.outer(t, inv_freq) - if self.scale is None: - self._cos_cached = torch.cos(freqs).to(dtype) - self._sin_cached = torch.sin(freqs).to(dtype) - else: - power = ( - torch.arange(seqlen, dtype=self.scale.dtype, device=self.scale.device) - - seqlen // 2 - ) / self.scale_base - scale = self.scale.to(device=power.device) ** rearrange(power, "s -> s 1") - # We want the multiplication by scale to happen in fp32 - self._cos_cached = (torch.cos(freqs) * scale).to(dtype) - self._sin_cached = (torch.sin(freqs) * scale).to(dtype) - self._cos_k_cached = (torch.cos(freqs) / scale).to(dtype) - self._sin_k_cached = (torch.sin(freqs) / scale).to(dtype) - - def forward( - self, - qkv: torch.Tensor, - kv: Optional[torch.Tensor] = None, - seqlen_offset: Union[int, torch.Tensor] = 0, - max_seqlen: Optional[int] = None, - num_heads_q: Optional[int] = None, - ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: - """ - qkv: (batch, seqlen, 3, nheads, headdim) or (batch, seqlen, num_heads_q + 2 * num_heads_k, headdim) - if kv is none, else it's just q of shape (batch, seqlen, nheads, headdim). - If qkv has shape (batch, seqlen, num_heads_q + 2 * num_heads_k, headdim) (e.g. MQA / GQA), - then num_heads_q must be provided. - kv: (batch, seqlen, 2, nheads, headdim) - seqlen_offset: (batch_size,) or int. Each sequence in x is shifted by this amount. - Most commonly used in inference when we have KV cache. - If it's a tensor of shape (batch_size,), then to update the cos / sin cache, one - should pass in max_seqlen, which will update the cos / sin cache up to that length. - Apply rotary embedding *inplace* to qkv and / or kv. - """ - seqlen = qkv.shape[1] - if max_seqlen is not None: - self._update_cos_sin_cache(max_seqlen, device=qkv.device, dtype=qkv.dtype) - elif isinstance(seqlen_offset, int): - self._update_cos_sin_cache(seqlen + seqlen_offset, device=qkv.device, dtype=qkv.dtype) - if kv is None: - return apply_rotary_emb_qkv_( - qkv, - self._cos_cached, - self._sin_cached, - self._cos_k_cached if self.scale is not None else None, - self._sin_k_cached if self.scale is not None else None, - interleaved=self.interleaved, - seqlen_offsets=seqlen_offset, - num_heads_q=num_heads_q, - ) - else: - q = qkv - q = apply_rotary_emb_func( - q, - self._cos_cached, - self._sin_cached, - interleaved=self.interleaved, - inplace=True, - seqlen_offsets=seqlen_offset, - ) - kv = apply_rotary_emb_kv_( - kv, - self._cos_cached if self.scale is None else self._cos_k_cached, - self._sin_cached if self.scale is None else self._sin_k_cached, - interleaved=self.interleaved, - seqlen_offsets=seqlen_offset, - ) - return q, kv diff --git a/build/torch211-cxx11-cu130-x86_64-linux/metadata.json b/build/torch211-cxx11-cu130-x86_64-linux/metadata.json deleted file mode 100644 index 6f92e389db28cb8794b3a36cea6454910bf83143..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu130-x86_64-linux/metadata.json +++ /dev/null @@ -1,42 +0,0 @@ -{ - "name": "flash-attn2", - "id": "_flash_attn2_cuda_f12afc9", - "version": 1, - "license": "BSD-3-Clause", - "python-depends": [], - "backend": { - "type": "cuda", - "archs": [ - "10.0", - "12.0", - "8.0", - "9.0" - ] - }, - "digest": { - "algorithm": "sha256", - "files": { - "__init__.py": "Cd9MEdd/WUYnBmfiQpC4SyPqAxrTilTjvsvswNju4zc=", - "_flash_attn2_cuda_f12afc9.abi3.so": "dRDocm3TsrXPRtJBM+wQRX5t4fBec3b5sKozg36YV10=", - "_ops.py": "GQbmTTeSOU+7Q2nI68Hz+1RvEogyfEFccx1ubUYw+2s=", - "bert_padding.py": "gF1EmsdJ+HpQ86MRQ4VxDw+Sb/RVISdQALdNnoByHlw=", - "flash_attn2/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY=", - "flash_attn_interface.py": "W7a9bHR2ly6QL3aKqsBqxWA7M7ZWYcZPmtsPJu6rAfY=", - "layers/__init__.py": "47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU=", - "layers/patch_embed.py": "H58CgME/qSOPTZLOG08wFgrQS1j34pvNwMPrkTj3Ek4=", - "layers/rotary.py": "rFIdNrZQFNawwLD6d7t5LE2m+M93wwDZ8EuExw+gfD8=", - "ops/__init__.py": "47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU=", - "ops/activations.py": "t5lzNg1In8LP6bKeTnyeMizwqjv27JGbJ6ylPdGvZYg=", - "ops/fused_dense.py": "BqT8VTSCGdLJAjh6WBMkCRFYrdVodPeoElWBvH6AIzA=", - "ops/layer_norm.py": "zr7NXIm+2mtEynTp1CS0fbFGI2Mqdp41dY4AfDWF6EQ=", - "ops/rms_norm.py": "XEnihcj0a4aSz4LO55m5iKGVn4HKTeKN8TIyHjuDgxI=", - "ops/triton/__init__.py": "AbpHGcgLb+kRsJGnwFEktk7uzpZOCcBY74+YBdrKVGs=", - "ops/triton/cross_entropy.py": "y56nmuOtpSNxM7sDZcUSgXKIsIEayfuvrU8bAVog/nw=", - "ops/triton/k_activations.py": "+Z3vIyO4JkqBMipKsPvhzmxljtBdIhJCsl/M+/ESqBo=", - "ops/triton/layer_norm.py": "BBQdJiqopbIS12viUqzx3SvOQCNbDSAYLJXhpyMGyC8=", - "ops/triton/linear.py": "OtRvKz8xdpl+7v3q/ZTaS9fdBt9XrzMyapgRr50uBbM=", - "ops/triton/mlp.py": "uWA+D7HyuRUsOT+4nhi3vvCDNZnNEXwYcHMQ02UREfQ=", - "ops/triton/rotary.py": "0YSrlbep0lJmzLpTiANs0slTFeGcQQqR+el6uzbSLvI=" - } - } -} \ No newline at end of file diff --git a/build/torch211-cxx11-cu130-x86_64-linux/metadata.json.sigstore b/build/torch211-cxx11-cu130-x86_64-linux/metadata.json.sigstore deleted file mode 100644 index 9adea6709ee074f9b37cecf6531e7af9afb946aa..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu130-x86_64-linux/metadata.json.sigstore +++ /dev/null @@ -1 +0,0 @@ -{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json","verificationMaterial":{"certificate":{"rawBytes":"MIIHTDCCBtKgAwIBAgIUX6ZnghPpKhD7tVc30KkJfxja6zIwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwNjI0MDEwMTIyWhcNMjYwNjI0MDExMTIyWjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAE7BJ5s3IfanTDEXxZUZoEoFs9nu3k1g+j7K4lPP2awrsvXFeI/DjT5Cw+FkVx0HS5zpNHDhNsCejUKlGglOp/36OCBfEwggXtMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUrE5CYjO6ylo8S0KWKUsE3WrC3TswHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoZjEyYWZjOTJjNDBiYmQ3NmQ5OTBkOTExZWJhYmYyOGU0OTBhN2JkZTATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoZjEyYWZjOTJjNDBiYmQ3NmQ5OTBkOTExZWJhYmYyOGU0OTBhN2JkZTAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoZjEyYWZjOTJjNDBiYmQ3NmQ5OTBkOTExZWJhYmYyOGU0OTBhN2JkZTAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKGYxMmFmYzkyYzQwYmJkNzZkOTkwZDkxMWViYWJmMjhlNDkwYTdiZGUwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMjgwMzYyNjg5MzUvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBiwYKKwYBBAHWeQIEAgR9BHsAeQB3AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABnvclYBwAAAQDAEgwRgIhAM3UT/FMzt+FdYEsh4DgKpHVcNb1WFuU03qUgyV7ExT2AiEAmPN2Lm6IESTD9caAaQrVzYpoJXP1DDwba9N5WSWs94YwCgYIKoZIzj0EAwMDaAAwZQIxAOapLP1p1FTk2060qCd7jB7q6t8YrQq51EMxNe0cAgELufm7d5rVdLOvhBhS1+uZuQIwW2h1BFAVFlWAmPfMxc5GslVCjwELBPAJG+OYRruW2X5cQ/Gzkxj9eYUv0lEzMZB5"},"tlogEntries":[{"logIndex":"1933772229","logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="},"kindVersion":{"kind":"hashedrekord","version":"0.0.1"},"integratedTime":"1782262882","inclusionPromise":{"signedEntryTimestamp":"MEUCIQCUAFSc/28JMiCkYZ5xu5/svMquWh08ZAlnAo1GDPkoQwIgGvJlQu9YNvYStkoCMTHLH4JLlQoVj4MLagzKEiWRjg4="},"inclusionProof":{"logIndex":"1811867967","rootHash":"9LMHECpLu3rg+650qhpPVxp4oWkdV/gp81gK29797ug=","treeSize":"1811867972","hashes":["qcFYAQ3AEGkFbHccRmGSeVWJn+XKNM0HKFEOkf+Bi7k=","IiSip5cbVxRP17mCQ9Sb8F+PC34lfbIr5xkxmz2LH1I=","Bqw0/SQNh6pnB9MxccrEdYBGv7+N7ceE1QFs6B4jVw8=","uvhHaYxyahbHQmnyV/6u9RYsxQaxkBtWsJhE4sdP+vE=","S6eN8brKiJ+xmpO98Z9Z06+TxjiOCeFMar7PV39sw+8=","edxgzYd42eXvAOlZZKQKqpwqeFzt3qNDebZlAz0JYEQ=","0P2BnIOxsuQi+xOcKT6E3cPsqK9mRZsSdD0lKdbql6w=","oTwvuS73/Mr8Eznxu3zK7mutJv/d67rW15lLY1XiQKs=","fqtytk2if0R+M3kWRpiPUrBEhf2z8y5Li6MeOuWmiks=","mzp9m+2At7GZdgfx0m7RwtO9YzltCZ6UhZBy4p+lUdM=","ipWWLDkqay5JAYe6P4Nrqsck/1C0C5uwqT9QbFsiY1E=","COKvK4TO+dkpqUDGbTbSKewajXzAonx1B1YaRb/ymLw=","dQlR+DXGDKXm5V1fupxaMSOwNZ9F1MybXyyxB96yj80=","JKJdxSZwKPO6QFzLTxHCrUyI5sDU/CwY57zDlA9ni6Q=","OgfnQCh7xr8xUTErIw//YAKaKMDwRJWiThvYExQvLJM=","S8Oe9ZicoJqbVuL1I0dIdTgVpLc82Cm2CMOcV83sKVs=","E1O2ilWlpFRLWGdbmJUMDZ2qbGu/eDZ12m6YjQ9/JPA=","nyv2TDlfkz0XfRwjY98zsbUbOkBZ+Tgz7PML0c2uN14=","OMfDh4k9YD0E+t4sjqJsh+vg8D1ueo1cC07exyxmr6s=","+vHKm4wNwGJ46ceyZ8mePD7Vn1SediStNHCStRoTSWQ=","5HwWjuGNDisetbiWM0cWIXM65TIM3EJlqk13ntCQlcU=","NDIRlqxO0IaGr94q86MIU7VCAN3yEEgOHS4loya5lu4=","daxmZaajRpZV+JxHiOYZhJBiSKN5ucqjh2WnGbHhirw=","DOCeoSMovIvLExkhIvisow9AuNXgeWs4ECkyR6EcqYU="],"checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n1811867972\n9LMHECpLu3rg+650qhpPVxp4oWkdV/gp81gK29797ug=\n\n— rekor.sigstore.dev wNI9ajBGAiEA5Y/xx9BY/vf8HclLN2KhGMWHvaojfJax1ukxP5BX17oCIQDiGhS5XbpGykVBgi7uNbc0heXFngI5vycSMYwxE1PJlg==\n"}},"canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiI4NjE5Y2Q0NmNjZjYwMzQzMTUyMWNhYjk0MzkxYjJjZmJlNjc1NDk4MWE2YmFmYWIzNmJlZWYyOGM4OGJkNDg0In19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FUUNIMlAyOUlpaFNPTDJuNVQ4Z2hPTjVSdEJKMTJ1cGtPUHUrUmZtdmZkOWJBQ0lRRGJad2pHYUJmTzA0Z2kreXhmU2RUbUI5TUN0NVdUdGdMQ1NRMlArNVJvWnc9PSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFVSRU5EUW5STFowRjNTVUpCWjBsVldEWmFibWRvVUhCTGFFUTNkRlpqTXpCTGEwcG1lR3BoTm5wSmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDVxU1RCTlJFVjNUVlJKZVZkb1kwNU5hbGwzVG1wSk1FMUVSWGhOVkVsNVYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVUzUWtvMWN6TkpabUZ1VkVSRldIaGFWVnB2Ulc5R2N6bHVkVE5yTVdjcmFqZExOR3dLVUZBeVlYZHljM1pZUm1WSkwwUnFWRFZEZHl0R2ExWjRNRWhUTlhwd1RraEVhRTV6UTJWcVZVdHNSMmRzVDNBdk16WlBRMEptUlhkbloxaDBUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlZ5UlRWRENsbHFUelo1Ykc4NFV6QkxWMHRWYzBVelYzSkRNMVJ6ZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOWFha1Y1V1ZkYWFrOVVTbXBPUkVKcFdXMVJNMDV0VVRWUFZFSnJDazlVUlhoYVYwcG9XVzFaZVU5SFZUQlBWRUpvVGpKS2ExcFVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMXBxUlhsWlYxcHFUMVJLYWs1RVFtbFpiVkV6VG0xUk5VOVVRbXRQVkVWNFdsZEthRmx0V1hrS1QwZFZNRTlVUW1oT01rcHJXbFJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMXBxUlhsWlYxcHFUMVJLYWs1RVFtbFpiVkV6VG0xUk5VOVVRbXNLVDFSRmVGcFhTbWhaYlZsNVQwZFZNRTlVUW1oT01rcHJXbFJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwZFplRTF0Um0wS1dYcHJlVmw2VVhkWmJVcHJUbnBhYTA5VWEzZGFSR3Q0VFZkV2FWbFhTbTFOYW1oc1RrUnJkMWxVWkdsYVIxVjNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOYW1kM1RYcFplVTVxWnpWTmVsVjJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwZDFsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTT1VKSWMwRUtaVkZDTTBGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtNTJZMnhaUW5kQlFVRlJSQXBCUldkM1VtZEphRUZOTTFWVUwwWk5lblFyUm1SWlJYTm9ORVJuUzNCSVZtTk9ZakZYUm5WVk1ETnhWV2Q1VmpkRmVGUXlRV2xGUVcxUVRqSk1iVFpKQ2tWVFZFUTVZMkZCWVZGeVZucFpjRzlLV0ZBeFJFUjNZbUU1VGpWWFUxZHpPVFJaZDBObldVbExiMXBKZW1vd1JVRjNUVVJoUVVGM1dsRkplRUZQWVhBS1RGQXhjREZHVkdzeU1EWXdjVU5rTjJwQ04zRTJkRGhaY2xGeE5URkZUWGhPWlRCalFXZEZUSFZtYlRka05YSldaRXhQZG1oQ2FGTXhLM1ZhZFZGSmR3cFhNbWd4UWtaQlZrWnNWMEZ0VUdaTmVHTTFSM05zVmtOcWQwVk1RbEJCU2tjclQxbFNjblZYTWxnMVkxRXZSM3ByZUdvNVpWbFZkakJzUlhwTldrSTFDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}],"timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyTADAgEAMIICwAYJKoZIhvcNAQcCoIICsTCCAq0CAQMxDTALBglghkgBZQMEAgEwgbcGCyqGSIb3DQEJEAEEoIGnBIGkMIGhAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgjUB1jr6Gt7e/bLnUtziWd21eQs14BVAECoc6RzQOcW0CFHdjfXJ0LOXAbvrWV17haauFulzOGA8yMDI2MDYyNDAxMDEyMlowAwIBAaAypDAwLjEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MRUwEwYDVQQDEwxzaWdzdG9yZS10c2GgADGCAdswggHXAgEBMFEwOTEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MSAwHgYDVQQDExdzaWdzdG9yZS10c2Etc2VsZnNpZ25lZAIUOhNULwyQYe68wUMvy4qOiyojiwwwCwYJYIZIAWUDBAIBoIH8MBoGCSqGSIb3DQEJAzENBgsqhkiG9w0BCRABBDAcBgkqhkiG9w0BCQUxDxcNMjYwNjI0MDEwMTIyWjAvBgkqhkiG9w0BCQQxIgQgjiYddk1X61Vnqu1NnGLa74RG0wopJCxEtDSG/7o/TA4wgY4GCyqGSIb3DQEJEAIvMX8wfTB7MHkEIIX5J7wHq2LKw7RDVsEO/IGyxog/2nq55thw2dE6zQW3MFUwPaQ7MDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAoGCCqGSM49BAMCBGcwZQIwDWVggVKc+cxuyp/W4sX1M0yxhvHh+dAVFJn2Tf0X2AFUTOHTJ5XulY9nZ6aaxv0oAjEAyN8ET1Q8/D35dfaUWrZ0StLn9okxBMZDQZxChddDgi5AT4zop/xn8zdFC0nGi2QG"}]}},"messageSignature":{"messageDigest":{"algorithm":"SHA2_256","digest":"hhnNRsz2A0MVIcq5Q5Gyz75nVJgaa6+rNr7vKMiL1IQ="},"signature":"MEQCH2P29IihSOL2n5T8ghON5RtBJ12upkOPu+Rfmvfd9bACIQDbZwjGaBfO04gi+yxfSdTmB9MCt5WTtgLCSQ2P+5RoZw=="}} \ No newline at end of file diff --git a/build/torch211-cxx11-cu130-x86_64-linux/ops/__init__.py b/build/torch211-cxx11-cu130-x86_64-linux/ops/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/build/torch211-cxx11-cu130-x86_64-linux/ops/activations.py b/build/torch211-cxx11-cu130-x86_64-linux/ops/activations.py deleted file mode 100644 index 7c09649fc41e12d5a360c5672825d8380bc7ec80..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu130-x86_64-linux/ops/activations.py +++ /dev/null @@ -1,135 +0,0 @@ -# Copied from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/model/layers/activations.py -import math - -import torch -import torch.nn as nn -import torch.nn.functional as F - -# 1/sqrt(2*pi)-> 0.3989423 -# 1/sqrt(2) -> 0.70710678 -# sqrt(2/pi) -> 0.79788456 - -# this function is tanh approximation of gelu -# actual gelu is: -# x * 0.5 * (1.0 + torch.erf(x * 0.70710678)) -@torch.jit.script -def bias_gelu(y, bias): - x = bias + y - return (x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))).to(dtype=y.dtype) - - -# gradient of tanh approximation of gelu -# gradient of actual gelu is: -# 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x) -@torch.jit.script -def bias_gelu_back(g, y, bias): - """Assume that y has shape (B, D) and bias has shape (D)""" - x = bias + y - tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)) - # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243 - ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( - 1 + tanh_out - ) - grad_y = ff * g - return grad_y.to(dtype=y.dtype), grad_y.sum(dim=(0), dtype=bias.dtype) - - -class GeLUFunction(torch.autograd.Function): - @staticmethod - # bias is an optional argument - def forward(ctx, input, bias): - ctx.save_for_backward(input, bias) - return bias_gelu(input, bias) - - @staticmethod - def backward(ctx, grad_output): - input, bias = ctx.saved_tensors - tmp = bias_gelu_back(grad_output, input, bias) - return tmp, tmp - - -bias_gelu_impl = GeLUFunction.apply - -# this function is tanh approximation of gelu -# actual gelu is: -# x * 0.5 * (1.0 + torch.erf(x * 0.70710678)) -@torch.jit.script -def gelu_fwd(x): - return (x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))).to(dtype=x.dtype) - - -# gradient of tanh approximation of gelu -# gradient of actual gelu is: -# 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x) -@torch.jit.script -def gelu_bwd(g, x): - tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)) - # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243 - ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( - 1 + tanh_out - ) - return (ff * g).to(dtype=x.dtype) - - -class FastGeLUFunction(torch.autograd.Function): - @staticmethod - # bias is an optional argument - def forward(ctx, input): - ctx.save_for_backward(input) - return gelu_fwd(input) - - @staticmethod - def backward(ctx, grad_output): - (input,) = ctx.saved_tensors - tmp = gelu_bwd(grad_output, input) - return tmp - - -fast_gelu_impl = FastGeLUFunction.apply - - -@torch.jit.script -def relu_bwd(g, x): - return torch.where(x >= 0, g, 0.0).to(dtype=x.dtype) - - -@torch.jit.script -def sqrelu_fwd(x): - r = F.relu(x) - return (r * r).to(dtype=x.dtype) - - -@torch.jit.script -def sqrelu_bwd(g, x): - return (2.0 * g * F.relu(x)).to(dtype=x.dtype) - - -swiglu_fwd_codestring = """ -template T swiglu_fwd(T x, T y) { - return float(x) * float(y) / (1.0f + ::exp(-float(x))); -} -""" -swiglu_bwd_codestring = """ -template void swiglu_bwd(T x, T y, T g, T& dx, T& dy) { - float x_sigmoid = 1.0f / (1.0f + ::exp(-float(x))); - dx = x_sigmoid * (1 + float(x) * (1.0f - x_sigmoid)) * float(g) * float(y); - dy = float(x) * x_sigmoid * float(g); -} -""" -swiglu_fwd = torch.cuda.jiterator._create_jit_fn(swiglu_fwd_codestring) -swiglu_bwd = torch.cuda.jiterator._create_multi_output_jit_fn(swiglu_bwd_codestring, num_outputs=2) - - -class SwiGLUFunction(torch.autograd.Function): - - @staticmethod - def forward(ctx, x, y): - ctx.save_for_backward(x, y) - return swiglu_fwd(x, y) - - @staticmethod - def backward(ctx, dout): - x, y = ctx.saved_tensors - return swiglu_bwd(x, y, dout) - -swiglu = SwiGLUFunction.apply diff --git a/build/torch211-cxx11-cu130-x86_64-linux/ops/fused_dense.py b/build/torch211-cxx11-cu130-x86_64-linux/ops/fused_dense.py deleted file mode 100644 index 6b4033d134e4093fe278f7b3f8c7d3128ce9f36d..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu130-x86_64-linux/ops/fused_dense.py +++ /dev/null @@ -1,688 +0,0 @@ -# Copyright (c) 2023, Tri Dao. -# Inspired by https://github.com/NVIDIA/apex/blob/master/apex/fused_dense/fused_dense.py -# We make it work with pytorch amp and with bfloat16. -# The TensorParallel linear modules are inspired by https://github.com/NVIDIA/apex/blob/master/apex/transformer/tensor_parallel/layers.py -from functools import partial -from typing import Optional - -# import fused_dense_cuda # from apex -import fused_dense_lib as fused_dense_cuda -import torch -import torch.nn as nn -import torch.nn.functional as F -from torch import Tensor -from torch.distributed import ProcessGroup - -from flash_attn.utils.torch import custom_fwd, custom_bwd -from flash_attn.ops.activations import gelu_bwd, relu_bwd, sqrelu_bwd, sqrelu_fwd -from flash_attn.utils.distributed import ( - all_gather_raw, - all_reduce, - all_reduce_raw, - reduce_scatter, - reduce_scatter_raw, -) - - -class FusedDenseFunc(torch.autograd.Function): - @staticmethod - @custom_fwd - def forward( - ctx, x, weight, bias, return_residual=False, process_group=None, sequence_parallel=True - ): - """ - If process_group is not None and sequence_parallel=True, we're doing Tensor Parallel - with sequence parallelism: we do an all_gather_raw of x before doing the matmul. - """ - ctx.compute_weight_gradient = weight.requires_grad - ctx.return_residual = return_residual - ctx.process_group = process_group - ctx.sequence_parallel = sequence_parallel - - if torch.is_autocast_enabled(): - x = x.to(dtype=torch.get_autocast_gpu_dtype()) - x = x.contiguous() - if process_group is not None and sequence_parallel: - # We want to kick off the all_gather early, before weight dtype conversion - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - else: - total_x = x - - if torch.is_autocast_enabled(): - weight = weight.to(dtype=torch.get_autocast_gpu_dtype()) - bias = bias.to(dtype=torch.get_autocast_gpu_dtype()) if bias is not None else None - weight = weight.contiguous() - if process_group is not None and sequence_parallel: - handle_x.wait() - batch_shape, n = total_x.shape[:-1], total_x.shape[-1] - batch_dim = batch_shape.numel() - # https://github.com/pytorch/pytorch/blob/5b51849b48a7dbccd297286cc0110def4706f9e7/aten/src/ATen/native/cuda/Blas.cpp#L174 - if min(batch_dim, n, *weight.shape) > 65535 * 32: - raise RuntimeError("fused_dense only supports matrix dims <= 2M") - output = F.linear(total_x, weight, bias) - if ctx.compute_weight_gradient: - ctx.save_for_backward(x, weight) - else: - ctx.save_for_backward(weight) - return output if not return_residual else (output, x) - - @staticmethod - @custom_bwd - def backward(ctx, grad_output, *args): - grad_output = grad_output.contiguous() - if ctx.return_residual: - (grad_input,) = args - grad_input = grad_input.contiguous() - process_group = ctx.process_group - sequence_parallel = ctx.sequence_parallel - if ctx.compute_weight_gradient: - x, weight = ctx.saved_tensors - if process_group is not None and sequence_parallel: - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - else: - total_x = x - else: - (weight,) = ctx.saved_tensors - total_x = None - batch_shape = grad_output.shape[:-1] - batch_dim = batch_shape.numel() - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - if ctx.needs_input_grad[0]: - if not ctx.return_residual: - grad_input = F.linear(grad_output, weight.t()) - else: - grad_input = torch.addmm( - grad_input.reshape(batch_dim, grad_input.shape[-1]), grad_output, weight - ) - grad_input = grad_input.reshape(*batch_shape, grad_input.shape[-1]) - if process_group is not None: - reduce_fn = reduce_scatter_raw if sequence_parallel else all_reduce_raw - grad_input, handle_grad_input = reduce_fn(grad_input, process_group, async_op=True) - else: - grad_input = None - if ctx.needs_input_grad[1]: - assert ctx.compute_weight_gradient - if process_group is not None and sequence_parallel: - handle_x.wait() - grad_weight, grad_bias = fused_dense_cuda.linear_bias_wgrad( - total_x.reshape(batch_dim, total_x.shape[-1]), grad_output, ctx.needs_input_grad[2] - ) - else: - grad_weight = None - grad_bias = grad_output if ctx.needs_input_grad[2] else None - if process_group is not None and ctx.needs_input_grad[0]: - handle_grad_input.wait() - return grad_input, grad_weight, grad_bias, None, None, None - - -def fused_dense_func( - x: Tensor, - weight: Tensor, - bias: Optional[Tensor] = None, - return_residual: bool = False, - process_group: Optional[ProcessGroup] = None, - sequence_parallel: bool = True, -): - dtype_eligible = x.dtype in [torch.float16, torch.bfloat16] or ( - x.dtype == torch.float32 and torch.is_autocast_enabled() - ) - if x.is_cuda and weight.is_cuda and (bias is None or bias.is_cuda) and dtype_eligible: - return FusedDenseFunc.apply( - x, weight, bias, return_residual, process_group, sequence_parallel - ) - else: - assert process_group is None - out = F.linear(x, weight, bias) - return out if not return_residual else (out, x) - - -class FusedDense(nn.Linear): - def __init__( - self, - in_features: int, - out_features: int, - bias: bool = True, - return_residual: bool = False, - device=None, - dtype=None, - ) -> None: - super().__init__(in_features, out_features, bias=bias, device=device, dtype=dtype) - self.return_residual = return_residual - - def forward(self, x, process_group=None): - """ - If process_group is not None, we're doing Tensor Parallel with sequence parallelism: - we do an all_gather of x before doing the matmul. - """ - return fused_dense_func( - x, - self.weight, - self.bias, - return_residual=self.return_residual, - process_group=process_group, - ) - - -class ColumnParallelLinear(nn.Linear): - def __init__( - self, - in_features: int, - out_features: int, - process_group: ProcessGroup, - bias: bool = True, - sequence_parallel=True, - multiple_of=1, - device=None, - dtype=None, - ) -> None: - world_size = torch.distributed.get_world_size(process_group) - if out_features % multiple_of: - raise ValueError(f"out_features ({out_features}) must be a multiple of {multiple_of}") - multiple = out_features // multiple_of - # We want to split @multiple across world_size, but it could be an uneven split - div = multiple // world_size - mod = multiple % world_size - # The first @mod ranks get @div + 1 copies, the rest get @div copies - local_multiple = div + int(torch.distributed.get_rank(process_group) < mod) - super().__init__( - in_features, local_multiple * multiple_of, bias=bias, device=device, dtype=dtype - ) - self.process_group = process_group - self.sequence_parallel = sequence_parallel - - def forward(self, x): - # If self.sequence_parallel is True, we're doing Tensor Parallel with sequence parallelism: - # we do an all_gather of x before doing the matmul. - # If not, then the input is already gathered. - return fused_dense_func( - x, - self.weight, - self.bias, - process_group=self.process_group, - sequence_parallel=self.sequence_parallel, - ) - - -class RowParallelLinear(nn.Linear): - def __init__( - self, - in_features: int, - out_features: int, - process_group: ProcessGroup, - bias: bool = True, - sequence_parallel=True, - multiple_of=1, - device=None, - dtype=None, - ) -> None: - world_size = torch.distributed.get_world_size(process_group) - rank = torch.distributed.get_rank(process_group) - if in_features % multiple_of: - raise ValueError(f"in_features ({in_features}) must be a multiple of {multiple_of}") - multiple = in_features // multiple_of - # We want to split @multiple across world_size, but it could be an uneven split - div = multiple // world_size - mod = multiple % world_size - # The first @mod ranks get @div + 1 copies, the rest get @div copies - local_multiple = div + int(torch.distributed.get_rank(process_group) < mod) - # Only rank 0 will have bias - super().__init__( - local_multiple * multiple_of, - out_features, - bias=bias and rank == 0, - device=device, - dtype=dtype, - ) - self.process_group = process_group - self.sequence_parallel = sequence_parallel - - def forward(self, x): - """ - We're doing Tensor Parallel with sequence parallelism: we do the matmul and then - a reduce_scatter of the result. - """ - out = fused_dense_func(x, self.weight, self.bias) - reduce_fn = reduce_scatter if self.sequence_parallel else all_reduce - return reduce_fn(out, self.process_group) - - -class FusedMLPFunc(torch.autograd.Function): - @staticmethod - @custom_fwd - def forward( - ctx, - x, - weight1, - bias1, - weight2, - bias2, - activation="gelu_approx", - save_pre_act=True, - return_residual=False, - checkpoint_lvl=0, - heuristic=0, - process_group=None, - sequence_parallel=True, - ): - """ - If process_group is not None and sequence_parallel=True, we're doing Tensor Parallel - with sequence parallelism: we do an all_gather of x before doing the matmul. - If sequence_parallel=False, then the input is already gathered. - - checkpoint_lvl: - 0: no recomputation in the bwd - 1: recompute gelu_out / relu_out in the bwd - 2: recompute pre_act and gelu_out / relu_out in the bwd - """ - assert -1 <= heuristic <= 4 - assert activation in ["gelu_approx", "relu", "sqrelu"] - if activation == "sqrelu": - assert heuristic == -1 - if not save_pre_act: - checkpoint_lvl = 2 - assert checkpoint_lvl in [0, 1, 2] - ctx.return_residual = return_residual - ctx.process_group = process_group - ctx.sequence_parallel = sequence_parallel - ctx.checkpoint_lvl = checkpoint_lvl - ctx.activation = activation - ctx.heuristic = heuristic - - if torch.is_autocast_enabled(): - x = x.to(dtype=torch.get_autocast_gpu_dtype()) - x = x.contiguous() - if process_group is not None and sequence_parallel: - # We want to kick off the all_gather early, before weight dtype conversion - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - else: - total_x = x - - if torch.is_autocast_enabled(): - dtype = torch.get_autocast_gpu_dtype() - weight1, weight2 = [a.to(dtype=dtype) for a in [weight1, weight2]] - bias1 = bias1.to(dtype=dtype) if bias1 is not None else None - bias2 = bias2.to(dtype=dtype) if bias2 is not None else None - weight1 = weight1.contiguous() - bias1 = bias1.contiguous() if bias1 is not None else None - weight2 = weight2.contiguous() - bias2 = bias2.contiguous() if bias2 is not None else None - if process_group is not None and sequence_parallel: - handle_x.wait() - batch_shape, n = total_x.shape[:-1], total_x.shape[-1] - batch_dim = batch_shape.numel() - # https://github.com/pytorch/pytorch/blob/5b51849b48a7dbccd297286cc0110def4706f9e7/aten/src/ATen/native/cuda/Blas.cpp#L174 - if min(batch_dim, n, *weight1.shape, *weight2.shape) > 65535 * 32: - raise RuntimeError("fused_dense only supports matrix dims <= 2M") - if heuristic == -1: - pre_act = F.linear(total_x, weight1, bias1) - activation_fn = ( - partial(F.gelu, approximate="tanh") - if activation == "gelu_approx" - else (sqrelu_fwd if activation == "sqrelu" else F.relu) - ) - with torch.jit.fuser("fuser2"): - output1 = activation_fn(pre_act) - # This is before adding bias1 - # pre_act = F.linear(total_x.reshape(batch_dim, n), weight1) - # with torch.jit.fuser('fuser2'): - # output1 = bias_gelu(pre_act, bias1) - else: - is_gelu = activation == "gelu_approx" - output1, *rest = fused_dense_cuda.linear_act_forward( - total_x.reshape(batch_dim, n), weight1, bias1, is_gelu, save_pre_act, heuristic - ) - if save_pre_act: - pre_act = rest[0] - output2 = F.linear(output1, weight2, bias2) - if checkpoint_lvl == 0 or (checkpoint_lvl == 1 and activation == "relu"): - # For RELU the pre_act is very small (just a bit-mask) so we just save it - ctx.save_for_backward(x, weight1, weight2, pre_act, output1) - elif checkpoint_lvl == 1: - ctx.save_for_backward(x, weight1, weight2, pre_act) - elif checkpoint_lvl == 2: - ctx.save_for_backward(x, weight1, weight2, bias1) - output2 = output2.reshape(*batch_shape, output2.shape[-1]) - return output2 if not return_residual else (output2, x) - - @staticmethod - @custom_bwd - def backward(ctx, grad_output, *args): - grad_output = grad_output.contiguous() - checkpoint_lvl = ctx.checkpoint_lvl - activation = ctx.activation - activation_fn = ( - partial(F.gelu, approximate="tanh") - if activation == "gelu_approx" - else (sqrelu_fwd if activation == "sqrelu" else F.relu) - ) - if ctx.return_residual: - (grad_input,) = args - grad_input = grad_input.contiguous() - process_group = ctx.process_group - sequence_parallel = ctx.sequence_parallel - x, weight1, weight2, *rest = ctx.saved_tensors - if process_group is None or not sequence_parallel: - total_x = x - batch_shape = grad_output.shape[:-1] - batch_dim = batch_shape.numel() - if checkpoint_lvl in [0, 1]: - if process_group is not None and sequence_parallel: - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - if checkpoint_lvl == 0 or (checkpoint_lvl == 1 and activation == "relu"): - pre_act, output1 = rest - elif checkpoint_lvl == 1: - (pre_act,) = rest - with torch.jit.fuser("fuser2"): - output1 = activation_fn(pre_act) - elif checkpoint_lvl == 2: - (bias1,) = rest - if process_group is not None and sequence_parallel: - total_x, _ = all_gather_raw(x, process_group) - if ctx.heuristic == -1: - pre_act = F.linear(total_x, weight1, bias1) - with torch.jit.fuser("fuser2"): - output1 = activation_fn(pre_act) - else: - output1, pre_act = fused_dense_cuda.linear_act_forward( - total_x.reshape(batch_dim, total_x.shape[-1]), - weight1, - bias1, - activation == "gelu_approx", - True, - ctx.heuristic, - ) - - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - output1 = output1.reshape(batch_dim, output1.shape[-1]) - pre_act = pre_act.reshape(batch_dim, pre_act.shape[-1]) - if ctx.needs_input_grad[3]: - grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad( - output1, grad_output, ctx.needs_input_grad[4] - ) - else: - grad_weight2 = None - grad_bias2 = grad_output if ctx.needs_input_grad[4] else None - if ctx.heuristic == -1: - # grad_pre_act = matmul_dgelu(grad_output, weight2, pre_act) - grad_output1 = F.linear(grad_output, weight2.t()) - activation_grad_fn = ( - gelu_bwd - if activation == "gelu_approx" - else (sqrelu_bwd if activation == "sqrelu" else relu_bwd) - ) - with torch.jit.fuser("fuser2"): - grad_pre_act = activation_grad_fn(grad_output1, pre_act) - else: - # The cublasLt epilogue has to compute both gelu/relu grad and bias grad, we can't - # just compute gelu/relu grad - grad_pre_act, grad_bias1 = fused_dense_cuda.bias_act_linear_dgrad_bgrad( - weight2, grad_output, pre_act, activation == "gelu_approx", ctx.heuristic - ) - if not ctx.needs_input_grad[2]: - grad_bias1 = None - if ctx.needs_input_grad[0]: - if not ctx.return_residual: - grad_input = F.linear(grad_pre_act, weight1.t()) - else: - grad_input = torch.addmm( - grad_input.reshape(batch_dim, grad_input.shape[-1]), grad_pre_act, weight1 - ) - grad_input = grad_input.reshape(*batch_shape, grad_input.shape[-1]) - if process_group is not None: - reduce_fn = reduce_scatter_raw if sequence_parallel else all_reduce_raw - grad_input, handle_grad_input = reduce_fn(grad_input, process_group, async_op=True) - else: - grad_input = None - if ctx.heuristic == -1: - if ctx.needs_input_grad[1]: - if process_group is not None and sequence_parallel and checkpoint_lvl != 2: - handle_x.wait() - grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_wgrad( - total_x.reshape(batch_dim, total_x.shape[-1]), - grad_pre_act, - ctx.needs_input_grad[2], - ) - else: - grad_weight1 = None - grad_bias1 = grad_pre_act if ctx.needs_input_grad[2] else None - else: - if ctx.needs_input_grad[1]: - if process_group is not None and sequence_parallel and checkpoint_lvl != 2: - handle_x.wait() - grad_weight1 = F.linear( - grad_pre_act.t(), total_x.reshape(batch_dim, total_x.shape[-1]).t() - ) - else: - grad_weight1 = None - if process_group is not None and ctx.needs_input_grad[0]: - handle_grad_input.wait() - return ( - grad_input, - grad_weight1, - grad_bias1, - grad_weight2, - grad_bias2, - None, - None, - None, - None, - None, - None, - None, - ) - - -def fused_mlp_func( - x: Tensor, - weight1: Tensor, - weight2: Tensor, - bias1: Optional[Tensor] = None, - bias2: Optional[Tensor] = None, - activation: str = "gelu_approx", - save_pre_act: bool = True, - return_residual: bool = False, - checkpoint_lvl: int = 0, - heuristic: int = 0, - process_group: Optional[ProcessGroup] = None, - sequence_parallel: bool = True, -): - assert activation in ["gelu_approx", "relu", "sqrelu"] - dtype_eligible = x.dtype in [torch.float16, torch.bfloat16] or ( - x.dtype == torch.float32 and torch.is_autocast_enabled() - ) - # If we save pre-activation, dimension must be divisible by 128 (relu) or 8 (gelu) - dim_eligible = not save_pre_act or (x.shape[-1] % (128 if activation == "relu" else 8) == 0) - if ( - x.is_cuda - and weight1.is_cuda - and weight2.is_cuda - and (bias1 is None or bias1.is_cuda) - and (bias2 is None or bias2.is_cuda) - and dtype_eligible - and dim_eligible - ): - return FusedMLPFunc.apply( - x, - weight1, - bias1, - weight2, - bias2, - activation, - save_pre_act, - return_residual, - checkpoint_lvl, - heuristic, - process_group, - sequence_parallel, - ) - else: - assert process_group is None - pre_act = F.linear(x, weight1, bias1) - activation_fn = ( - partial(F.gelu, approximate="tanh") - if activation == "gelu_approx" - else partial(F.relu, inplace=True) - ) - output1 = activation_fn(pre_act) - output2 = F.linear(output1, weight2, bias2) - return output2 if not return_residual else (output2, x) - - -class FusedMLP(nn.Module): - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - bias1=True, - bias2=True, - activation="gelu_approx", - return_residual=False, - checkpoint_lvl=0, - heuristic="auto", - device=None, - dtype=None, - ): - """ - If process_group is not None, we're doing Tensor Parallel with sequence parallelism: - we do an all_gather of x before doing the matmul, gelu, then matmul. - Finally we do a reduce_scatter of the output. - - checkpoint_lvl (increasing lvl means slower but more memory saving): - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute pre_act and gelu_out in the bwd - heuristic: - -1: don't fuse gemm + gelu (separate kernel) - 0..4: use this heuristic for the algo section in the fused gemm + gelu - 'auto': heuristic will be picked automatically: - For CUDA >= 11.8, we set heuristic=0 for both fp16 and bf16 for best perf. - For CUDA <= 11.7, we set heuristic=1 for fp16 and heuristic=-1 for bf16. - For H100, we set heuristic=-1 for both fp16 and bf16 as the fused cuBlasLt implementation - is slower than the unfused version. - return_residual: whether to return the input x along with the output. This is for - performance reason: for post-norm architecture, returning the input allows us - to fuse the backward of nn.Linear with the residual connection. - """ - assert checkpoint_lvl in [0, 1, 2] - assert activation in ["gelu_approx", "relu", "sqrelu"] - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features * 4 - self.activation = activation - self.return_residual = return_residual - self.checkpoint_lvl = checkpoint_lvl - self.heuristic = heuristic if activation != "sqrelu" else -1 - self.fc1 = nn.Linear(in_features, hidden_features, bias=bias1, **factory_kwargs) - self.fc2 = nn.Linear(hidden_features, out_features, bias=bias2, **factory_kwargs) - - def forward(self, x, process_group=None): - dtype = x.dtype if not torch.is_autocast_enabled() else torch.get_autocast_gpu_dtype() - if self.heuristic == "auto": - if self.activation == "gelu_approx": - if torch.cuda.get_device_capability("cuda") == (9, 0): - heuristic = -1 - else: - cuda_ver = tuple(map(int, torch.version.cuda.split("."))) - heuristic = 0 if cuda_ver >= (11, 8) else (1 if dtype == torch.float16 else -1) - else: - heuristic = 0 - else: - heuristic = self.heuristic - out = fused_mlp_func( - x, - self.fc1.weight, - self.fc2.weight, - self.fc1.bias, - self.fc2.bias, - activation=self.activation, - save_pre_act=self.training, - return_residual=self.return_residual, - checkpoint_lvl=self.checkpoint_lvl, - heuristic=heuristic, - process_group=process_group, - ) - if self.return_residual: - out, x = out - if process_group is not None: - out = reduce_scatter(out, process_group) - return out if not self.return_residual else (out, x) - - -class ParallelFusedMLP(nn.Module): - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - activation="gelu_approx", - process_group: ProcessGroup = None, - bias1=True, - bias2=True, - sequence_parallel=True, - checkpoint_lvl=0, - heuristic="auto", - device=None, - dtype=None, - ): - """ - process_group is required. We're doing Tensor Parallel with sequence parallelism: - we do an all_gather of x before doing the matmul, gelu, then matmul. - Finally we do a reduce_scatter of the output. - - checkpoint_lvl (increasing lvl means slower but more memory saving): - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute pre_act and gelu_out in the bwd - heuristic: - -1: don't fuse gemm + gelu (separate kernel) - 0..4: use this heuristic for the algo section in the fused gemm + gelu - 'auto': heuristic will be picked automatically: - For CUDA >= 11.8, we set heuristic=0 for both fp16 and bf16 for best perf. - For CUDA <= 11.7, we set heuristic=1 for fp16 and heuristic=-1 for bf16. - """ - assert checkpoint_lvl in [0, 1, 2] - assert activation in ["gelu_approx", "relu", "sqrelu"] - assert process_group is not None - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features * 4 - self.activation = activation - self.process_group = process_group - self.sequence_parallel = sequence_parallel - self.checkpoint_lvl = checkpoint_lvl - self.heuristic = heuristic if activation != "sqrelu" else -1 - self.fc1 = ColumnParallelLinear( - in_features, hidden_features, process_group, bias=bias1, **factory_kwargs - ) - self.fc2 = RowParallelLinear( - hidden_features, out_features, process_group, bias=bias2, **factory_kwargs - ) - - def forward(self, x): - dtype = x.dtype if not torch.is_autocast_enabled() else torch.get_autocast_gpu_dtype() - if self.heuristic == "auto": - if self.activation == "gelu_approx": - cuda_ver = tuple(map(int, torch.version.cuda.split("."))) - heuristic = 0 if cuda_ver >= (11, 8) else (1 if dtype == torch.float16 else -1) - else: - heuristic = 0 - else: - heuristic = self.heuristic - out = fused_mlp_func( - x, - self.fc1.weight, - self.fc2.weight, - self.fc1.bias, - self.fc2.bias, - activation=self.activation, - save_pre_act=self.training, - checkpoint_lvl=self.checkpoint_lvl, - heuristic=heuristic, - process_group=self.process_group, - sequence_parallel=self.sequence_parallel, - ) - reduce_fn = reduce_scatter if self.sequence_parallel else all_reduce - return reduce_fn(out, self.process_group) diff --git a/build/torch211-cxx11-cu130-x86_64-linux/ops/layer_norm.py b/build/torch211-cxx11-cu130-x86_64-linux/ops/layer_norm.py deleted file mode 100644 index 4b6cd798fd02844ef9cd3897f8ab95e490e638bf..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu130-x86_64-linux/ops/layer_norm.py +++ /dev/null @@ -1,800 +0,0 @@ -# Copyright (c) 2022, Tri Dao. -# Adapted from https://github.com/NVIDIA/apex/blob/master/apex/contrib/layer_norm/layer_norm.py - -import dropout_layer_norm -import torch -from torch.nn import init - - -def maybe_align(x, alignment_in_bytes=16): - """Assume that x already has last dim divisible by alignment_in_bytes""" - # TD [2023-07-04] I'm not 100% sure that clone will align the memory - # https://discuss.pytorch.org/t/how-to-ensure-that-tensor-data-ptr-is-aligned-to-16-bytes/183440 - return x if x.data_ptr() % alignment_in_bytes == 0 else x.clone() - - -def _dropout_add_layer_norm_forward( - x0, - residual, - gamma, - beta, - rowscale, - colscale, - dropout_p, - epsilon, - residual_in_fp32=False, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes""" - hidden_size = gamma.numel() - x0mat = x0.view((-1, hidden_size)) - residualmat = residual.view((-1, hidden_size)) if residual is not None else None - rowscale = rowscale.view(-1) if rowscale is not None else None - zmat, xmat, dmask, mu, rsigma = dropout_layer_norm.dropout_add_ln_fwd( - x0mat, - residualmat, - gamma, - beta, - rowscale, - colscale, - None, - None, - dropout_p, - epsilon, - 1.0, - 0, - None, - residual_in_fp32, - is_rms_norm, - ) - # dmask is None if dropout_p == 0.0 - # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype - return zmat, xmat if xmat is not None else x0mat, dmask, mu, rsigma - - -def _dropout_add_layer_norm_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - rowscale, - colscale, - dropout_p, - has_residual, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes - dx == None means that it was a post-norm architecture - (x = drop(x0) + residual was not returned in the fwd). - x0 must not be None if we have colscale. - """ - hidden_size = gamma.numel() - xmat = x.view((-1, hidden_size)) - dzmat = dz.view(xmat.shape) - dxmat = dx.view(xmat.shape) if dx is not None else None - x0mat = x0.view((-1, hidden_size)) if x0 is not None else None - rowscale = rowscale.view(-1) if rowscale is not None else None - if colscale is not None: - assert x0 is not None, "x0 is required to compute the gradient of colscale" - dx0mat, dresidualmat, dgamma, dbeta, _, _, *rest = dropout_layer_norm.dropout_add_ln_bwd( - dzmat, - dxmat, - xmat, - x0mat, - dmask, - mu, - rsigma, - gamma, - rowscale, - colscale, - None, - None, - dropout_p, - 1.0, - 0, - has_residual, - is_rms_norm, - ) - # dresidualmat is None if not has_residual - if colscale is None: - return dx0mat, dresidualmat, dgamma, dbeta - else: - dcolscale = rest[0] - return dx0mat, dresidualmat, dgamma, dbeta, dcolscale - - -def _dropout_add_layer_norm_subset_forward( - x0, - residual, - gamma, - beta, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32=False, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes""" - hidden_size = gamma.numel() - x0mat = x0.view((-1, hidden_size)) - residualmat = residual.view((-1, hidden_size)) if residual is not None else None - x0_subset = x0_subset.view(-1) if x0_subset is not None else None - out_subset = out_subset.view(-1) if out_subset is not None else None - zmat, xmat, dmask, mu, rsigma = dropout_layer_norm.dropout_add_ln_fwd( - x0mat, - residualmat, - gamma, - beta, - None, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - None, - residual_in_fp32, - is_rms_norm, - ) - # dmask is None if dropout_p == 0.0 - # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype - return zmat, xmat if xmat is not None else x0mat, dmask, mu, rsigma - - -def _dropout_add_layer_norm_subset_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - colscale, - x0_subset, - out_subset, - dropout_p, - rowscale_const, - x0_numrows, - has_residual, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes - dx == None means that it was a post-norm architecture - (x = drop(x0) + residual was not returned in the fwd). - x0 must not be None if we have colscale. - """ - hidden_size = gamma.numel() - xmat = x.view((-1, hidden_size)) - dzmat = dz.view(-1, hidden_size) - dxmat = dx.view(xmat.shape) if dx is not None else None - x0mat = x0.view((-1, hidden_size)) if x0 is not None else None - x0_subset = x0_subset.view(-1) if x0_subset is not None else None - out_subset = out_subset.view(-1) if out_subset is not None else None - if colscale is not None: - assert x0 is not None, "x0 is required to compute the gradient of colscale" - dx0mat, dresidualmat, dgamma, dbeta, _, _, *rest = dropout_layer_norm.dropout_add_ln_bwd( - dzmat, - dxmat, - xmat, - x0mat, - dmask, - mu, - rsigma, - gamma, - None, - colscale, - x0_subset, - out_subset, - dropout_p, - rowscale_const, - x0_numrows, - has_residual, - is_rms_norm, - ) - # dresidualmat is None if not has_residual - if colscale is None: - return dx0mat, dresidualmat, dgamma, dbeta - else: - dcolscale = rest[0] - return dx0mat, dresidualmat, dgamma, dbeta, dcolscale - - -def _dropout_add_layer_norm_parallel_residual_forward( - x0, - x1, - residual, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - residual_in_fp32=False, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes""" - hidden_size = gamma0.numel() - x0mat = x0.view((-1, hidden_size)) - x1mat = x1.view((-1, hidden_size)) if x1 is not None else None - residualmat = residual.view((-1, hidden_size)) if residual is not None else None - ( - z0mat, - z1mat, - xmat, - dmask0, - dmask1, - mu, - rsigma, - ) = dropout_layer_norm.dropout_add_ln_parallel_residual_fwd( - x0mat, - x1mat, - residualmat, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - None, - residual_in_fp32, - is_rms_norm, - ) - # dmask0 and dmask1 are None if dropout_p == 0.0 - # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype - return z0mat, z1mat, xmat if xmat is not None else x0mat, dmask0, dmask1, mu, rsigma - - -def _dropout_add_layer_norm_parallel_residual_backward( - dz0, - dz1, - dx, - x, - dmask0, - dmask1, - mu, - rsigma, - gamma0, - gamma1, - dropout_p, - has_x1, - has_residual, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes - dx == None means that it was a post-norm architecture - (x = drop(x0) + residual was not returned in the fwd). - """ - hidden_size = gamma0.numel() - xmat = x.view((-1, hidden_size)) - dz0mat = dz0.view(xmat.shape) - dz1mat = dz1.view(xmat.shape) if dz1 is not None else None - dxmat = dx.view(xmat.shape) if dx is not None else None - ( - dx0mat, - dx1mat, - dresidualmat, - dgamma0, - dbeta0, - dgamma1, - dbeta1, - *rest, - ) = dropout_layer_norm.dropout_add_ln_parallel_residual_bwd( - dz0mat, - dz1mat, - dxmat, - xmat, - dmask0, - dmask1, - mu, - rsigma, - gamma0, - gamma1, - dropout_p, - has_x1, - has_residual, - is_rms_norm, - ) - # dresidualmat is None if not has_residual - return dx0mat, dx1mat, dresidualmat, dgamma0, dbeta0, dgamma1, dbeta1 - - -class DropoutAddLayerNormFn(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x0, - residual, - gamma, - beta, - rowscale, - colscale, - dropout_p, - epsilon, - residual_in_fp32=False, - prenorm=False, - is_rms_norm=False, - return_dmask=False, - ): - x0 = maybe_align(x0.contiguous(), 16) - residual = maybe_align(residual.contiguous(), 16) if residual is not None else None - gamma = maybe_align(gamma.contiguous(), 16) - beta = maybe_align(beta.contiguous(), 16) if beta is not None else None - rowscale = maybe_align(rowscale.contiguous(), 16) if rowscale is not None else None - colscale = maybe_align(colscale.contiguous(), 16) if colscale is not None else None - zmat, xmat, dmask, mu, rsigma = _dropout_add_layer_norm_forward( - x0, - residual, - gamma, - beta, - rowscale, - colscale, - dropout_p, - epsilon, - residual_in_fp32, - is_rms_norm, - ) - # Only need to save x0 if we need to compute gradient wrt colscale - x0_saved = x0 if colscale is not None else None - ctx.save_for_backward( - xmat.view(x0.shape), x0_saved, dmask, gamma, mu, rsigma, rowscale, colscale - ) - ctx.prenorm = prenorm - ctx.dropout_p = dropout_p - ctx.has_residual = residual is not None - ctx.is_rms_norm = is_rms_norm - ctx.has_beta = beta is not None - if not return_dmask: - return ( - zmat.view(x0.shape) if not prenorm else (zmat.view(x0.shape), xmat.view(x0.shape)) - ) - else: - dmask = ( - dmask.view(x0.shape) - if dropout_p > 0.0 - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - ctx.mark_non_differentiable(dmask) - return ( - (zmat.view(x0.shape), dmask) - if not prenorm - else (zmat.view(x0.shape), xmat.view(x0.shape), dmask) - ) - - @staticmethod - def backward(ctx, dz, *args): - # assert dz.is_contiguous() - dz = maybe_align(dz.contiguous(), 16) # this happens! - dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None - x, x0, dmask, gamma, mu, rsigma, rowscale, colscale = ctx.saved_tensors - # x0 is None if colscale is None - dropout_p = ctx.dropout_p - has_residual = ctx.has_residual - dx0mat, dresidualmat, dgamma, dbeta, *rest = _dropout_add_layer_norm_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - rowscale, - colscale, - dropout_p, - has_residual, - ctx.is_rms_norm, - ) - dx0 = dx0mat.view(x.shape) - dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None - dcolscale = rest[0] if colscale is not None else None - return ( - dx0, - dresidual, - dgamma, - dbeta if ctx.has_beta else None, - None, - dcolscale, - None, - None, - None, - None, - None, - None, - ) - - -class DropoutAddLayerNormSubsetFn(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x0, - residual, - gamma, - beta, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32=False, - prenorm=False, - is_rms_norm=False, - return_dmask=False, - ): - x0 = maybe_align(x0.contiguous(), 16) - residual = maybe_align(residual.contiguous(), 16) if residual is not None else None - gamma = maybe_align(gamma.contiguous(), 16) - beta = maybe_align(beta.contiguous(), 16) if beta is not None else None - colscale = maybe_align(colscale.contiguous(), 16) if colscale is not None else None - zmat, xmat, dmask, mu, rsigma = _dropout_add_layer_norm_subset_forward( - x0, - residual, - gamma, - beta, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32, - is_rms_norm, - ) - # Only need to save x0 if we need to compute gradient wrt colscale - x0_saved = x0 if colscale is not None else None - x_shape = (-1, *x0.shape[1:]) - ctx.save_for_backward( - xmat.view(x_shape), x0_saved, dmask, gamma, mu, rsigma, colscale, x0_subset, out_subset - ) - ctx.prenorm = prenorm - ctx.dropout_p = dropout_p - ctx.rowscale_const = rowscale_const - ctx.x0_numrows = x0.shape[:-1].numel() - ctx.has_residual = residual is not None - ctx.is_rms_norm = is_rms_norm - ctx.has_beta = beta is not None - z_shape = (-1, *x0.shape[1:]) - if not return_dmask: - return zmat.view(z_shape) if not prenorm else (zmat.view(z_shape), xmat.view(x0.shape)) - else: - z = zmat.view(z_shape) - dmask = ( - dmask.view(x0.shape) - if dropout_p > 0.0 - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - ctx.mark_non_differentiable(dmask) - return (z, dmask) if not prenorm else (z, xmat.view(x_shape), dmask) - - @staticmethod - def backward(ctx, dz, *args): - # assert dz.is_contiguous() - dz = maybe_align(dz.contiguous(), 16) # this happens! - dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None - x, x0, dmask, gamma, mu, rsigma, colscale, x0_subset, out_subset = ctx.saved_tensors - # x0 is None if colscale is None - dropout_p = ctx.dropout_p - has_residual = ctx.has_residual - dx0mat, dresidualmat, dgamma, dbeta, *rest = _dropout_add_layer_norm_subset_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - colscale, - x0_subset, - out_subset, - dropout_p, - ctx.rowscale_const, - ctx.x0_numrows, - has_residual, - ctx.is_rms_norm, - ) - dx0 = dx0mat.view(-1, *x.shape[1:]) - dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None - dcolscale = rest[0] if colscale is not None else None - return ( - dx0, - dresidual, - dgamma, - dbeta if ctx.has_beta else None, - dcolscale, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) - - -class DropoutAddLayerNormParallelResidualFn(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x0, - x1, - residual, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - residual_in_fp32=False, - prenorm=False, - is_rms_norm=False, - return_dmask=False, - ): - x0 = maybe_align(x0.contiguous(), 16) - x1 = maybe_align(x1.contiguous(), 16) if x1 is not None else None - residual = maybe_align(residual.contiguous(), 16) if residual is not None else None - gamma0 = maybe_align(gamma0.contiguous(), 16) - beta0 = maybe_align(beta0.contiguous(), 16) if beta0 is not None else None - gamma1 = maybe_align(gamma1.contiguous(), 16) if gamma1 is not None else None - beta1 = maybe_align(beta1.contiguous(), 16) if beta1 is not None else None - ( - z0mat, - z1mat, - xmat, - dmask0, - dmask1, - mu, - rsigma, - ) = _dropout_add_layer_norm_parallel_residual_forward( - x0, - x1, - residual, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - residual_in_fp32, - is_rms_norm, - ) - ctx.save_for_backward(xmat.view(x0.shape), dmask0, dmask1, gamma0, gamma1, mu, rsigma) - ctx.prenorm = prenorm - ctx.dropout_p = dropout_p - ctx.has_x1 = x1 is not None - ctx.has_residual = residual is not None - ctx.is_rms_norm = is_rms_norm - ctx.has_beta = beta0 is not None - z = (z0mat.view(x0.shape), z1mat.view(x0.shape) if z1mat is not None else None) - if not return_dmask: - return z if not prenorm else (*z, xmat.view(x0.shape)) - else: - dmask0 = ( - dmask0.view(x0.shape) - if dropout_p > 0.0 - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - dmask1 = ( - dmask1.view(x0.shape) - if dropout_p > 0.0 and x1 is not None - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - ctx.mark_non_differentiable(dmask0) - ctx.mark_non_differentiable(dmask1) - return ( - (*z, dmask0, dmask1) if not prenorm else (*z, xmat.view(x0.shape), dmask0, dmask1) - ) - - @staticmethod - def backward(ctx, dz0, dz1, *args): - dz0 = maybe_align(dz0.contiguous(), 16) # this happens! - dz1 = maybe_align(dz1.contiguous(), 16) if dz1 is not None else None - dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None - x, dmask0, dmask1, gamma0, gamma1, mu, rsigma = ctx.saved_tensors - dropout_p = ctx.dropout_p - has_x1 = ctx.has_x1 - has_residual = ctx.has_residual - ( - dx0mat, - dx1mat, - dresidualmat, - dgamma0, - dbeta0, - dgamma1, - dbeta1, - ) = _dropout_add_layer_norm_parallel_residual_backward( - dz0, - dz1, - dx, - x, - dmask0, - dmask1, - mu, - rsigma, - gamma0, - gamma1, - dropout_p, - has_x1, - has_residual, - ctx.is_rms_norm, - ) - dx0 = dx0mat.view(x.shape) - dx1 = dx1mat.view(x.shape) if dx1mat is not None else None - dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None - return ( - dx0, - dx1, - dresidual, - dgamma0, - dbeta0 if ctx.has_beta else None, - dgamma1, - dbeta1 if ctx.has_beta else None, - None, - None, - None, - None, - None, - None, - ) - - -def layer_norm(x, weight, bias, epsilon): - return DropoutAddLayerNormFn.apply(x, None, weight, bias, None, None, 0.0, epsilon, False) - - -def dropout_add_layer_norm( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - rowscale=None, - layerscale=None, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormFn.apply( - x0, - residual, - weight, - bias, - rowscale, - layerscale, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - False, - return_dropout_mask, - ) - - -def dropout_add_layer_norm_subset( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - layerscale=None, - x0_subset=None, - out_subset=None, - rowscale_const=1.0, - out_numrows=0, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormSubsetFn.apply( - x0, - residual, - weight, - bias, - layerscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32, - prenorm, - False, - return_dropout_mask, - ) - - -def dropout_add_layer_norm_parallel_residual( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormParallelResidualFn.apply( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - False, - return_dropout_mask, - ) - - -class DropoutAddLayerNorm(torch.nn.Module): - def __init__( - self, - hidden_size, - prenorm=False, - p=0.0, - eps=1e-5, - residual_in_fp32=False, - device=None, - dtype=None, - ): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.prenorm = prenorm - self.p = p - self.eps = eps - self.residual_in_fp32 = residual_in_fp32 - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.bias = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.reset_parameters() - - def reset_parameters(self): - init.ones_(self.weight) - init.zeros_(self.bias) - - def forward(self, x0, residual=None): - return dropout_add_layer_norm( - x0, - residual, - self.weight, - self.bias, - self.p if self.training else 0.0, - self.eps, - prenorm=self.prenorm, - residual_in_fp32=self.residual_in_fp32, - ) diff --git a/build/torch211-cxx11-cu130-x86_64-linux/ops/rms_norm.py b/build/torch211-cxx11-cu130-x86_64-linux/ops/rms_norm.py deleted file mode 100644 index 068348d61290e3839dd082b540d898578ba1e8e2..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu130-x86_64-linux/ops/rms_norm.py +++ /dev/null @@ -1,174 +0,0 @@ -# Copyright (c) 2022, Tri Dao. -# Adapted from https://github.com/NVIDIA/apex/blob/master/apex/contrib/layer_norm/layer_norm.py - -import torch -from torch.nn import init - -from flash_attn.ops.layer_norm import ( - DropoutAddLayerNormFn, - DropoutAddLayerNormParallelResidualFn, - DropoutAddLayerNormSubsetFn, -) - - -def rms_norm(x, weight, epsilon): - return DropoutAddLayerNormFn.apply( - x, None, weight, None, None, None, 0.0, epsilon, False, False, True - ) - - -def dropout_add_rms_norm( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - rowscale=None, - layerscale=None, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormFn.apply( - x0, - residual, - weight, - bias, - rowscale, - layerscale, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - True, - return_dropout_mask, - ) - - -def dropout_add_rms_norm_subset( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - layerscale=None, - x0_subset=None, - out_subset=None, - rowscale_const=1.0, - out_numrows=0, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormSubsetFn.apply( - x0, - residual, - weight, - bias, - layerscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32, - prenorm, - True, - return_dropout_mask, - ) - - -def dropout_add_rms_norm_parallel_residual( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormParallelResidualFn.apply( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - True, - return_dropout_mask, - ) - - -class RMSNorm(torch.nn.Module): - def __init__(self, hidden_size, eps=1e-5, device=None, dtype=None): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.eps = eps - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.register_parameter("bias", None) - self.reset_parameters() - - def reset_parameters(self): - init.ones_(self.weight) - - def forward(self, x): - return rms_norm(x, self.weight, self.eps) - - -class DropoutAddRMSNorm(torch.nn.Module): - def __init__( - self, - hidden_size, - prenorm=False, - p=0.0, - eps=1e-5, - residual_in_fp32=False, - device=None, - dtype=None, - ): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.prenorm = prenorm - self.p = p - self.eps = eps - self.residual_in_fp32 = residual_in_fp32 - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.register_parameter("bias", None) - self.reset_parameters() - - def reset_parameters(self): - init.ones_(self.weight) - - def forward(self, x0, residual=None): - return dropout_add_rms_norm( - x0, - residual, - self.weight, - None, - self.p if self.training else 0.0, - self.eps, - prenorm=self.prenorm, - residual_in_fp32=self.residual_in_fp32, - ) diff --git a/build/torch211-cxx11-cu130-x86_64-linux/ops/triton/__init__.py b/build/torch211-cxx11-cu130-x86_64-linux/ops/triton/__init__.py deleted file mode 100644 index 8b137891791fe96927ad78e64b0aad7bded08bdc..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu130-x86_64-linux/ops/triton/__init__.py +++ /dev/null @@ -1 +0,0 @@ - diff --git a/build/torch211-cxx11-cu130-x86_64-linux/ops/triton/cross_entropy.py b/build/torch211-cxx11-cu130-x86_64-linux/ops/triton/cross_entropy.py deleted file mode 100644 index 1b5a415b73f236f3e05fb14b9141959559e18526..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu130-x86_64-linux/ops/triton/cross_entropy.py +++ /dev/null @@ -1,330 +0,0 @@ -# Copyright (c) 2023, Tri Dao. - -from typing import Tuple, Optional, Union - -import torch -import torch.nn.functional as F - -import triton -import triton.language as tl - -# `all_gather_into_tensor` and `reduce_scatter_tensor` are new placeholders for -# `_all_gather_base` and `_reduce_scatter_base`. They require the most recent -# version of PyTorch. The following 2 lines are for backward compatibility with -# older PyTorch. -if "all_gather_into_tensor" not in dir(torch.distributed): - torch.distributed.all_gather_into_tensor = torch.distributed._all_gather_base - - -@triton.heuristics( - { - "HAS_SMOOTHING": lambda args: args["smoothing"] > 0.0, - } -) -@triton.jit -def cross_entropy_fwd_kernel( - loss_ptr, # data ptrs - lse_ptr, - z_loss_ptr, - logits_ptr, - labels_ptr, - smoothing, - logit_scale, - lse_square_scale, - ignore_index, - total_classes, - class_start_idx, # Useful for tensor parallel when each rank only has a subset of classes - n_cols, # shapes - logits_row_stride, # strides - BLOCK_SIZE: tl.constexpr, - HAS_SMOOTHING: tl.constexpr, - # if SPLIT (e.g. tensor parallel), don't include the LSE in the loss since it's not the final LSE - SPLIT: tl.constexpr, - PRECOMPUTED_LSE: tl.constexpr, # If LSE is already computed (also no smoothing and logit_scale == 1.0) -): - row_idx = tl.program_id(0) - logits_ptr = logits_ptr + row_idx * logits_row_stride.to(tl.int64) - sum_logits = 0.0 # For smoothing - if not PRECOMPUTED_LSE: - # Statistics for online softmax - m_i = -float("inf") - l_i = 0.0 - for col_offset in range(0, n_cols, BLOCK_SIZE): - cols = col_offset + tl.arange(0, BLOCK_SIZE) - logits = tl.load(logits_ptr + cols, mask=cols < n_cols, other=-float("inf")).to( - tl.float32 - ) * logit_scale - if HAS_SMOOTHING: - sum_logits += tl.sum(tl.where(cols < n_cols, logits, 0.0)) - m_i_new = tl.maximum(m_i, tl.max(logits)) - l_i = tl.exp(m_i - m_i_new) * l_i + tl.sum(tl.exp(logits - m_i_new)) - m_i = m_i_new - lse = tl.log(l_i) + m_i - tl.store(lse_ptr + row_idx, lse) - else: - lse = tl.load(lse_ptr + row_idx) - label_idx = tl.load(labels_ptr + row_idx) - if label_idx == ignore_index: - loss = 0.0 - z_loss = 0.0 - else: - label_idx -= class_start_idx - if label_idx >= 0 and label_idx < n_cols: - logits_label = tl.load(logits_ptr + label_idx) * logit_scale - if HAS_SMOOTHING: - loss = ( - (lse if not SPLIT else 0.0) - - smoothing * sum_logits / total_classes - - (1 - smoothing) * logits_label - ) - else: - loss = (lse if not SPLIT else 0.0) - logits_label - else: - # If label is out of bounds, we set the CE loss to 0.0. But we still want the smoothing loss - if HAS_SMOOTHING: - loss = smoothing * ((lse if not SPLIT else 0.0) - sum_logits / total_classes) - else: - loss = 0.0 - if not SPLIT: - z_loss = lse_square_scale * lse * lse - loss += z_loss - else: - z_loss = 0.0 - tl.store(loss_ptr + row_idx, loss) - if not SPLIT: - tl.store(z_loss_ptr + row_idx, z_loss) - - -@triton.heuristics( - { - "HAS_SMOOTHING": lambda args: args["smoothing"] > 0.0, - } -) -@triton.jit -def cross_entropy_bwd_kernel( - dlogits_ptr, # data ptrs - dloss_ptr, - logits_ptr, - lse_ptr, - labels_ptr, - smoothing, - logit_scale, - lse_square_scale, - ignore_index, - total_classes, - class_start_idx, # Useful for tensor parallel when each rank only has a subset of classes - n_cols, # shapes - logits_row_stride, # strides - dlogits_row_stride, - dloss_row_stride, - BLOCK_SIZE: tl.constexpr, - HAS_SMOOTHING: tl.constexpr, -): - row_idx = tl.program_id(0) - col_block_idx = tl.program_id(1) - logits_ptr = logits_ptr + row_idx * logits_row_stride.to(tl.int64) - dlogits_ptr = dlogits_ptr + row_idx * dlogits_row_stride.to(tl.int64) - col_offsets = col_block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - label_idx = tl.load(labels_ptr + row_idx) - if label_idx != ignore_index: - dloss = tl.load(dloss_ptr + row_idx * dloss_row_stride) - else: - dloss = 0.0 - logits = tl.load(logits_ptr + col_offsets, mask=col_offsets < n_cols, other=-float("inf")).to( - tl.float32 - ) * logit_scale - lse = tl.load(lse_ptr + row_idx) - probs = tl.exp(logits - lse) - probs += 2.0 * lse_square_scale * lse * probs - label_idx -= class_start_idx - if HAS_SMOOTHING: - smooth_positive = 1.0 - smoothing - smooth_negative = smoothing / total_classes - probs = tl.where(col_offsets == label_idx, probs - smooth_positive, probs) - smooth_negative - else: - probs = tl.where(col_offsets == label_idx, probs - 1.0, probs) - tl.store(dlogits_ptr + col_offsets, (dloss * logit_scale) * probs, mask=col_offsets < n_cols) - - -class CrossEntropyLoss(torch.autograd.Function): - - @staticmethod - def forward( - ctx, - logits, - labels, - precomputed_lse=None, - smoothing=0.0, - logit_scale=1.0, - lse_square_scale=0.0, - ignore_index=-100, - inplace_backward=False, - process_group=None, - ): - # For some reason Triton generates wrong code when labels has dtype long and its address - # is not aligned to 16 bytes. The ld.global.b64 seems to load the wrong label index. - if labels.dtype == torch.long and labels.data_ptr() % 16 != 0: - labels = F.pad(labels, (0, 1))[..., :-1] - assert labels.data_ptr() % 16 == 0 - assert logit_scale > 0.0 - n_rows, n_cols = logits.shape - assert labels.shape == (n_rows,) - world_size = 1 if process_group is None else torch.distributed.get_world_size(process_group) - total_classes = world_size * n_cols - rank = 0 if process_group is None else torch.distributed.get_rank(process_group) - class_start_idx = rank * n_cols - use_precomputed_lse = precomputed_lse is not None and logit_scale == 1.0 and smoothing == 0.0 - - if logits.stride(-1) != 1: - logits = logits.contiguous() - MAX_BLOCK_SIZE = 16 * 1024 - BLOCK_SIZE = min(triton.next_power_of_2(n_cols), MAX_BLOCK_SIZE) - num_warps = ( - 4 - if BLOCK_SIZE < 2048 - else (8 if BLOCK_SIZE < 8192 else (16 if BLOCK_SIZE < 128 * 1024 else 32)) - ) - losses = torch.empty(n_rows, dtype=torch.float, device=logits.device) - if use_precomputed_lse: - assert precomputed_lse.shape == (n_rows,) - lse = precomputed_lse.contiguous() - else: - lse = torch.empty(n_rows, dtype=torch.float, device=logits.device) - z_losses = torch.empty(n_rows, dtype=torch.float, device=logits.device) - # Need this, otherwise Triton tries to launch from cuda:0 and we get - # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?) - with torch.cuda.device(logits.device.index): - cross_entropy_fwd_kernel[(n_rows,)]( - losses, # data ptrs - lse, - z_losses, - logits, - labels, - smoothing, - logit_scale, - lse_square_scale, - ignore_index, - total_classes, - class_start_idx, - n_cols, # shapes - logits.stride(0), # strides - BLOCK_SIZE=BLOCK_SIZE, # constants - SPLIT=world_size > 1, - PRECOMPUTED_LSE=use_precomputed_lse, - num_warps=num_warps, - ) - - if world_size > 1: - # If there's no smoothing, if labels are in the vocab of this partition, losses contains - # - predicted logit, and 0 otherwise. - # If there's smoothing=0.1, for labels in the vocab of this partition, losses contains - # -0.9 * predicted logit - 0.1 * sum logit / total_classes. - # For labels not in the vocab of this partition, losses contains - # -0.1 * sum logit / total_classes. - if world_size > 1: - lse_allgather = torch.empty(world_size, n_rows, dtype=lse.dtype, device=lse.device) - torch.distributed.all_gather_into_tensor(lse_allgather, lse, group=process_group) - handle_losses = torch.distributed.all_reduce( - losses, op=torch.distributed.ReduceOp.SUM, group=process_group, async_op=True - ) - lse = torch.logsumexp(lse_allgather, dim=0) - handle_losses.wait() - # After the allreduce, if there's no smoothing, the total losses are - predicted_logit, - # we just have to add the (global) lse. - # If there's smoothing=0.1, the total losses are - # -0.9 * predicted_logit - 0.1 * sum logit / total_classes. - # Again, we just have to add the (global) lse. - losses += lse - if lse_square_scale != 0.0: - z_losses = lse_square_scale * lse.square() - z_losses.masked_fill_(labels == ignore_index, 0.0) - losses += z_losses - else: - z_losses = torch.zeros_like(losses) - losses.masked_fill_(labels == ignore_index, 0.0) - - ctx.save_for_backward(logits, lse, labels) - ctx.mark_non_differentiable(z_losses) - ctx.smoothing = smoothing - ctx.logit_scale = logit_scale - ctx.lse_square_scale = lse_square_scale - ctx.ignore_index = ignore_index - ctx.total_classes = total_classes - ctx.class_start_idx = class_start_idx - ctx.inplace_backward = inplace_backward - return losses, z_losses - - @staticmethod - def backward(ctx, grad_losses, grad_z_losses): - del grad_z_losses # z_losses are only for logging. - - logits, lse, labels = ctx.saved_tensors - dlogits = logits if ctx.inplace_backward else torch.empty_like(logits) - n_rows, n_cols = logits.shape - BLOCK_SIZE = min(triton.next_power_of_2(n_cols), 4 * 1024) - num_warps = 4 if BLOCK_SIZE < 2048 else (8 if BLOCK_SIZE < 8192 else 16) - grid = lambda META: (n_rows, triton.cdiv(n_cols, META["BLOCK_SIZE"])) # noqa - # Need this, otherwise Triton tries to launch from cuda:0 and we get - # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?) - with torch.cuda.device(logits.device.index): - cross_entropy_bwd_kernel[grid]( - dlogits, # data ptrs - grad_losses, - logits, - lse, - labels, - ctx.smoothing, - ctx.logit_scale, - ctx.lse_square_scale, - ctx.ignore_index, - ctx.total_classes, - ctx.class_start_idx, - n_cols, # shapes - logits.stride(0), # strides - dlogits.stride(0), - grad_losses.stride(0), - BLOCK_SIZE=BLOCK_SIZE, # constants - num_warps=num_warps, - ) - return dlogits, None, None, None, None, None, None, None, None, None - - -def cross_entropy_loss( - logits: torch.Tensor, - labels: torch.Tensor, - precomputed_lse: Optional[torch.Tensor] = None, - label_smoothing: float = 0.0, - logit_scale: float = 1.0, - lse_square_scale: float = 0.0, - ignore_index=-100, - inplace_backward: bool = False, - process_group=None, -) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Arguments: - logits: (batch, vocab_size) - labels: (batch,) - label_smoothing: float - logit_scale: float. Multiply logits by this scale before calculating the loss. - lse_square_scale: float. If > 0, we add lse_square_scale * lse(logits) ^ 2 to the loss. - This is also referred to as "z-loss". - ignore_index: int. If labels == ignore_index, the loss is set to 0.0. - inplace_backward: bool. If True, we do the backward pass in-place by modifying the logits. - This saves memory. - process_group: if not None, we're doing Tensor Parallel: each process is responsible for - one part of the vocab. The loss will be aggregated across processes. - Returns: - losses: (batch,), float - z_losses: (batch,), float - """ - return CrossEntropyLoss.apply( - logits, - labels, - precomputed_lse, - label_smoothing, - logit_scale, - lse_square_scale, - ignore_index, - inplace_backward, - process_group, - ) diff --git a/build/torch211-cxx11-cu130-x86_64-linux/ops/triton/k_activations.py b/build/torch211-cxx11-cu130-x86_64-linux/ops/triton/k_activations.py deleted file mode 100644 index efb83c358eb4a85d069ee340a3c83f418f9a805b..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu130-x86_64-linux/ops/triton/k_activations.py +++ /dev/null @@ -1,162 +0,0 @@ -# Adapted from https://github.com/facebookresearch/xformers/blob/main/xformers/triton/k_activations.py -# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved. -# -# This source code is licensed under the BSD license found in the -# LICENSE file in the root directory of this source tree. - -import math -from enum import Enum -from typing import Optional - -import triton -import triton.language as tl - -_sqrt2pi = math.sqrt(2.0 / math.pi) -_sqrt1_2 = math.sqrt(1.0 / 2) -_gaussian_pdf_normalization = 1.0 / math.sqrt(2 * math.pi) - - -class Activation(str, Enum): - SquaredReLU = "squared_relu" - GeLU = "gelu" - GeLUApprox = "gelu_approx" - LeakyReLU = "leaky_relu" - ReLU = "relu" - - -def get_triton_activation_kernel(activation: Optional[Activation]): - return ( - { - Activation.ReLU: relu, - Activation.LeakyReLU: leaky_relu, - Activation.GeLU: gelu, - Activation.GeLUApprox: gelu_approx, - Activation.SquaredReLU: squared_relu, - }[activation] - if activation - else None - ) - - -def get_triton_activation_bwd_kernel(activation: Optional[Activation]): - return ( - { - Activation.ReLU: relu_grad, - Activation.LeakyReLU: leaky_relu_grad, - Activation.GeLU: gelu_grad, - Activation.GeLUApprox: gelu_approx_grad, - Activation.SquaredReLU: squared_relu_grad, - }[activation] - if activation - else None - ) - - -@triton.jit -def tanh(x): - # Tanh is just a scaled sigmoid - return 2 * tl.sigmoid(2 * x) - 1 - - -@triton.jit -def cosh(x): - exp_x = tl.exp(x) - return (exp_x + 1.0 / exp_x) * 0.5 - - -# a Triton implementation of the most used activations -# See for instance http://arxiv.org/abs/1606.08415 for an overview - -# ReLU -@triton.jit -def relu(x): - """ - ReLU_ activation function - - .. _ReLU: https://pytorch.org/docs/stable/generated/torch.nn.ReLU.html - """ - zero = 0.0 - return tl.where(x >= 0, x, zero.to(x.dtype)) - - -@triton.jit -def relu_grad(x): - # ReLU is different from other activations - # in that it does not require the input to retrospectively compute its gradient - # here the input is the downstream gradient, and we return the upstream gradient directly - zero = 0.0 - one = 1.0 - return tl.where(x >= 0, one.to(x.dtype), zero.to(x.dtype)) - - -@triton.jit -def squared_relu(x): - """ - Squared ReLU activation, as proposed in the Primer_ paper. - - .. _Primer: https://arxiv.org/abs/2109.08668 - """ - x_ = relu(x) - return (x_ * x_).to(x.dtype) - - -@triton.jit -def squared_relu_grad(x): - return tl.where(x >= 0, 2.0 * x, 0.0) - - -# Leaky ReLU -@triton.jit -def leaky_relu(x): - """ - LeakyReLU_ activation - - .. _LeakyReLU: https://pytorch.org/docs/stable/generated/torch.nn.LeakyReLU.html - """ - scale = 0.01 + 0.0 - scale = scale.to(x.dtype) - return tl.where(x >= 0, x, scale * x) - - -@triton.jit -def leaky_relu_grad(x): - min_grad = 0.01 - max_grad = 1 - - min_grad = min_grad.to(x.dtype) - max_grad = max_grad.to(x.dtype) - - return tl.where(x >= 0, max_grad, min_grad) - - -@triton.jit -def gelu(x): - """Gaussian Error Linear Unit (GELU)""" - return x * 0.5 * (1.0 + tl.libdevice.erf(x * _sqrt1_2)) - - -@triton.jit -def gelu_grad(x): - cdf = 0.5 * (1.0 + tl.libdevice.erf(x * _sqrt1_2)) - pdf = tl.exp(-0.5 * x * x) * _gaussian_pdf_normalization - return cdf + x * pdf - - -@triton.jit -def gelu_approx(x): - """ - GeLU_ activation - Gaussian error linear unit, with tanh approximation - - .. _GeLU: https://arxiv.org/pdf/1606.08415.pdf - """ - return 0.5 * x * (1.0 + tanh(_sqrt2pi * x * (1.0 + 0.044715 * x * x))) - - -@triton.jit -def gelu_approx_grad(x): - # CREDITS: Fast implementation proposed in - # https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/fused_bias_gelu.py#L30 - tanh_out = tanh(0.79788456 * x * (1 + 0.044715 * x * x)) - return 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( - 1 + tanh_out - ) diff --git a/build/torch211-cxx11-cu130-x86_64-linux/ops/triton/layer_norm.py b/build/torch211-cxx11-cu130-x86_64-linux/ops/triton/layer_norm.py deleted file mode 100644 index 5e6f0b3697b612e6a19031bbc55ac76ad233dc85..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu130-x86_64-linux/ops/triton/layer_norm.py +++ /dev/null @@ -1,1254 +0,0 @@ -# Copyright (c) 2024, Tri Dao. -# Implement dropout + residual + layer_norm / rms_norm. - -# Based on the Triton LayerNorm tutorial: https://triton-lang.org/main/getting-started/tutorials/05-layer-norm.html -# For the backward pass, we keep weight_grad and bias_grad in registers and accumulate. -# This is faster for dimensions up to 8k, but after that it's much slower due to register spilling. -# The models we train have hidden dim up to 8k anyway (e.g. Llama 70B), so this is fine. - -import math -from typing import Optional, List - -import torch -import torch.nn.functional as F -from torch import Tensor - -import triton -import triton.language as tl - -from flash_attn.utils.torch import custom_fwd, custom_bwd -from flash_attn.utils.library import triton_op - -from ..._ops import add_op_namespace_prefix - - -def maybe_contiguous_lastdim(x): - return x.contiguous() if x is not None and x.stride(-1) != 1 else x - - -def maybe_contiguous(x): - return x.contiguous() if x is not None else None - - -def triton_autotune_configs(): - # Return configs with a valid warp count for the current device - configs = [] - # Maximum threads per block is architecture-dependent in theory, but in reality all are 1024 - max_threads_per_block = 1024 - # Default to warp size 32 if not defined by device - warp_size = getattr(torch.cuda.get_device_properties(torch.cuda.current_device()), "warp_size", 32) - # Autotune for warp counts which are powers of 2 and do not exceed thread per block limit - return [triton.Config({}, num_warps=warp_count) for warp_count in [1, 2, 4, 8, 16, 32] - if warp_count * warp_size <= max_threads_per_block] - # return [triton.Config({}, num_warps=8)] - - -def layer_norm_ref( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - zero_centered_weight=False, - dropout_mask=None, - dropout_mask1=None, - upcast=False, -): - dtype = x.dtype - if upcast: - x = x.float() - weight = weight.float() - bias = bias.float() if bias is not None else None - residual = residual.float() if residual is not None else residual - x1 = x1.float() if x1 is not None else None - weight1 = weight1.float() if weight1 is not None else None - bias1 = bias1.float() if bias1 is not None else None - if zero_centered_weight: - weight = weight + 1.0 - if weight1 is not None: - weight1 = weight1 + 1.0 - if x1 is not None: - assert rowscale is None, "rowscale is not supported with parallel LayerNorm" - if rowscale is not None: - x = x * rowscale[..., None] - if dropout_p > 0.0: - if dropout_mask is not None: - x = x.masked_fill(~dropout_mask, 0.0) / (1.0 - dropout_p) - else: - x = F.dropout(x, p=dropout_p) - if x1 is not None: - if dropout_mask1 is not None: - x1 = x1.masked_fill(~dropout_mask1, 0.0) / (1.0 - dropout_p) - else: - x1 = F.dropout(x1, p=dropout_p) - if x1 is not None: - x = x + x1 - if residual is not None: - x = (x + residual).to(x.dtype) - out = F.layer_norm(x.to(weight.dtype), x.shape[-1:], weight=weight, bias=bias, eps=eps).to( - dtype - ) - if weight1 is None: - return out if not prenorm else (out, x) - else: - out1 = F.layer_norm( - x.to(weight1.dtype), x.shape[-1:], weight=weight1, bias=bias1, eps=eps - ).to(dtype) - return (out, out1) if not prenorm else (out, out1, x) - - -def rms_norm_ref( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - zero_centered_weight=False, - dropout_mask=None, - dropout_mask1=None, - upcast=False, -): - dtype = x.dtype - if upcast: - x = x.float() - weight = weight.float() - bias = bias.float() if bias is not None else None - residual = residual.float() if residual is not None else residual - x1 = x1.float() if x1 is not None else None - weight1 = weight1.float() if weight1 is not None else None - bias1 = bias1.float() if bias1 is not None else None - if zero_centered_weight: - weight = weight + 1.0 - if weight1 is not None: - weight1 = weight1 + 1.0 - if x1 is not None: - assert rowscale is None, "rowscale is not supported with parallel LayerNorm" - if rowscale is not None: - x = x * rowscale[..., None] - if dropout_p > 0.0: - if dropout_mask is not None: - x = x.masked_fill(~dropout_mask, 0.0) / (1.0 - dropout_p) - else: - x = F.dropout(x, p=dropout_p) - if x1 is not None: - if dropout_mask1 is not None: - x1 = x1.masked_fill(~dropout_mask1, 0.0) / (1.0 - dropout_p) - else: - x1 = F.dropout(x1, p=dropout_p) - if x1 is not None: - x = x + x1 - if residual is not None: - x = (x + residual).to(x.dtype) - rstd = 1 / torch.sqrt((x.square()).mean(dim=-1, keepdim=True) + eps) - out = ((x * rstd * weight) + bias if bias is not None else (x * rstd * weight)).to(dtype) - if weight1 is None: - return out if not prenorm else (out, x) - else: - out1 = ((x * rstd * weight1) + bias1 if bias1 is not None else (x * rstd * weight1)).to( - dtype - ) - return (out, out1) if not prenorm else (out, out1, x) - - -@triton.autotune( - configs=triton_autotune_configs(), - key=["N", "HAS_RESIDUAL", "STORE_RESIDUAL_OUT", "IS_RMS_NORM", "HAS_BIAS", "HAS_X1", "HAS_W1", "HAS_B1"], -) -# torch compile doesn't like triton.heuristics, so we set these manually when calling the kernel -# @triton.heuristics({"HAS_BIAS": lambda args: args["B"] is not None}) -# @triton.heuristics({"HAS_RESIDUAL": lambda args: args["RESIDUAL"] is not None}) -# @triton.heuristics({"HAS_X1": lambda args: args["X1"] is not None}) -# @triton.heuristics({"HAS_W1": lambda args: args["W1"] is not None}) -# @triton.heuristics({"HAS_B1": lambda args: args["B1"] is not None}) -@triton.jit -def _layer_norm_fwd_1pass_kernel( - X, # pointer to the input - Y, # pointer to the output - W, # pointer to the weights - B, # pointer to the biases - RESIDUAL, # pointer to the residual - X1, - W1, - B1, - Y1, - RESIDUAL_OUT, # pointer to the residual - ROWSCALE, - SEEDS, # Dropout seeds for each row - DROPOUT_MASK, - DROPOUT_MASK1, - Mean, # pointer to the mean - Rstd, # pointer to the 1/std - stride_x_row, # how much to increase the pointer when moving by 1 row - stride_y_row, - stride_res_row, - stride_res_out_row, - stride_x1_row, - stride_y1_row, - M, # number of rows in X - N, # number of columns in X - eps, # epsilon to avoid division by zero - dropout_p, # Dropout probability - zero_centered_weight, # If true, add 1.0 to the weight - IS_RMS_NORM: tl.constexpr, - BLOCK_N: tl.constexpr, - HAS_RESIDUAL: tl.constexpr, - STORE_RESIDUAL_OUT: tl.constexpr, - HAS_BIAS: tl.constexpr, - HAS_DROPOUT: tl.constexpr, - STORE_DROPOUT_MASK: tl.constexpr, - HAS_ROWSCALE: tl.constexpr, - HAS_X1: tl.constexpr, - HAS_W1: tl.constexpr, - HAS_B1: tl.constexpr, -): - # Map the program id to the row of X and Y it should compute. - row = tl.program_id(0) - X += row * stride_x_row - Y += row * stride_y_row - if HAS_RESIDUAL: - RESIDUAL += row * stride_res_row - if STORE_RESIDUAL_OUT: - RESIDUAL_OUT += row * stride_res_out_row - if HAS_X1: - X1 += row * stride_x1_row - if HAS_W1: - Y1 += row * stride_y1_row - # Compute mean and variance - cols = tl.arange(0, BLOCK_N) - x = tl.load(X + cols, mask=cols < N, other=0.0).to(tl.float32) - if HAS_ROWSCALE: - rowscale = tl.load(ROWSCALE + row).to(tl.float32) - x *= rowscale - if HAS_DROPOUT: - # Compute dropout mask - # 7 rounds is good enough, and reduces register pressure - keep_mask = tl.rand(tl.load(SEEDS + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - x = tl.where(keep_mask, x / (1.0 - dropout_p), 0.0) - if STORE_DROPOUT_MASK: - tl.store(DROPOUT_MASK + row * N + cols, keep_mask, mask=cols < N) - if HAS_X1: - x1 = tl.load(X1 + cols, mask=cols < N, other=0.0).to(tl.float32) - if HAS_ROWSCALE: - rowscale = tl.load(ROWSCALE + M + row).to(tl.float32) - x1 *= rowscale - if HAS_DROPOUT: - # Compute dropout mask - # 7 rounds is good enough, and reduces register pressure - keep_mask = ( - tl.rand(tl.load(SEEDS + M + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - ) - x1 = tl.where(keep_mask, x1 / (1.0 - dropout_p), 0.0) - if STORE_DROPOUT_MASK: - tl.store(DROPOUT_MASK1 + row * N + cols, keep_mask, mask=cols < N) - x += x1 - if HAS_RESIDUAL: - residual = tl.load(RESIDUAL + cols, mask=cols < N, other=0.0).to(tl.float32) - x += residual - if STORE_RESIDUAL_OUT: - tl.store(RESIDUAL_OUT + cols, x, mask=cols < N) - if not IS_RMS_NORM: - mean = tl.sum(x, axis=0) / N - tl.store(Mean + row, mean) - xbar = tl.where(cols < N, x - mean, 0.0) - var = tl.sum(xbar * xbar, axis=0) / N - else: - xbar = tl.where(cols < N, x, 0.0) - var = tl.sum(xbar * xbar, axis=0) / N - rstd = 1 / tl.sqrt(var + eps) - tl.store(Rstd + row, rstd) - # Normalize and apply linear transformation - mask = cols < N - w = tl.load(W + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w += 1.0 - if HAS_BIAS: - b = tl.load(B + cols, mask=mask).to(tl.float32) - x_hat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd - y = x_hat * w + b if HAS_BIAS else x_hat * w - # Write output - tl.store(Y + cols, y, mask=mask) - if HAS_W1: - w1 = tl.load(W1 + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w1 += 1.0 - if HAS_B1: - b1 = tl.load(B1 + cols, mask=mask).to(tl.float32) - y1 = x_hat * w1 + b1 if HAS_B1 else x_hat * w1 - tl.store(Y1 + cols, y1, mask=mask) - - -def _layer_norm_fwd( - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - residual: Optional[Tensor] = None, - x1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - out_dtype: Optional[torch.dtype] = None, - residual_dtype: Optional[torch.dtype] = None, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - return_dropout_mask: bool = False, - out: Optional[Tensor] = None, - residual_out: Optional[Tensor] = None -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - # Need to wrap to handle the case where residual_out is a alias of x, which makes torch.library - # and torch.compile unhappy. Also allocate memory for out and residual_out if they are None - # so that _layer_norm_fwd_impl doesn't have to return them. - if out is None: - out = torch.empty_like(x, dtype=x.dtype if out_dtype is None else out_dtype) - if residual is not None: - residual_dtype = residual.dtype - if residual_out is None and ( - residual is not None - or (residual_dtype is not None and residual_dtype != x.dtype) - or dropout_p > 0.0 - or rowscale is not None - or x1 is not None - ): - residual_out = torch.empty_like( - x, dtype=residual_dtype if residual_dtype is not None else x.dtype - ) - else: - residual_out = None - y1, mean, rstd, seeds, dropout_mask, dropout_mask1 = _layer_norm_fwd_impl( - x, - weight, - bias, - eps, - out, - residual=residual, - x1=x1, - weight1=weight1, - bias1=bias1, - dropout_p=dropout_p, - rowscale=rowscale, - zero_centered_weight=zero_centered_weight, - is_rms_norm=is_rms_norm, - return_dropout_mask=return_dropout_mask, - residual_out=residual_out, - ) - # residual_out is None if residual is None and residual_dtype == input_dtype and dropout_p == 0.0 - if residual_out is None: - residual_out = x - return out, y1, mean, rstd, residual_out, seeds, dropout_mask, dropout_mask1 - - -# [2025-04-28] torch.library.triton_op ignores the schema argument, but here we need the schema -# since we're returning a tuple of tensors -@triton_op(add_op_namespace_prefix("layer_norm_fwd_impl"), mutates_args={"out", "residual_out"}, - schema="(Tensor x, Tensor weight, Tensor bias, float eps, Tensor(a!) out, Tensor? residual, Tensor? x1, Tensor? weight1, Tensor? bias1, float dropout_p, Tensor? rowscale, bool zero_centered_weight, bool is_rms_norm, bool return_dropout_mask, Tensor(a!)? residual_out) -> (Tensor y1, Tensor mean, Tensor rstd, Tensor seeds, Tensor dropout_mask, Tensor dropout_mask1)") -def _layer_norm_fwd_impl( - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - out: Tensor, - residual: Optional[Tensor] = None, - x1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - return_dropout_mask: bool = False, - residual_out: Optional[Tensor] = None -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - M, N = x.shape - assert x.stride(-1) == 1 - if residual is not None: - assert residual.stride(-1) == 1 - assert residual.shape == (M, N) - assert weight.shape == (N,) - assert weight.stride(-1) == 1 - if bias is not None: - assert bias.stride(-1) == 1 - assert bias.shape == (N,) - if x1 is not None: - assert x1.shape == x.shape - assert rowscale is None - assert x1.stride(-1) == 1 - if weight1 is not None: - assert weight1.shape == (N,) - assert weight1.stride(-1) == 1 - if bias1 is not None: - assert bias1.shape == (N,) - assert bias1.stride(-1) == 1 - if rowscale is not None: - assert rowscale.is_contiguous() - assert rowscale.shape == (M,) - assert out.shape == x.shape - assert out.stride(-1) == 1 - if residual_out is not None: - assert residual_out.shape == x.shape - assert residual_out.stride(-1) == 1 - if weight1 is not None: - y1 = torch.empty_like(out) - assert y1.stride(-1) == 1 - else: - y1 = None - mean = torch.empty((M,), dtype=torch.float32, device=x.device) if not is_rms_norm else None - rstd = torch.empty((M,), dtype=torch.float32, device=x.device) - if dropout_p > 0.0: - seeds = torch.randint( - 2**32, (M if x1 is None else 2 * M,), device=x.device, dtype=torch.int64 - ) - else: - seeds = None - if return_dropout_mask and dropout_p > 0.0: - dropout_mask = torch.empty(M, N, device=x.device, dtype=torch.bool) - if x1 is not None: - dropout_mask1 = torch.empty(M, N, device=x.device, dtype=torch.bool) - else: - dropout_mask1 = None - else: - dropout_mask, dropout_mask1 = None, None - # Less than 64KB per feature: enqueue fused kernel - MAX_FUSED_SIZE = 65536 // x.element_size() - BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) - if N > BLOCK_N: - raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") - with torch.cuda.device(x.device.index): - torch.library.wrap_triton(_layer_norm_fwd_1pass_kernel)[(M,)]( - x, - out, - weight, - bias, - residual, - x1, - weight1, - bias1, - y1, - residual_out, - rowscale, - seeds, - dropout_mask, - dropout_mask1, - mean, - rstd, - x.stride(0), - out.stride(0), - residual.stride(0) if residual is not None else 0, - residual_out.stride(0) if residual_out is not None else 0, - x1.stride(0) if x1 is not None else 0, - y1.stride(0) if y1 is not None else 0, - M, - N, - eps, - dropout_p, - # Passing bool make torch inductor very unhappy since it then tries to compare to int_max - int(zero_centered_weight), - is_rms_norm, - BLOCK_N, - residual is not None, - residual_out is not None, - bias is not None, - dropout_p > 0.0, - dropout_mask is not None, - rowscale is not None, - HAS_X1=x1 is not None, - HAS_W1=weight1 is not None, - HAS_B1=bias1 is not None, - ) - return y1, mean, rstd, seeds, dropout_mask, dropout_mask1 - - -@triton.autotune( - configs=triton_autotune_configs(), - key=["N", "HAS_DRESIDUAL", "STORE_DRESIDUAL", "IS_RMS_NORM", "HAS_BIAS", "HAS_DROPOUT"], -) -# torch compile doesn't like triton.heuristics, so we set these manually when calling the kernel -# @triton.heuristics({"HAS_BIAS": lambda args: args["B"] is not None}) -# @triton.heuristics({"HAS_DRESIDUAL": lambda args: args["DRESIDUAL"] is not None}) -# @triton.heuristics({"STORE_DRESIDUAL": lambda args: args["DRESIDUAL_IN"] is not None}) -# @triton.heuristics({"HAS_ROWSCALE": lambda args: args["ROWSCALE"] is not None}) -# @triton.heuristics({"HAS_DY1": lambda args: args["DY1"] is not None}) -# @triton.heuristics({"HAS_DX1": lambda args: args["DX1"] is not None}) -# @triton.heuristics({"HAS_B1": lambda args: args["DB1"] is not None}) -# @triton.heuristics({"RECOMPUTE_OUTPUT": lambda args: args["Y"] is not None}) -@triton.jit -def _layer_norm_bwd_kernel( - X, # pointer to the input - W, # pointer to the weights - B, # pointer to the biases - Y, # pointer to the output to be recomputed - DY, # pointer to the output gradient - DX, # pointer to the input gradient - DW, # pointer to the partial sum of weights gradient - DB, # pointer to the partial sum of biases gradient - DRESIDUAL, - W1, - DY1, - DX1, - DW1, - DB1, - DRESIDUAL_IN, - ROWSCALE, - SEEDS, - Mean, # pointer to the mean - Rstd, # pointer to the 1/std - stride_x_row, # how much to increase the pointer when moving by 1 row - stride_y_row, - stride_dy_row, - stride_dx_row, - stride_dres_row, - stride_dy1_row, - stride_dx1_row, - stride_dres_in_row, - M, # number of rows in X - N, # number of columns in X - eps, # epsilon to avoid division by zero - dropout_p, - zero_centered_weight, - rows_per_program, - IS_RMS_NORM: tl.constexpr, - BLOCK_N: tl.constexpr, - HAS_DRESIDUAL: tl.constexpr, - STORE_DRESIDUAL: tl.constexpr, - HAS_BIAS: tl.constexpr, - HAS_DROPOUT: tl.constexpr, - HAS_ROWSCALE: tl.constexpr, - HAS_DY1: tl.constexpr, - HAS_DX1: tl.constexpr, - HAS_B1: tl.constexpr, - RECOMPUTE_OUTPUT: tl.constexpr, -): - # Map the program id to the elements of X, DX, and DY it should compute. - row_block_id = tl.program_id(0) - row_start = row_block_id * rows_per_program - # Do not early exit if row_start >= M, because we need to write DW and DB - cols = tl.arange(0, BLOCK_N) - mask = cols < N - X += row_start * stride_x_row - if HAS_DRESIDUAL: - DRESIDUAL += row_start * stride_dres_row - if STORE_DRESIDUAL: - DRESIDUAL_IN += row_start * stride_dres_in_row - DY += row_start * stride_dy_row - DX += row_start * stride_dx_row - if HAS_DY1: - DY1 += row_start * stride_dy1_row - if HAS_DX1: - DX1 += row_start * stride_dx1_row - if RECOMPUTE_OUTPUT: - Y += row_start * stride_y_row - w = tl.load(W + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w += 1.0 - if RECOMPUTE_OUTPUT and HAS_BIAS: - b = tl.load(B + cols, mask=mask, other=0.0).to(tl.float32) - if HAS_DY1: - w1 = tl.load(W1 + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w1 += 1.0 - dw = tl.zeros((BLOCK_N,), dtype=tl.float32) - if HAS_BIAS: - db = tl.zeros((BLOCK_N,), dtype=tl.float32) - if HAS_DY1: - dw1 = tl.zeros((BLOCK_N,), dtype=tl.float32) - if HAS_B1: - db1 = tl.zeros((BLOCK_N,), dtype=tl.float32) - row_end = min((row_block_id + 1) * rows_per_program, M) - for row in range(row_start, row_end): - # Load data to SRAM - x = tl.load(X + cols, mask=mask, other=0).to(tl.float32) - dy = tl.load(DY + cols, mask=mask, other=0).to(tl.float32) - if HAS_DY1: - dy1 = tl.load(DY1 + cols, mask=mask, other=0).to(tl.float32) - if not IS_RMS_NORM: - mean = tl.load(Mean + row) - rstd = tl.load(Rstd + row) - # Compute dx - xhat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd - xhat = tl.where(mask, xhat, 0.0) - if RECOMPUTE_OUTPUT: - y = xhat * w + b if HAS_BIAS else xhat * w - tl.store(Y + cols, y, mask=mask) - wdy = w * dy - dw += dy * xhat - if HAS_BIAS: - db += dy - if HAS_DY1: - wdy += w1 * dy1 - dw1 += dy1 * xhat - if HAS_B1: - db1 += dy1 - if not IS_RMS_NORM: - c1 = tl.sum(xhat * wdy, axis=0) / N - c2 = tl.sum(wdy, axis=0) / N - dx = (wdy - (xhat * c1 + c2)) * rstd - else: - c1 = tl.sum(xhat * wdy, axis=0) / N - dx = (wdy - xhat * c1) * rstd - if HAS_DRESIDUAL: - dres = tl.load(DRESIDUAL + cols, mask=mask, other=0).to(tl.float32) - dx += dres - # Write dx - if STORE_DRESIDUAL: - tl.store(DRESIDUAL_IN + cols, dx, mask=mask) - if HAS_DX1: - if HAS_DROPOUT: - keep_mask = ( - tl.rand(tl.load(SEEDS + M + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - ) - dx1 = tl.where(keep_mask, dx / (1.0 - dropout_p), 0.0) - else: - dx1 = dx - tl.store(DX1 + cols, dx1, mask=mask) - if HAS_DROPOUT: - keep_mask = tl.rand(tl.load(SEEDS + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - dx = tl.where(keep_mask, dx / (1.0 - dropout_p), 0.0) - if HAS_ROWSCALE: - rowscale = tl.load(ROWSCALE + row).to(tl.float32) - dx *= rowscale - tl.store(DX + cols, dx, mask=mask) - - X += stride_x_row - if HAS_DRESIDUAL: - DRESIDUAL += stride_dres_row - if STORE_DRESIDUAL: - DRESIDUAL_IN += stride_dres_in_row - if RECOMPUTE_OUTPUT: - Y += stride_y_row - DY += stride_dy_row - DX += stride_dx_row - if HAS_DY1: - DY1 += stride_dy1_row - if HAS_DX1: - DX1 += stride_dx1_row - tl.store(DW + row_block_id * N + cols, dw, mask=mask) - if HAS_BIAS: - tl.store(DB + row_block_id * N + cols, db, mask=mask) - if HAS_DY1: - tl.store(DW1 + row_block_id * N + cols, dw1, mask=mask) - if HAS_B1: - tl.store(DB1 + row_block_id * N + cols, db1, mask=mask) - - -def _layer_norm_bwd( - dy: Tensor, - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - mean: Tensor, - rstd: Tensor, - dresidual: Optional[Tensor] = None, - dy1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - seeds: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - has_residual: bool = False, - has_x1: bool = False, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - x_dtype: Optional[torch.dtype] = None, - recompute_output: bool = False, -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - # Need to wrap to handle the case where dresidual_in or dx1 are aliases of x, - # which makes torch.library unhappy - dx, dw, db, dresidual_in, dx1, dw1, db1, y = _layer_norm_bwd_impl( - dy, - x, - weight, - bias, - eps, - mean, - rstd, - dresidual, - dy1, - weight1, - bias1, - seeds, - dropout_p, - rowscale, - has_residual, - has_x1, - zero_centered_weight, - is_rms_norm, - x_dtype=x_dtype, - recompute_output=recompute_output, - ) - # Don't need to compute dresidual_in separately in this case - if has_residual and dx.dtype == x.dtype and dropout_p == 0.0 and rowscale is None: - dresidual_in = dx - if has_x1 and dropout_p == 0.0: - dx1 = dx - return dx, dw, db, dresidual_in, dx1, dw1, db1, y - - - -@triton_op(add_op_namespace_prefix("layer_norm_bwd_impl"), mutates_args={}, - schema="(Tensor dy, Tensor x, Tensor weight, Tensor bias, float eps, Tensor mean, Tensor rstd, Tensor? dresidual, Tensor? dy1, Tensor? weight1, Tensor? bias1, Tensor? seeds, float dropout_p, Tensor? rowscale, bool has_residual, bool has_x1, bool zero_centered_weight, bool is_rms_norm, ScalarType? x_dtype, bool recompute_output) -> (Tensor dx, Tensor dw, Tensor db, Tensor dresidual_in, Tensor dx1, Tensor dw1, Tensor db1, Tensor y)", - allow_decomposition=False, # Don't let torch.compile trace inside - ) -def _layer_norm_bwd_impl( - dy: Tensor, - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - mean: Tensor, - rstd: Tensor, - dresidual: Optional[Tensor] = None, - dy1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - seeds: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - has_residual: bool = False, - has_x1: bool = False, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - x_dtype: Optional[torch.dtype] = None, - recompute_output: bool = False, -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - M, N = x.shape - assert x.stride(-1) == 1 - dy = maybe_contiguous_lastdim(dy) - assert dy.stride(-1) == 1 - assert dy.shape == (M, N) - if dresidual is not None: - dresidual = maybe_contiguous_lastdim(dresidual) - assert dresidual.stride(-1) == 1 - assert dresidual.shape == (M, N) - assert weight.shape == (N,) - assert weight.stride(-1) == 1 - if bias is not None: - assert bias.stride(-1) == 1 - assert bias.shape == (N,) - if dy1 is not None: - dy1 = maybe_contiguous_lastdim(dy1) - assert weight1 is not None - assert dy1.shape == dy.shape - assert dy1.stride(-1) == 1 - if weight1 is not None: - assert weight1.shape == (N,) - assert weight1.stride(-1) == 1 - if bias1 is not None: - assert bias1.shape == (N,) - assert bias1.stride(-1) == 1 - if seeds is not None: - assert seeds.is_contiguous() - assert seeds.shape == (M if not has_x1 else M * 2,) - if rowscale is not None: - assert rowscale.is_contiguous() - assert rowscale.shape == (M,) - # allocate output - dx = ( - torch.empty_like(x) - if x_dtype is None - else torch.empty(M, N, dtype=x_dtype, device=x.device) - ) - dresidual_in = ( - torch.empty_like(x) - if has_residual - and (dx.dtype != x.dtype or dropout_p > 0.0 or rowscale is not None or has_x1) - else None - ) - dx1 = torch.empty_like(dx) if (has_x1 and dropout_p > 0.0) else None - y = torch.empty(M, N, dtype=dy.dtype, device=dy.device) if recompute_output else None - if recompute_output: - assert weight1 is None, "recompute_output is not supported with parallel LayerNorm" - - # Less than 64KB per feature: enqueue fused kernel - MAX_FUSED_SIZE = 65536 // x.element_size() - BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) - if N > BLOCK_N: - raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") - # Increasing the multiple (e.g. 8) will allow more thread blocks to be launched and hide the - # latency of the gmem reads/writes, but will increase the time of summing up dw / db. - sm_count = torch.cuda.get_device_properties(x.device).multi_processor_count * 8 - _dw = torch.empty((sm_count, N), dtype=torch.float32, device=weight.device) - _db = ( - torch.empty((sm_count, N), dtype=torch.float32, device=bias.device) - if bias is not None - else None - ) - _dw1 = torch.empty_like(_dw) if weight1 is not None else None - _db1 = torch.empty_like(_db) if bias1 is not None else None - rows_per_program = math.ceil(M / sm_count) - grid = (sm_count,) - with torch.cuda.device(x.device.index): - torch.library.wrap_triton(_layer_norm_bwd_kernel)[grid]( - x, - weight, - bias, - y, - dy, - dx, - _dw, - _db, - dresidual, - weight1, - dy1, - dx1, - _dw1, - _db1, - dresidual_in, - rowscale, - seeds, - mean, - rstd, - x.stride(0), - 0 if not recompute_output else y.stride(0), - dy.stride(0), - dx.stride(0), - dresidual.stride(0) if dresidual is not None else 0, - dy1.stride(0) if dy1 is not None else 0, - dx1.stride(0) if dx1 is not None else 0, - dresidual_in.stride(0) if dresidual_in is not None else 0, - M, - N, - eps, - dropout_p, - # Passing bool make torch inductor very unhappy since it then tries to compare to int_max - int(zero_centered_weight), - rows_per_program, - is_rms_norm, - BLOCK_N, - dresidual is not None, - dresidual_in is not None, - bias is not None, - dropout_p > 0.0, - HAS_ROWSCALE=rowscale is not None, - HAS_DY1=dy1 is not None, - HAS_DX1=dx1 is not None, - HAS_B1=bias1 is not None, - RECOMPUTE_OUTPUT=y is not None, - ) - dw = _dw.sum(0).to(weight.dtype) - db = _db.sum(0).to(bias.dtype) if bias is not None else None - dw1 = _dw1.sum(0).to(weight1.dtype) if weight1 is not None else None - db1 = _db1.sum(0).to(bias1.dtype) if bias1 is not None else None - # dresidual_in and dx1 could be None, the wrapper will handle assigning them from dx - return dx, dw, db, dresidual_in, dx1, dw1, db1, y - - -class LayerNormFn(torch.autograd.Function): - - @staticmethod - def forward( - ctx, - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - residual_in_fp32=False, - zero_centered_weight=False, - is_rms_norm=False, - return_dropout_mask=False, - out_dtype=None, - out=None, - residual_out=None - ): - x_shape_og = x.shape - # reshape input data into 2D tensor - x = maybe_contiguous_lastdim(x.reshape(-1, x.shape[-1])) - if residual is not None: - assert residual.shape == x_shape_og - residual = maybe_contiguous_lastdim(residual.reshape(-1, residual.shape[-1])) - if x1 is not None: - assert x1.shape == x_shape_og - assert rowscale is None, "rowscale is not supported with parallel LayerNorm" - x1 = maybe_contiguous_lastdim(x1.reshape(-1, x1.shape[-1])) - weight = weight.contiguous() - bias = maybe_contiguous(bias) - weight1 = maybe_contiguous(weight1) - bias1 = maybe_contiguous(bias1) - if rowscale is not None: - rowscale = rowscale.reshape(-1).contiguous() - residual_dtype = ( - residual.dtype - if residual is not None - else (torch.float32 if residual_in_fp32 else None) - ) - if out is not None: - out = out.reshape(-1, out.shape[-1]) - if residual_out is not None: - residual_out = residual_out.reshape(-1, residual_out.shape[-1]) - y, y1, mean, rstd, residual_out, seeds, dropout_mask, dropout_mask1 = _layer_norm_fwd( - x, - weight, - bias, - eps, - residual, - x1, - weight1, - bias1, - dropout_p=dropout_p, - rowscale=rowscale, - out_dtype=out_dtype, - residual_dtype=residual_dtype, - zero_centered_weight=zero_centered_weight, - is_rms_norm=is_rms_norm, - return_dropout_mask=return_dropout_mask, - out=out, - residual_out=residual_out, - ) - ctx.save_for_backward( - residual_out, weight, bias, weight1, bias1, rowscale, seeds, mean, rstd - ) - ctx.x_shape_og = x_shape_og - ctx.eps = eps - ctx.dropout_p = dropout_p - ctx.is_rms_norm = is_rms_norm - ctx.has_residual = residual is not None - ctx.has_x1 = x1 is not None - ctx.prenorm = prenorm - ctx.x_dtype = x.dtype - ctx.zero_centered_weight = zero_centered_weight - y = y.reshape(x_shape_og) - y1 = y1.reshape(x_shape_og) if y1 is not None else None - residual_out = residual_out.reshape(x_shape_og) if residual_out is not None else None - dropout_mask = dropout_mask.reshape(x_shape_og) if dropout_mask is not None else None - dropout_mask1 = dropout_mask1.reshape(x_shape_og) if dropout_mask1 is not None else None - if not return_dropout_mask: - if weight1 is None: - return y if not prenorm else (y, residual_out) - else: - return (y, y1) if not prenorm else (y, y1, residual_out) - else: - if weight1 is None: - return ( - (y, dropout_mask, dropout_mask1) - if not prenorm - else (y, residual_out, dropout_mask, dropout_mask1) - ) - else: - return ( - (y, y1, dropout_mask, dropout_mask1) - if not prenorm - else (y, y1, residual_out, dropout_mask, dropout_mask1) - ) - - @staticmethod - def backward(ctx, dy, *args): - x, weight, bias, weight1, bias1, rowscale, seeds, mean, rstd = ctx.saved_tensors - dy = dy.reshape(-1, dy.shape[-1]) - if weight1 is not None: - dy1, args = args[0], args[1:] - dy1 = dy1.reshape(-1, dy1.shape[-1]) - assert dy1.shape == x.shape - else: - dy1 = None - if ctx.prenorm: - dresidual = args[0] - dresidual = dresidual.reshape(-1, dresidual.shape[-1]) - assert dresidual.shape == x.shape - else: - dresidual = None - dx, dw, db, dresidual_in, dx1, dw1, db1, _ = _layer_norm_bwd( - dy, - x, - weight, - bias, - ctx.eps, - mean, - rstd, - dresidual, - dy1, - weight1, - bias1, - seeds, - ctx.dropout_p, - rowscale, - ctx.has_residual, - ctx.has_x1, - ctx.zero_centered_weight, - ctx.is_rms_norm, - x_dtype=ctx.x_dtype, - recompute_output=False, - ) - return ( - dx.reshape(ctx.x_shape_og), - dw, - db, - dresidual_in.reshape(ctx.x_shape_og) if ctx.has_residual else None, - dx1.reshape(ctx.x_shape_og) if dx1 is not None else None, - dw1, - db1, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) - - -def layer_norm_fn( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - residual_in_fp32=False, - zero_centered_weight=False, - is_rms_norm=False, - return_dropout_mask=False, - out_dtype=None, - out=None, - residual_out=None -): - return LayerNormFn.apply( - x, - weight, - bias, - residual, - x1, - weight1, - bias1, - eps, - dropout_p, - rowscale, - prenorm, - residual_in_fp32, - zero_centered_weight, - is_rms_norm, - return_dropout_mask, - out_dtype, - out, - residual_out - ) - - -def rms_norm_fn( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - residual_in_fp32=False, - zero_centered_weight=False, - return_dropout_mask=False, - out_dtype=None, - out=None, - residual_out=None -): - return LayerNormFn.apply( - x, - weight, - bias, - residual, - x1, - weight1, - bias1, - eps, - dropout_p, - rowscale, - prenorm, - residual_in_fp32, - zero_centered_weight, - True, - return_dropout_mask, - out_dtype, - out, - residual_out - ) - - -class RMSNorm(torch.nn.Module): - - def __init__(self, hidden_size, eps=1e-5, dropout_p=0.0, zero_centered_weight=False, - device=None, dtype=None): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.eps = eps - if dropout_p > 0.0: - self.drop = torch.nn.Dropout(dropout_p) - else: - self.drop = None - self.zero_centered_weight = zero_centered_weight - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.register_parameter("bias", None) - self.reset_parameters() - - def reset_parameters(self): - if not self.zero_centered_weight: - torch.nn.init.ones_(self.weight) - else: - torch.nn.init.zeros_(self.weight) - - def forward(self, x, residual=None, prenorm=False, residual_in_fp32=False): - return rms_norm_fn( - x, - self.weight, - self.bias, - residual=residual, - eps=self.eps, - dropout_p=self.drop.p if self.drop is not None and self.training else 0.0, - prenorm=prenorm, - residual_in_fp32=residual_in_fp32, - zero_centered_weight=self.zero_centered_weight, - ) - - -class LayerNormLinearFn(torch.autograd.Function): - - @staticmethod - @custom_fwd - def forward( - ctx, - x, - norm_weight, - norm_bias, - linear_weight, - linear_bias, - residual=None, - eps=1e-6, - prenorm=False, - residual_in_fp32=False, - is_rms_norm=False, - ): - x_shape_og = x.shape - # reshape input data into 2D tensor - x = maybe_contiguous_lastdim(x.reshape(-1, x.shape[-1])) - if residual is not None: - assert residual.shape == x_shape_og - residual = maybe_contiguous_lastdim(residual.reshape(-1, residual.shape[-1])) - norm_weight = norm_weight.contiguous() - norm_bias = maybe_contiguous(norm_bias) - residual_dtype = ( - residual.dtype - if residual is not None - else (torch.float32 if residual_in_fp32 else None) - ) - y, _, mean, rstd, residual_out, *rest = _layer_norm_fwd( - x, - norm_weight, - norm_bias, - eps, - residual, - out_dtype=None if not torch.is_autocast_enabled() else torch.get_autocast_dtype("cuda"), - residual_dtype=residual_dtype, - is_rms_norm=is_rms_norm, - ) - y = y.reshape(x_shape_og) - dtype = torch.get_autocast_dtype("cuda") if torch.is_autocast_enabled() else y.dtype - linear_weight = linear_weight.to(dtype) - linear_bias = linear_bias.to(dtype) if linear_bias is not None else None - out = F.linear(y.to(linear_weight.dtype), linear_weight, linear_bias) - # We don't store y, will be recomputed in the backward pass to save memory - ctx.save_for_backward(residual_out, norm_weight, norm_bias, linear_weight, mean, rstd) - ctx.x_shape_og = x_shape_og - ctx.eps = eps - ctx.is_rms_norm = is_rms_norm - ctx.has_residual = residual is not None - ctx.prenorm = prenorm - ctx.x_dtype = x.dtype - ctx.linear_bias_is_none = linear_bias is None - return out if not prenorm else (out, residual_out.reshape(x_shape_og)) - - @staticmethod - @custom_bwd - def backward(ctx, dout, *args): - x, norm_weight, norm_bias, linear_weight, mean, rstd = ctx.saved_tensors - dout = dout.reshape(-1, dout.shape[-1]) - dy = F.linear(dout, linear_weight.t()) - dlinear_bias = None if ctx.linear_bias_is_none else dout.sum(0) - dy = maybe_contiguous_lastdim(dy) - assert dy.shape == x.shape - if ctx.prenorm: - dresidual = args[0] - dresidual = maybe_contiguous_lastdim(dresidual.reshape(-1, dresidual.shape[-1])) - assert dresidual.shape == x.shape - else: - dresidual = None - dx, dnorm_weight, dnorm_bias, dresidual_in, _, _, _, y = _layer_norm_bwd( - dy, - x, - norm_weight, - norm_bias, - ctx.eps, - mean, - rstd, - dresidual=dresidual, - has_residual=ctx.has_residual, - is_rms_norm=ctx.is_rms_norm, - x_dtype=ctx.x_dtype, - recompute_output=True, - ) - dlinear_weight = torch.einsum("bo,bi->oi", dout, y) - return ( - dx.reshape(ctx.x_shape_og), - dnorm_weight, - dnorm_bias, - dlinear_weight, - dlinear_bias, - dresidual_in.reshape(ctx.x_shape_og) if ctx.has_residual else None, - None, - None, - None, - None, - ) - - -def layer_norm_linear_fn( - x, - norm_weight, - norm_bias, - linear_weight, - linear_bias, - residual=None, - eps=1e-6, - prenorm=False, - residual_in_fp32=False, - is_rms_norm=False, -): - return LayerNormLinearFn.apply( - x, - norm_weight, - norm_bias, - linear_weight, - linear_bias, - residual, - eps, - prenorm, - residual_in_fp32, - is_rms_norm, - ) diff --git a/build/torch211-cxx11-cu130-x86_64-linux/ops/triton/linear.py b/build/torch211-cxx11-cu130-x86_64-linux/ops/triton/linear.py deleted file mode 100644 index a8966dbc345ab0e593df0124451ee7be3dae131a..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu130-x86_64-linux/ops/triton/linear.py +++ /dev/null @@ -1,594 +0,0 @@ -# Adapted from https://github.com/ELS-RD/kernl/blob/main/src/kernl/implementations/linear_layer.py -# and https://github.com/openai/triton/blob/master/python/triton/ops/matmul.py -from typing import Optional - -import torch -import triton -import triton.language as tl -from triton.ops.matmul_perf_model import early_config_prune, estimate_matmul_time - -from flash_attn.ops.triton.k_activations import ( - gelu, - gelu_approx, - gelu_approx_grad, - gelu_grad, - squared_relu, - squared_relu_grad, -) - -# CREDITS: Initially inspired by the Triton tutorial on matrix multiplications - - -def init_to_zero(name): - return lambda nargs: nargs[name].zero_() - - -def get_configs_io_bound(): - configs = [] - for num_stages in [2, 3, 4, 5, 6]: - for block_m in [16, 32]: - for block_k in [32, 64]: - for block_n in [32, 64, 128, 256]: - num_warps = 2 if block_n <= 64 else 4 - configs.append( - triton.Config( - { - "BLOCK_M": block_m, - "BLOCK_N": block_n, - "BLOCK_K": block_k, - "SPLIT_K": 1, - }, - num_stages=num_stages, - num_warps=num_warps, - ) - ) - # split_k not used - # for split_k in [2, 4, 8, 16]: - # configs.append(triton.Config( - # {'BLOCK_M': block_m, 'BLOCK_N': block_n, 'BLOCK_K': block_k, 'SPLIT_K': split_k}, - # num_stages=num_stages, num_warps=num_warps, pre_hook=init_to_zero('C'))) - return configs - - -@triton.autotune( - configs=[ - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - # good for int8 - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=4, - num_warps=4, - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - ] - + get_configs_io_bound(), - key=["CACHE_KEY_M", "CACHE_KEY_N", "CACHE_KEY_K"], - prune_configs_by={ - "early_config_prune": early_config_prune, - "perf_model": estimate_matmul_time, - "top_k": 10, - }, -) -@triton.heuristics( - { - "EVEN_K": lambda args: args["K"] % (args["BLOCK_K"] * args["SPLIT_K"]) == 0, - } -) -@triton.jit -def kernel_fwd( - C, # Pointers to matrices - ACT_INPUT, - A, - B, - bias, - # Matrix dimensions - M, - N, - K, - CACHE_KEY_M, - CACHE_KEY_N, - CACHE_KEY_K, - # The stride variables represent how much to increase the ptr by when moving by 1 - # element in a particular dimension. E.g. stride_am is how much to increase a_ptr - # by to get the element one row down (A has M rows) - stride_cm, - # stride_cn, # Assume that stride_cn == 1 - stride_am, - stride_ak, - stride_bn, - stride_bk, - # Meta-parameters - BLOCK_M: tl.constexpr, - GROUP_M: tl.constexpr, - BLOCK_N: tl.constexpr, - BLOCK_K: tl.constexpr, - # split k not used, not performant with activation, kept because early_config_prune is expecting it - SPLIT_K: tl.constexpr, - EVEN_K: tl.constexpr, - A_ROWMAJOR: tl.constexpr, - B_COLMAJOR: tl.constexpr, - BIAS: tl.constexpr, - SAVE_ACT_INPUT: tl.constexpr, - ACTIVATION: tl.constexpr, -): - - """ - Kernel for computing Out = activation(A x W + C) - - Input has shape (M, K) - - Weight has shape (K, N) - - Bias has shape (N,) - - Output has shape (M, N) - - ActInputs (optional) has shape (M, N) - 'ActInputs' optionally saves the A x W + C intermediate for backward computations - This kernel will consolidate over K - """ - - pid = tl.program_id(axis=0) - - grid_m = (M + BLOCK_M - 1) // BLOCK_M - grid_n = (N + BLOCK_N - 1) // BLOCK_N - # re-order program ID for better L2 performance - width = GROUP_M * grid_n - group_id = pid // width - group_size = min(grid_m - group_id * GROUP_M, GROUP_M) - pid_m = group_id * GROUP_M + (pid % group_size) - pid_n = (pid % width) // (group_size) - - # now compute the block that each program will go through - # rm (resp. rn) denotes a range of indices - # for rows (resp. col) of C - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - # trick to avoid masking on M and N axis - ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M) - rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N) - rk = tl.arange(0, BLOCK_K) - - if A_ROWMAJOR: - A = A + (ram[:, None] * stride_am + rk[None, :]) - else: - A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak) - if B_COLMAJOR: - B = B + (rk[:, None] + rbn[None, :] * stride_bn) - else: - B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn) - - acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) - - for k in range(K, 0, -BLOCK_K): - if EVEN_K: - a = tl.load(A) - b = tl.load(B) - else: - a = tl.load(A, mask=rk[None, :] < k, other=0.0) - b = tl.load(B, mask=rk[:, None] < k, other=0.0) - acc += tl.dot(a, b) - - if A_ROWMAJOR: - A += BLOCK_K - else: - A += BLOCK_K * stride_ak - if B_COLMAJOR: - B += BLOCK_K - else: - B += BLOCK_K * stride_bk - - # Putting bias after the matmul (instead of before) is faster, idk why - if BIAS: - bias = tl.load(bias + rn, mask=rn < N, other=0.0).to(tl.float32) - acc += bias[None, :] - - # optional: save the activation inputs - if SAVE_ACT_INPUT: - # act_in_ptrs = ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :] * stride_cn - act_in_ptrs = ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :] - tl.store(act_in_ptrs, acc) - - # optional: fused activation (while the data is in shared memory) - if ACTIVATION == "gelu": - acc = gelu(acc) - elif ACTIVATION == "gelu_approx": - acc = gelu_approx(acc) - elif ACTIVATION == "squared_relu": - acc = squared_relu(acc) - # rematerialize rm and rn to save registers - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - - # write back result - # C = C + rm[:, None] * stride_cm + rn[None, :] * stride_cn - C = C + rm[:, None] * stride_cm + rn[None, :] - mask = (rm < M)[:, None] & (rn < N)[None, :] - tl.store(C, acc) - - -def triton_linear_act( - x: torch.Tensor, - weight: torch.Tensor, - bias: Optional[torch.Tensor] = None, - activation: str = "id", - save_act_input: bool = False, -) -> torch.Tensor: - """ - Compute e = activation(x @ weight.T + bias). - This wrapper kicks the `kernel_fwd` Triton kernel - :param x: input tensor - :param weight: weight matrix - :param bias: an optional bias tensor - :param activation: Activation name. Needs to be a Triton kernel. - :param act_input: an optional tensor to save the activation inputs (for backward) - :return: result tensor - """ - # if torch.is_autocast_enabled(): - # dtype = torch.get_autocast_gpu_dtype() - # x, weight, bias = [a.to(dtype=dtype) for a in [x, weight, bias]] - - assert activation in ["id", "gelu", "gelu_approx", "squared_relu"] - - batch_shape, n = x.shape[:-1], x.shape[-1] - batch_dim = batch_shape.numel() - x_reshaped = x.reshape(batch_dim, n) - - if x_reshaped.stride(0) > 1 and x_reshaped.stride(1) > 1: - x_reshaped = x_reshaped.contiguous() - if weight.stride(0) > 1 and weight.stride(1) > 1: - weight = weight.contiguous() - bias = bias.contiguous() if bias is not None else None - - assert ( - x.dtype == weight.dtype - ), f"Input and weight must have the same dtype, got {x.dtype} and {weight.dtype}" - if bias is not None: - assert ( - x.dtype == bias.dtype - ), f"Input and bias must have the same dtype, got {x.dtype} and {bias.dtype}" - assert ( - x_reshaped.shape[1] == weight.shape[1] - ), f"Incompatible dimensions: {x_reshaped.shape} - {weight.shape}" - - assert ( - bias is None or bias.shape[0] == weight.shape[0] - ), "Incompatible dimensions in between weight and bias" - - M, K = x_reshaped.shape - N, K = weight.shape - - output = torch.empty((M, N), device=x.device, dtype=x.dtype) - act_input = torch.empty_like(output) if save_act_input else None - - # 1D launch kernel where each block gets its own program. - grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),) # noqa - - kernel_fwd[grid]( - output, - act_input, - x_reshaped, - weight, # data ptrs - bias if bias is not None else x, # auto skip bias if not present - M, # shapes - N, - K, - M // 32, # key for triton cache (limit number of compilations) - N // 32, - K // 32, - stride_cm=output.stride(0), # strides - # stride_cn=output.stride(1), - stride_am=x_reshaped.stride(0), - stride_ak=x_reshaped.stride(1), - stride_bk=weight.stride(1), - stride_bn=weight.stride(0), - BIAS=bias is not None, # optional fused bias - SAVE_ACT_INPUT=save_act_input, # optional save activation inputs - ACTIVATION=activation, # optional fused activation - A_ROWMAJOR=x_reshaped.stride(1) == 1, - B_COLMAJOR=weight.stride(1) == 1, - GROUP_M=8, # speed optimization: group the programs - ) - - if not save_act_input: - return output.reshape(*batch_shape, output.shape[-1]) - else: - return ( - output.reshape(*batch_shape, output.shape[-1]), - act_input.reshape(*batch_shape, act_input.shape[-1]), - ) - - -@triton.autotune( - configs=[ - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - # good for int8 - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=4, - num_warps=4, - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - ] - + get_configs_io_bound(), - key=["CACHE_KEY_M", "CACHE_KEY_N", "CACHE_KEY_K"], - prune_configs_by={ - "early_config_prune": early_config_prune, - "perf_model": estimate_matmul_time, - "top_k": 10, - }, -) -@triton.heuristics( - { - "EVEN_K": lambda args: args["K"] % (args["BLOCK_K"] * args["SPLIT_K"]) == 0, - } -) -@triton.jit -def kernel_bwd( - C, # Pointers to matrices - ACT_INPUT, - A, - B, - # Matrix dimensions - M, - N, - K, - CACHE_KEY_M, - CACHE_KEY_N, - CACHE_KEY_K, - # The stride variables represent how much to increase the ptr by when moving by 1 - # element in a particular dimension. E.g. stride_am is how much to increase a_ptr - # by to get the element one row down (A has M rows) - stride_cm, - # stride_cn, # Assume that stride_cn == 1 - stride_am, - stride_ak, - stride_bk, - stride_bn, - # Meta-parameters - BLOCK_M: tl.constexpr, - GROUP_M: tl.constexpr, - BLOCK_N: tl.constexpr, - BLOCK_K: tl.constexpr, - # split k not used, not performant with activation, kept because early_config_prune is expecting it - SPLIT_K: tl.constexpr, - EVEN_K: tl.constexpr, - ACTIVATION: tl.constexpr, -): - - """ - Kernel for computing Out = activation(A x W + C) - - Input has shape (M, K) - - Weight has shape (K, N) - - Output has shape (M, N) - - ActInputs (optional) has shape (M, N) - 'ActInputs' optionally saves the A x W + C intermediate for backward computations - This kernel will consolidate over K - """ - - pid = tl.program_id(axis=0) - - grid_m = (M + BLOCK_M - 1) // BLOCK_M - grid_n = (N + BLOCK_N - 1) // BLOCK_N - # re-order program ID for better L2 performance - width = GROUP_M * grid_n - group_id = pid // width - group_size = min(grid_m - group_id * GROUP_M, GROUP_M) - pid_m = group_id * GROUP_M + (pid % group_size) - pid_n = (pid % width) // (group_size) - - # now compute the block that each program will go through - # rm (resp. rn) denotes a range of indices - # for rows (resp. col) of C - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - # trick to avoid masking on M and N axis - ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M) - rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N) - rk = tl.arange(0, BLOCK_K) - - A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak) - B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn) - - acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) - - for k in range(K, 0, -BLOCK_K): - if EVEN_K: - a = tl.load(A) - b = tl.load(B) - else: - a = tl.load(A, mask=rk[None, :] < k, other=0.0) - b = tl.load(B, mask=rk[:, None] < k, other=0.0) - acc += tl.dot(a, b) - - A += BLOCK_K * stride_ak - B += BLOCK_K * stride_bk - - # optional: fused activation (while the data is in shared memory) - if ACTIVATION != "id": - act_in_ptrs = ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :] - act_input = tl.load(act_in_ptrs).to(acc.dtype) - if ACTIVATION == "gelu": - acc *= gelu_grad(act_input) - elif ACTIVATION == "gelu_approx": - acc *= gelu_approx_grad(act_input) - elif ACTIVATION == "squared_relu": - acc *= squared_relu_grad(act_input) - - # rematerialize rm and rn to save registers - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - - # write back result - C = C + rm[:, None] * stride_cm + rn[None, :] - mask = (rm < M)[:, None] & (rn < N)[None, :] - tl.store(C, acc, mask=mask) - - -def triton_dgrad_act( - grad_output: torch.Tensor, - weight: torch.Tensor, - activation: str = "id", - act_input: Optional[torch.Tensor] = None, -) -> torch.Tensor: - """ - Compute e = activation(grad_output @ weight + bias). - This wrapper kicks the `kernel_fwd` Triton kernel - :param grad_output: input tensor - :param weight: weight matrix - :param activation: Activation name. Needs to be a Triton kernel. - :param act_input: an optional tensor to save the activation inputs (for backward) - :return: result tensor - """ - assert activation in ["id", "gelu", "gelu_approx", "squared_relu"] - - batch_shape, n = grad_output.shape[:-1], grad_output.shape[-1] - batch_dim = batch_shape.numel() - grad_output_reshaped = grad_output.reshape(batch_dim, n) - - if grad_output_reshaped.stride(0) > 1 and grad_output_reshaped.stride(1) > 1: - grad_output_reshaped = grad_output_reshaped.contiguous() - if weight.stride(0) > 1 and weight.stride(1) > 1: - weight = weight.contiguous() - - assert ( - grad_output.dtype == weight.dtype - ), f"grad_output and weight must have the same dtype, got {grad_output.dtype} and {weight.dtype}" - assert ( - grad_output_reshaped.shape[1] == weight.shape[0] - ), f"Incompatible dimensions: {grad_output_reshaped.shape} - {weight.shape}" - if activation != "id": - assert act_input is not None, f"act_input is required for activation {activation}" - - # M, N, K in bwd are different from M, N, K in fwd - M, K = grad_output_reshaped.shape - K, N = weight.shape - - grad_input = torch.empty((M, N), device=grad_output.device, dtype=grad_output.dtype) - - # 1D launch kernel where each block gets its own program. - grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),) # noqa - - kernel_bwd[grid]( - grad_input, - act_input, - grad_output_reshaped, - weight, # data ptrs - M, # shapes - N, - K, - M // 32, # key for triton cache (limit number of compilations) - N // 32, - K // 32, - stride_cm=grad_input.stride(0), # strides - # stride_cn=grad_input.stride(1), - stride_am=grad_output_reshaped.stride(0), - stride_ak=grad_output_reshaped.stride(1), - stride_bk=weight.stride(0), - stride_bn=weight.stride(1), - ACTIVATION=activation, # optional fused activation - GROUP_M=8, # speed optimization: group the programs - ) - - return grad_input.reshape(*batch_shape, grad_input.shape[-1]) diff --git a/build/torch211-cxx11-cu130-x86_64-linux/ops/triton/mlp.py b/build/torch211-cxx11-cu130-x86_64-linux/ops/triton/mlp.py deleted file mode 100644 index 059f4f8a5e174c1f4824e43d313fca18eaa799b8..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu130-x86_64-linux/ops/triton/mlp.py +++ /dev/null @@ -1,149 +0,0 @@ -# The triton fused matmul + sqrelu is faster for fp16 but slower for bf16, compared -# to naive implementation. -import fused_dense_lib as fused_dense_cuda -import torch -import torch.nn as nn -import torch.nn.functional as F - -from flash_attn.utils.torch import custom_fwd, custom_bwd -from flash_attn.ops.activations import sqrelu_bwd, sqrelu_fwd -from flash_attn.ops.triton.linear import triton_dgrad_act, triton_linear_act - - -class FusedDenseSqreluDenseFunc(torch.autograd.Function): - @staticmethod - @custom_fwd - def forward(ctx, x, weight1, bias1, weight2, bias2, checkpoint_lvl=0): - """checkpoint_lvl: - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute act_input and gelu_out in the bwd - """ - if torch.is_autocast_enabled(): - dtype = torch.get_autocast_gpu_dtype() - x, weight1, bias1, weight2, bias2 = [ - a.to(dtype=dtype) for a in [x, weight1, bias1, weight2, bias2] - ] - is_bf16 = x.dtype == torch.bfloat16 - assert checkpoint_lvl in [0, 1, 2] - x = x.contiguous() - weight1 = weight1.contiguous() - bias1 = bias1.contiguous() - weight2 = weight2.contiguous() - bias2 = bias2.contiguous() - batch_shape, n = x.shape[:-1], x.shape[-1] - batch_dim = batch_shape.numel() - if is_bf16: - act_input = fused_dense_cuda.linear_bias_forward( - x.reshape(batch_dim, n), weight1, bias1 - ) - output1 = sqrelu_fwd(act_input) - else: - save_act_input = checkpoint_lvl != 2 - result = triton_linear_act( - x.reshape(batch_dim, n), - weight1, - bias1, - activation="squared_relu", - save_act_input=save_act_input, - ) - if save_act_input: - output1, act_input = result - else: - output1 = result - output2 = fused_dense_cuda.linear_bias_forward(output1, weight2, bias2) - ctx.checkpoint_lvl = checkpoint_lvl - if checkpoint_lvl == 0: - ctx.save_for_backward(x, weight1, bias1, weight2, act_input, output1) - elif checkpoint_lvl == 1: - ctx.save_for_backward(x, weight1, bias1, weight2, act_input) - elif checkpoint_lvl == 2: - ctx.save_for_backward(x, weight1, bias1, weight2) - return output2.reshape(*batch_shape, output2.shape[-1]) - - @staticmethod - @custom_bwd - def backward(ctx, grad_output): - grad_output = grad_output.contiguous() - checkpoint_lvl = ctx.checkpoint_lvl - x, weight1, bias1, weight2, *rest = ctx.saved_tensors - batch_shape, n = x.shape[:-1], x.shape[-1] - batch_dim = batch_shape.numel() - is_bf16 = x.dtype == torch.bfloat16 - if checkpoint_lvl == 0: - act_input, output1 = rest - elif checkpoint_lvl == 1: - (act_input,) = rest - output1 = sqrelu_fwd(act_input) - elif checkpoint_lvl == 2: - if is_bf16: - act_input = fused_dense_cuda.linear_bias_forward( - x.reshape(batch_dim, n), weight1, bias1 - ) - output1 = sqrelu_fwd(act_input) - else: - output1, act_input = triton_linear_act( - x.reshape(batch_dim, n), - weight1, - bias1, - activation="squared_relu", - save_act_input=True, - ) - - if is_bf16: - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad(output1, grad_output) - grad_output1 = grad_output @ weight2 - grad_act_input = sqrelu_bwd(grad_output1, act_input) - grad_input, grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_backward( - x.reshape(batch_dim, n), weight1, grad_act_input - ) - else: - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad(output1, grad_output) - grad_act_input = triton_dgrad_act( - grad_output, weight2, activation="squared_relu", act_input=act_input - ) - grad_input, grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_backward( - x.reshape(batch_dim, n), weight1, grad_act_input - ) - return grad_input.reshape_as(x), grad_weight1, grad_bias1, grad_weight2, grad_bias2, None - - -fused_dense_sqrelu_dense_function = FusedDenseSqreluDenseFunc.apply - - -class FusedDenseSqreluDense(nn.Module): - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - bias1=True, - bias2=True, - checkpoint_lvl=0, - device=None, - dtype=None, - ): - """ - checkpoint_lvl (increasing lvl means slower but more memory saving): - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute gelu_in and gelu_out in the bwd - """ - assert checkpoint_lvl in [0, 1, 2] - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features * 4 - assert bias1 == True, "DenseSqreluDense module without bias is currently not supported" - assert bias2 == True, "DenseSqreluDense module without bias is currently not supported" - self.checkpoint_lvl = checkpoint_lvl - self.fc1 = nn.Linear(in_features, hidden_features, bias=bias1, **factory_kwargs) - self.fc2 = nn.Linear(hidden_features, out_features, bias=bias2, **factory_kwargs) - - def forward(self, x): - assert x.is_cuda - return fused_dense_sqrelu_dense_function( - x, self.fc1.weight, self.fc1.bias, self.fc2.weight, self.fc2.bias, self.checkpoint_lvl - ) diff --git a/build/torch211-cxx11-cu130-x86_64-linux/ops/triton/rotary.py b/build/torch211-cxx11-cu130-x86_64-linux/ops/triton/rotary.py deleted file mode 100644 index 6ad3ab412a084270120167ad32b67ad1720f2c37..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-cu130-x86_64-linux/ops/triton/rotary.py +++ /dev/null @@ -1,186 +0,0 @@ -# Copyright (c) 2025, Tri Dao. -# As of 2025-04-23, we require triton >= 3.0 - -from typing import Optional, Union - -import torch - -import triton -import triton.language as tl - - -@triton.jit -def rotary_kernel( - OUT, # Pointers to matrices - X, - COS, - SIN, - CU_SEQLENS, - SEQLEN_OFFSETS, # this could be int or a pointer - # Matrix dimensions - seqlen, - nheads, - seqlen_ro, - # strides - stride_out_batch, - stride_out_seqlen, - stride_out_nheads, - stride_out_headdim, - stride_x_batch, - stride_x_seqlen, - stride_x_nheads, - stride_x_headdim, - # Meta-parameters - # We want ROTARY_DIM to be constexpr, otherwise the triton compiler doesn't know that - # the mask is constant every 8 elements, and it will generate LDG.16 instead of LDG.128 - ROTARY_DIM: tl.constexpr, - IS_SEQLEN_OFFSETS_TENSOR: tl.constexpr, - IS_VARLEN: tl.constexpr, - INTERLEAVED: tl.constexpr, - CONJUGATE: tl.constexpr, - BLOCK_H: tl.constexpr, - BLOCK_M: tl.constexpr, -): - BLOCK_K: tl.constexpr = triton.next_power_of_2(ROTARY_DIM) - ROTARY_DIM_HALF = ROTARY_DIM // 2 - pid_head = tl.program_id(axis=0) - pid_m = tl.program_id(axis=1) - pid_batch = tl.program_id(axis=2) - - if not IS_VARLEN: - X = X + pid_batch * stride_x_batch - OUT = OUT + pid_batch * stride_out_batch - else: - start_idx = tl.load(CU_SEQLENS + pid_batch) - seqlen = tl.load(CU_SEQLENS + pid_batch + 1) - start_idx - X = X + start_idx * stride_x_seqlen - OUT = OUT + start_idx * stride_out_seqlen - - if pid_m * BLOCK_M >= seqlen: - return - - rh = pid_head * BLOCK_H + tl.arange(0, BLOCK_H) - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - if not IS_SEQLEN_OFFSETS_TENSOR: - rm_cs = rm + SEQLEN_OFFSETS - else: - rm_cs = rm + tl.load(SEQLEN_OFFSETS + pid_batch) - - rk_half = tl.arange(0, BLOCK_K // 2) - COS = COS + (rm_cs[:, None] * ROTARY_DIM_HALF + rk_half[None, :]) - SIN = SIN + (rm_cs[:, None] * ROTARY_DIM_HALF + rk_half[None, :]) - mask_cs = (rm_cs[:, None] < seqlen_ro) & (rk_half[None, :] < ROTARY_DIM_HALF) - cos = tl.load(COS, mask=mask_cs, other=1.0).to(tl.float32) - sin = tl.load(SIN, mask=mask_cs, other=0.0).to(tl.float32) - if CONJUGATE: - sin = -sin - - if not INTERLEAVED: - # Load the 1st and 2nd halves of X, do calculation, then store to 1st and 2nd halves of OUT - X = X + (rh[:, None, None] * stride_x_nheads + rm[None, :, None] * stride_x_seqlen + rk_half[None, None, :] * stride_x_headdim) - OUT = OUT + (rh[:, None, None] * stride_out_nheads + rm[None, :, None] * stride_out_seqlen + rk_half[None, None, :] * stride_out_headdim) - mask = (rh[:, None, None] < nheads) & (rm[None, :, None] < seqlen) & (rk_half[None, None, :] < ROTARY_DIM_HALF) - x0 = tl.load(X, mask=mask, other=0.0).to(tl.float32) - x1 = tl.load(X + ROTARY_DIM_HALF * stride_x_headdim, mask=mask, other=0.0,).to(tl.float32) - o0 = x0 * cos - x1 * sin - o1 = x0 * sin + x1 * cos - tl.store(OUT, o0, mask=mask) - tl.store(OUT + ROTARY_DIM_HALF * stride_out_headdim, o1, mask=mask) - else: - rk = tl.arange(0, BLOCK_K) - X = X + (rh[:, None, None] * stride_x_nheads + rm[None, :, None] * stride_x_seqlen + rk[None, None, :] * stride_x_headdim) - OUT = OUT + (rh[:, None, None] * stride_out_nheads + rm[None, :, None] * stride_out_seqlen + rk[None, None, :] * stride_out_headdim) - mask = (rh[:, None, None] < nheads) & (rm[None, :, None] < seqlen) & (rk[None, None, :] < ROTARY_DIM) - x = tl.load(X, mask=mask, other=0.0).to(tl.float32) - x0, x1 = tl.split(tl.reshape(x, [BLOCK_H, BLOCK_M, BLOCK_K // 2, 2])) - o0 = x0 * cos - x1 * sin - o1 = x0 * sin + x1 * cos - o = tl.reshape(tl.join(o0, o1), [BLOCK_H, BLOCK_M, BLOCK_K]) - tl.store(OUT, o, mask=mask) - - -def apply_rotary( - x: torch.Tensor, - cos: torch.Tensor, - sin: torch.Tensor, - seqlen_offsets: Union[int, torch.Tensor] = 0, - cu_seqlens: Optional[torch.Tensor] = None, - max_seqlen: Optional[int] = None, - interleaved=False, - inplace=False, - conjugate=False, -) -> torch.Tensor: - """ - Arguments: - x: (batch, seqlen, nheads, headdim) if cu_seqlens is None - else (total_seqlen, nheads, headdim). - cos: (seqlen_ro, rotary_dim / 2) - sin: (seqlen_ro, rotary_dim / 2) - seqlen_offsets: integer or integer tensor of size (batch,) - cu_seqlens: (batch + 1,) or None - max_seqlen: int - Returns: - y: (batch, seqlen, nheads, headdim) - """ - is_varlen = cu_seqlens is not None - if not is_varlen: - batch, seqlen, nheads, headdim = x.shape - else: - assert max_seqlen is not None, "If cu_seqlens is passed in, then max_seqlen must be passed" - total_seqlen, nheads, headdim = x.shape - batch_p_1 = cu_seqlens.shape[0] - batch = batch_p_1 - 1 - seqlen = max_seqlen - seqlen_ro, rotary_dim = cos.shape - assert sin.shape == cos.shape - rotary_dim *= 2 - assert rotary_dim <= headdim, "rotary_dim must be <= headdim" - assert headdim <= 256, "Only support headdim <= 256" - assert seqlen_ro >= seqlen, "seqlen_ro must be >= seqlen" - - cos, sin = cos.contiguous(), sin.contiguous() - if isinstance(seqlen_offsets, torch.Tensor): - assert seqlen_offsets.shape == (batch,) - assert seqlen_offsets.dtype in [torch.int32, torch.int64] - seqlen_offsets = seqlen_offsets.contiguous() - else: - assert seqlen_offsets + seqlen <= seqlen_ro - - output = torch.empty_like(x) if not inplace else x - if rotary_dim < headdim and not inplace: - output[..., rotary_dim:].copy_(x[..., rotary_dim:]) - - grid = lambda META: (triton.cdiv(nheads, META["BLOCK_H"]), triton.cdiv(seqlen, META["BLOCK_M"]), batch) # noqa - BLOCK_M = 8 if rotary_dim <= 128 else 4 - - # Need this, otherwise Triton tries to launch from cuda:0 and we get - # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?) - device_ctx = torch.cuda.device(x.device.index) if x.device.type == 'cuda' else torch.xpu.device(x.device.index) - with device_ctx: - torch.library.wrap_triton(rotary_kernel)[grid]( - output, # data ptrs - x, - cos, - sin, - cu_seqlens, - seqlen_offsets, - seqlen, # shapes - nheads, - seqlen_ro, - output.stride(0) if not is_varlen else 0, # batch_strides if not varlen else 0 - output.stride(-3), # seqlen_stride or total_seqlen_stride - output.stride(-2), # nheads_stride - output.stride(-1), # headdim_stride - x.stride(0) if not is_varlen else 0, # batch_strides if not varlen else 0 - x.stride(-3), # seqlen stride or total_seqlen_stride - x.stride(-2), # nheads stride - x.stride(-1), # headdim stride - rotary_dim, - isinstance(seqlen_offsets, torch.Tensor), - is_varlen, - interleaved, - conjugate, - BLOCK_M=BLOCK_M, - BLOCK_H=2, - ) - return output diff --git a/build/torch211-cxx11-xpu20253-x86_64-linux/__init__.py b/build/torch211-cxx11-xpu20253-x86_64-linux/__init__.py deleted file mode 100644 index ecc2f9d896b6c93f90b0a1499856dc0612177422..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-xpu20253-x86_64-linux/__init__.py +++ /dev/null @@ -1,393 +0,0 @@ -from typing import Optional, List -import torch -from ._ops import ops as flash_attn_ops -from .flash_attn_interface import ( - flash_attn_func, - flash_attn_kvpacked_func, - flash_attn_qkvpacked_func, - flash_attn_varlen_func, - flash_attn_varlen_kvpacked_func, - flash_attn_varlen_qkvpacked_func, - flash_attn_with_kvcache, -) - - -def fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - return_softmax: bool = False, - gen: Optional[torch.Generator] = None, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention. - - Args: - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Optional output tensor, same shape as q - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - return_softmax: Whether to return softmax weights - gen: Optional random number generator - - Returns: - List of tensors: [output, softmax_lse, (softmax if return_softmax)] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.fwd( - q, - k, - v, - out, - alibi_slopes, - p_dropout, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - gen, - ) - - -def varlen_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - out: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - block_table: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - max_seqlen_q: int = 0, - max_seqlen_k: int = 0, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - zero_tensors: bool = False, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - return_softmax: bool = False, - gen: Optional[torch.Generator] = None, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention with variable sequence lengths. - - Args: - q: Query tensor of shape [total_q, num_heads, head_size] - k: Key tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - v: Value tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1] - cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1] - out: Optional output tensor of shape [total_q, num_heads, head_size] - seqused_k: Optional tensor specifying how many keys to use per batch element [batch_size] - leftpad_k: Optional left padding for keys of shape [batch_size] - block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq] - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - max_seqlen_q: Maximum sequence length for queries - max_seqlen_k: Maximum sequence length for keys - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - zero_tensors: Whether to zero tensors before computation - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - return_softmax: Whether to return softmax weights - gen: Optional random number generator - - Returns: - List of tensors: [output, softmax_lse, (softmax if return_softmax)] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.varlen_fwd( - q, - k, - v, - out, - cu_seqlens_q, - cu_seqlens_k, - seqused_k, - leftpad_k, - block_table, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - p_dropout, - softmax_scale, - zero_tensors, - is_causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - gen, - ) - - -def bwd( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - p_dropout: float = 0.0, - 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, - gen: Optional[torch.Generator] = None, - rng_state: Optional[torch.Tensor] = None, -) -> List[torch.Tensor]: - """ - Backward pass for multi-head attention. - - Args: - dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size] - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size] - softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q] - dq: Optional gradient tensor for queries, same shape as q - dk: Optional gradient tensor for keys, same shape as k - dv: Optional gradient tensor for values, same shape as v - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - deterministic: Whether to use deterministic algorithms - gen: Optional random number generator - rng_state: Optional RNG state from forward pass - - Returns: - List of tensors: [dq, dk, dv] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - alibi_slopes, - p_dropout, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - deterministic, - gen, - rng_state, - ) - - -def varlen_bwd( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - max_seqlen_q: int = 0, - max_seqlen_k: int = 0, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - zero_tensors: bool = False, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - deterministic: bool = False, - gen: Optional[torch.Generator] = None, - rng_state: Optional[torch.Tensor] = None, -) -> List[torch.Tensor]: - """ - Backward pass for multi-head attention with variable sequence lengths. - - Args: - dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size] - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size] - softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q] - cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1] - cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1] - dq: Optional gradient tensor for queries, same shape as q - dk: Optional gradient tensor for keys, same shape as k - dv: Optional gradient tensor for values, same shape as v - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - max_seqlen_q: Maximum sequence length for queries - max_seqlen_k: Maximum sequence length for keys - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - zero_tensors: Whether to zero tensors before computation - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - deterministic: Whether to use deterministic algorithms - gen: Optional random number generator - rng_state: Optional RNG state from forward pass - - Returns: - List of tensors: [dq, dk, dv] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.varlen_bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - p_dropout, - softmax_scale, - zero_tensors, - is_causal, - window_size_left, - window_size_right, - softcap, - deterministic, - gen, - rng_state, - ) - - -def fwd_kvcache( - q: torch.Tensor, - kcache: torch.Tensor, - vcache: torch.Tensor, - k: Optional[torch.Tensor] = None, - v: Optional[torch.Tensor] = None, - seqlens_k: Optional[torch.Tensor] = None, - rotary_cos: Optional[torch.Tensor] = None, - rotary_sin: Optional[torch.Tensor] = None, - cache_batch_idx: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - block_table: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - out: 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, - is_rotary_interleaved: bool = False, - num_splits: int = 1, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention with KV cache. - - Args: - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - kcache: Key cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - vcache: Value cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - k: Optional new keys tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size] - v: Optional new values tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size] - seqlens_k: Optional sequence lengths for keys of shape [batch_size] - rotary_cos: Optional rotary cosine tensor of shape [seqlen_ro, rotary_dim/2] - rotary_sin: Optional rotary sine tensor of shape [seqlen_ro, rotary_dim/2] - cache_batch_idx: Optional indices to index into the KV cache - leftpad_k: Optional left padding for keys of shape [batch_size] - block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq] - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - out: Optional output tensor, same shape as q - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - is_rotary_interleaved: Whether rotary embeddings are interleaved - num_splits: Number of splits for computation - - Returns: - List of tensors: [output, softmax_lse] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.fwd_kvcache( - q, - kcache, - vcache, - k, - v, - seqlens_k, - rotary_cos, - rotary_sin, - cache_batch_idx, - leftpad_k, - block_table, - alibi_slopes, - out, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - is_rotary_interleaved, - num_splits, - ) diff --git a/build/torch211-cxx11-xpu20253-x86_64-linux/_flash_attn2_xpu_f12afc9.abi3.so b/build/torch211-cxx11-xpu20253-x86_64-linux/_flash_attn2_xpu_f12afc9.abi3.so deleted file mode 100644 index 00e192f6d5a4d3442d5d1f5e0c0ba019b39b5616..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-xpu20253-x86_64-linux/_flash_attn2_xpu_f12afc9.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:56c6bc9987537f70b10f8d8cca00c0d3f3ff3aefdb9f2f32f01e232da786a7bf -size 82767552 diff --git a/build/torch211-cxx11-xpu20253-x86_64-linux/_ops.py b/build/torch211-cxx11-xpu20253-x86_64-linux/_ops.py deleted file mode 100644 index ad32616257d5cef8c756337dded818dda050fc0c..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-xpu20253-x86_64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _flash_attn2_xpu_f12afc9 -ops = torch.ops._flash_attn2_xpu_f12afc9 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_flash_attn2_xpu_f12afc9::{op_name}" diff --git a/build/torch211-cxx11-xpu20253-x86_64-linux/bert_padding.py b/build/torch211-cxx11-xpu20253-x86_64-linux/bert_padding.py deleted file mode 100644 index 3c2d35159a014a9d03aabead9e52e009168696ea..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-xpu20253-x86_64-linux/bert_padding.py +++ /dev/null @@ -1,218 +0,0 @@ -# Adapted from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/padding.py - -import torch -import torch.nn.functional as F -from einops import rearrange, repeat - - -class IndexFirstAxis(torch.autograd.Function): - @staticmethod - def forward(ctx, input, indices): - ctx.save_for_backward(indices) - assert input.ndim >= 2 - ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] - second_dim = other_shape.numel() - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - # return input[indices] - return torch.gather( - rearrange(input, "b ... -> b (...)"), 0, repeat(indices, "z -> z d", d=second_dim) - ).reshape(-1, *other_shape) - - @staticmethod - def backward(ctx, grad_output): - (indices,) = ctx.saved_tensors - assert grad_output.ndim >= 2 - other_shape = grad_output.shape[1:] - grad_output = rearrange(grad_output, "b ... -> b (...)") - grad_input = torch.zeros( - [ctx.first_axis_dim, grad_output.shape[1]], - device=grad_output.device, - dtype=grad_output.dtype, - ) - # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. - # grad_input[indices] = grad_output - grad_input.scatter_(0, repeat(indices, "z -> z d", d=grad_output.shape[1]), grad_output) - return grad_input.reshape(ctx.first_axis_dim, *other_shape), None - - -index_first_axis = IndexFirstAxis.apply - - -class IndexPutFirstAxis(torch.autograd.Function): - @staticmethod - def forward(ctx, values, indices, first_axis_dim): - ctx.save_for_backward(indices) - assert indices.ndim == 1 - assert values.ndim >= 2 - output = torch.zeros( - first_axis_dim, *values.shape[1:], device=values.device, dtype=values.dtype - ) - # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. - output[indices] = values - # output.scatter_(0, repeat(indices, 'z -> z d', d=values.shape[1]), values) - return output - - @staticmethod - def backward(ctx, grad_output): - (indices,) = ctx.saved_tensors - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - grad_values = grad_output[indices] - # grad_values = torch.gather(grad_output, 0, repeat(indices, 'z -> z d', d=grad_output.shape[1])) - return grad_values, None, None - - -index_put_first_axis = IndexPutFirstAxis.apply - - -class IndexFirstAxisResidual(torch.autograd.Function): - @staticmethod - def forward(ctx, input, indices): - ctx.save_for_backward(indices) - assert input.ndim >= 2 - ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] - second_dim = other_shape.numel() - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - output = input[indices] - # We don't want to reshape input (b ... -> b (...)) since it could change the channel_last - # memory format to channel_first. In other words, input might not be contiguous. - # If we don't detach, Pytorch complains about output being a view and is being modified inplace - return output, input.detach() - - @staticmethod - def backward(ctx, grad_output, grad_residual): - (indices,) = ctx.saved_tensors - assert grad_output.ndim >= 2 - other_shape = grad_output.shape[1:] - assert grad_residual.shape[1:] == other_shape - grad_input = grad_residual - # grad_input[indices] += grad_output - indices = indices.reshape(indices.shape[0], *((1,) * (grad_output.ndim - 1))) - indices = indices.expand_as(grad_output) - grad_input.scatter_add_(0, indices, grad_output) - return grad_input.reshape(ctx.first_axis_dim, *other_shape), None - - -index_first_axis_residual = IndexFirstAxisResidual.apply - - -def unpad_input(hidden_states, attention_mask, unused_mask=None): - """ - Arguments: - hidden_states: (batch, seqlen, ...) - attention_mask: (batch, seqlen), bool / int, 1 means valid and 0 means not valid. - unused_mask: (batch, seqlen), bool / int, 1 means the element is allocated but unused. - Return: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens selected in attention_mask + unused_mask. - indices: (total_nnz), the indices of masked tokens from the flattened input sequence. - cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. - max_seqlen_in_batch: int - seqused: (batch), returns the number of tokens selected in attention_mask + unused_mask. - """ - all_masks = (attention_mask + unused_mask) if unused_mask is not None else attention_mask - seqlens_in_batch = all_masks.sum(dim=-1, dtype=torch.int32) - used_seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) - indices = torch.nonzero(all_masks.flatten(), as_tuple=False).flatten() - max_seqlen_in_batch = seqlens_in_batch.max().item() - cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) - # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the - # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim - # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to - # index with integer indices. Moreover, torch's index is a bit slower than it needs to be, - # so we write custom forward and backward to make it a bit faster. - return ( - index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices), - indices, - cu_seqlens, - max_seqlen_in_batch, - used_seqlens_in_batch, - ) - - -def unpad_input_for_concatenated_sequences(hidden_states, attention_mask_in_length): - """ - Supports concatenating short samples in one sequence. The attention_mask_in_length is utilized to mask other short samples. It helps efficient training of variant lengths-based samples (e.g., the supervised fine-tuning task in large language model). - The motivation for this function is explained [here](https://github.com/Dao-AILab/flash-attention/issues/432#issuecomment-1668822286). - - For example, if batch = 3 and seqlen = 6, the attention_mask_in_length is: - ``` - [ - [2, 3, 0, 0, 0, 0], - [3, 2, 0, 0, 0, 0], - [6, 0, 0, 0, 0, 0] - ] - ``` - , which refers to the 3D-attention mask: - ``` - [ - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [0, 0, 1, 0, 0, 0], - [0, 0, 1, 1, 0, 0], - [0, 0, 1, 1, 1, 0], - [0, 0, 0, 0, 0, 1] - ], - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [0, 0, 0, 1, 0, 0], - [0, 0, 0, 1, 1, 0], - [0, 0, 0, 0, 0, 1] - ], - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [1, 1, 1, 1, 0, 0], - [1, 1, 1, 1, 1, 0], - [1, 1, 1, 1, 1, 1] - ] - ] - ```. - - Arguments: - hidden_states: (batch, seqlen, ...) - attention_mask_in_length: (batch, seqlen), int, a nonzero number (e.g., 1, 2, 3, etc.) means length of concatenated sequence in b-th batch, and 0 means none. - Return: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. - indices: (total_nnz), the indices of non-masked tokens from the flattened input sequence. - cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. - max_seqlen_in_batch: int - """ - length = attention_mask_in_length.sum(dim=-1) - seqlen = attention_mask_in_length.size(-1) - attention_mask_2d = torch.arange(seqlen, device=length.device, dtype=length.dtype).expand(len(length), seqlen) < length.unsqueeze(1) - real_indices_idx = torch.nonzero(attention_mask_in_length.flatten(), as_tuple=False).flatten() - seqlens_in_batch = attention_mask_in_length.flatten()[real_indices_idx] - indices = torch.nonzero(attention_mask_2d.flatten(), as_tuple=False).flatten() - max_seqlen_in_batch = seqlens_in_batch.max().item() - cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) - # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the - # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim - # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to - # index with integer indices. Moreover, torch's index is a bit slower than it needs to be, - # so we write custom forward and backward to make it a bit faster. - return ( - index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices), - indices, - cu_seqlens, - max_seqlen_in_batch, - ) - - -def pad_input(hidden_states, indices, batch, seqlen): - """ - Arguments: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. - indices: (total_nnz), the indices that represent the non-masked tokens of the original padded input sequence. - batch: int, batch size for the padded sequence. - seqlen: int, maximum sequence length for the padded sequence. - Return: - hidden_states: (batch, seqlen, ...) - """ - dim = hidden_states.shape[-1] - # output = torch.zeros((batch * seqlen), dim, device=hidden_states.device, dtype=hidden_states.dtype) - # output[indices] = hidden_states - output = index_put_first_axis(hidden_states, indices, batch * seqlen) - return rearrange(output, "(b s) ... -> b s ...", b=batch) diff --git a/build/torch211-cxx11-xpu20253-x86_64-linux/flash_attn2/__init__.py b/build/torch211-cxx11-xpu20253-x86_64-linux/flash_attn2/__init__.py deleted file mode 100644 index a9b2672c1cd85b74c1b3ded0fc0b2100e1aeac23..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-xpu20253-x86_64-linux/flash_attn2/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -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-xpu20253-x86_64-linux/flash_attn_interface.py b/build/torch211-cxx11-xpu20253-x86_64-linux/flash_attn_interface.py deleted file mode 100644 index b014e569dc9812a9c840d1043f2cebe767b4d477..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-xpu20253-x86_64-linux/flash_attn_interface.py +++ /dev/null @@ -1,1631 +0,0 @@ -# Copyright (c) 2023, Tri Dao. - -from typing import Optional, Sequence, Tuple, Union - -import torch -import torch.nn as nn -import os - -# # isort: off -# # We need to import the CUDA kernels after importing torch -# USE_TRITON_ROCM = os.getenv("FLASH_ATTENTION_TRITON_AMD_ENABLE", "FALSE") == "TRUE" -# if USE_TRITON_ROCM: -# from .flash_attn_triton_amd import interface_fa as flash_attn -# else: -# import flash_attn_2_cuda as flash_attn - - -from . import _ops -from ._ops import ops as flash_attn -from ._ops import add_op_namespace_prefix - -# # isort: on - -def maybe_contiguous(x): - return x.contiguous() if x is not None and x.stride(-1) != 1 else x - - -def _get_device(): - if torch.xpu.is_available(): - return "xpu" - elif torch.cuda.is_available(): - return "cuda" - else: - return "cpu" - - -def _get_block_size_n(device, head_dim, is_dropout, is_causal): - assert head_dim <= 256 - - if device.type == "xpu": - if head_dim <= 96: - return 64 - elif head_dim <= 128: - return 32 - elif head_dim <= 256: - return 64 - else: - return 32 - - # This should match the block sizes in the CUDA kernel - major, minor = torch.cuda.get_device_capability(device) - is_sm8x = major == 8 and minor > 0 # Only include sm86 and sm89, exclude sm80 (A100) - is_sm80 = major == 8 and minor == 0 - is_sm90 = major == 9 and minor == 0 - if head_dim <= 32: - return 128 - if head_dim <= 64: - return 128 if not is_dropout else 64 - elif head_dim <= 96: - return 64 - elif head_dim <= 128: - if is_sm8x: - return 64 if (not is_dropout and is_causal) else 32 - else: - return 64 if not is_dropout else 32 - elif head_dim <= 192: - return 64 - elif head_dim <= 224: - return 64 - elif head_dim <= 256: - return 64 - - -def round_multiple(x, m): - return (x + m - 1) // m * m - - -# torch.compile() support is only enabled for pytorch >= 2.4 -# The reason for this is that we are using the new custom_op and register_fake -# APIs, which support inplace modification of inputs in the function itself -if torch.__version__ >= "2.4.0": - _torch_custom_op_wrapper = torch.library.custom_op - _torch_register_fake_wrapper = torch.library.register_fake -else: - def noop_custom_op_wrapper(name, fn=None, /, *, mutates_args, device_types=None, schema=None): - def wrap(func): - return func - if fn is None: - return wrap - return fn - def noop_register_fake_wrapper(op, fn=None, /, *, lib=None, _stacklevel=1): - def wrap(func): - return func - if fn is None: - return wrap - return fn - _torch_custom_op_wrapper = noop_custom_op_wrapper - _torch_register_fake_wrapper = noop_register_fake_wrapper - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_forward"), mutates_args=(), device_types=_get_device()) -def _flash_attn_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - return_softmax: bool -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - out, softmax_lse, S_dmask, rng_state = flash_attn.fwd( - q, - k, - v, - None, - alibi_slopes, - dropout_p, - softmax_scale, - causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - None, - ) - return out, softmax_lse, S_dmask, rng_state - - -@_torch_register_fake_wrapper(add_op_namespace_prefix("_flash_attn_forward")) -def _flash_attn_forward_fake( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - return_softmax: bool -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - batch_size, seqlen_q, num_heads, head_size = q.shape - seqlen_k = k.shape[1] - out = torch.empty_like(q) - softmax_lse = torch.empty((batch_size, num_heads, seqlen_q), dtype=torch.float32, device=q.device, layout=q.layout) - p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) - if return_softmax: - p = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128), round_multiple(seqlen_k, 128)), dtype=q.dtype, device=q.device, layout=q.layout) - rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) - - return out, softmax_lse, p, rng_state - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_forward = _ops.ops._flash_attn_forward -else: - _wrapped_flash_attn_forward = _flash_attn_forward - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_varlen_forward"), mutates_args=(), device_types=_get_device()) -def _flash_attn_varlen_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - alibi_slopes: Optional[torch.Tensor] = None, - return_softmax: bool = False, - block_table: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - out, softmax_lse, S_dmask, rng_state = flash_attn.varlen_fwd( - q, - k, - v, - None, - cu_seqlens_q, - cu_seqlens_k, - seqused_k, - leftpad_k, - block_table, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - zero_tensors, - causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - None, - ) - # if out.isnan().any() or softmax_lse.isnan().any(): - # breakpoint() - return out, softmax_lse, S_dmask, rng_state - - -@_torch_register_fake_wrapper(add_op_namespace_prefix("_flash_attn_varlen_forward")) -def _flash_attn_varlen_forward_fake( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - alibi_slopes: Optional[torch.Tensor] = None, - return_softmax: bool = False, - block_table: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - paged_kv = block_table is not None - batch_size = cu_seqlens_q.numel() - 1 - total_q, num_heads, _ = q.shape - - out = torch.empty_like(q) - softmax_lse = torch.empty((num_heads, total_q), dtype=torch.float32, device=q.device, layout=q.layout) - p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) - seqlen_q_rounded = round_multiple(max_seqlen_q, 128) - seqlen_k_rounded = round_multiple(max_seqlen_k, 128) - if return_softmax: - p = torch.empty((batch_size, num_heads, seqlen_q_rounded, seqlen_k_rounded), dtype=q.dtype, device=q.device, layout=q.layout) - rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) - return out, softmax_lse, p, rng_state - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_varlen_forward = _ops.ops._flash_attn_varlen_forward -else: - _wrapped_flash_attn_varlen_forward = _flash_attn_varlen_forward - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_backward"), mutates_args=("dq", "dk", "dv"), device_types=_get_device()) -def _flash_attn_backward( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, -) -> 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)] - ( - dq, - dk, - dv, - softmax_d, - ) = flash_attn.bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - alibi_slopes, - dropout_p, - softmax_scale, - causal, - window_size_left, - window_size_right, - softcap, - deterministic, - None, - rng_state, - ) - return softmax_d - - -@_torch_register_fake_wrapper(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, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, -) -> torch.Tensor: - dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] - if dq is None: - dq = torch.empty_like(q) - if dk is None: - dk = torch.empty_like(k) - if dv is None: - dv = torch.empty_like(v) - batch_size, seqlen_q, num_heads, _ = q.shape - softmax_d = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128)), device=q.device, dtype=torch.float32) - - return softmax_d - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_backward = _ops.ops._flash_attn_backward -else: - _wrapped_flash_attn_backward = _flash_attn_backward - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_varlen_backward"), mutates_args=("dq", "dk", "dv"), device_types=_get_device()) -def _flash_attn_varlen_backward( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> 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)] - ( - dq, - dk, - dv, - softmax_d, - ) = flash_attn.varlen_bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - zero_tensors, - causal, - window_size_left, - window_size_right, - softcap, - deterministic, - None, - rng_state, - ) - # if dk.isnan().any() or dk.isnan().any() or dv.isnan().any() or softmax_d.isnan().any(): - # breakpoint() - return softmax_d - - -@_torch_register_fake_wrapper(add_op_namespace_prefix("_flash_attn_varlen_backward")) -def _flash_attn_varlen_backward_fake( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> torch.Tensor: - dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] - batch_size = cu_seqlens_q.numel() - 1 - total_q, num_heads, _ = q.shape - - if dq is None: - dq = torch.empty_like(q) - if dk is None: - dk = torch.empty_like(k) - if dv is None: - dv = torch.empty_like(v) - softmax_d = torch.empty((num_heads, total_q + 128 * batch_size), device=q.device, dtype=torch.float32) - - return softmax_d - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_varlen_backward = _ops.ops._flash_attn_varlen_backward -else: - _wrapped_flash_attn_varlen_backward = _flash_attn_varlen_backward - - -class FlashAttnQKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and qkv.requires_grad - if softmax_scale is None: - softmax_scale = qkv.shape[-1] ** (-0.5) - q, k, v = qkv[:, :, 0].detach(), qkv[:, :, 1].detach(), qkv[:, :, 2].detach() - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) - dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dqkv[:, :, 0], - dqkv[:, :, 1], - dqkv[:, :, 2], - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dqkv = dqkv[..., : dout.shape[-1]] # We could have padded the head dimension - return dqkv, None, None, None, None, None, None, None, None, None - - -class FlashAttnVarlenQKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - cu_seqlens, - max_seqlen, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and qkv.requires_grad - if softmax_scale is None: - softmax_scale = qkv.shape[-1] ** (-0.5) - q, k, v = qkv[:, 0].detach(), qkv[:, 1].detach(), qkv[:, 2].detach() - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens, - cu_seqlens, - max_seqlen, - max_seqlen, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=None, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens, rng_state) - ctx.dropout_p = dropout_p - ctx.max_seqlen = max_seqlen - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens, rng_state = ctx.saved_tensors - qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) - dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dqkv[:, 0], - dqkv[:, 1], - dqkv[:, 2], - cu_seqlens, - cu_seqlens, - ctx.max_seqlen, - ctx.max_seqlen, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - 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 FlashAttnKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - kv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, kv] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - k, v = kv[:, :, 0].detach(), kv[:, :, 1].detach() - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - dq = torch.empty_like(q) - kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) - dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dkv[:, :, 0], - dkv[:, :, 1], - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dkv = dkv[..., : dout.shape[-1]] - return dq, dkv, None, None, None, None, None, None, None, None, None - - -class FlashAttnVarlenKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, kv] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - k, v = kv[:, 0].detach(), kv[:, 1].detach() - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=None, - ) - if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state - ) - ctx.dropout_p = dropout_p - 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.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors - dq = torch.empty_like(q) - kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) - dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dkv[:, 0], - dkv[:, 1], - cu_seqlens_q, - cu_seqlens_k, - ctx.max_seqlen_q, - ctx.max_seqlen_k, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dkv = dkv[..., : dout.shape[-1]] - return dq, dkv, None, None, None, None, None, None, None, None, None, None, None, None, None - - -class FlashAttnFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - k, - v, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, k, v] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dk = dk[..., : dout.shape[-1]] - dv = dv[..., : dout.shape[-1]] - return dq, dk, dv, 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, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - block_table, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, k, v] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=block_table, - ) - if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state - ) - ctx.dropout_p = dropout_p - 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.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors - dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - ctx.max_seqlen_q, - ctx.max_seqlen_k, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dk = dk[..., : dout.shape[-1]] - dv = dv[..., : dout.shape[-1]] - return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None - - -def flash_attn_qkvpacked_func( - qkv, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # <=0.0 means deactivate - alibi_slopes=None, - deterministic=False, - 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, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_kvpacked_func( - q, - kv, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, -): - """dropout_p should be set to 0.0 during evaluation - If 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 K, V. - 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) - kv: (batch_size, seqlen, 2, 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. - softcap: float. Anything > 0 activates softcapping 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). - 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 FlashAttnKVPackedFunc.apply( - q, - kv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_func( - q, - k, - v, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - 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). - 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 FlashAttnFunc.apply( - q, - k, - v, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_qkvpacked_func( - qkv, - cu_seqlens, - max_seqlen, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - 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_varlen_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_varlen_kvpacked_func and flash_attn_varlen_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: (total, 3, nheads, headdim), where total = total number of tokens in the batch. - cu_seqlens: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into qkv. - max_seqlen: int. Maximum sequence length in the batch. - 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenQKVPackedFunc.apply( - qkv, - cu_seqlens, - max_seqlen, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_kvpacked_func( - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, -): - """dropout_p should be set to 0.0 during evaluation - If 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 K, V. - 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: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch. - kv: (total_k, 2, nheads_k, headdim), where total_k = total number of key tokens in the batch. - cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into q. - cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into kv. - max_seqlen_q: int. Maximum query sequence length in the batch. - max_seqlen_k: int. Maximum key sequence length in the batch. - 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 + 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenKVPackedFunc.apply( - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_func( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, - block_table=None, -): - """dropout_p should be set to 0.0 during evaluation - Supports multi-query and grouped-query attention (MQA/GQA) by passing in K, V 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: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch. - k: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. - v: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. - cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into q. - cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into kv. - max_seqlen_q: int. Maximum query sequence length in the batch. - max_seqlen_k: int. Maximum key sequence length in the batch. - 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 + 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenFunc.apply( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - block_table, - False if q.device.type == "cpu" else torch.is_grad_enabled(), - ) - - -def flash_attn_with_kvcache( - q, - k_cache, - v_cache, - k=None, - v=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, - block_table: Optional[torch.Tensor] = None, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - rotary_interleaved=True, - alibi_slopes=None, - num_splits=0, - 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 block_table, - or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_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) if there's no block_table, - or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_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). Similar to k. - 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. - block_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). - 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. - 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" - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-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) - cache_batch_idx = maybe_contiguous(cache_batch_idx) - block_table = maybe_contiguous(block_table) - out, softmax_lse = flash_attn.fwd_kvcache( - q, - k_cache, - v_cache, - k, - v, - cache_seqlens, - rotary_cos, - rotary_sin, - cache_batch_idx, - cache_leftpad, - block_table, - alibi_slopes, - None, - softmax_scale, - causal, - window_size[0], - window_size[1], - softcap, - rotary_interleaved, - num_splits, - ) - return (out, softmax_lse) if return_softmax_lse else out diff --git a/build/torch211-cxx11-xpu20253-x86_64-linux/layers/__init__.py b/build/torch211-cxx11-xpu20253-x86_64-linux/layers/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/build/torch211-cxx11-xpu20253-x86_64-linux/layers/patch_embed.py b/build/torch211-cxx11-xpu20253-x86_64-linux/layers/patch_embed.py deleted file mode 100644 index 05562f8e8bcdb58e947c6f402a49eacd2d031871..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-xpu20253-x86_64-linux/layers/patch_embed.py +++ /dev/null @@ -1,67 +0,0 @@ -# We use the same API as https://github.com/rwightman/pytorch-image-models/blob/v0.6.11/timm/models/layers/patch_embed.py -# But we use nn.Linear instead of Conv2d and it's about 8x faster. - -from functools import partial - -import torch.nn as nn -from einops import rearrange -from torch import _assert -from torch.nn.modules.utils import _pair - -try: - from flash_attn.ops.fused_dense import FusedDense -except ImportError: - FusedDense = None - - -class PatchEmbed(nn.Module): - """2D Image to Patch Embedding""" - - def __init__( - self, - img_size=224, - patch_size=16, - in_chans=3, - embed_dim=768, - norm_layer=None, - flatten=True, - bias=True, - fused_bias_fc=False, - ): - super().__init__() - img_size = _pair(img_size) - patch_size = _pair(patch_size) - self.img_size = img_size - self.patch_size = patch_size - self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1]) - self.num_patches = self.grid_size[0] * self.grid_size[1] - self.flatten = flatten - if fused_bias_fc and FusedDense is None: - raise ImportError("fused_dense is not installed") - - linear_cls = nn.Linear if not fused_bias_fc or not bias else FusedDense - self.proj = linear_cls(in_chans * patch_size[0] * patch_size[1], embed_dim, bias=bias) - self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() - - def forward(self, x): - _, _, H, W = x.shape - _assert( - H == self.img_size[0], - f"Input image height ({H}) doesn't match model ({self.img_size[0]}).", - ) - _assert( - W == self.img_size[1], - f"Input image width ({W}) doesn't match model ({self.img_size[1]}).", - ) - x = self.proj( - rearrange( - x, - "b c (h p1) (w p2) -> b h w (c p1 p2)", - p1=self.patch_size[0], - p2=self.patch_size[1], - ) - ) - if self.flatten: - x = rearrange(x, "b h w c -> b (h w) c") - x = self.norm(x) - return x diff --git a/build/torch211-cxx11-xpu20253-x86_64-linux/layers/rotary.py b/build/torch211-cxx11-xpu20253-x86_64-linux/layers/rotary.py deleted file mode 100644 index d1bfc21fc7de1dd287e8f382847b194a48075981..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-xpu20253-x86_64-linux/layers/rotary.py +++ /dev/null @@ -1,483 +0,0 @@ -# Copyright (c) 2025, Tri Dao - -import math -from functools import partial -from typing import Optional, Tuple, Union - -import torch -from torch import Tensor - -from einops import rearrange, repeat -# from flash_attn.ops.triton.rotary import apply_rotary -from ..ops.triton.rotary import apply_rotary - - -def rotate_half(x, interleaved=False): - if not interleaved: - x1, x2 = x.chunk(2, dim=-1) - return torch.cat((-x2, x1), dim=-1) - else: - x1, x2 = x[..., ::2], x[..., 1::2] - return rearrange(torch.stack((-x2, x1), dim=-1), "... d two -> ... (d two)", two=2) - - -def apply_rotary_emb_torch(x, cos, sin, interleaved=False): - """ - x: (batch_size, seqlen, nheads, headdim) - cos, sin: (seqlen, rotary_dim / 2) or (batch_size, seqlen, rotary_dim / 2) - """ - ro_dim = cos.shape[-1] * 2 - assert ro_dim <= x.shape[-1] - cos = repeat(cos, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)") - sin = repeat(sin, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)") - return torch.cat( - [x[..., :ro_dim] * cos + rotate_half(x[..., :ro_dim], interleaved) * sin, x[..., ro_dim:]], - dim=-1, - ) - - -class ApplyRotaryEmb(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x, - cos, - sin, - interleaved=False, - inplace=False, - seqlen_offsets: Union[int, Tensor] = 0, - cu_seqlens: Optional[Tensor] = None, - max_seqlen: Optional[int] = None, - ): - out = apply_rotary( - x, - cos, - sin, - seqlen_offsets=seqlen_offsets, - cu_seqlens=cu_seqlens, - max_seqlen=max_seqlen, - interleaved=interleaved, - inplace=inplace, - ) - if isinstance(seqlen_offsets, int): - ctx.save_for_backward(cos, sin, cu_seqlens) # Can't save int with save_for_backward - ctx.seqlen_offsets = seqlen_offsets - else: - ctx.save_for_backward(cos, sin, cu_seqlens, seqlen_offsets) - ctx.seqlen_offsets = None - ctx.interleaved = interleaved - ctx.inplace = inplace - ctx.max_seqlen = max_seqlen - return out if not inplace else x - - @staticmethod - def backward(ctx, do): - seqlen_offsets = ctx.seqlen_offsets - if seqlen_offsets is None: - cos, sin, cu_seqlens, seqlen_offsets = ctx.saved_tensors - else: - cos, sin, cu_seqlens = ctx.saved_tensors - dx = apply_rotary( - do, - cos, - sin, - seqlen_offsets=seqlen_offsets, - cu_seqlens=cu_seqlens, - max_seqlen=ctx.max_seqlen, - interleaved=ctx.interleaved, - inplace=ctx.inplace, - conjugate=True, - ) - return dx, None, None, None, None, None, None, None - - -def apply_rotary_emb( - x, - cos, - sin, - interleaved=False, - inplace=False, - seqlen_offsets: Union[int, Tensor] = 0, - cu_seqlens: Optional[Tensor] = None, - max_seqlen: Optional[int] = None, -): - """ - Arguments: - x: (batch_size, seqlen, nheads, headdim) if cu_seqlens is None - else (total_seqlen, nheads, headdim) - cos, sin: (seqlen_rotary, rotary_dim / 2) - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead - of 1st half and 2nd half (GPT-NeoX style). - inplace: if True, apply rotary embedding in-place. - seqlen_offsets: (batch_size,) or int. Each sequence in x is shifted by this amount. - Most commonly used in inference when we have KV cache. - cu_seqlens: (batch + 1,) or None - max_seqlen: int - Return: - out: (batch_size, seqlen, nheads, headdim) if cu_seqlens is None - else (total_seqlen, nheads, headdim) - rotary_dim must be <= headdim - Apply rotary embedding to the first rotary_dim of x. - """ - return ApplyRotaryEmb.apply( - x, cos, sin, interleaved, inplace, seqlen_offsets, cu_seqlens, max_seqlen - ) - - -# For backward compatibility -apply_rotary_emb_func = apply_rotary_emb - - -def _apply_rotary_emb_qkv( - qkv, - cos, - sin, - cos_k=None, - sin_k=None, - interleaved=False, - inplace=False, - conjugate=False, - seqlen_offsets: Union[int, Tensor] = 0, - num_heads_q: Optional[int] = None, -): - apply_rotary_fn = partial( - apply_rotary, - interleaved=interleaved, - inplace=inplace, - conjugate=conjugate, - seqlen_offsets=seqlen_offsets - ) - if cos_k is None and sin_k is None and qkv.is_contiguous(): - # Call 1 kernel instead of 2 kernels - # We need qkv to be contiguous so that when we reshape to combine (3, nheads) - # dimensions, we get the same tensor - if qkv.dim() == 5: - batch, seqlen, three, nheads, headdim = qkv.shape - assert three == 3 - # qk = rearrange(qkv[:, :, :2], "b s t h d -> b s (t h) d") - qk = qkv[:, :, :2].reshape(batch, seqlen, -1, headdim) - qk = apply_rotary_fn(qk, cos, sin) - else: - assert qkv.dim() == 4 - assert num_heads_q is not None - num_heads_k = (qkv.shape[2] - num_heads_q) // 2 - assert qkv.shape[2] == num_heads_q + 2 * num_heads_k - qk = qkv[:, :, :num_heads_q + num_heads_k] - qk = apply_rotary_fn(qk, cos, sin) - if not inplace: - if qkv.dim() == 5: - qkv = torch.cat([rearrange(qk, "b s (t h) d -> b s t h d", t=2), qkv[:, :, 2:]], dim=2) - else: - qkv = torch.cat([qk, qkv[:, :, num_heads_q + num_heads_k :]], dim=2) - else: - cos_k = cos if cos_k is None else cos_k - sin_k = sin if sin_k is None else sin_k - if qkv.dim() == 5: - batch, seqlen, three, nheads, headdim = qkv.shape - assert three == 3 - q, k = qkv[:, :, 0], qkv[:, :, 1] - else: - assert qkv.dim() == 4 - assert num_heads_q is not None - num_heads_k = (qkv.shape[2] - num_heads_q) // 2 - assert qkv.shape[2] == num_heads_q + 2 * num_heads_k - q, k = qkv[:, :, :num_heads_q], qkv[:, :, num_heads_q : num_heads_q + num_heads_k] - q = apply_rotary_fn(q, cos, sin) - k = apply_rotary_fn(k, cos_k, sin_k) - if not inplace: - if qkv.dim() == 5: - qkv = torch.stack([q, k, qkv[:, :, 2]], dim=2) - else: - qkv = torch.cat([q, k, qkv[:, :, num_heads_q + num_heads_k:]], dim=2) - return qkv - - -class ApplyRotaryEmbQKV_(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - cos, - sin, - cos_k=None, - sin_k=None, - interleaved=False, - seqlen_offsets: Union[int, torch.Tensor] = 0, - num_heads_q: Optional[int] = None, - ): - # apply_rotary_emb_qkv_inplace( - qkv = _apply_rotary_emb_qkv( - qkv, cos, sin, cos_k, sin_k, interleaved=interleaved, inplace=True, - seqlen_offsets=seqlen_offsets, num_heads_q=num_heads_q, - ) - if isinstance(seqlen_offsets, int): - ctx.save_for_backward(cos, sin, cos_k, sin_k) - ctx.seqlen_offsets = seqlen_offsets - else: - ctx.save_for_backward(cos, sin, cos_k, sin_k, seqlen_offsets) - ctx.seqlen_offsets = None - ctx.interleaved = interleaved - ctx.num_heads_q = num_heads_q - return qkv - - @staticmethod - def backward(ctx, dqkv): - seqlen_offsets = ctx.seqlen_offsets - if seqlen_offsets is None: - cos, sin, cos_k, sin_k, seqlen_offsets = ctx.saved_tensors - else: - cos, sin, cos_k, sin_k = ctx.saved_tensors - dqkv = _apply_rotary_emb_qkv( - dqkv, cos, sin, cos_k, sin_k, interleaved=ctx.interleaved, inplace=True, - seqlen_offsets=seqlen_offsets, num_heads_q=ctx.num_heads_q, conjugate=True, - ) - return dqkv, None, None, None, None, None, None, None - - -def apply_rotary_emb_qkv_( - qkv, - cos, - sin, - cos_k=None, - sin_k=None, - interleaved=False, - seqlen_offsets: Union[int, torch.Tensor] = 0, - num_heads_q: Optional[int] = None, -): - """ - Arguments: - qkv: (batch_size, seqlen, 3, nheads, headdim) or (batch_size, seqlen, num_heads_q + 2 * num_heads_k, headdim). - If qkv has shape (batch_size, seqlen, num_heads_q + 2 * num_heads_k, headdim) (e.g. MQA / GQA), - then num_heads_q must be provided. - cos, sin: (seqlen, rotary_dim / 2) - cos_k, sin_k: (seqlen, rotary_dim / 2), optional - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead of - 1st half and 2nd half (GPT-NeoX style). - seqlen_offsets: (batch_size,) or int. Each sequence in Q and K is shifted by this amount. - Most commonly used in inference when we have KV cache. - Return: - qkv: (batch_size, seqlen, 3, nheads, headdim) or (batch_size, seqlen, num_heads_q + 2 * num_heads_k, headdim) - rotary_dim must be <= headdim - Apply rotary embedding *inplace* to the first rotary_dim of Q and K. - """ - return ApplyRotaryEmbQKV_.apply( - qkv, cos, sin, cos_k, sin_k, interleaved, seqlen_offsets, num_heads_q - ) - - -class ApplyRotaryEmbKV_(torch.autograd.Function): - - @staticmethod - def forward(ctx, kv, cos, sin, interleaved=False, seqlen_offsets: Union[int, torch.Tensor] = 0): - batch, seqlen, two, nheads, headdim = kv.shape - assert two == 2 - k = kv[:, :, 0] - apply_rotary( - k, cos, sin, seqlen_offsets=seqlen_offsets, interleaved=interleaved, inplace=True - ) - if isinstance(seqlen_offsets, int): - ctx.save_for_backward(cos, sin) # Can't save int with save_for_backward - ctx.seqlen_offsets = seqlen_offsets - else: - ctx.save_for_backward(cos, sin, seqlen_offsets) - ctx.seqlen_offsets = None - ctx.interleaved = interleaved - return kv - - @staticmethod - def backward(ctx, dkv): - seqlen_offsets = ctx.seqlen_offsets - if seqlen_offsets is None: - cos, sin, seqlen_offsets = ctx.saved_tensors - else: - cos, sin = ctx.saved_tensors - apply_rotary( - dkv[:, :, 0], - cos, - sin, - seqlen_offsets=seqlen_offsets, - interleaved=ctx.interleaved, - inplace=True, - conjugate=True, - ) - return dkv, None, None, None, None - - -apply_rotary_emb_kv_ = ApplyRotaryEmbKV_.apply - - -def apply_rotary_emb_kv_( - kv, - cos, - sin, - interleaved=False, - seqlen_offsets: Union[int, torch.Tensor] = 0, -): - """ - Arguments: - kv: (batch_size, seqlen, 2, nheads, headdim) - cos, sin: (seqlen, rotary_dim / 2) - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead of - 1st half and 2nd half (GPT-NeoX style). - seqlen_offsets: (batch_size,) or int. Each sequence in Q and K is shifted by this amount. - Most commonly used in inference when we have KV cache. - Return: - kv: (batch_size, seqlen, 2, nheads, headdim) - rotary_dim must be <= headdim - Apply rotary embedding *inplace* to the first rotary_dim of K. - """ - return ApplyRotaryEmbKV_.apply(kv, cos, sin, interleaved, seqlen_offsets) - - -class RotaryEmbedding(torch.nn.Module): - """ - The rotary position embeddings from RoFormer_ (Su et. al). - A crucial insight from the method is that the query and keys are - transformed by rotation matrices which depend on the relative positions. - - Other implementations are available in the Rotary Transformer repo_ and in - GPT-NeoX_, GPT-NeoX was an inspiration - - .. _RoFormer: https://arxiv.org/abs/2104.09864 - .. _repo: https://github.com/ZhuiyiTechnology/roformer - .. _GPT-NeoX: https://github.com/EleutherAI/gpt-neox - - If scale_base is not None, this implements XPos (Sun et al., https://arxiv.org/abs/2212.10554). - A recommended value for scale_base is 512: https://github.com/HazyResearch/flash-attention/issues/96 - Reference: https://github.com/sunyt32/torchscale/blob/main/torchscale/component/xpos_relative_position.py - """ - - def __init__( - self, - dim: int, - base=10000.0, - interleaved=False, - scale_base=None, - device=None, - ): - """ - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead - of 1st half and 2nd half (GPT-NeoX style). - """ - super().__init__() - self.dim = dim - self.base = float(base) - # Generate and save the inverse frequency buffer (non trainable) - inv_freq = self._compute_inv_freq(device) - self.register_buffer("inv_freq", inv_freq, persistent=False) - self.interleaved = interleaved - self.scale_base = scale_base - scale = ( - (torch.arange(0, dim, 2, device=device, dtype=torch.float32) + 0.4 * dim) / (1.4 * dim) - if scale_base is not None - else None - ) - self.register_buffer("scale", scale, persistent=False) - - self._seq_len_cached = 0 - self._cos_cached = None - self._sin_cached = None - self._cos_k_cached = None - self._sin_k_cached = None - - def _compute_inv_freq(self, device=None): - return 1.0 / ( - self.base - ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim) - ) - - def _update_cos_sin_cache(self, seqlen, device=None, dtype=None): - # Reset the tables if the sequence length has changed, - # if we're on a new device (possibly due to tracing for instance), - # or if we're switching from inference mode to training - if ( - seqlen > self._seq_len_cached - or self._cos_cached is None - or self._cos_cached.device != device - or self._cos_cached.dtype != dtype - or (self.training and self._cos_cached.is_inference()) - ): - self._seq_len_cached = seqlen - # We want fp32 here, not self.inv_freq.dtype, since the model could be loaded in bf16 - # And the output of arange can be quite large, so bf16 would lose a lot of precision. - t = torch.arange(seqlen, device=device, dtype=torch.float32) - # We want fp32 here as well since inv_freq will be multiplied with t, and the output - # will be large. Having it in bf16 will lose a lot of precision and cause the - # cos & sin output to change significantly. - # We want to recompute self.inv_freq if it was not loaded in fp32 - if self.inv_freq.dtype != torch.float32: - inv_freq = self._compute_inv_freq(device=device) - else: - inv_freq = self.inv_freq - # Don't do einsum, it converts fp32 to bf16 under AMP - # freqs = torch.einsum("i,j->ij", t, self.inv_freq) - freqs = torch.outer(t, inv_freq) - if self.scale is None: - self._cos_cached = torch.cos(freqs).to(dtype) - self._sin_cached = torch.sin(freqs).to(dtype) - else: - power = ( - torch.arange(seqlen, dtype=self.scale.dtype, device=self.scale.device) - - seqlen // 2 - ) / self.scale_base - scale = self.scale.to(device=power.device) ** rearrange(power, "s -> s 1") - # We want the multiplication by scale to happen in fp32 - self._cos_cached = (torch.cos(freqs) * scale).to(dtype) - self._sin_cached = (torch.sin(freqs) * scale).to(dtype) - self._cos_k_cached = (torch.cos(freqs) / scale).to(dtype) - self._sin_k_cached = (torch.sin(freqs) / scale).to(dtype) - - def forward( - self, - qkv: torch.Tensor, - kv: Optional[torch.Tensor] = None, - seqlen_offset: Union[int, torch.Tensor] = 0, - max_seqlen: Optional[int] = None, - num_heads_q: Optional[int] = None, - ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: - """ - qkv: (batch, seqlen, 3, nheads, headdim) or (batch, seqlen, num_heads_q + 2 * num_heads_k, headdim) - if kv is none, else it's just q of shape (batch, seqlen, nheads, headdim). - If qkv has shape (batch, seqlen, num_heads_q + 2 * num_heads_k, headdim) (e.g. MQA / GQA), - then num_heads_q must be provided. - kv: (batch, seqlen, 2, nheads, headdim) - seqlen_offset: (batch_size,) or int. Each sequence in x is shifted by this amount. - Most commonly used in inference when we have KV cache. - If it's a tensor of shape (batch_size,), then to update the cos / sin cache, one - should pass in max_seqlen, which will update the cos / sin cache up to that length. - Apply rotary embedding *inplace* to qkv and / or kv. - """ - seqlen = qkv.shape[1] - if max_seqlen is not None: - self._update_cos_sin_cache(max_seqlen, device=qkv.device, dtype=qkv.dtype) - elif isinstance(seqlen_offset, int): - self._update_cos_sin_cache(seqlen + seqlen_offset, device=qkv.device, dtype=qkv.dtype) - if kv is None: - return apply_rotary_emb_qkv_( - qkv, - self._cos_cached, - self._sin_cached, - self._cos_k_cached if self.scale is not None else None, - self._sin_k_cached if self.scale is not None else None, - interleaved=self.interleaved, - seqlen_offsets=seqlen_offset, - num_heads_q=num_heads_q, - ) - else: - q = qkv - q = apply_rotary_emb_func( - q, - self._cos_cached, - self._sin_cached, - interleaved=self.interleaved, - inplace=True, - seqlen_offsets=seqlen_offset, - ) - kv = apply_rotary_emb_kv_( - kv, - self._cos_cached if self.scale is None else self._cos_k_cached, - self._sin_cached if self.scale is None else self._sin_k_cached, - interleaved=self.interleaved, - seqlen_offsets=seqlen_offset, - ) - return q, kv diff --git a/build/torch211-cxx11-xpu20253-x86_64-linux/metadata.json b/build/torch211-cxx11-xpu20253-x86_64-linux/metadata.json deleted file mode 100644 index 92873a6f76b3571e5dc2d66b68d6de7a93225c42..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-xpu20253-x86_64-linux/metadata.json +++ /dev/null @@ -1,36 +0,0 @@ -{ - "name": "flash-attn2", - "id": "_flash_attn2_xpu_f12afc9", - "version": 1, - "license": "BSD-3-Clause", - "python-depends": [], - "backend": { - "type": "xpu" - }, - "digest": { - "algorithm": "sha256", - "files": { - "__init__.py": "Cd9MEdd/WUYnBmfiQpC4SyPqAxrTilTjvsvswNju4zc=", - "_flash_attn2_xpu_f12afc9.abi3.so": "Vsa8mYdTf3CxD42MygDA0/P/Ou/bny8y8B4jLaeGp78=", - "_ops.py": "cmwaGTsubLtetqeo95vdShoO62KcV0DrdeNH34KtCyk=", - "bert_padding.py": "gF1EmsdJ+HpQ86MRQ4VxDw+Sb/RVISdQALdNnoByHlw=", - "flash_attn2/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY=", - "flash_attn_interface.py": "W7a9bHR2ly6QL3aKqsBqxWA7M7ZWYcZPmtsPJu6rAfY=", - "layers/__init__.py": "47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU=", - "layers/patch_embed.py": "H58CgME/qSOPTZLOG08wFgrQS1j34pvNwMPrkTj3Ek4=", - "layers/rotary.py": "rFIdNrZQFNawwLD6d7t5LE2m+M93wwDZ8EuExw+gfD8=", - "ops/__init__.py": "47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU=", - "ops/activations.py": "t5lzNg1In8LP6bKeTnyeMizwqjv27JGbJ6ylPdGvZYg=", - "ops/fused_dense.py": "BqT8VTSCGdLJAjh6WBMkCRFYrdVodPeoElWBvH6AIzA=", - "ops/layer_norm.py": "zr7NXIm+2mtEynTp1CS0fbFGI2Mqdp41dY4AfDWF6EQ=", - "ops/rms_norm.py": "XEnihcj0a4aSz4LO55m5iKGVn4HKTeKN8TIyHjuDgxI=", - "ops/triton/__init__.py": "AbpHGcgLb+kRsJGnwFEktk7uzpZOCcBY74+YBdrKVGs=", - "ops/triton/cross_entropy.py": "y56nmuOtpSNxM7sDZcUSgXKIsIEayfuvrU8bAVog/nw=", - "ops/triton/k_activations.py": "+Z3vIyO4JkqBMipKsPvhzmxljtBdIhJCsl/M+/ESqBo=", - "ops/triton/layer_norm.py": "BBQdJiqopbIS12viUqzx3SvOQCNbDSAYLJXhpyMGyC8=", - "ops/triton/linear.py": "OtRvKz8xdpl+7v3q/ZTaS9fdBt9XrzMyapgRr50uBbM=", - "ops/triton/mlp.py": "uWA+D7HyuRUsOT+4nhi3vvCDNZnNEXwYcHMQ02UREfQ=", - "ops/triton/rotary.py": "0YSrlbep0lJmzLpTiANs0slTFeGcQQqR+el6uzbSLvI=" - } - } -} \ No newline at end of file diff --git a/build/torch211-cxx11-xpu20253-x86_64-linux/metadata.json.sigstore b/build/torch211-cxx11-xpu20253-x86_64-linux/metadata.json.sigstore deleted file mode 100644 index edd65259635b546cb002e20585fa63a9bb7672ac..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-xpu20253-x86_64-linux/metadata.json.sigstore +++ /dev/null @@ -1 +0,0 @@ -{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json","verificationMaterial":{"certificate":{"rawBytes":"MIIHSjCCBs+gAwIBAgIUQHOVUD+ToW9T2L+QQgiY/WnvsFMwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwNjIzMTgyMDI0WhcNMjYwNjIzMTgzMDI0WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEpQ3zkVVjKVgymQzbJOlrc2QsXJzuAcRN07QTJTzuJTNTy+c23ESb0JLXyF5pos7GaWWmnkOFkvIfNIpRcWZqs6OCBe4wggXqMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUOH1IWfwK03NemEkI0k7F7FufmIEwHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoZjEyYWZjOTJjNDBiYmQ3NmQ5OTBkOTExZWJhYmYyOGU0OTBhN2JkZTATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoZjEyYWZjOTJjNDBiYmQ3NmQ5OTBkOTExZWJhYmYyOGU0OTBhN2JkZTAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoZjEyYWZjOTJjNDBiYmQ3NmQ5OTBkOTExZWJhYmYyOGU0OTBhN2JkZTAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKGYxMmFmYzkyYzQwYmJkNzZkOTkwZDkxMWViYWJmMjhlNDkwYTdiZGUwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMjgwMzYyNjg5MzUvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBiAYKKwYBBAHWeQIEAgR6BHgAdgB0AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABnvW2RmUAAAQDAEUwQwIgSmciocMm02ZPe2XT/MvODB41tk6i8X+LxyLDMMBViFMCHzvbP6LCRVmx3wmFCbNzTBkTH60DrHiSNwwCrvTwomQwCgYIKoZIzj0EAwMDaQAwZgIxAINJS+M81K1ZUvmMBqDcqfdq9RJcIcfQv/8ivMvkBpq7+TZyNc9U3LpwukV22n+3ZAIxAISGXCdFg495JGxyuXUnMBTwweebTGNpc98Ns9XL4qqaooDWIVdLOVgwaraTlAb3+Q=="},"tlogEntries":[{"logIndex":"1929966264","logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="},"kindVersion":{"kind":"hashedrekord","version":"0.0.1"},"integratedTime":"1782238824","inclusionPromise":{"signedEntryTimestamp":"MEYCIQDFywAtNZdQZuR45CA6x55RSmWyOUXCUgffw3gLPLAcxAIhALo8AG93dyweru4be9A2VKIPAZKZjw201SwZsu9noYy6"},"inclusionProof":{"logIndex":"1808062002","rootHash":"YYgXjax87Z42EwwFoE6cT6cjl4Jr2r2RvqVXObZl69s=","treeSize":"1808062045","hashes":["EK4WhDG1gOMRsOFTggkuPXKr5OETgUFy1dV0vhOK4O4=","xslKR9vDihRC6ueQF0jiU0tRG9DJOLfYZBzOzZCU0eI=","RxXbC01J6YtWGzIggaplB7eihXYt0pqFMQNjziSQuk4=","NOFwrcf3DOzRML/zvmapb9zLoIqzDhOPzjfP0x/PeYI=","Qn7q0H1JCoNlGaEd940w0Hj9EgzgY9vthZZDs9DOpNw=","uUReT35PLLmVIxxE2SMm7H9Hh5LqTPpQZ+7F2D+0yyc=","z/jLVVhuq1FK09hB5SAFyl8+hsl06onfgCkqgkAOcis=","jgKZMZUAOznhZow0dOq9j+6kBbYqJbpAwbHuyR0vtw0=","43kco2WsynkkNcu9FNHzcL7ZH810rbB/a5HS62l4P9Q=","Rf4AP9k6JckOtBvSb8IZhPtOdgQbDwTqzHj/Ob7GoR0=","2azvF0ZQXjI0TTkEYlWo40jJTIJKrIaRQSKDw70qaz8=","MXnb5KZ/vcuCLPtNBIq6MAywaatte0rS/H2lJyqLfIY=","QdxFjkGuHUV4B069H6Stxfn75M+rztEbA+7TnHdszN8=","nyv2TDlfkz0XfRwjY98zsbUbOkBZ+Tgz7PML0c2uN14=","OMfDh4k9YD0E+t4sjqJsh+vg8D1ueo1cC07exyxmr6s=","+vHKm4wNwGJ46ceyZ8mePD7Vn1SediStNHCStRoTSWQ=","5HwWjuGNDisetbiWM0cWIXM65TIM3EJlqk13ntCQlcU=","NDIRlqxO0IaGr94q86MIU7VCAN3yEEgOHS4loya5lu4=","daxmZaajRpZV+JxHiOYZhJBiSKN5ucqjh2WnGbHhirw=","DOCeoSMovIvLExkhIvisow9AuNXgeWs4ECkyR6EcqYU="],"checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n1808062045\nYYgXjax87Z42EwwFoE6cT6cjl4Jr2r2RvqVXObZl69s=\n\n— rekor.sigstore.dev wNI9ajBFAiEAk0pd6Fk+H57TJ5GR9L6lBCpYm6BC3rWTFgvJDV7WfmgCIDu4uOrAMYKV6foQH20j8yUn/Lw0xxM3pBISnipOb3gY\n"}},"canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiI3YWM3ZjJjMWE1NzI4Y2M1OTk4OTZkNjc4NzUyMDVkYWVkZjIwMWY4YjE3N2E5ZjhkNzZmODIyODNjZDczNDdlIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FUUNJRTBFc29nVkszWnB6QnZBWVRrdVVnMTRsRTl4c25LbGlDMlNVWWFycnFDWEFpQkhlL1pENmlCbkowMy9yOW0ybFQ2ZS8wYkxZNkFtNUswbnVab3BWTU40cEE9PSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRha05EUW5NclowRjNTVUpCWjBsVlVVaFBWbFZFSzFSdlZ6bFVNa3dyVVZGbmFWa3ZWMjUyYzBaTmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDVxU1hwTlZHZDVUVVJKTUZkb1kwNU5hbGwzVG1wSmVrMVVaM3BOUkVrd1YycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZ3VVRONmExWldha3RXWjNsdFVYcGlTazlzY21NeVVYTllTbnAxUVdOU1RqQTNVVlFLU2xSNmRVcFVUbFI1SzJNeU0wVlRZakJLVEZoNVJqVndiM00zUjJGWFYyMXVhMDlHYTNaSlprNUpjRkpqVjFweGN6WlBRMEpsTkhkbloxaHhUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlZQU0RGSkNsZG1kMHN3TTA1bGJVVnJTVEJyTjBZM1JuVm1iVWxGZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOWFha1Y1V1ZkYWFrOVVTbXBPUkVKcFdXMVJNMDV0VVRWUFZFSnJDazlVUlhoYVYwcG9XVzFaZVU5SFZUQlBWRUpvVGpKS2ExcFVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMXBxUlhsWlYxcHFUMVJLYWs1RVFtbFpiVkV6VG0xUk5VOVVRbXRQVkVWNFdsZEthRmx0V1hrS1QwZFZNRTlVUW1oT01rcHJXbFJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMXBxUlhsWlYxcHFUMVJLYWs1RVFtbFpiVkV6VG0xUk5VOVVRbXNLVDFSRmVGcFhTbWhaYlZsNVQwZFZNRTlVUW1oT01rcHJXbFJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwZFplRTF0Um0wS1dYcHJlVmw2VVhkWmJVcHJUbnBhYTA5VWEzZGFSR3Q0VFZkV2FWbFhTbTFOYW1oc1RrUnJkMWxVWkdsYVIxVjNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOYW1kM1RYcFplVTVxWnpWTmVsVjJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwUVZsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTTmtKSVowRUtaR2RDTUVGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtNTJWekpTYlZWQlFVRlJSQXBCUlZWM1VYZEpaMU50WTJsdlkwMXRNREphVUdVeVdGUXZUWFpQUkVJME1YUnJObWs0V0N0TWVIbE1SRTFOUWxacFJrMURTSHAyWWxBMlRFTlNWbTE0Q2pOM2JVWkRZazU2VkVKclZFZzJNRVJ5U0dsVFRuZDNRM0oyVkhkdmJWRjNRMmRaU1V0dldrbDZhakJGUVhkTlJHRlJRWGRhWjBsNFFVbE9TbE1yVFRnS01Vc3hXbFYyYlUxQ2NVUmpjV1prY1RsU1NtTkpZMlpSZGk4NGFYWk5kbXRDY0hFM0sxUmFlVTVqT1ZVelRIQjNkV3RXTWpKdUt6TmFRVWw0UVVsVFJ3cFlRMlJHWnpRNU5VcEhlSGwxV0ZWdVRVSlVkM2RsWldKVVIwNXdZems0VG5NNVdFdzBjWEZoYjI5RVYwbFdaRXhQVm1kM1lYSmhWR3hCWWpNclVUMDlDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}],"timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyjADAgEAMIICwQYJKoZIhvcNAQcCoIICsjCCAq4CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgRxIHq5Tmqq1j3edyvcmcfcWJ1tqW2qS3BQ7+dreB8SoCFQDkm2vzaf4Cip4fuwa5jel3YPRj5hgPMjAyNjA2MjMxODIwMjRaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHbMIIB1wIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MDYyMzE4MjAyNFowLwYJKoZIhvcNAQkEMSIEIPMnsVyzhPLOQhrqWwvCXMHbtneH2MM5Wm1qtjQYkt3AMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRnMGUCMQChBlm8p2nE8nHoRkd+WHekWd8bP1DgtDNA9QZGvYdQov+6llANZFsxGVGkjrxGiDACMEON0NvxSTKwuPCDblY14JOfHE81ojhc6ttsBqvRhfcSNpDeG6Yt+V/RsHLZ6zohEA=="}]}},"messageSignature":{"messageDigest":{"algorithm":"SHA2_256","digest":"esfywaVyjMWZiW1nh1IF2u3yAfixd6n412+CKDzXNH4="},"signature":"MEQCIE0EsogVK3ZpzBvAYTkuUg14lE9xsnKliC2SUYarrqCXAiBHe/ZD6iBnJ03/r9m2lT6e/0bLY6Am5K0nuZopVMN4pA=="}} \ No newline at end of file diff --git a/build/torch211-cxx11-xpu20253-x86_64-linux/ops/__init__.py b/build/torch211-cxx11-xpu20253-x86_64-linux/ops/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/build/torch211-cxx11-xpu20253-x86_64-linux/ops/activations.py b/build/torch211-cxx11-xpu20253-x86_64-linux/ops/activations.py deleted file mode 100644 index 7c09649fc41e12d5a360c5672825d8380bc7ec80..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-xpu20253-x86_64-linux/ops/activations.py +++ /dev/null @@ -1,135 +0,0 @@ -# Copied from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/model/layers/activations.py -import math - -import torch -import torch.nn as nn -import torch.nn.functional as F - -# 1/sqrt(2*pi)-> 0.3989423 -# 1/sqrt(2) -> 0.70710678 -# sqrt(2/pi) -> 0.79788456 - -# this function is tanh approximation of gelu -# actual gelu is: -# x * 0.5 * (1.0 + torch.erf(x * 0.70710678)) -@torch.jit.script -def bias_gelu(y, bias): - x = bias + y - return (x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))).to(dtype=y.dtype) - - -# gradient of tanh approximation of gelu -# gradient of actual gelu is: -# 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x) -@torch.jit.script -def bias_gelu_back(g, y, bias): - """Assume that y has shape (B, D) and bias has shape (D)""" - x = bias + y - tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)) - # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243 - ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( - 1 + tanh_out - ) - grad_y = ff * g - return grad_y.to(dtype=y.dtype), grad_y.sum(dim=(0), dtype=bias.dtype) - - -class GeLUFunction(torch.autograd.Function): - @staticmethod - # bias is an optional argument - def forward(ctx, input, bias): - ctx.save_for_backward(input, bias) - return bias_gelu(input, bias) - - @staticmethod - def backward(ctx, grad_output): - input, bias = ctx.saved_tensors - tmp = bias_gelu_back(grad_output, input, bias) - return tmp, tmp - - -bias_gelu_impl = GeLUFunction.apply - -# this function is tanh approximation of gelu -# actual gelu is: -# x * 0.5 * (1.0 + torch.erf(x * 0.70710678)) -@torch.jit.script -def gelu_fwd(x): - return (x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))).to(dtype=x.dtype) - - -# gradient of tanh approximation of gelu -# gradient of actual gelu is: -# 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x) -@torch.jit.script -def gelu_bwd(g, x): - tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)) - # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243 - ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( - 1 + tanh_out - ) - return (ff * g).to(dtype=x.dtype) - - -class FastGeLUFunction(torch.autograd.Function): - @staticmethod - # bias is an optional argument - def forward(ctx, input): - ctx.save_for_backward(input) - return gelu_fwd(input) - - @staticmethod - def backward(ctx, grad_output): - (input,) = ctx.saved_tensors - tmp = gelu_bwd(grad_output, input) - return tmp - - -fast_gelu_impl = FastGeLUFunction.apply - - -@torch.jit.script -def relu_bwd(g, x): - return torch.where(x >= 0, g, 0.0).to(dtype=x.dtype) - - -@torch.jit.script -def sqrelu_fwd(x): - r = F.relu(x) - return (r * r).to(dtype=x.dtype) - - -@torch.jit.script -def sqrelu_bwd(g, x): - return (2.0 * g * F.relu(x)).to(dtype=x.dtype) - - -swiglu_fwd_codestring = """ -template T swiglu_fwd(T x, T y) { - return float(x) * float(y) / (1.0f + ::exp(-float(x))); -} -""" -swiglu_bwd_codestring = """ -template void swiglu_bwd(T x, T y, T g, T& dx, T& dy) { - float x_sigmoid = 1.0f / (1.0f + ::exp(-float(x))); - dx = x_sigmoid * (1 + float(x) * (1.0f - x_sigmoid)) * float(g) * float(y); - dy = float(x) * x_sigmoid * float(g); -} -""" -swiglu_fwd = torch.cuda.jiterator._create_jit_fn(swiglu_fwd_codestring) -swiglu_bwd = torch.cuda.jiterator._create_multi_output_jit_fn(swiglu_bwd_codestring, num_outputs=2) - - -class SwiGLUFunction(torch.autograd.Function): - - @staticmethod - def forward(ctx, x, y): - ctx.save_for_backward(x, y) - return swiglu_fwd(x, y) - - @staticmethod - def backward(ctx, dout): - x, y = ctx.saved_tensors - return swiglu_bwd(x, y, dout) - -swiglu = SwiGLUFunction.apply diff --git a/build/torch211-cxx11-xpu20253-x86_64-linux/ops/fused_dense.py b/build/torch211-cxx11-xpu20253-x86_64-linux/ops/fused_dense.py deleted file mode 100644 index 6b4033d134e4093fe278f7b3f8c7d3128ce9f36d..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-xpu20253-x86_64-linux/ops/fused_dense.py +++ /dev/null @@ -1,688 +0,0 @@ -# Copyright (c) 2023, Tri Dao. -# Inspired by https://github.com/NVIDIA/apex/blob/master/apex/fused_dense/fused_dense.py -# We make it work with pytorch amp and with bfloat16. -# The TensorParallel linear modules are inspired by https://github.com/NVIDIA/apex/blob/master/apex/transformer/tensor_parallel/layers.py -from functools import partial -from typing import Optional - -# import fused_dense_cuda # from apex -import fused_dense_lib as fused_dense_cuda -import torch -import torch.nn as nn -import torch.nn.functional as F -from torch import Tensor -from torch.distributed import ProcessGroup - -from flash_attn.utils.torch import custom_fwd, custom_bwd -from flash_attn.ops.activations import gelu_bwd, relu_bwd, sqrelu_bwd, sqrelu_fwd -from flash_attn.utils.distributed import ( - all_gather_raw, - all_reduce, - all_reduce_raw, - reduce_scatter, - reduce_scatter_raw, -) - - -class FusedDenseFunc(torch.autograd.Function): - @staticmethod - @custom_fwd - def forward( - ctx, x, weight, bias, return_residual=False, process_group=None, sequence_parallel=True - ): - """ - If process_group is not None and sequence_parallel=True, we're doing Tensor Parallel - with sequence parallelism: we do an all_gather_raw of x before doing the matmul. - """ - ctx.compute_weight_gradient = weight.requires_grad - ctx.return_residual = return_residual - ctx.process_group = process_group - ctx.sequence_parallel = sequence_parallel - - if torch.is_autocast_enabled(): - x = x.to(dtype=torch.get_autocast_gpu_dtype()) - x = x.contiguous() - if process_group is not None and sequence_parallel: - # We want to kick off the all_gather early, before weight dtype conversion - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - else: - total_x = x - - if torch.is_autocast_enabled(): - weight = weight.to(dtype=torch.get_autocast_gpu_dtype()) - bias = bias.to(dtype=torch.get_autocast_gpu_dtype()) if bias is not None else None - weight = weight.contiguous() - if process_group is not None and sequence_parallel: - handle_x.wait() - batch_shape, n = total_x.shape[:-1], total_x.shape[-1] - batch_dim = batch_shape.numel() - # https://github.com/pytorch/pytorch/blob/5b51849b48a7dbccd297286cc0110def4706f9e7/aten/src/ATen/native/cuda/Blas.cpp#L174 - if min(batch_dim, n, *weight.shape) > 65535 * 32: - raise RuntimeError("fused_dense only supports matrix dims <= 2M") - output = F.linear(total_x, weight, bias) - if ctx.compute_weight_gradient: - ctx.save_for_backward(x, weight) - else: - ctx.save_for_backward(weight) - return output if not return_residual else (output, x) - - @staticmethod - @custom_bwd - def backward(ctx, grad_output, *args): - grad_output = grad_output.contiguous() - if ctx.return_residual: - (grad_input,) = args - grad_input = grad_input.contiguous() - process_group = ctx.process_group - sequence_parallel = ctx.sequence_parallel - if ctx.compute_weight_gradient: - x, weight = ctx.saved_tensors - if process_group is not None and sequence_parallel: - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - else: - total_x = x - else: - (weight,) = ctx.saved_tensors - total_x = None - batch_shape = grad_output.shape[:-1] - batch_dim = batch_shape.numel() - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - if ctx.needs_input_grad[0]: - if not ctx.return_residual: - grad_input = F.linear(grad_output, weight.t()) - else: - grad_input = torch.addmm( - grad_input.reshape(batch_dim, grad_input.shape[-1]), grad_output, weight - ) - grad_input = grad_input.reshape(*batch_shape, grad_input.shape[-1]) - if process_group is not None: - reduce_fn = reduce_scatter_raw if sequence_parallel else all_reduce_raw - grad_input, handle_grad_input = reduce_fn(grad_input, process_group, async_op=True) - else: - grad_input = None - if ctx.needs_input_grad[1]: - assert ctx.compute_weight_gradient - if process_group is not None and sequence_parallel: - handle_x.wait() - grad_weight, grad_bias = fused_dense_cuda.linear_bias_wgrad( - total_x.reshape(batch_dim, total_x.shape[-1]), grad_output, ctx.needs_input_grad[2] - ) - else: - grad_weight = None - grad_bias = grad_output if ctx.needs_input_grad[2] else None - if process_group is not None and ctx.needs_input_grad[0]: - handle_grad_input.wait() - return grad_input, grad_weight, grad_bias, None, None, None - - -def fused_dense_func( - x: Tensor, - weight: Tensor, - bias: Optional[Tensor] = None, - return_residual: bool = False, - process_group: Optional[ProcessGroup] = None, - sequence_parallel: bool = True, -): - dtype_eligible = x.dtype in [torch.float16, torch.bfloat16] or ( - x.dtype == torch.float32 and torch.is_autocast_enabled() - ) - if x.is_cuda and weight.is_cuda and (bias is None or bias.is_cuda) and dtype_eligible: - return FusedDenseFunc.apply( - x, weight, bias, return_residual, process_group, sequence_parallel - ) - else: - assert process_group is None - out = F.linear(x, weight, bias) - return out if not return_residual else (out, x) - - -class FusedDense(nn.Linear): - def __init__( - self, - in_features: int, - out_features: int, - bias: bool = True, - return_residual: bool = False, - device=None, - dtype=None, - ) -> None: - super().__init__(in_features, out_features, bias=bias, device=device, dtype=dtype) - self.return_residual = return_residual - - def forward(self, x, process_group=None): - """ - If process_group is not None, we're doing Tensor Parallel with sequence parallelism: - we do an all_gather of x before doing the matmul. - """ - return fused_dense_func( - x, - self.weight, - self.bias, - return_residual=self.return_residual, - process_group=process_group, - ) - - -class ColumnParallelLinear(nn.Linear): - def __init__( - self, - in_features: int, - out_features: int, - process_group: ProcessGroup, - bias: bool = True, - sequence_parallel=True, - multiple_of=1, - device=None, - dtype=None, - ) -> None: - world_size = torch.distributed.get_world_size(process_group) - if out_features % multiple_of: - raise ValueError(f"out_features ({out_features}) must be a multiple of {multiple_of}") - multiple = out_features // multiple_of - # We want to split @multiple across world_size, but it could be an uneven split - div = multiple // world_size - mod = multiple % world_size - # The first @mod ranks get @div + 1 copies, the rest get @div copies - local_multiple = div + int(torch.distributed.get_rank(process_group) < mod) - super().__init__( - in_features, local_multiple * multiple_of, bias=bias, device=device, dtype=dtype - ) - self.process_group = process_group - self.sequence_parallel = sequence_parallel - - def forward(self, x): - # If self.sequence_parallel is True, we're doing Tensor Parallel with sequence parallelism: - # we do an all_gather of x before doing the matmul. - # If not, then the input is already gathered. - return fused_dense_func( - x, - self.weight, - self.bias, - process_group=self.process_group, - sequence_parallel=self.sequence_parallel, - ) - - -class RowParallelLinear(nn.Linear): - def __init__( - self, - in_features: int, - out_features: int, - process_group: ProcessGroup, - bias: bool = True, - sequence_parallel=True, - multiple_of=1, - device=None, - dtype=None, - ) -> None: - world_size = torch.distributed.get_world_size(process_group) - rank = torch.distributed.get_rank(process_group) - if in_features % multiple_of: - raise ValueError(f"in_features ({in_features}) must be a multiple of {multiple_of}") - multiple = in_features // multiple_of - # We want to split @multiple across world_size, but it could be an uneven split - div = multiple // world_size - mod = multiple % world_size - # The first @mod ranks get @div + 1 copies, the rest get @div copies - local_multiple = div + int(torch.distributed.get_rank(process_group) < mod) - # Only rank 0 will have bias - super().__init__( - local_multiple * multiple_of, - out_features, - bias=bias and rank == 0, - device=device, - dtype=dtype, - ) - self.process_group = process_group - self.sequence_parallel = sequence_parallel - - def forward(self, x): - """ - We're doing Tensor Parallel with sequence parallelism: we do the matmul and then - a reduce_scatter of the result. - """ - out = fused_dense_func(x, self.weight, self.bias) - reduce_fn = reduce_scatter if self.sequence_parallel else all_reduce - return reduce_fn(out, self.process_group) - - -class FusedMLPFunc(torch.autograd.Function): - @staticmethod - @custom_fwd - def forward( - ctx, - x, - weight1, - bias1, - weight2, - bias2, - activation="gelu_approx", - save_pre_act=True, - return_residual=False, - checkpoint_lvl=0, - heuristic=0, - process_group=None, - sequence_parallel=True, - ): - """ - If process_group is not None and sequence_parallel=True, we're doing Tensor Parallel - with sequence parallelism: we do an all_gather of x before doing the matmul. - If sequence_parallel=False, then the input is already gathered. - - checkpoint_lvl: - 0: no recomputation in the bwd - 1: recompute gelu_out / relu_out in the bwd - 2: recompute pre_act and gelu_out / relu_out in the bwd - """ - assert -1 <= heuristic <= 4 - assert activation in ["gelu_approx", "relu", "sqrelu"] - if activation == "sqrelu": - assert heuristic == -1 - if not save_pre_act: - checkpoint_lvl = 2 - assert checkpoint_lvl in [0, 1, 2] - ctx.return_residual = return_residual - ctx.process_group = process_group - ctx.sequence_parallel = sequence_parallel - ctx.checkpoint_lvl = checkpoint_lvl - ctx.activation = activation - ctx.heuristic = heuristic - - if torch.is_autocast_enabled(): - x = x.to(dtype=torch.get_autocast_gpu_dtype()) - x = x.contiguous() - if process_group is not None and sequence_parallel: - # We want to kick off the all_gather early, before weight dtype conversion - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - else: - total_x = x - - if torch.is_autocast_enabled(): - dtype = torch.get_autocast_gpu_dtype() - weight1, weight2 = [a.to(dtype=dtype) for a in [weight1, weight2]] - bias1 = bias1.to(dtype=dtype) if bias1 is not None else None - bias2 = bias2.to(dtype=dtype) if bias2 is not None else None - weight1 = weight1.contiguous() - bias1 = bias1.contiguous() if bias1 is not None else None - weight2 = weight2.contiguous() - bias2 = bias2.contiguous() if bias2 is not None else None - if process_group is not None and sequence_parallel: - handle_x.wait() - batch_shape, n = total_x.shape[:-1], total_x.shape[-1] - batch_dim = batch_shape.numel() - # https://github.com/pytorch/pytorch/blob/5b51849b48a7dbccd297286cc0110def4706f9e7/aten/src/ATen/native/cuda/Blas.cpp#L174 - if min(batch_dim, n, *weight1.shape, *weight2.shape) > 65535 * 32: - raise RuntimeError("fused_dense only supports matrix dims <= 2M") - if heuristic == -1: - pre_act = F.linear(total_x, weight1, bias1) - activation_fn = ( - partial(F.gelu, approximate="tanh") - if activation == "gelu_approx" - else (sqrelu_fwd if activation == "sqrelu" else F.relu) - ) - with torch.jit.fuser("fuser2"): - output1 = activation_fn(pre_act) - # This is before adding bias1 - # pre_act = F.linear(total_x.reshape(batch_dim, n), weight1) - # with torch.jit.fuser('fuser2'): - # output1 = bias_gelu(pre_act, bias1) - else: - is_gelu = activation == "gelu_approx" - output1, *rest = fused_dense_cuda.linear_act_forward( - total_x.reshape(batch_dim, n), weight1, bias1, is_gelu, save_pre_act, heuristic - ) - if save_pre_act: - pre_act = rest[0] - output2 = F.linear(output1, weight2, bias2) - if checkpoint_lvl == 0 or (checkpoint_lvl == 1 and activation == "relu"): - # For RELU the pre_act is very small (just a bit-mask) so we just save it - ctx.save_for_backward(x, weight1, weight2, pre_act, output1) - elif checkpoint_lvl == 1: - ctx.save_for_backward(x, weight1, weight2, pre_act) - elif checkpoint_lvl == 2: - ctx.save_for_backward(x, weight1, weight2, bias1) - output2 = output2.reshape(*batch_shape, output2.shape[-1]) - return output2 if not return_residual else (output2, x) - - @staticmethod - @custom_bwd - def backward(ctx, grad_output, *args): - grad_output = grad_output.contiguous() - checkpoint_lvl = ctx.checkpoint_lvl - activation = ctx.activation - activation_fn = ( - partial(F.gelu, approximate="tanh") - if activation == "gelu_approx" - else (sqrelu_fwd if activation == "sqrelu" else F.relu) - ) - if ctx.return_residual: - (grad_input,) = args - grad_input = grad_input.contiguous() - process_group = ctx.process_group - sequence_parallel = ctx.sequence_parallel - x, weight1, weight2, *rest = ctx.saved_tensors - if process_group is None or not sequence_parallel: - total_x = x - batch_shape = grad_output.shape[:-1] - batch_dim = batch_shape.numel() - if checkpoint_lvl in [0, 1]: - if process_group is not None and sequence_parallel: - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - if checkpoint_lvl == 0 or (checkpoint_lvl == 1 and activation == "relu"): - pre_act, output1 = rest - elif checkpoint_lvl == 1: - (pre_act,) = rest - with torch.jit.fuser("fuser2"): - output1 = activation_fn(pre_act) - elif checkpoint_lvl == 2: - (bias1,) = rest - if process_group is not None and sequence_parallel: - total_x, _ = all_gather_raw(x, process_group) - if ctx.heuristic == -1: - pre_act = F.linear(total_x, weight1, bias1) - with torch.jit.fuser("fuser2"): - output1 = activation_fn(pre_act) - else: - output1, pre_act = fused_dense_cuda.linear_act_forward( - total_x.reshape(batch_dim, total_x.shape[-1]), - weight1, - bias1, - activation == "gelu_approx", - True, - ctx.heuristic, - ) - - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - output1 = output1.reshape(batch_dim, output1.shape[-1]) - pre_act = pre_act.reshape(batch_dim, pre_act.shape[-1]) - if ctx.needs_input_grad[3]: - grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad( - output1, grad_output, ctx.needs_input_grad[4] - ) - else: - grad_weight2 = None - grad_bias2 = grad_output if ctx.needs_input_grad[4] else None - if ctx.heuristic == -1: - # grad_pre_act = matmul_dgelu(grad_output, weight2, pre_act) - grad_output1 = F.linear(grad_output, weight2.t()) - activation_grad_fn = ( - gelu_bwd - if activation == "gelu_approx" - else (sqrelu_bwd if activation == "sqrelu" else relu_bwd) - ) - with torch.jit.fuser("fuser2"): - grad_pre_act = activation_grad_fn(grad_output1, pre_act) - else: - # The cublasLt epilogue has to compute both gelu/relu grad and bias grad, we can't - # just compute gelu/relu grad - grad_pre_act, grad_bias1 = fused_dense_cuda.bias_act_linear_dgrad_bgrad( - weight2, grad_output, pre_act, activation == "gelu_approx", ctx.heuristic - ) - if not ctx.needs_input_grad[2]: - grad_bias1 = None - if ctx.needs_input_grad[0]: - if not ctx.return_residual: - grad_input = F.linear(grad_pre_act, weight1.t()) - else: - grad_input = torch.addmm( - grad_input.reshape(batch_dim, grad_input.shape[-1]), grad_pre_act, weight1 - ) - grad_input = grad_input.reshape(*batch_shape, grad_input.shape[-1]) - if process_group is not None: - reduce_fn = reduce_scatter_raw if sequence_parallel else all_reduce_raw - grad_input, handle_grad_input = reduce_fn(grad_input, process_group, async_op=True) - else: - grad_input = None - if ctx.heuristic == -1: - if ctx.needs_input_grad[1]: - if process_group is not None and sequence_parallel and checkpoint_lvl != 2: - handle_x.wait() - grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_wgrad( - total_x.reshape(batch_dim, total_x.shape[-1]), - grad_pre_act, - ctx.needs_input_grad[2], - ) - else: - grad_weight1 = None - grad_bias1 = grad_pre_act if ctx.needs_input_grad[2] else None - else: - if ctx.needs_input_grad[1]: - if process_group is not None and sequence_parallel and checkpoint_lvl != 2: - handle_x.wait() - grad_weight1 = F.linear( - grad_pre_act.t(), total_x.reshape(batch_dim, total_x.shape[-1]).t() - ) - else: - grad_weight1 = None - if process_group is not None and ctx.needs_input_grad[0]: - handle_grad_input.wait() - return ( - grad_input, - grad_weight1, - grad_bias1, - grad_weight2, - grad_bias2, - None, - None, - None, - None, - None, - None, - None, - ) - - -def fused_mlp_func( - x: Tensor, - weight1: Tensor, - weight2: Tensor, - bias1: Optional[Tensor] = None, - bias2: Optional[Tensor] = None, - activation: str = "gelu_approx", - save_pre_act: bool = True, - return_residual: bool = False, - checkpoint_lvl: int = 0, - heuristic: int = 0, - process_group: Optional[ProcessGroup] = None, - sequence_parallel: bool = True, -): - assert activation in ["gelu_approx", "relu", "sqrelu"] - dtype_eligible = x.dtype in [torch.float16, torch.bfloat16] or ( - x.dtype == torch.float32 and torch.is_autocast_enabled() - ) - # If we save pre-activation, dimension must be divisible by 128 (relu) or 8 (gelu) - dim_eligible = not save_pre_act or (x.shape[-1] % (128 if activation == "relu" else 8) == 0) - if ( - x.is_cuda - and weight1.is_cuda - and weight2.is_cuda - and (bias1 is None or bias1.is_cuda) - and (bias2 is None or bias2.is_cuda) - and dtype_eligible - and dim_eligible - ): - return FusedMLPFunc.apply( - x, - weight1, - bias1, - weight2, - bias2, - activation, - save_pre_act, - return_residual, - checkpoint_lvl, - heuristic, - process_group, - sequence_parallel, - ) - else: - assert process_group is None - pre_act = F.linear(x, weight1, bias1) - activation_fn = ( - partial(F.gelu, approximate="tanh") - if activation == "gelu_approx" - else partial(F.relu, inplace=True) - ) - output1 = activation_fn(pre_act) - output2 = F.linear(output1, weight2, bias2) - return output2 if not return_residual else (output2, x) - - -class FusedMLP(nn.Module): - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - bias1=True, - bias2=True, - activation="gelu_approx", - return_residual=False, - checkpoint_lvl=0, - heuristic="auto", - device=None, - dtype=None, - ): - """ - If process_group is not None, we're doing Tensor Parallel with sequence parallelism: - we do an all_gather of x before doing the matmul, gelu, then matmul. - Finally we do a reduce_scatter of the output. - - checkpoint_lvl (increasing lvl means slower but more memory saving): - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute pre_act and gelu_out in the bwd - heuristic: - -1: don't fuse gemm + gelu (separate kernel) - 0..4: use this heuristic for the algo section in the fused gemm + gelu - 'auto': heuristic will be picked automatically: - For CUDA >= 11.8, we set heuristic=0 for both fp16 and bf16 for best perf. - For CUDA <= 11.7, we set heuristic=1 for fp16 and heuristic=-1 for bf16. - For H100, we set heuristic=-1 for both fp16 and bf16 as the fused cuBlasLt implementation - is slower than the unfused version. - return_residual: whether to return the input x along with the output. This is for - performance reason: for post-norm architecture, returning the input allows us - to fuse the backward of nn.Linear with the residual connection. - """ - assert checkpoint_lvl in [0, 1, 2] - assert activation in ["gelu_approx", "relu", "sqrelu"] - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features * 4 - self.activation = activation - self.return_residual = return_residual - self.checkpoint_lvl = checkpoint_lvl - self.heuristic = heuristic if activation != "sqrelu" else -1 - self.fc1 = nn.Linear(in_features, hidden_features, bias=bias1, **factory_kwargs) - self.fc2 = nn.Linear(hidden_features, out_features, bias=bias2, **factory_kwargs) - - def forward(self, x, process_group=None): - dtype = x.dtype if not torch.is_autocast_enabled() else torch.get_autocast_gpu_dtype() - if self.heuristic == "auto": - if self.activation == "gelu_approx": - if torch.cuda.get_device_capability("cuda") == (9, 0): - heuristic = -1 - else: - cuda_ver = tuple(map(int, torch.version.cuda.split("."))) - heuristic = 0 if cuda_ver >= (11, 8) else (1 if dtype == torch.float16 else -1) - else: - heuristic = 0 - else: - heuristic = self.heuristic - out = fused_mlp_func( - x, - self.fc1.weight, - self.fc2.weight, - self.fc1.bias, - self.fc2.bias, - activation=self.activation, - save_pre_act=self.training, - return_residual=self.return_residual, - checkpoint_lvl=self.checkpoint_lvl, - heuristic=heuristic, - process_group=process_group, - ) - if self.return_residual: - out, x = out - if process_group is not None: - out = reduce_scatter(out, process_group) - return out if not self.return_residual else (out, x) - - -class ParallelFusedMLP(nn.Module): - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - activation="gelu_approx", - process_group: ProcessGroup = None, - bias1=True, - bias2=True, - sequence_parallel=True, - checkpoint_lvl=0, - heuristic="auto", - device=None, - dtype=None, - ): - """ - process_group is required. We're doing Tensor Parallel with sequence parallelism: - we do an all_gather of x before doing the matmul, gelu, then matmul. - Finally we do a reduce_scatter of the output. - - checkpoint_lvl (increasing lvl means slower but more memory saving): - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute pre_act and gelu_out in the bwd - heuristic: - -1: don't fuse gemm + gelu (separate kernel) - 0..4: use this heuristic for the algo section in the fused gemm + gelu - 'auto': heuristic will be picked automatically: - For CUDA >= 11.8, we set heuristic=0 for both fp16 and bf16 for best perf. - For CUDA <= 11.7, we set heuristic=1 for fp16 and heuristic=-1 for bf16. - """ - assert checkpoint_lvl in [0, 1, 2] - assert activation in ["gelu_approx", "relu", "sqrelu"] - assert process_group is not None - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features * 4 - self.activation = activation - self.process_group = process_group - self.sequence_parallel = sequence_parallel - self.checkpoint_lvl = checkpoint_lvl - self.heuristic = heuristic if activation != "sqrelu" else -1 - self.fc1 = ColumnParallelLinear( - in_features, hidden_features, process_group, bias=bias1, **factory_kwargs - ) - self.fc2 = RowParallelLinear( - hidden_features, out_features, process_group, bias=bias2, **factory_kwargs - ) - - def forward(self, x): - dtype = x.dtype if not torch.is_autocast_enabled() else torch.get_autocast_gpu_dtype() - if self.heuristic == "auto": - if self.activation == "gelu_approx": - cuda_ver = tuple(map(int, torch.version.cuda.split("."))) - heuristic = 0 if cuda_ver >= (11, 8) else (1 if dtype == torch.float16 else -1) - else: - heuristic = 0 - else: - heuristic = self.heuristic - out = fused_mlp_func( - x, - self.fc1.weight, - self.fc2.weight, - self.fc1.bias, - self.fc2.bias, - activation=self.activation, - save_pre_act=self.training, - checkpoint_lvl=self.checkpoint_lvl, - heuristic=heuristic, - process_group=self.process_group, - sequence_parallel=self.sequence_parallel, - ) - reduce_fn = reduce_scatter if self.sequence_parallel else all_reduce - return reduce_fn(out, self.process_group) diff --git a/build/torch211-cxx11-xpu20253-x86_64-linux/ops/layer_norm.py b/build/torch211-cxx11-xpu20253-x86_64-linux/ops/layer_norm.py deleted file mode 100644 index 4b6cd798fd02844ef9cd3897f8ab95e490e638bf..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-xpu20253-x86_64-linux/ops/layer_norm.py +++ /dev/null @@ -1,800 +0,0 @@ -# Copyright (c) 2022, Tri Dao. -# Adapted from https://github.com/NVIDIA/apex/blob/master/apex/contrib/layer_norm/layer_norm.py - -import dropout_layer_norm -import torch -from torch.nn import init - - -def maybe_align(x, alignment_in_bytes=16): - """Assume that x already has last dim divisible by alignment_in_bytes""" - # TD [2023-07-04] I'm not 100% sure that clone will align the memory - # https://discuss.pytorch.org/t/how-to-ensure-that-tensor-data-ptr-is-aligned-to-16-bytes/183440 - return x if x.data_ptr() % alignment_in_bytes == 0 else x.clone() - - -def _dropout_add_layer_norm_forward( - x0, - residual, - gamma, - beta, - rowscale, - colscale, - dropout_p, - epsilon, - residual_in_fp32=False, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes""" - hidden_size = gamma.numel() - x0mat = x0.view((-1, hidden_size)) - residualmat = residual.view((-1, hidden_size)) if residual is not None else None - rowscale = rowscale.view(-1) if rowscale is not None else None - zmat, xmat, dmask, mu, rsigma = dropout_layer_norm.dropout_add_ln_fwd( - x0mat, - residualmat, - gamma, - beta, - rowscale, - colscale, - None, - None, - dropout_p, - epsilon, - 1.0, - 0, - None, - residual_in_fp32, - is_rms_norm, - ) - # dmask is None if dropout_p == 0.0 - # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype - return zmat, xmat if xmat is not None else x0mat, dmask, mu, rsigma - - -def _dropout_add_layer_norm_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - rowscale, - colscale, - dropout_p, - has_residual, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes - dx == None means that it was a post-norm architecture - (x = drop(x0) + residual was not returned in the fwd). - x0 must not be None if we have colscale. - """ - hidden_size = gamma.numel() - xmat = x.view((-1, hidden_size)) - dzmat = dz.view(xmat.shape) - dxmat = dx.view(xmat.shape) if dx is not None else None - x0mat = x0.view((-1, hidden_size)) if x0 is not None else None - rowscale = rowscale.view(-1) if rowscale is not None else None - if colscale is not None: - assert x0 is not None, "x0 is required to compute the gradient of colscale" - dx0mat, dresidualmat, dgamma, dbeta, _, _, *rest = dropout_layer_norm.dropout_add_ln_bwd( - dzmat, - dxmat, - xmat, - x0mat, - dmask, - mu, - rsigma, - gamma, - rowscale, - colscale, - None, - None, - dropout_p, - 1.0, - 0, - has_residual, - is_rms_norm, - ) - # dresidualmat is None if not has_residual - if colscale is None: - return dx0mat, dresidualmat, dgamma, dbeta - else: - dcolscale = rest[0] - return dx0mat, dresidualmat, dgamma, dbeta, dcolscale - - -def _dropout_add_layer_norm_subset_forward( - x0, - residual, - gamma, - beta, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32=False, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes""" - hidden_size = gamma.numel() - x0mat = x0.view((-1, hidden_size)) - residualmat = residual.view((-1, hidden_size)) if residual is not None else None - x0_subset = x0_subset.view(-1) if x0_subset is not None else None - out_subset = out_subset.view(-1) if out_subset is not None else None - zmat, xmat, dmask, mu, rsigma = dropout_layer_norm.dropout_add_ln_fwd( - x0mat, - residualmat, - gamma, - beta, - None, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - None, - residual_in_fp32, - is_rms_norm, - ) - # dmask is None if dropout_p == 0.0 - # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype - return zmat, xmat if xmat is not None else x0mat, dmask, mu, rsigma - - -def _dropout_add_layer_norm_subset_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - colscale, - x0_subset, - out_subset, - dropout_p, - rowscale_const, - x0_numrows, - has_residual, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes - dx == None means that it was a post-norm architecture - (x = drop(x0) + residual was not returned in the fwd). - x0 must not be None if we have colscale. - """ - hidden_size = gamma.numel() - xmat = x.view((-1, hidden_size)) - dzmat = dz.view(-1, hidden_size) - dxmat = dx.view(xmat.shape) if dx is not None else None - x0mat = x0.view((-1, hidden_size)) if x0 is not None else None - x0_subset = x0_subset.view(-1) if x0_subset is not None else None - out_subset = out_subset.view(-1) if out_subset is not None else None - if colscale is not None: - assert x0 is not None, "x0 is required to compute the gradient of colscale" - dx0mat, dresidualmat, dgamma, dbeta, _, _, *rest = dropout_layer_norm.dropout_add_ln_bwd( - dzmat, - dxmat, - xmat, - x0mat, - dmask, - mu, - rsigma, - gamma, - None, - colscale, - x0_subset, - out_subset, - dropout_p, - rowscale_const, - x0_numrows, - has_residual, - is_rms_norm, - ) - # dresidualmat is None if not has_residual - if colscale is None: - return dx0mat, dresidualmat, dgamma, dbeta - else: - dcolscale = rest[0] - return dx0mat, dresidualmat, dgamma, dbeta, dcolscale - - -def _dropout_add_layer_norm_parallel_residual_forward( - x0, - x1, - residual, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - residual_in_fp32=False, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes""" - hidden_size = gamma0.numel() - x0mat = x0.view((-1, hidden_size)) - x1mat = x1.view((-1, hidden_size)) if x1 is not None else None - residualmat = residual.view((-1, hidden_size)) if residual is not None else None - ( - z0mat, - z1mat, - xmat, - dmask0, - dmask1, - mu, - rsigma, - ) = dropout_layer_norm.dropout_add_ln_parallel_residual_fwd( - x0mat, - x1mat, - residualmat, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - None, - residual_in_fp32, - is_rms_norm, - ) - # dmask0 and dmask1 are None if dropout_p == 0.0 - # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype - return z0mat, z1mat, xmat if xmat is not None else x0mat, dmask0, dmask1, mu, rsigma - - -def _dropout_add_layer_norm_parallel_residual_backward( - dz0, - dz1, - dx, - x, - dmask0, - dmask1, - mu, - rsigma, - gamma0, - gamma1, - dropout_p, - has_x1, - has_residual, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes - dx == None means that it was a post-norm architecture - (x = drop(x0) + residual was not returned in the fwd). - """ - hidden_size = gamma0.numel() - xmat = x.view((-1, hidden_size)) - dz0mat = dz0.view(xmat.shape) - dz1mat = dz1.view(xmat.shape) if dz1 is not None else None - dxmat = dx.view(xmat.shape) if dx is not None else None - ( - dx0mat, - dx1mat, - dresidualmat, - dgamma0, - dbeta0, - dgamma1, - dbeta1, - *rest, - ) = dropout_layer_norm.dropout_add_ln_parallel_residual_bwd( - dz0mat, - dz1mat, - dxmat, - xmat, - dmask0, - dmask1, - mu, - rsigma, - gamma0, - gamma1, - dropout_p, - has_x1, - has_residual, - is_rms_norm, - ) - # dresidualmat is None if not has_residual - return dx0mat, dx1mat, dresidualmat, dgamma0, dbeta0, dgamma1, dbeta1 - - -class DropoutAddLayerNormFn(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x0, - residual, - gamma, - beta, - rowscale, - colscale, - dropout_p, - epsilon, - residual_in_fp32=False, - prenorm=False, - is_rms_norm=False, - return_dmask=False, - ): - x0 = maybe_align(x0.contiguous(), 16) - residual = maybe_align(residual.contiguous(), 16) if residual is not None else None - gamma = maybe_align(gamma.contiguous(), 16) - beta = maybe_align(beta.contiguous(), 16) if beta is not None else None - rowscale = maybe_align(rowscale.contiguous(), 16) if rowscale is not None else None - colscale = maybe_align(colscale.contiguous(), 16) if colscale is not None else None - zmat, xmat, dmask, mu, rsigma = _dropout_add_layer_norm_forward( - x0, - residual, - gamma, - beta, - rowscale, - colscale, - dropout_p, - epsilon, - residual_in_fp32, - is_rms_norm, - ) - # Only need to save x0 if we need to compute gradient wrt colscale - x0_saved = x0 if colscale is not None else None - ctx.save_for_backward( - xmat.view(x0.shape), x0_saved, dmask, gamma, mu, rsigma, rowscale, colscale - ) - ctx.prenorm = prenorm - ctx.dropout_p = dropout_p - ctx.has_residual = residual is not None - ctx.is_rms_norm = is_rms_norm - ctx.has_beta = beta is not None - if not return_dmask: - return ( - zmat.view(x0.shape) if not prenorm else (zmat.view(x0.shape), xmat.view(x0.shape)) - ) - else: - dmask = ( - dmask.view(x0.shape) - if dropout_p > 0.0 - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - ctx.mark_non_differentiable(dmask) - return ( - (zmat.view(x0.shape), dmask) - if not prenorm - else (zmat.view(x0.shape), xmat.view(x0.shape), dmask) - ) - - @staticmethod - def backward(ctx, dz, *args): - # assert dz.is_contiguous() - dz = maybe_align(dz.contiguous(), 16) # this happens! - dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None - x, x0, dmask, gamma, mu, rsigma, rowscale, colscale = ctx.saved_tensors - # x0 is None if colscale is None - dropout_p = ctx.dropout_p - has_residual = ctx.has_residual - dx0mat, dresidualmat, dgamma, dbeta, *rest = _dropout_add_layer_norm_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - rowscale, - colscale, - dropout_p, - has_residual, - ctx.is_rms_norm, - ) - dx0 = dx0mat.view(x.shape) - dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None - dcolscale = rest[0] if colscale is not None else None - return ( - dx0, - dresidual, - dgamma, - dbeta if ctx.has_beta else None, - None, - dcolscale, - None, - None, - None, - None, - None, - None, - ) - - -class DropoutAddLayerNormSubsetFn(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x0, - residual, - gamma, - beta, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32=False, - prenorm=False, - is_rms_norm=False, - return_dmask=False, - ): - x0 = maybe_align(x0.contiguous(), 16) - residual = maybe_align(residual.contiguous(), 16) if residual is not None else None - gamma = maybe_align(gamma.contiguous(), 16) - beta = maybe_align(beta.contiguous(), 16) if beta is not None else None - colscale = maybe_align(colscale.contiguous(), 16) if colscale is not None else None - zmat, xmat, dmask, mu, rsigma = _dropout_add_layer_norm_subset_forward( - x0, - residual, - gamma, - beta, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32, - is_rms_norm, - ) - # Only need to save x0 if we need to compute gradient wrt colscale - x0_saved = x0 if colscale is not None else None - x_shape = (-1, *x0.shape[1:]) - ctx.save_for_backward( - xmat.view(x_shape), x0_saved, dmask, gamma, mu, rsigma, colscale, x0_subset, out_subset - ) - ctx.prenorm = prenorm - ctx.dropout_p = dropout_p - ctx.rowscale_const = rowscale_const - ctx.x0_numrows = x0.shape[:-1].numel() - ctx.has_residual = residual is not None - ctx.is_rms_norm = is_rms_norm - ctx.has_beta = beta is not None - z_shape = (-1, *x0.shape[1:]) - if not return_dmask: - return zmat.view(z_shape) if not prenorm else (zmat.view(z_shape), xmat.view(x0.shape)) - else: - z = zmat.view(z_shape) - dmask = ( - dmask.view(x0.shape) - if dropout_p > 0.0 - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - ctx.mark_non_differentiable(dmask) - return (z, dmask) if not prenorm else (z, xmat.view(x_shape), dmask) - - @staticmethod - def backward(ctx, dz, *args): - # assert dz.is_contiguous() - dz = maybe_align(dz.contiguous(), 16) # this happens! - dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None - x, x0, dmask, gamma, mu, rsigma, colscale, x0_subset, out_subset = ctx.saved_tensors - # x0 is None if colscale is None - dropout_p = ctx.dropout_p - has_residual = ctx.has_residual - dx0mat, dresidualmat, dgamma, dbeta, *rest = _dropout_add_layer_norm_subset_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - colscale, - x0_subset, - out_subset, - dropout_p, - ctx.rowscale_const, - ctx.x0_numrows, - has_residual, - ctx.is_rms_norm, - ) - dx0 = dx0mat.view(-1, *x.shape[1:]) - dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None - dcolscale = rest[0] if colscale is not None else None - return ( - dx0, - dresidual, - dgamma, - dbeta if ctx.has_beta else None, - dcolscale, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) - - -class DropoutAddLayerNormParallelResidualFn(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x0, - x1, - residual, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - residual_in_fp32=False, - prenorm=False, - is_rms_norm=False, - return_dmask=False, - ): - x0 = maybe_align(x0.contiguous(), 16) - x1 = maybe_align(x1.contiguous(), 16) if x1 is not None else None - residual = maybe_align(residual.contiguous(), 16) if residual is not None else None - gamma0 = maybe_align(gamma0.contiguous(), 16) - beta0 = maybe_align(beta0.contiguous(), 16) if beta0 is not None else None - gamma1 = maybe_align(gamma1.contiguous(), 16) if gamma1 is not None else None - beta1 = maybe_align(beta1.contiguous(), 16) if beta1 is not None else None - ( - z0mat, - z1mat, - xmat, - dmask0, - dmask1, - mu, - rsigma, - ) = _dropout_add_layer_norm_parallel_residual_forward( - x0, - x1, - residual, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - residual_in_fp32, - is_rms_norm, - ) - ctx.save_for_backward(xmat.view(x0.shape), dmask0, dmask1, gamma0, gamma1, mu, rsigma) - ctx.prenorm = prenorm - ctx.dropout_p = dropout_p - ctx.has_x1 = x1 is not None - ctx.has_residual = residual is not None - ctx.is_rms_norm = is_rms_norm - ctx.has_beta = beta0 is not None - z = (z0mat.view(x0.shape), z1mat.view(x0.shape) if z1mat is not None else None) - if not return_dmask: - return z if not prenorm else (*z, xmat.view(x0.shape)) - else: - dmask0 = ( - dmask0.view(x0.shape) - if dropout_p > 0.0 - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - dmask1 = ( - dmask1.view(x0.shape) - if dropout_p > 0.0 and x1 is not None - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - ctx.mark_non_differentiable(dmask0) - ctx.mark_non_differentiable(dmask1) - return ( - (*z, dmask0, dmask1) if not prenorm else (*z, xmat.view(x0.shape), dmask0, dmask1) - ) - - @staticmethod - def backward(ctx, dz0, dz1, *args): - dz0 = maybe_align(dz0.contiguous(), 16) # this happens! - dz1 = maybe_align(dz1.contiguous(), 16) if dz1 is not None else None - dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None - x, dmask0, dmask1, gamma0, gamma1, mu, rsigma = ctx.saved_tensors - dropout_p = ctx.dropout_p - has_x1 = ctx.has_x1 - has_residual = ctx.has_residual - ( - dx0mat, - dx1mat, - dresidualmat, - dgamma0, - dbeta0, - dgamma1, - dbeta1, - ) = _dropout_add_layer_norm_parallel_residual_backward( - dz0, - dz1, - dx, - x, - dmask0, - dmask1, - mu, - rsigma, - gamma0, - gamma1, - dropout_p, - has_x1, - has_residual, - ctx.is_rms_norm, - ) - dx0 = dx0mat.view(x.shape) - dx1 = dx1mat.view(x.shape) if dx1mat is not None else None - dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None - return ( - dx0, - dx1, - dresidual, - dgamma0, - dbeta0 if ctx.has_beta else None, - dgamma1, - dbeta1 if ctx.has_beta else None, - None, - None, - None, - None, - None, - None, - ) - - -def layer_norm(x, weight, bias, epsilon): - return DropoutAddLayerNormFn.apply(x, None, weight, bias, None, None, 0.0, epsilon, False) - - -def dropout_add_layer_norm( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - rowscale=None, - layerscale=None, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormFn.apply( - x0, - residual, - weight, - bias, - rowscale, - layerscale, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - False, - return_dropout_mask, - ) - - -def dropout_add_layer_norm_subset( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - layerscale=None, - x0_subset=None, - out_subset=None, - rowscale_const=1.0, - out_numrows=0, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormSubsetFn.apply( - x0, - residual, - weight, - bias, - layerscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32, - prenorm, - False, - return_dropout_mask, - ) - - -def dropout_add_layer_norm_parallel_residual( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormParallelResidualFn.apply( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - False, - return_dropout_mask, - ) - - -class DropoutAddLayerNorm(torch.nn.Module): - def __init__( - self, - hidden_size, - prenorm=False, - p=0.0, - eps=1e-5, - residual_in_fp32=False, - device=None, - dtype=None, - ): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.prenorm = prenorm - self.p = p - self.eps = eps - self.residual_in_fp32 = residual_in_fp32 - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.bias = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.reset_parameters() - - def reset_parameters(self): - init.ones_(self.weight) - init.zeros_(self.bias) - - def forward(self, x0, residual=None): - return dropout_add_layer_norm( - x0, - residual, - self.weight, - self.bias, - self.p if self.training else 0.0, - self.eps, - prenorm=self.prenorm, - residual_in_fp32=self.residual_in_fp32, - ) diff --git a/build/torch211-cxx11-xpu20253-x86_64-linux/ops/rms_norm.py b/build/torch211-cxx11-xpu20253-x86_64-linux/ops/rms_norm.py deleted file mode 100644 index 068348d61290e3839dd082b540d898578ba1e8e2..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-xpu20253-x86_64-linux/ops/rms_norm.py +++ /dev/null @@ -1,174 +0,0 @@ -# Copyright (c) 2022, Tri Dao. -# Adapted from https://github.com/NVIDIA/apex/blob/master/apex/contrib/layer_norm/layer_norm.py - -import torch -from torch.nn import init - -from flash_attn.ops.layer_norm import ( - DropoutAddLayerNormFn, - DropoutAddLayerNormParallelResidualFn, - DropoutAddLayerNormSubsetFn, -) - - -def rms_norm(x, weight, epsilon): - return DropoutAddLayerNormFn.apply( - x, None, weight, None, None, None, 0.0, epsilon, False, False, True - ) - - -def dropout_add_rms_norm( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - rowscale=None, - layerscale=None, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormFn.apply( - x0, - residual, - weight, - bias, - rowscale, - layerscale, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - True, - return_dropout_mask, - ) - - -def dropout_add_rms_norm_subset( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - layerscale=None, - x0_subset=None, - out_subset=None, - rowscale_const=1.0, - out_numrows=0, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormSubsetFn.apply( - x0, - residual, - weight, - bias, - layerscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32, - prenorm, - True, - return_dropout_mask, - ) - - -def dropout_add_rms_norm_parallel_residual( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormParallelResidualFn.apply( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - True, - return_dropout_mask, - ) - - -class RMSNorm(torch.nn.Module): - def __init__(self, hidden_size, eps=1e-5, device=None, dtype=None): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.eps = eps - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.register_parameter("bias", None) - self.reset_parameters() - - def reset_parameters(self): - init.ones_(self.weight) - - def forward(self, x): - return rms_norm(x, self.weight, self.eps) - - -class DropoutAddRMSNorm(torch.nn.Module): - def __init__( - self, - hidden_size, - prenorm=False, - p=0.0, - eps=1e-5, - residual_in_fp32=False, - device=None, - dtype=None, - ): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.prenorm = prenorm - self.p = p - self.eps = eps - self.residual_in_fp32 = residual_in_fp32 - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.register_parameter("bias", None) - self.reset_parameters() - - def reset_parameters(self): - init.ones_(self.weight) - - def forward(self, x0, residual=None): - return dropout_add_rms_norm( - x0, - residual, - self.weight, - None, - self.p if self.training else 0.0, - self.eps, - prenorm=self.prenorm, - residual_in_fp32=self.residual_in_fp32, - ) diff --git a/build/torch211-cxx11-xpu20253-x86_64-linux/ops/triton/__init__.py b/build/torch211-cxx11-xpu20253-x86_64-linux/ops/triton/__init__.py deleted file mode 100644 index 8b137891791fe96927ad78e64b0aad7bded08bdc..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-xpu20253-x86_64-linux/ops/triton/__init__.py +++ /dev/null @@ -1 +0,0 @@ - diff --git a/build/torch211-cxx11-xpu20253-x86_64-linux/ops/triton/cross_entropy.py b/build/torch211-cxx11-xpu20253-x86_64-linux/ops/triton/cross_entropy.py deleted file mode 100644 index 1b5a415b73f236f3e05fb14b9141959559e18526..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-xpu20253-x86_64-linux/ops/triton/cross_entropy.py +++ /dev/null @@ -1,330 +0,0 @@ -# Copyright (c) 2023, Tri Dao. - -from typing import Tuple, Optional, Union - -import torch -import torch.nn.functional as F - -import triton -import triton.language as tl - -# `all_gather_into_tensor` and `reduce_scatter_tensor` are new placeholders for -# `_all_gather_base` and `_reduce_scatter_base`. They require the most recent -# version of PyTorch. The following 2 lines are for backward compatibility with -# older PyTorch. -if "all_gather_into_tensor" not in dir(torch.distributed): - torch.distributed.all_gather_into_tensor = torch.distributed._all_gather_base - - -@triton.heuristics( - { - "HAS_SMOOTHING": lambda args: args["smoothing"] > 0.0, - } -) -@triton.jit -def cross_entropy_fwd_kernel( - loss_ptr, # data ptrs - lse_ptr, - z_loss_ptr, - logits_ptr, - labels_ptr, - smoothing, - logit_scale, - lse_square_scale, - ignore_index, - total_classes, - class_start_idx, # Useful for tensor parallel when each rank only has a subset of classes - n_cols, # shapes - logits_row_stride, # strides - BLOCK_SIZE: tl.constexpr, - HAS_SMOOTHING: tl.constexpr, - # if SPLIT (e.g. tensor parallel), don't include the LSE in the loss since it's not the final LSE - SPLIT: tl.constexpr, - PRECOMPUTED_LSE: tl.constexpr, # If LSE is already computed (also no smoothing and logit_scale == 1.0) -): - row_idx = tl.program_id(0) - logits_ptr = logits_ptr + row_idx * logits_row_stride.to(tl.int64) - sum_logits = 0.0 # For smoothing - if not PRECOMPUTED_LSE: - # Statistics for online softmax - m_i = -float("inf") - l_i = 0.0 - for col_offset in range(0, n_cols, BLOCK_SIZE): - cols = col_offset + tl.arange(0, BLOCK_SIZE) - logits = tl.load(logits_ptr + cols, mask=cols < n_cols, other=-float("inf")).to( - tl.float32 - ) * logit_scale - if HAS_SMOOTHING: - sum_logits += tl.sum(tl.where(cols < n_cols, logits, 0.0)) - m_i_new = tl.maximum(m_i, tl.max(logits)) - l_i = tl.exp(m_i - m_i_new) * l_i + tl.sum(tl.exp(logits - m_i_new)) - m_i = m_i_new - lse = tl.log(l_i) + m_i - tl.store(lse_ptr + row_idx, lse) - else: - lse = tl.load(lse_ptr + row_idx) - label_idx = tl.load(labels_ptr + row_idx) - if label_idx == ignore_index: - loss = 0.0 - z_loss = 0.0 - else: - label_idx -= class_start_idx - if label_idx >= 0 and label_idx < n_cols: - logits_label = tl.load(logits_ptr + label_idx) * logit_scale - if HAS_SMOOTHING: - loss = ( - (lse if not SPLIT else 0.0) - - smoothing * sum_logits / total_classes - - (1 - smoothing) * logits_label - ) - else: - loss = (lse if not SPLIT else 0.0) - logits_label - else: - # If label is out of bounds, we set the CE loss to 0.0. But we still want the smoothing loss - if HAS_SMOOTHING: - loss = smoothing * ((lse if not SPLIT else 0.0) - sum_logits / total_classes) - else: - loss = 0.0 - if not SPLIT: - z_loss = lse_square_scale * lse * lse - loss += z_loss - else: - z_loss = 0.0 - tl.store(loss_ptr + row_idx, loss) - if not SPLIT: - tl.store(z_loss_ptr + row_idx, z_loss) - - -@triton.heuristics( - { - "HAS_SMOOTHING": lambda args: args["smoothing"] > 0.0, - } -) -@triton.jit -def cross_entropy_bwd_kernel( - dlogits_ptr, # data ptrs - dloss_ptr, - logits_ptr, - lse_ptr, - labels_ptr, - smoothing, - logit_scale, - lse_square_scale, - ignore_index, - total_classes, - class_start_idx, # Useful for tensor parallel when each rank only has a subset of classes - n_cols, # shapes - logits_row_stride, # strides - dlogits_row_stride, - dloss_row_stride, - BLOCK_SIZE: tl.constexpr, - HAS_SMOOTHING: tl.constexpr, -): - row_idx = tl.program_id(0) - col_block_idx = tl.program_id(1) - logits_ptr = logits_ptr + row_idx * logits_row_stride.to(tl.int64) - dlogits_ptr = dlogits_ptr + row_idx * dlogits_row_stride.to(tl.int64) - col_offsets = col_block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - label_idx = tl.load(labels_ptr + row_idx) - if label_idx != ignore_index: - dloss = tl.load(dloss_ptr + row_idx * dloss_row_stride) - else: - dloss = 0.0 - logits = tl.load(logits_ptr + col_offsets, mask=col_offsets < n_cols, other=-float("inf")).to( - tl.float32 - ) * logit_scale - lse = tl.load(lse_ptr + row_idx) - probs = tl.exp(logits - lse) - probs += 2.0 * lse_square_scale * lse * probs - label_idx -= class_start_idx - if HAS_SMOOTHING: - smooth_positive = 1.0 - smoothing - smooth_negative = smoothing / total_classes - probs = tl.where(col_offsets == label_idx, probs - smooth_positive, probs) - smooth_negative - else: - probs = tl.where(col_offsets == label_idx, probs - 1.0, probs) - tl.store(dlogits_ptr + col_offsets, (dloss * logit_scale) * probs, mask=col_offsets < n_cols) - - -class CrossEntropyLoss(torch.autograd.Function): - - @staticmethod - def forward( - ctx, - logits, - labels, - precomputed_lse=None, - smoothing=0.0, - logit_scale=1.0, - lse_square_scale=0.0, - ignore_index=-100, - inplace_backward=False, - process_group=None, - ): - # For some reason Triton generates wrong code when labels has dtype long and its address - # is not aligned to 16 bytes. The ld.global.b64 seems to load the wrong label index. - if labels.dtype == torch.long and labels.data_ptr() % 16 != 0: - labels = F.pad(labels, (0, 1))[..., :-1] - assert labels.data_ptr() % 16 == 0 - assert logit_scale > 0.0 - n_rows, n_cols = logits.shape - assert labels.shape == (n_rows,) - world_size = 1 if process_group is None else torch.distributed.get_world_size(process_group) - total_classes = world_size * n_cols - rank = 0 if process_group is None else torch.distributed.get_rank(process_group) - class_start_idx = rank * n_cols - use_precomputed_lse = precomputed_lse is not None and logit_scale == 1.0 and smoothing == 0.0 - - if logits.stride(-1) != 1: - logits = logits.contiguous() - MAX_BLOCK_SIZE = 16 * 1024 - BLOCK_SIZE = min(triton.next_power_of_2(n_cols), MAX_BLOCK_SIZE) - num_warps = ( - 4 - if BLOCK_SIZE < 2048 - else (8 if BLOCK_SIZE < 8192 else (16 if BLOCK_SIZE < 128 * 1024 else 32)) - ) - losses = torch.empty(n_rows, dtype=torch.float, device=logits.device) - if use_precomputed_lse: - assert precomputed_lse.shape == (n_rows,) - lse = precomputed_lse.contiguous() - else: - lse = torch.empty(n_rows, dtype=torch.float, device=logits.device) - z_losses = torch.empty(n_rows, dtype=torch.float, device=logits.device) - # Need this, otherwise Triton tries to launch from cuda:0 and we get - # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?) - with torch.cuda.device(logits.device.index): - cross_entropy_fwd_kernel[(n_rows,)]( - losses, # data ptrs - lse, - z_losses, - logits, - labels, - smoothing, - logit_scale, - lse_square_scale, - ignore_index, - total_classes, - class_start_idx, - n_cols, # shapes - logits.stride(0), # strides - BLOCK_SIZE=BLOCK_SIZE, # constants - SPLIT=world_size > 1, - PRECOMPUTED_LSE=use_precomputed_lse, - num_warps=num_warps, - ) - - if world_size > 1: - # If there's no smoothing, if labels are in the vocab of this partition, losses contains - # - predicted logit, and 0 otherwise. - # If there's smoothing=0.1, for labels in the vocab of this partition, losses contains - # -0.9 * predicted logit - 0.1 * sum logit / total_classes. - # For labels not in the vocab of this partition, losses contains - # -0.1 * sum logit / total_classes. - if world_size > 1: - lse_allgather = torch.empty(world_size, n_rows, dtype=lse.dtype, device=lse.device) - torch.distributed.all_gather_into_tensor(lse_allgather, lse, group=process_group) - handle_losses = torch.distributed.all_reduce( - losses, op=torch.distributed.ReduceOp.SUM, group=process_group, async_op=True - ) - lse = torch.logsumexp(lse_allgather, dim=0) - handle_losses.wait() - # After the allreduce, if there's no smoothing, the total losses are - predicted_logit, - # we just have to add the (global) lse. - # If there's smoothing=0.1, the total losses are - # -0.9 * predicted_logit - 0.1 * sum logit / total_classes. - # Again, we just have to add the (global) lse. - losses += lse - if lse_square_scale != 0.0: - z_losses = lse_square_scale * lse.square() - z_losses.masked_fill_(labels == ignore_index, 0.0) - losses += z_losses - else: - z_losses = torch.zeros_like(losses) - losses.masked_fill_(labels == ignore_index, 0.0) - - ctx.save_for_backward(logits, lse, labels) - ctx.mark_non_differentiable(z_losses) - ctx.smoothing = smoothing - ctx.logit_scale = logit_scale - ctx.lse_square_scale = lse_square_scale - ctx.ignore_index = ignore_index - ctx.total_classes = total_classes - ctx.class_start_idx = class_start_idx - ctx.inplace_backward = inplace_backward - return losses, z_losses - - @staticmethod - def backward(ctx, grad_losses, grad_z_losses): - del grad_z_losses # z_losses are only for logging. - - logits, lse, labels = ctx.saved_tensors - dlogits = logits if ctx.inplace_backward else torch.empty_like(logits) - n_rows, n_cols = logits.shape - BLOCK_SIZE = min(triton.next_power_of_2(n_cols), 4 * 1024) - num_warps = 4 if BLOCK_SIZE < 2048 else (8 if BLOCK_SIZE < 8192 else 16) - grid = lambda META: (n_rows, triton.cdiv(n_cols, META["BLOCK_SIZE"])) # noqa - # Need this, otherwise Triton tries to launch from cuda:0 and we get - # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?) - with torch.cuda.device(logits.device.index): - cross_entropy_bwd_kernel[grid]( - dlogits, # data ptrs - grad_losses, - logits, - lse, - labels, - ctx.smoothing, - ctx.logit_scale, - ctx.lse_square_scale, - ctx.ignore_index, - ctx.total_classes, - ctx.class_start_idx, - n_cols, # shapes - logits.stride(0), # strides - dlogits.stride(0), - grad_losses.stride(0), - BLOCK_SIZE=BLOCK_SIZE, # constants - num_warps=num_warps, - ) - return dlogits, None, None, None, None, None, None, None, None, None - - -def cross_entropy_loss( - logits: torch.Tensor, - labels: torch.Tensor, - precomputed_lse: Optional[torch.Tensor] = None, - label_smoothing: float = 0.0, - logit_scale: float = 1.0, - lse_square_scale: float = 0.0, - ignore_index=-100, - inplace_backward: bool = False, - process_group=None, -) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Arguments: - logits: (batch, vocab_size) - labels: (batch,) - label_smoothing: float - logit_scale: float. Multiply logits by this scale before calculating the loss. - lse_square_scale: float. If > 0, we add lse_square_scale * lse(logits) ^ 2 to the loss. - This is also referred to as "z-loss". - ignore_index: int. If labels == ignore_index, the loss is set to 0.0. - inplace_backward: bool. If True, we do the backward pass in-place by modifying the logits. - This saves memory. - process_group: if not None, we're doing Tensor Parallel: each process is responsible for - one part of the vocab. The loss will be aggregated across processes. - Returns: - losses: (batch,), float - z_losses: (batch,), float - """ - return CrossEntropyLoss.apply( - logits, - labels, - precomputed_lse, - label_smoothing, - logit_scale, - lse_square_scale, - ignore_index, - inplace_backward, - process_group, - ) diff --git a/build/torch211-cxx11-xpu20253-x86_64-linux/ops/triton/k_activations.py b/build/torch211-cxx11-xpu20253-x86_64-linux/ops/triton/k_activations.py deleted file mode 100644 index efb83c358eb4a85d069ee340a3c83f418f9a805b..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-xpu20253-x86_64-linux/ops/triton/k_activations.py +++ /dev/null @@ -1,162 +0,0 @@ -# Adapted from https://github.com/facebookresearch/xformers/blob/main/xformers/triton/k_activations.py -# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved. -# -# This source code is licensed under the BSD license found in the -# LICENSE file in the root directory of this source tree. - -import math -from enum import Enum -from typing import Optional - -import triton -import triton.language as tl - -_sqrt2pi = math.sqrt(2.0 / math.pi) -_sqrt1_2 = math.sqrt(1.0 / 2) -_gaussian_pdf_normalization = 1.0 / math.sqrt(2 * math.pi) - - -class Activation(str, Enum): - SquaredReLU = "squared_relu" - GeLU = "gelu" - GeLUApprox = "gelu_approx" - LeakyReLU = "leaky_relu" - ReLU = "relu" - - -def get_triton_activation_kernel(activation: Optional[Activation]): - return ( - { - Activation.ReLU: relu, - Activation.LeakyReLU: leaky_relu, - Activation.GeLU: gelu, - Activation.GeLUApprox: gelu_approx, - Activation.SquaredReLU: squared_relu, - }[activation] - if activation - else None - ) - - -def get_triton_activation_bwd_kernel(activation: Optional[Activation]): - return ( - { - Activation.ReLU: relu_grad, - Activation.LeakyReLU: leaky_relu_grad, - Activation.GeLU: gelu_grad, - Activation.GeLUApprox: gelu_approx_grad, - Activation.SquaredReLU: squared_relu_grad, - }[activation] - if activation - else None - ) - - -@triton.jit -def tanh(x): - # Tanh is just a scaled sigmoid - return 2 * tl.sigmoid(2 * x) - 1 - - -@triton.jit -def cosh(x): - exp_x = tl.exp(x) - return (exp_x + 1.0 / exp_x) * 0.5 - - -# a Triton implementation of the most used activations -# See for instance http://arxiv.org/abs/1606.08415 for an overview - -# ReLU -@triton.jit -def relu(x): - """ - ReLU_ activation function - - .. _ReLU: https://pytorch.org/docs/stable/generated/torch.nn.ReLU.html - """ - zero = 0.0 - return tl.where(x >= 0, x, zero.to(x.dtype)) - - -@triton.jit -def relu_grad(x): - # ReLU is different from other activations - # in that it does not require the input to retrospectively compute its gradient - # here the input is the downstream gradient, and we return the upstream gradient directly - zero = 0.0 - one = 1.0 - return tl.where(x >= 0, one.to(x.dtype), zero.to(x.dtype)) - - -@triton.jit -def squared_relu(x): - """ - Squared ReLU activation, as proposed in the Primer_ paper. - - .. _Primer: https://arxiv.org/abs/2109.08668 - """ - x_ = relu(x) - return (x_ * x_).to(x.dtype) - - -@triton.jit -def squared_relu_grad(x): - return tl.where(x >= 0, 2.0 * x, 0.0) - - -# Leaky ReLU -@triton.jit -def leaky_relu(x): - """ - LeakyReLU_ activation - - .. _LeakyReLU: https://pytorch.org/docs/stable/generated/torch.nn.LeakyReLU.html - """ - scale = 0.01 + 0.0 - scale = scale.to(x.dtype) - return tl.where(x >= 0, x, scale * x) - - -@triton.jit -def leaky_relu_grad(x): - min_grad = 0.01 - max_grad = 1 - - min_grad = min_grad.to(x.dtype) - max_grad = max_grad.to(x.dtype) - - return tl.where(x >= 0, max_grad, min_grad) - - -@triton.jit -def gelu(x): - """Gaussian Error Linear Unit (GELU)""" - return x * 0.5 * (1.0 + tl.libdevice.erf(x * _sqrt1_2)) - - -@triton.jit -def gelu_grad(x): - cdf = 0.5 * (1.0 + tl.libdevice.erf(x * _sqrt1_2)) - pdf = tl.exp(-0.5 * x * x) * _gaussian_pdf_normalization - return cdf + x * pdf - - -@triton.jit -def gelu_approx(x): - """ - GeLU_ activation - Gaussian error linear unit, with tanh approximation - - .. _GeLU: https://arxiv.org/pdf/1606.08415.pdf - """ - return 0.5 * x * (1.0 + tanh(_sqrt2pi * x * (1.0 + 0.044715 * x * x))) - - -@triton.jit -def gelu_approx_grad(x): - # CREDITS: Fast implementation proposed in - # https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/fused_bias_gelu.py#L30 - tanh_out = tanh(0.79788456 * x * (1 + 0.044715 * x * x)) - return 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( - 1 + tanh_out - ) diff --git a/build/torch211-cxx11-xpu20253-x86_64-linux/ops/triton/layer_norm.py b/build/torch211-cxx11-xpu20253-x86_64-linux/ops/triton/layer_norm.py deleted file mode 100644 index 5e6f0b3697b612e6a19031bbc55ac76ad233dc85..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-xpu20253-x86_64-linux/ops/triton/layer_norm.py +++ /dev/null @@ -1,1254 +0,0 @@ -# Copyright (c) 2024, Tri Dao. -# Implement dropout + residual + layer_norm / rms_norm. - -# Based on the Triton LayerNorm tutorial: https://triton-lang.org/main/getting-started/tutorials/05-layer-norm.html -# For the backward pass, we keep weight_grad and bias_grad in registers and accumulate. -# This is faster for dimensions up to 8k, but after that it's much slower due to register spilling. -# The models we train have hidden dim up to 8k anyway (e.g. Llama 70B), so this is fine. - -import math -from typing import Optional, List - -import torch -import torch.nn.functional as F -from torch import Tensor - -import triton -import triton.language as tl - -from flash_attn.utils.torch import custom_fwd, custom_bwd -from flash_attn.utils.library import triton_op - -from ..._ops import add_op_namespace_prefix - - -def maybe_contiguous_lastdim(x): - return x.contiguous() if x is not None and x.stride(-1) != 1 else x - - -def maybe_contiguous(x): - return x.contiguous() if x is not None else None - - -def triton_autotune_configs(): - # Return configs with a valid warp count for the current device - configs = [] - # Maximum threads per block is architecture-dependent in theory, but in reality all are 1024 - max_threads_per_block = 1024 - # Default to warp size 32 if not defined by device - warp_size = getattr(torch.cuda.get_device_properties(torch.cuda.current_device()), "warp_size", 32) - # Autotune for warp counts which are powers of 2 and do not exceed thread per block limit - return [triton.Config({}, num_warps=warp_count) for warp_count in [1, 2, 4, 8, 16, 32] - if warp_count * warp_size <= max_threads_per_block] - # return [triton.Config({}, num_warps=8)] - - -def layer_norm_ref( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - zero_centered_weight=False, - dropout_mask=None, - dropout_mask1=None, - upcast=False, -): - dtype = x.dtype - if upcast: - x = x.float() - weight = weight.float() - bias = bias.float() if bias is not None else None - residual = residual.float() if residual is not None else residual - x1 = x1.float() if x1 is not None else None - weight1 = weight1.float() if weight1 is not None else None - bias1 = bias1.float() if bias1 is not None else None - if zero_centered_weight: - weight = weight + 1.0 - if weight1 is not None: - weight1 = weight1 + 1.0 - if x1 is not None: - assert rowscale is None, "rowscale is not supported with parallel LayerNorm" - if rowscale is not None: - x = x * rowscale[..., None] - if dropout_p > 0.0: - if dropout_mask is not None: - x = x.masked_fill(~dropout_mask, 0.0) / (1.0 - dropout_p) - else: - x = F.dropout(x, p=dropout_p) - if x1 is not None: - if dropout_mask1 is not None: - x1 = x1.masked_fill(~dropout_mask1, 0.0) / (1.0 - dropout_p) - else: - x1 = F.dropout(x1, p=dropout_p) - if x1 is not None: - x = x + x1 - if residual is not None: - x = (x + residual).to(x.dtype) - out = F.layer_norm(x.to(weight.dtype), x.shape[-1:], weight=weight, bias=bias, eps=eps).to( - dtype - ) - if weight1 is None: - return out if not prenorm else (out, x) - else: - out1 = F.layer_norm( - x.to(weight1.dtype), x.shape[-1:], weight=weight1, bias=bias1, eps=eps - ).to(dtype) - return (out, out1) if not prenorm else (out, out1, x) - - -def rms_norm_ref( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - zero_centered_weight=False, - dropout_mask=None, - dropout_mask1=None, - upcast=False, -): - dtype = x.dtype - if upcast: - x = x.float() - weight = weight.float() - bias = bias.float() if bias is not None else None - residual = residual.float() if residual is not None else residual - x1 = x1.float() if x1 is not None else None - weight1 = weight1.float() if weight1 is not None else None - bias1 = bias1.float() if bias1 is not None else None - if zero_centered_weight: - weight = weight + 1.0 - if weight1 is not None: - weight1 = weight1 + 1.0 - if x1 is not None: - assert rowscale is None, "rowscale is not supported with parallel LayerNorm" - if rowscale is not None: - x = x * rowscale[..., None] - if dropout_p > 0.0: - if dropout_mask is not None: - x = x.masked_fill(~dropout_mask, 0.0) / (1.0 - dropout_p) - else: - x = F.dropout(x, p=dropout_p) - if x1 is not None: - if dropout_mask1 is not None: - x1 = x1.masked_fill(~dropout_mask1, 0.0) / (1.0 - dropout_p) - else: - x1 = F.dropout(x1, p=dropout_p) - if x1 is not None: - x = x + x1 - if residual is not None: - x = (x + residual).to(x.dtype) - rstd = 1 / torch.sqrt((x.square()).mean(dim=-1, keepdim=True) + eps) - out = ((x * rstd * weight) + bias if bias is not None else (x * rstd * weight)).to(dtype) - if weight1 is None: - return out if not prenorm else (out, x) - else: - out1 = ((x * rstd * weight1) + bias1 if bias1 is not None else (x * rstd * weight1)).to( - dtype - ) - return (out, out1) if not prenorm else (out, out1, x) - - -@triton.autotune( - configs=triton_autotune_configs(), - key=["N", "HAS_RESIDUAL", "STORE_RESIDUAL_OUT", "IS_RMS_NORM", "HAS_BIAS", "HAS_X1", "HAS_W1", "HAS_B1"], -) -# torch compile doesn't like triton.heuristics, so we set these manually when calling the kernel -# @triton.heuristics({"HAS_BIAS": lambda args: args["B"] is not None}) -# @triton.heuristics({"HAS_RESIDUAL": lambda args: args["RESIDUAL"] is not None}) -# @triton.heuristics({"HAS_X1": lambda args: args["X1"] is not None}) -# @triton.heuristics({"HAS_W1": lambda args: args["W1"] is not None}) -# @triton.heuristics({"HAS_B1": lambda args: args["B1"] is not None}) -@triton.jit -def _layer_norm_fwd_1pass_kernel( - X, # pointer to the input - Y, # pointer to the output - W, # pointer to the weights - B, # pointer to the biases - RESIDUAL, # pointer to the residual - X1, - W1, - B1, - Y1, - RESIDUAL_OUT, # pointer to the residual - ROWSCALE, - SEEDS, # Dropout seeds for each row - DROPOUT_MASK, - DROPOUT_MASK1, - Mean, # pointer to the mean - Rstd, # pointer to the 1/std - stride_x_row, # how much to increase the pointer when moving by 1 row - stride_y_row, - stride_res_row, - stride_res_out_row, - stride_x1_row, - stride_y1_row, - M, # number of rows in X - N, # number of columns in X - eps, # epsilon to avoid division by zero - dropout_p, # Dropout probability - zero_centered_weight, # If true, add 1.0 to the weight - IS_RMS_NORM: tl.constexpr, - BLOCK_N: tl.constexpr, - HAS_RESIDUAL: tl.constexpr, - STORE_RESIDUAL_OUT: tl.constexpr, - HAS_BIAS: tl.constexpr, - HAS_DROPOUT: tl.constexpr, - STORE_DROPOUT_MASK: tl.constexpr, - HAS_ROWSCALE: tl.constexpr, - HAS_X1: tl.constexpr, - HAS_W1: tl.constexpr, - HAS_B1: tl.constexpr, -): - # Map the program id to the row of X and Y it should compute. - row = tl.program_id(0) - X += row * stride_x_row - Y += row * stride_y_row - if HAS_RESIDUAL: - RESIDUAL += row * stride_res_row - if STORE_RESIDUAL_OUT: - RESIDUAL_OUT += row * stride_res_out_row - if HAS_X1: - X1 += row * stride_x1_row - if HAS_W1: - Y1 += row * stride_y1_row - # Compute mean and variance - cols = tl.arange(0, BLOCK_N) - x = tl.load(X + cols, mask=cols < N, other=0.0).to(tl.float32) - if HAS_ROWSCALE: - rowscale = tl.load(ROWSCALE + row).to(tl.float32) - x *= rowscale - if HAS_DROPOUT: - # Compute dropout mask - # 7 rounds is good enough, and reduces register pressure - keep_mask = tl.rand(tl.load(SEEDS + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - x = tl.where(keep_mask, x / (1.0 - dropout_p), 0.0) - if STORE_DROPOUT_MASK: - tl.store(DROPOUT_MASK + row * N + cols, keep_mask, mask=cols < N) - if HAS_X1: - x1 = tl.load(X1 + cols, mask=cols < N, other=0.0).to(tl.float32) - if HAS_ROWSCALE: - rowscale = tl.load(ROWSCALE + M + row).to(tl.float32) - x1 *= rowscale - if HAS_DROPOUT: - # Compute dropout mask - # 7 rounds is good enough, and reduces register pressure - keep_mask = ( - tl.rand(tl.load(SEEDS + M + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - ) - x1 = tl.where(keep_mask, x1 / (1.0 - dropout_p), 0.0) - if STORE_DROPOUT_MASK: - tl.store(DROPOUT_MASK1 + row * N + cols, keep_mask, mask=cols < N) - x += x1 - if HAS_RESIDUAL: - residual = tl.load(RESIDUAL + cols, mask=cols < N, other=0.0).to(tl.float32) - x += residual - if STORE_RESIDUAL_OUT: - tl.store(RESIDUAL_OUT + cols, x, mask=cols < N) - if not IS_RMS_NORM: - mean = tl.sum(x, axis=0) / N - tl.store(Mean + row, mean) - xbar = tl.where(cols < N, x - mean, 0.0) - var = tl.sum(xbar * xbar, axis=0) / N - else: - xbar = tl.where(cols < N, x, 0.0) - var = tl.sum(xbar * xbar, axis=0) / N - rstd = 1 / tl.sqrt(var + eps) - tl.store(Rstd + row, rstd) - # Normalize and apply linear transformation - mask = cols < N - w = tl.load(W + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w += 1.0 - if HAS_BIAS: - b = tl.load(B + cols, mask=mask).to(tl.float32) - x_hat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd - y = x_hat * w + b if HAS_BIAS else x_hat * w - # Write output - tl.store(Y + cols, y, mask=mask) - if HAS_W1: - w1 = tl.load(W1 + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w1 += 1.0 - if HAS_B1: - b1 = tl.load(B1 + cols, mask=mask).to(tl.float32) - y1 = x_hat * w1 + b1 if HAS_B1 else x_hat * w1 - tl.store(Y1 + cols, y1, mask=mask) - - -def _layer_norm_fwd( - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - residual: Optional[Tensor] = None, - x1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - out_dtype: Optional[torch.dtype] = None, - residual_dtype: Optional[torch.dtype] = None, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - return_dropout_mask: bool = False, - out: Optional[Tensor] = None, - residual_out: Optional[Tensor] = None -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - # Need to wrap to handle the case where residual_out is a alias of x, which makes torch.library - # and torch.compile unhappy. Also allocate memory for out and residual_out if they are None - # so that _layer_norm_fwd_impl doesn't have to return them. - if out is None: - out = torch.empty_like(x, dtype=x.dtype if out_dtype is None else out_dtype) - if residual is not None: - residual_dtype = residual.dtype - if residual_out is None and ( - residual is not None - or (residual_dtype is not None and residual_dtype != x.dtype) - or dropout_p > 0.0 - or rowscale is not None - or x1 is not None - ): - residual_out = torch.empty_like( - x, dtype=residual_dtype if residual_dtype is not None else x.dtype - ) - else: - residual_out = None - y1, mean, rstd, seeds, dropout_mask, dropout_mask1 = _layer_norm_fwd_impl( - x, - weight, - bias, - eps, - out, - residual=residual, - x1=x1, - weight1=weight1, - bias1=bias1, - dropout_p=dropout_p, - rowscale=rowscale, - zero_centered_weight=zero_centered_weight, - is_rms_norm=is_rms_norm, - return_dropout_mask=return_dropout_mask, - residual_out=residual_out, - ) - # residual_out is None if residual is None and residual_dtype == input_dtype and dropout_p == 0.0 - if residual_out is None: - residual_out = x - return out, y1, mean, rstd, residual_out, seeds, dropout_mask, dropout_mask1 - - -# [2025-04-28] torch.library.triton_op ignores the schema argument, but here we need the schema -# since we're returning a tuple of tensors -@triton_op(add_op_namespace_prefix("layer_norm_fwd_impl"), mutates_args={"out", "residual_out"}, - schema="(Tensor x, Tensor weight, Tensor bias, float eps, Tensor(a!) out, Tensor? residual, Tensor? x1, Tensor? weight1, Tensor? bias1, float dropout_p, Tensor? rowscale, bool zero_centered_weight, bool is_rms_norm, bool return_dropout_mask, Tensor(a!)? residual_out) -> (Tensor y1, Tensor mean, Tensor rstd, Tensor seeds, Tensor dropout_mask, Tensor dropout_mask1)") -def _layer_norm_fwd_impl( - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - out: Tensor, - residual: Optional[Tensor] = None, - x1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - return_dropout_mask: bool = False, - residual_out: Optional[Tensor] = None -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - M, N = x.shape - assert x.stride(-1) == 1 - if residual is not None: - assert residual.stride(-1) == 1 - assert residual.shape == (M, N) - assert weight.shape == (N,) - assert weight.stride(-1) == 1 - if bias is not None: - assert bias.stride(-1) == 1 - assert bias.shape == (N,) - if x1 is not None: - assert x1.shape == x.shape - assert rowscale is None - assert x1.stride(-1) == 1 - if weight1 is not None: - assert weight1.shape == (N,) - assert weight1.stride(-1) == 1 - if bias1 is not None: - assert bias1.shape == (N,) - assert bias1.stride(-1) == 1 - if rowscale is not None: - assert rowscale.is_contiguous() - assert rowscale.shape == (M,) - assert out.shape == x.shape - assert out.stride(-1) == 1 - if residual_out is not None: - assert residual_out.shape == x.shape - assert residual_out.stride(-1) == 1 - if weight1 is not None: - y1 = torch.empty_like(out) - assert y1.stride(-1) == 1 - else: - y1 = None - mean = torch.empty((M,), dtype=torch.float32, device=x.device) if not is_rms_norm else None - rstd = torch.empty((M,), dtype=torch.float32, device=x.device) - if dropout_p > 0.0: - seeds = torch.randint( - 2**32, (M if x1 is None else 2 * M,), device=x.device, dtype=torch.int64 - ) - else: - seeds = None - if return_dropout_mask and dropout_p > 0.0: - dropout_mask = torch.empty(M, N, device=x.device, dtype=torch.bool) - if x1 is not None: - dropout_mask1 = torch.empty(M, N, device=x.device, dtype=torch.bool) - else: - dropout_mask1 = None - else: - dropout_mask, dropout_mask1 = None, None - # Less than 64KB per feature: enqueue fused kernel - MAX_FUSED_SIZE = 65536 // x.element_size() - BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) - if N > BLOCK_N: - raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") - with torch.cuda.device(x.device.index): - torch.library.wrap_triton(_layer_norm_fwd_1pass_kernel)[(M,)]( - x, - out, - weight, - bias, - residual, - x1, - weight1, - bias1, - y1, - residual_out, - rowscale, - seeds, - dropout_mask, - dropout_mask1, - mean, - rstd, - x.stride(0), - out.stride(0), - residual.stride(0) if residual is not None else 0, - residual_out.stride(0) if residual_out is not None else 0, - x1.stride(0) if x1 is not None else 0, - y1.stride(0) if y1 is not None else 0, - M, - N, - eps, - dropout_p, - # Passing bool make torch inductor very unhappy since it then tries to compare to int_max - int(zero_centered_weight), - is_rms_norm, - BLOCK_N, - residual is not None, - residual_out is not None, - bias is not None, - dropout_p > 0.0, - dropout_mask is not None, - rowscale is not None, - HAS_X1=x1 is not None, - HAS_W1=weight1 is not None, - HAS_B1=bias1 is not None, - ) - return y1, mean, rstd, seeds, dropout_mask, dropout_mask1 - - -@triton.autotune( - configs=triton_autotune_configs(), - key=["N", "HAS_DRESIDUAL", "STORE_DRESIDUAL", "IS_RMS_NORM", "HAS_BIAS", "HAS_DROPOUT"], -) -# torch compile doesn't like triton.heuristics, so we set these manually when calling the kernel -# @triton.heuristics({"HAS_BIAS": lambda args: args["B"] is not None}) -# @triton.heuristics({"HAS_DRESIDUAL": lambda args: args["DRESIDUAL"] is not None}) -# @triton.heuristics({"STORE_DRESIDUAL": lambda args: args["DRESIDUAL_IN"] is not None}) -# @triton.heuristics({"HAS_ROWSCALE": lambda args: args["ROWSCALE"] is not None}) -# @triton.heuristics({"HAS_DY1": lambda args: args["DY1"] is not None}) -# @triton.heuristics({"HAS_DX1": lambda args: args["DX1"] is not None}) -# @triton.heuristics({"HAS_B1": lambda args: args["DB1"] is not None}) -# @triton.heuristics({"RECOMPUTE_OUTPUT": lambda args: args["Y"] is not None}) -@triton.jit -def _layer_norm_bwd_kernel( - X, # pointer to the input - W, # pointer to the weights - B, # pointer to the biases - Y, # pointer to the output to be recomputed - DY, # pointer to the output gradient - DX, # pointer to the input gradient - DW, # pointer to the partial sum of weights gradient - DB, # pointer to the partial sum of biases gradient - DRESIDUAL, - W1, - DY1, - DX1, - DW1, - DB1, - DRESIDUAL_IN, - ROWSCALE, - SEEDS, - Mean, # pointer to the mean - Rstd, # pointer to the 1/std - stride_x_row, # how much to increase the pointer when moving by 1 row - stride_y_row, - stride_dy_row, - stride_dx_row, - stride_dres_row, - stride_dy1_row, - stride_dx1_row, - stride_dres_in_row, - M, # number of rows in X - N, # number of columns in X - eps, # epsilon to avoid division by zero - dropout_p, - zero_centered_weight, - rows_per_program, - IS_RMS_NORM: tl.constexpr, - BLOCK_N: tl.constexpr, - HAS_DRESIDUAL: tl.constexpr, - STORE_DRESIDUAL: tl.constexpr, - HAS_BIAS: tl.constexpr, - HAS_DROPOUT: tl.constexpr, - HAS_ROWSCALE: tl.constexpr, - HAS_DY1: tl.constexpr, - HAS_DX1: tl.constexpr, - HAS_B1: tl.constexpr, - RECOMPUTE_OUTPUT: tl.constexpr, -): - # Map the program id to the elements of X, DX, and DY it should compute. - row_block_id = tl.program_id(0) - row_start = row_block_id * rows_per_program - # Do not early exit if row_start >= M, because we need to write DW and DB - cols = tl.arange(0, BLOCK_N) - mask = cols < N - X += row_start * stride_x_row - if HAS_DRESIDUAL: - DRESIDUAL += row_start * stride_dres_row - if STORE_DRESIDUAL: - DRESIDUAL_IN += row_start * stride_dres_in_row - DY += row_start * stride_dy_row - DX += row_start * stride_dx_row - if HAS_DY1: - DY1 += row_start * stride_dy1_row - if HAS_DX1: - DX1 += row_start * stride_dx1_row - if RECOMPUTE_OUTPUT: - Y += row_start * stride_y_row - w = tl.load(W + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w += 1.0 - if RECOMPUTE_OUTPUT and HAS_BIAS: - b = tl.load(B + cols, mask=mask, other=0.0).to(tl.float32) - if HAS_DY1: - w1 = tl.load(W1 + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w1 += 1.0 - dw = tl.zeros((BLOCK_N,), dtype=tl.float32) - if HAS_BIAS: - db = tl.zeros((BLOCK_N,), dtype=tl.float32) - if HAS_DY1: - dw1 = tl.zeros((BLOCK_N,), dtype=tl.float32) - if HAS_B1: - db1 = tl.zeros((BLOCK_N,), dtype=tl.float32) - row_end = min((row_block_id + 1) * rows_per_program, M) - for row in range(row_start, row_end): - # Load data to SRAM - x = tl.load(X + cols, mask=mask, other=0).to(tl.float32) - dy = tl.load(DY + cols, mask=mask, other=0).to(tl.float32) - if HAS_DY1: - dy1 = tl.load(DY1 + cols, mask=mask, other=0).to(tl.float32) - if not IS_RMS_NORM: - mean = tl.load(Mean + row) - rstd = tl.load(Rstd + row) - # Compute dx - xhat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd - xhat = tl.where(mask, xhat, 0.0) - if RECOMPUTE_OUTPUT: - y = xhat * w + b if HAS_BIAS else xhat * w - tl.store(Y + cols, y, mask=mask) - wdy = w * dy - dw += dy * xhat - if HAS_BIAS: - db += dy - if HAS_DY1: - wdy += w1 * dy1 - dw1 += dy1 * xhat - if HAS_B1: - db1 += dy1 - if not IS_RMS_NORM: - c1 = tl.sum(xhat * wdy, axis=0) / N - c2 = tl.sum(wdy, axis=0) / N - dx = (wdy - (xhat * c1 + c2)) * rstd - else: - c1 = tl.sum(xhat * wdy, axis=0) / N - dx = (wdy - xhat * c1) * rstd - if HAS_DRESIDUAL: - dres = tl.load(DRESIDUAL + cols, mask=mask, other=0).to(tl.float32) - dx += dres - # Write dx - if STORE_DRESIDUAL: - tl.store(DRESIDUAL_IN + cols, dx, mask=mask) - if HAS_DX1: - if HAS_DROPOUT: - keep_mask = ( - tl.rand(tl.load(SEEDS + M + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - ) - dx1 = tl.where(keep_mask, dx / (1.0 - dropout_p), 0.0) - else: - dx1 = dx - tl.store(DX1 + cols, dx1, mask=mask) - if HAS_DROPOUT: - keep_mask = tl.rand(tl.load(SEEDS + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - dx = tl.where(keep_mask, dx / (1.0 - dropout_p), 0.0) - if HAS_ROWSCALE: - rowscale = tl.load(ROWSCALE + row).to(tl.float32) - dx *= rowscale - tl.store(DX + cols, dx, mask=mask) - - X += stride_x_row - if HAS_DRESIDUAL: - DRESIDUAL += stride_dres_row - if STORE_DRESIDUAL: - DRESIDUAL_IN += stride_dres_in_row - if RECOMPUTE_OUTPUT: - Y += stride_y_row - DY += stride_dy_row - DX += stride_dx_row - if HAS_DY1: - DY1 += stride_dy1_row - if HAS_DX1: - DX1 += stride_dx1_row - tl.store(DW + row_block_id * N + cols, dw, mask=mask) - if HAS_BIAS: - tl.store(DB + row_block_id * N + cols, db, mask=mask) - if HAS_DY1: - tl.store(DW1 + row_block_id * N + cols, dw1, mask=mask) - if HAS_B1: - tl.store(DB1 + row_block_id * N + cols, db1, mask=mask) - - -def _layer_norm_bwd( - dy: Tensor, - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - mean: Tensor, - rstd: Tensor, - dresidual: Optional[Tensor] = None, - dy1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - seeds: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - has_residual: bool = False, - has_x1: bool = False, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - x_dtype: Optional[torch.dtype] = None, - recompute_output: bool = False, -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - # Need to wrap to handle the case where dresidual_in or dx1 are aliases of x, - # which makes torch.library unhappy - dx, dw, db, dresidual_in, dx1, dw1, db1, y = _layer_norm_bwd_impl( - dy, - x, - weight, - bias, - eps, - mean, - rstd, - dresidual, - dy1, - weight1, - bias1, - seeds, - dropout_p, - rowscale, - has_residual, - has_x1, - zero_centered_weight, - is_rms_norm, - x_dtype=x_dtype, - recompute_output=recompute_output, - ) - # Don't need to compute dresidual_in separately in this case - if has_residual and dx.dtype == x.dtype and dropout_p == 0.0 and rowscale is None: - dresidual_in = dx - if has_x1 and dropout_p == 0.0: - dx1 = dx - return dx, dw, db, dresidual_in, dx1, dw1, db1, y - - - -@triton_op(add_op_namespace_prefix("layer_norm_bwd_impl"), mutates_args={}, - schema="(Tensor dy, Tensor x, Tensor weight, Tensor bias, float eps, Tensor mean, Tensor rstd, Tensor? dresidual, Tensor? dy1, Tensor? weight1, Tensor? bias1, Tensor? seeds, float dropout_p, Tensor? rowscale, bool has_residual, bool has_x1, bool zero_centered_weight, bool is_rms_norm, ScalarType? x_dtype, bool recompute_output) -> (Tensor dx, Tensor dw, Tensor db, Tensor dresidual_in, Tensor dx1, Tensor dw1, Tensor db1, Tensor y)", - allow_decomposition=False, # Don't let torch.compile trace inside - ) -def _layer_norm_bwd_impl( - dy: Tensor, - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - mean: Tensor, - rstd: Tensor, - dresidual: Optional[Tensor] = None, - dy1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - seeds: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - has_residual: bool = False, - has_x1: bool = False, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - x_dtype: Optional[torch.dtype] = None, - recompute_output: bool = False, -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - M, N = x.shape - assert x.stride(-1) == 1 - dy = maybe_contiguous_lastdim(dy) - assert dy.stride(-1) == 1 - assert dy.shape == (M, N) - if dresidual is not None: - dresidual = maybe_contiguous_lastdim(dresidual) - assert dresidual.stride(-1) == 1 - assert dresidual.shape == (M, N) - assert weight.shape == (N,) - assert weight.stride(-1) == 1 - if bias is not None: - assert bias.stride(-1) == 1 - assert bias.shape == (N,) - if dy1 is not None: - dy1 = maybe_contiguous_lastdim(dy1) - assert weight1 is not None - assert dy1.shape == dy.shape - assert dy1.stride(-1) == 1 - if weight1 is not None: - assert weight1.shape == (N,) - assert weight1.stride(-1) == 1 - if bias1 is not None: - assert bias1.shape == (N,) - assert bias1.stride(-1) == 1 - if seeds is not None: - assert seeds.is_contiguous() - assert seeds.shape == (M if not has_x1 else M * 2,) - if rowscale is not None: - assert rowscale.is_contiguous() - assert rowscale.shape == (M,) - # allocate output - dx = ( - torch.empty_like(x) - if x_dtype is None - else torch.empty(M, N, dtype=x_dtype, device=x.device) - ) - dresidual_in = ( - torch.empty_like(x) - if has_residual - and (dx.dtype != x.dtype or dropout_p > 0.0 or rowscale is not None or has_x1) - else None - ) - dx1 = torch.empty_like(dx) if (has_x1 and dropout_p > 0.0) else None - y = torch.empty(M, N, dtype=dy.dtype, device=dy.device) if recompute_output else None - if recompute_output: - assert weight1 is None, "recompute_output is not supported with parallel LayerNorm" - - # Less than 64KB per feature: enqueue fused kernel - MAX_FUSED_SIZE = 65536 // x.element_size() - BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) - if N > BLOCK_N: - raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") - # Increasing the multiple (e.g. 8) will allow more thread blocks to be launched and hide the - # latency of the gmem reads/writes, but will increase the time of summing up dw / db. - sm_count = torch.cuda.get_device_properties(x.device).multi_processor_count * 8 - _dw = torch.empty((sm_count, N), dtype=torch.float32, device=weight.device) - _db = ( - torch.empty((sm_count, N), dtype=torch.float32, device=bias.device) - if bias is not None - else None - ) - _dw1 = torch.empty_like(_dw) if weight1 is not None else None - _db1 = torch.empty_like(_db) if bias1 is not None else None - rows_per_program = math.ceil(M / sm_count) - grid = (sm_count,) - with torch.cuda.device(x.device.index): - torch.library.wrap_triton(_layer_norm_bwd_kernel)[grid]( - x, - weight, - bias, - y, - dy, - dx, - _dw, - _db, - dresidual, - weight1, - dy1, - dx1, - _dw1, - _db1, - dresidual_in, - rowscale, - seeds, - mean, - rstd, - x.stride(0), - 0 if not recompute_output else y.stride(0), - dy.stride(0), - dx.stride(0), - dresidual.stride(0) if dresidual is not None else 0, - dy1.stride(0) if dy1 is not None else 0, - dx1.stride(0) if dx1 is not None else 0, - dresidual_in.stride(0) if dresidual_in is not None else 0, - M, - N, - eps, - dropout_p, - # Passing bool make torch inductor very unhappy since it then tries to compare to int_max - int(zero_centered_weight), - rows_per_program, - is_rms_norm, - BLOCK_N, - dresidual is not None, - dresidual_in is not None, - bias is not None, - dropout_p > 0.0, - HAS_ROWSCALE=rowscale is not None, - HAS_DY1=dy1 is not None, - HAS_DX1=dx1 is not None, - HAS_B1=bias1 is not None, - RECOMPUTE_OUTPUT=y is not None, - ) - dw = _dw.sum(0).to(weight.dtype) - db = _db.sum(0).to(bias.dtype) if bias is not None else None - dw1 = _dw1.sum(0).to(weight1.dtype) if weight1 is not None else None - db1 = _db1.sum(0).to(bias1.dtype) if bias1 is not None else None - # dresidual_in and dx1 could be None, the wrapper will handle assigning them from dx - return dx, dw, db, dresidual_in, dx1, dw1, db1, y - - -class LayerNormFn(torch.autograd.Function): - - @staticmethod - def forward( - ctx, - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - residual_in_fp32=False, - zero_centered_weight=False, - is_rms_norm=False, - return_dropout_mask=False, - out_dtype=None, - out=None, - residual_out=None - ): - x_shape_og = x.shape - # reshape input data into 2D tensor - x = maybe_contiguous_lastdim(x.reshape(-1, x.shape[-1])) - if residual is not None: - assert residual.shape == x_shape_og - residual = maybe_contiguous_lastdim(residual.reshape(-1, residual.shape[-1])) - if x1 is not None: - assert x1.shape == x_shape_og - assert rowscale is None, "rowscale is not supported with parallel LayerNorm" - x1 = maybe_contiguous_lastdim(x1.reshape(-1, x1.shape[-1])) - weight = weight.contiguous() - bias = maybe_contiguous(bias) - weight1 = maybe_contiguous(weight1) - bias1 = maybe_contiguous(bias1) - if rowscale is not None: - rowscale = rowscale.reshape(-1).contiguous() - residual_dtype = ( - residual.dtype - if residual is not None - else (torch.float32 if residual_in_fp32 else None) - ) - if out is not None: - out = out.reshape(-1, out.shape[-1]) - if residual_out is not None: - residual_out = residual_out.reshape(-1, residual_out.shape[-1]) - y, y1, mean, rstd, residual_out, seeds, dropout_mask, dropout_mask1 = _layer_norm_fwd( - x, - weight, - bias, - eps, - residual, - x1, - weight1, - bias1, - dropout_p=dropout_p, - rowscale=rowscale, - out_dtype=out_dtype, - residual_dtype=residual_dtype, - zero_centered_weight=zero_centered_weight, - is_rms_norm=is_rms_norm, - return_dropout_mask=return_dropout_mask, - out=out, - residual_out=residual_out, - ) - ctx.save_for_backward( - residual_out, weight, bias, weight1, bias1, rowscale, seeds, mean, rstd - ) - ctx.x_shape_og = x_shape_og - ctx.eps = eps - ctx.dropout_p = dropout_p - ctx.is_rms_norm = is_rms_norm - ctx.has_residual = residual is not None - ctx.has_x1 = x1 is not None - ctx.prenorm = prenorm - ctx.x_dtype = x.dtype - ctx.zero_centered_weight = zero_centered_weight - y = y.reshape(x_shape_og) - y1 = y1.reshape(x_shape_og) if y1 is not None else None - residual_out = residual_out.reshape(x_shape_og) if residual_out is not None else None - dropout_mask = dropout_mask.reshape(x_shape_og) if dropout_mask is not None else None - dropout_mask1 = dropout_mask1.reshape(x_shape_og) if dropout_mask1 is not None else None - if not return_dropout_mask: - if weight1 is None: - return y if not prenorm else (y, residual_out) - else: - return (y, y1) if not prenorm else (y, y1, residual_out) - else: - if weight1 is None: - return ( - (y, dropout_mask, dropout_mask1) - if not prenorm - else (y, residual_out, dropout_mask, dropout_mask1) - ) - else: - return ( - (y, y1, dropout_mask, dropout_mask1) - if not prenorm - else (y, y1, residual_out, dropout_mask, dropout_mask1) - ) - - @staticmethod - def backward(ctx, dy, *args): - x, weight, bias, weight1, bias1, rowscale, seeds, mean, rstd = ctx.saved_tensors - dy = dy.reshape(-1, dy.shape[-1]) - if weight1 is not None: - dy1, args = args[0], args[1:] - dy1 = dy1.reshape(-1, dy1.shape[-1]) - assert dy1.shape == x.shape - else: - dy1 = None - if ctx.prenorm: - dresidual = args[0] - dresidual = dresidual.reshape(-1, dresidual.shape[-1]) - assert dresidual.shape == x.shape - else: - dresidual = None - dx, dw, db, dresidual_in, dx1, dw1, db1, _ = _layer_norm_bwd( - dy, - x, - weight, - bias, - ctx.eps, - mean, - rstd, - dresidual, - dy1, - weight1, - bias1, - seeds, - ctx.dropout_p, - rowscale, - ctx.has_residual, - ctx.has_x1, - ctx.zero_centered_weight, - ctx.is_rms_norm, - x_dtype=ctx.x_dtype, - recompute_output=False, - ) - return ( - dx.reshape(ctx.x_shape_og), - dw, - db, - dresidual_in.reshape(ctx.x_shape_og) if ctx.has_residual else None, - dx1.reshape(ctx.x_shape_og) if dx1 is not None else None, - dw1, - db1, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) - - -def layer_norm_fn( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - residual_in_fp32=False, - zero_centered_weight=False, - is_rms_norm=False, - return_dropout_mask=False, - out_dtype=None, - out=None, - residual_out=None -): - return LayerNormFn.apply( - x, - weight, - bias, - residual, - x1, - weight1, - bias1, - eps, - dropout_p, - rowscale, - prenorm, - residual_in_fp32, - zero_centered_weight, - is_rms_norm, - return_dropout_mask, - out_dtype, - out, - residual_out - ) - - -def rms_norm_fn( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - residual_in_fp32=False, - zero_centered_weight=False, - return_dropout_mask=False, - out_dtype=None, - out=None, - residual_out=None -): - return LayerNormFn.apply( - x, - weight, - bias, - residual, - x1, - weight1, - bias1, - eps, - dropout_p, - rowscale, - prenorm, - residual_in_fp32, - zero_centered_weight, - True, - return_dropout_mask, - out_dtype, - out, - residual_out - ) - - -class RMSNorm(torch.nn.Module): - - def __init__(self, hidden_size, eps=1e-5, dropout_p=0.0, zero_centered_weight=False, - device=None, dtype=None): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.eps = eps - if dropout_p > 0.0: - self.drop = torch.nn.Dropout(dropout_p) - else: - self.drop = None - self.zero_centered_weight = zero_centered_weight - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.register_parameter("bias", None) - self.reset_parameters() - - def reset_parameters(self): - if not self.zero_centered_weight: - torch.nn.init.ones_(self.weight) - else: - torch.nn.init.zeros_(self.weight) - - def forward(self, x, residual=None, prenorm=False, residual_in_fp32=False): - return rms_norm_fn( - x, - self.weight, - self.bias, - residual=residual, - eps=self.eps, - dropout_p=self.drop.p if self.drop is not None and self.training else 0.0, - prenorm=prenorm, - residual_in_fp32=residual_in_fp32, - zero_centered_weight=self.zero_centered_weight, - ) - - -class LayerNormLinearFn(torch.autograd.Function): - - @staticmethod - @custom_fwd - def forward( - ctx, - x, - norm_weight, - norm_bias, - linear_weight, - linear_bias, - residual=None, - eps=1e-6, - prenorm=False, - residual_in_fp32=False, - is_rms_norm=False, - ): - x_shape_og = x.shape - # reshape input data into 2D tensor - x = maybe_contiguous_lastdim(x.reshape(-1, x.shape[-1])) - if residual is not None: - assert residual.shape == x_shape_og - residual = maybe_contiguous_lastdim(residual.reshape(-1, residual.shape[-1])) - norm_weight = norm_weight.contiguous() - norm_bias = maybe_contiguous(norm_bias) - residual_dtype = ( - residual.dtype - if residual is not None - else (torch.float32 if residual_in_fp32 else None) - ) - y, _, mean, rstd, residual_out, *rest = _layer_norm_fwd( - x, - norm_weight, - norm_bias, - eps, - residual, - out_dtype=None if not torch.is_autocast_enabled() else torch.get_autocast_dtype("cuda"), - residual_dtype=residual_dtype, - is_rms_norm=is_rms_norm, - ) - y = y.reshape(x_shape_og) - dtype = torch.get_autocast_dtype("cuda") if torch.is_autocast_enabled() else y.dtype - linear_weight = linear_weight.to(dtype) - linear_bias = linear_bias.to(dtype) if linear_bias is not None else None - out = F.linear(y.to(linear_weight.dtype), linear_weight, linear_bias) - # We don't store y, will be recomputed in the backward pass to save memory - ctx.save_for_backward(residual_out, norm_weight, norm_bias, linear_weight, mean, rstd) - ctx.x_shape_og = x_shape_og - ctx.eps = eps - ctx.is_rms_norm = is_rms_norm - ctx.has_residual = residual is not None - ctx.prenorm = prenorm - ctx.x_dtype = x.dtype - ctx.linear_bias_is_none = linear_bias is None - return out if not prenorm else (out, residual_out.reshape(x_shape_og)) - - @staticmethod - @custom_bwd - def backward(ctx, dout, *args): - x, norm_weight, norm_bias, linear_weight, mean, rstd = ctx.saved_tensors - dout = dout.reshape(-1, dout.shape[-1]) - dy = F.linear(dout, linear_weight.t()) - dlinear_bias = None if ctx.linear_bias_is_none else dout.sum(0) - dy = maybe_contiguous_lastdim(dy) - assert dy.shape == x.shape - if ctx.prenorm: - dresidual = args[0] - dresidual = maybe_contiguous_lastdim(dresidual.reshape(-1, dresidual.shape[-1])) - assert dresidual.shape == x.shape - else: - dresidual = None - dx, dnorm_weight, dnorm_bias, dresidual_in, _, _, _, y = _layer_norm_bwd( - dy, - x, - norm_weight, - norm_bias, - ctx.eps, - mean, - rstd, - dresidual=dresidual, - has_residual=ctx.has_residual, - is_rms_norm=ctx.is_rms_norm, - x_dtype=ctx.x_dtype, - recompute_output=True, - ) - dlinear_weight = torch.einsum("bo,bi->oi", dout, y) - return ( - dx.reshape(ctx.x_shape_og), - dnorm_weight, - dnorm_bias, - dlinear_weight, - dlinear_bias, - dresidual_in.reshape(ctx.x_shape_og) if ctx.has_residual else None, - None, - None, - None, - None, - ) - - -def layer_norm_linear_fn( - x, - norm_weight, - norm_bias, - linear_weight, - linear_bias, - residual=None, - eps=1e-6, - prenorm=False, - residual_in_fp32=False, - is_rms_norm=False, -): - return LayerNormLinearFn.apply( - x, - norm_weight, - norm_bias, - linear_weight, - linear_bias, - residual, - eps, - prenorm, - residual_in_fp32, - is_rms_norm, - ) diff --git a/build/torch211-cxx11-xpu20253-x86_64-linux/ops/triton/linear.py b/build/torch211-cxx11-xpu20253-x86_64-linux/ops/triton/linear.py deleted file mode 100644 index a8966dbc345ab0e593df0124451ee7be3dae131a..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-xpu20253-x86_64-linux/ops/triton/linear.py +++ /dev/null @@ -1,594 +0,0 @@ -# Adapted from https://github.com/ELS-RD/kernl/blob/main/src/kernl/implementations/linear_layer.py -# and https://github.com/openai/triton/blob/master/python/triton/ops/matmul.py -from typing import Optional - -import torch -import triton -import triton.language as tl -from triton.ops.matmul_perf_model import early_config_prune, estimate_matmul_time - -from flash_attn.ops.triton.k_activations import ( - gelu, - gelu_approx, - gelu_approx_grad, - gelu_grad, - squared_relu, - squared_relu_grad, -) - -# CREDITS: Initially inspired by the Triton tutorial on matrix multiplications - - -def init_to_zero(name): - return lambda nargs: nargs[name].zero_() - - -def get_configs_io_bound(): - configs = [] - for num_stages in [2, 3, 4, 5, 6]: - for block_m in [16, 32]: - for block_k in [32, 64]: - for block_n in [32, 64, 128, 256]: - num_warps = 2 if block_n <= 64 else 4 - configs.append( - triton.Config( - { - "BLOCK_M": block_m, - "BLOCK_N": block_n, - "BLOCK_K": block_k, - "SPLIT_K": 1, - }, - num_stages=num_stages, - num_warps=num_warps, - ) - ) - # split_k not used - # for split_k in [2, 4, 8, 16]: - # configs.append(triton.Config( - # {'BLOCK_M': block_m, 'BLOCK_N': block_n, 'BLOCK_K': block_k, 'SPLIT_K': split_k}, - # num_stages=num_stages, num_warps=num_warps, pre_hook=init_to_zero('C'))) - return configs - - -@triton.autotune( - configs=[ - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - # good for int8 - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=4, - num_warps=4, - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - ] - + get_configs_io_bound(), - key=["CACHE_KEY_M", "CACHE_KEY_N", "CACHE_KEY_K"], - prune_configs_by={ - "early_config_prune": early_config_prune, - "perf_model": estimate_matmul_time, - "top_k": 10, - }, -) -@triton.heuristics( - { - "EVEN_K": lambda args: args["K"] % (args["BLOCK_K"] * args["SPLIT_K"]) == 0, - } -) -@triton.jit -def kernel_fwd( - C, # Pointers to matrices - ACT_INPUT, - A, - B, - bias, - # Matrix dimensions - M, - N, - K, - CACHE_KEY_M, - CACHE_KEY_N, - CACHE_KEY_K, - # The stride variables represent how much to increase the ptr by when moving by 1 - # element in a particular dimension. E.g. stride_am is how much to increase a_ptr - # by to get the element one row down (A has M rows) - stride_cm, - # stride_cn, # Assume that stride_cn == 1 - stride_am, - stride_ak, - stride_bn, - stride_bk, - # Meta-parameters - BLOCK_M: tl.constexpr, - GROUP_M: tl.constexpr, - BLOCK_N: tl.constexpr, - BLOCK_K: tl.constexpr, - # split k not used, not performant with activation, kept because early_config_prune is expecting it - SPLIT_K: tl.constexpr, - EVEN_K: tl.constexpr, - A_ROWMAJOR: tl.constexpr, - B_COLMAJOR: tl.constexpr, - BIAS: tl.constexpr, - SAVE_ACT_INPUT: tl.constexpr, - ACTIVATION: tl.constexpr, -): - - """ - Kernel for computing Out = activation(A x W + C) - - Input has shape (M, K) - - Weight has shape (K, N) - - Bias has shape (N,) - - Output has shape (M, N) - - ActInputs (optional) has shape (M, N) - 'ActInputs' optionally saves the A x W + C intermediate for backward computations - This kernel will consolidate over K - """ - - pid = tl.program_id(axis=0) - - grid_m = (M + BLOCK_M - 1) // BLOCK_M - grid_n = (N + BLOCK_N - 1) // BLOCK_N - # re-order program ID for better L2 performance - width = GROUP_M * grid_n - group_id = pid // width - group_size = min(grid_m - group_id * GROUP_M, GROUP_M) - pid_m = group_id * GROUP_M + (pid % group_size) - pid_n = (pid % width) // (group_size) - - # now compute the block that each program will go through - # rm (resp. rn) denotes a range of indices - # for rows (resp. col) of C - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - # trick to avoid masking on M and N axis - ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M) - rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N) - rk = tl.arange(0, BLOCK_K) - - if A_ROWMAJOR: - A = A + (ram[:, None] * stride_am + rk[None, :]) - else: - A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak) - if B_COLMAJOR: - B = B + (rk[:, None] + rbn[None, :] * stride_bn) - else: - B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn) - - acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) - - for k in range(K, 0, -BLOCK_K): - if EVEN_K: - a = tl.load(A) - b = tl.load(B) - else: - a = tl.load(A, mask=rk[None, :] < k, other=0.0) - b = tl.load(B, mask=rk[:, None] < k, other=0.0) - acc += tl.dot(a, b) - - if A_ROWMAJOR: - A += BLOCK_K - else: - A += BLOCK_K * stride_ak - if B_COLMAJOR: - B += BLOCK_K - else: - B += BLOCK_K * stride_bk - - # Putting bias after the matmul (instead of before) is faster, idk why - if BIAS: - bias = tl.load(bias + rn, mask=rn < N, other=0.0).to(tl.float32) - acc += bias[None, :] - - # optional: save the activation inputs - if SAVE_ACT_INPUT: - # act_in_ptrs = ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :] * stride_cn - act_in_ptrs = ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :] - tl.store(act_in_ptrs, acc) - - # optional: fused activation (while the data is in shared memory) - if ACTIVATION == "gelu": - acc = gelu(acc) - elif ACTIVATION == "gelu_approx": - acc = gelu_approx(acc) - elif ACTIVATION == "squared_relu": - acc = squared_relu(acc) - # rematerialize rm and rn to save registers - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - - # write back result - # C = C + rm[:, None] * stride_cm + rn[None, :] * stride_cn - C = C + rm[:, None] * stride_cm + rn[None, :] - mask = (rm < M)[:, None] & (rn < N)[None, :] - tl.store(C, acc) - - -def triton_linear_act( - x: torch.Tensor, - weight: torch.Tensor, - bias: Optional[torch.Tensor] = None, - activation: str = "id", - save_act_input: bool = False, -) -> torch.Tensor: - """ - Compute e = activation(x @ weight.T + bias). - This wrapper kicks the `kernel_fwd` Triton kernel - :param x: input tensor - :param weight: weight matrix - :param bias: an optional bias tensor - :param activation: Activation name. Needs to be a Triton kernel. - :param act_input: an optional tensor to save the activation inputs (for backward) - :return: result tensor - """ - # if torch.is_autocast_enabled(): - # dtype = torch.get_autocast_gpu_dtype() - # x, weight, bias = [a.to(dtype=dtype) for a in [x, weight, bias]] - - assert activation in ["id", "gelu", "gelu_approx", "squared_relu"] - - batch_shape, n = x.shape[:-1], x.shape[-1] - batch_dim = batch_shape.numel() - x_reshaped = x.reshape(batch_dim, n) - - if x_reshaped.stride(0) > 1 and x_reshaped.stride(1) > 1: - x_reshaped = x_reshaped.contiguous() - if weight.stride(0) > 1 and weight.stride(1) > 1: - weight = weight.contiguous() - bias = bias.contiguous() if bias is not None else None - - assert ( - x.dtype == weight.dtype - ), f"Input and weight must have the same dtype, got {x.dtype} and {weight.dtype}" - if bias is not None: - assert ( - x.dtype == bias.dtype - ), f"Input and bias must have the same dtype, got {x.dtype} and {bias.dtype}" - assert ( - x_reshaped.shape[1] == weight.shape[1] - ), f"Incompatible dimensions: {x_reshaped.shape} - {weight.shape}" - - assert ( - bias is None or bias.shape[0] == weight.shape[0] - ), "Incompatible dimensions in between weight and bias" - - M, K = x_reshaped.shape - N, K = weight.shape - - output = torch.empty((M, N), device=x.device, dtype=x.dtype) - act_input = torch.empty_like(output) if save_act_input else None - - # 1D launch kernel where each block gets its own program. - grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),) # noqa - - kernel_fwd[grid]( - output, - act_input, - x_reshaped, - weight, # data ptrs - bias if bias is not None else x, # auto skip bias if not present - M, # shapes - N, - K, - M // 32, # key for triton cache (limit number of compilations) - N // 32, - K // 32, - stride_cm=output.stride(0), # strides - # stride_cn=output.stride(1), - stride_am=x_reshaped.stride(0), - stride_ak=x_reshaped.stride(1), - stride_bk=weight.stride(1), - stride_bn=weight.stride(0), - BIAS=bias is not None, # optional fused bias - SAVE_ACT_INPUT=save_act_input, # optional save activation inputs - ACTIVATION=activation, # optional fused activation - A_ROWMAJOR=x_reshaped.stride(1) == 1, - B_COLMAJOR=weight.stride(1) == 1, - GROUP_M=8, # speed optimization: group the programs - ) - - if not save_act_input: - return output.reshape(*batch_shape, output.shape[-1]) - else: - return ( - output.reshape(*batch_shape, output.shape[-1]), - act_input.reshape(*batch_shape, act_input.shape[-1]), - ) - - -@triton.autotune( - configs=[ - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - # good for int8 - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=4, - num_warps=4, - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - ] - + get_configs_io_bound(), - key=["CACHE_KEY_M", "CACHE_KEY_N", "CACHE_KEY_K"], - prune_configs_by={ - "early_config_prune": early_config_prune, - "perf_model": estimate_matmul_time, - "top_k": 10, - }, -) -@triton.heuristics( - { - "EVEN_K": lambda args: args["K"] % (args["BLOCK_K"] * args["SPLIT_K"]) == 0, - } -) -@triton.jit -def kernel_bwd( - C, # Pointers to matrices - ACT_INPUT, - A, - B, - # Matrix dimensions - M, - N, - K, - CACHE_KEY_M, - CACHE_KEY_N, - CACHE_KEY_K, - # The stride variables represent how much to increase the ptr by when moving by 1 - # element in a particular dimension. E.g. stride_am is how much to increase a_ptr - # by to get the element one row down (A has M rows) - stride_cm, - # stride_cn, # Assume that stride_cn == 1 - stride_am, - stride_ak, - stride_bk, - stride_bn, - # Meta-parameters - BLOCK_M: tl.constexpr, - GROUP_M: tl.constexpr, - BLOCK_N: tl.constexpr, - BLOCK_K: tl.constexpr, - # split k not used, not performant with activation, kept because early_config_prune is expecting it - SPLIT_K: tl.constexpr, - EVEN_K: tl.constexpr, - ACTIVATION: tl.constexpr, -): - - """ - Kernel for computing Out = activation(A x W + C) - - Input has shape (M, K) - - Weight has shape (K, N) - - Output has shape (M, N) - - ActInputs (optional) has shape (M, N) - 'ActInputs' optionally saves the A x W + C intermediate for backward computations - This kernel will consolidate over K - """ - - pid = tl.program_id(axis=0) - - grid_m = (M + BLOCK_M - 1) // BLOCK_M - grid_n = (N + BLOCK_N - 1) // BLOCK_N - # re-order program ID for better L2 performance - width = GROUP_M * grid_n - group_id = pid // width - group_size = min(grid_m - group_id * GROUP_M, GROUP_M) - pid_m = group_id * GROUP_M + (pid % group_size) - pid_n = (pid % width) // (group_size) - - # now compute the block that each program will go through - # rm (resp. rn) denotes a range of indices - # for rows (resp. col) of C - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - # trick to avoid masking on M and N axis - ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M) - rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N) - rk = tl.arange(0, BLOCK_K) - - A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak) - B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn) - - acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) - - for k in range(K, 0, -BLOCK_K): - if EVEN_K: - a = tl.load(A) - b = tl.load(B) - else: - a = tl.load(A, mask=rk[None, :] < k, other=0.0) - b = tl.load(B, mask=rk[:, None] < k, other=0.0) - acc += tl.dot(a, b) - - A += BLOCK_K * stride_ak - B += BLOCK_K * stride_bk - - # optional: fused activation (while the data is in shared memory) - if ACTIVATION != "id": - act_in_ptrs = ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :] - act_input = tl.load(act_in_ptrs).to(acc.dtype) - if ACTIVATION == "gelu": - acc *= gelu_grad(act_input) - elif ACTIVATION == "gelu_approx": - acc *= gelu_approx_grad(act_input) - elif ACTIVATION == "squared_relu": - acc *= squared_relu_grad(act_input) - - # rematerialize rm and rn to save registers - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - - # write back result - C = C + rm[:, None] * stride_cm + rn[None, :] - mask = (rm < M)[:, None] & (rn < N)[None, :] - tl.store(C, acc, mask=mask) - - -def triton_dgrad_act( - grad_output: torch.Tensor, - weight: torch.Tensor, - activation: str = "id", - act_input: Optional[torch.Tensor] = None, -) -> torch.Tensor: - """ - Compute e = activation(grad_output @ weight + bias). - This wrapper kicks the `kernel_fwd` Triton kernel - :param grad_output: input tensor - :param weight: weight matrix - :param activation: Activation name. Needs to be a Triton kernel. - :param act_input: an optional tensor to save the activation inputs (for backward) - :return: result tensor - """ - assert activation in ["id", "gelu", "gelu_approx", "squared_relu"] - - batch_shape, n = grad_output.shape[:-1], grad_output.shape[-1] - batch_dim = batch_shape.numel() - grad_output_reshaped = grad_output.reshape(batch_dim, n) - - if grad_output_reshaped.stride(0) > 1 and grad_output_reshaped.stride(1) > 1: - grad_output_reshaped = grad_output_reshaped.contiguous() - if weight.stride(0) > 1 and weight.stride(1) > 1: - weight = weight.contiguous() - - assert ( - grad_output.dtype == weight.dtype - ), f"grad_output and weight must have the same dtype, got {grad_output.dtype} and {weight.dtype}" - assert ( - grad_output_reshaped.shape[1] == weight.shape[0] - ), f"Incompatible dimensions: {grad_output_reshaped.shape} - {weight.shape}" - if activation != "id": - assert act_input is not None, f"act_input is required for activation {activation}" - - # M, N, K in bwd are different from M, N, K in fwd - M, K = grad_output_reshaped.shape - K, N = weight.shape - - grad_input = torch.empty((M, N), device=grad_output.device, dtype=grad_output.dtype) - - # 1D launch kernel where each block gets its own program. - grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),) # noqa - - kernel_bwd[grid]( - grad_input, - act_input, - grad_output_reshaped, - weight, # data ptrs - M, # shapes - N, - K, - M // 32, # key for triton cache (limit number of compilations) - N // 32, - K // 32, - stride_cm=grad_input.stride(0), # strides - # stride_cn=grad_input.stride(1), - stride_am=grad_output_reshaped.stride(0), - stride_ak=grad_output_reshaped.stride(1), - stride_bk=weight.stride(0), - stride_bn=weight.stride(1), - ACTIVATION=activation, # optional fused activation - GROUP_M=8, # speed optimization: group the programs - ) - - return grad_input.reshape(*batch_shape, grad_input.shape[-1]) diff --git a/build/torch211-cxx11-xpu20253-x86_64-linux/ops/triton/mlp.py b/build/torch211-cxx11-xpu20253-x86_64-linux/ops/triton/mlp.py deleted file mode 100644 index 059f4f8a5e174c1f4824e43d313fca18eaa799b8..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-xpu20253-x86_64-linux/ops/triton/mlp.py +++ /dev/null @@ -1,149 +0,0 @@ -# The triton fused matmul + sqrelu is faster for fp16 but slower for bf16, compared -# to naive implementation. -import fused_dense_lib as fused_dense_cuda -import torch -import torch.nn as nn -import torch.nn.functional as F - -from flash_attn.utils.torch import custom_fwd, custom_bwd -from flash_attn.ops.activations import sqrelu_bwd, sqrelu_fwd -from flash_attn.ops.triton.linear import triton_dgrad_act, triton_linear_act - - -class FusedDenseSqreluDenseFunc(torch.autograd.Function): - @staticmethod - @custom_fwd - def forward(ctx, x, weight1, bias1, weight2, bias2, checkpoint_lvl=0): - """checkpoint_lvl: - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute act_input and gelu_out in the bwd - """ - if torch.is_autocast_enabled(): - dtype = torch.get_autocast_gpu_dtype() - x, weight1, bias1, weight2, bias2 = [ - a.to(dtype=dtype) for a in [x, weight1, bias1, weight2, bias2] - ] - is_bf16 = x.dtype == torch.bfloat16 - assert checkpoint_lvl in [0, 1, 2] - x = x.contiguous() - weight1 = weight1.contiguous() - bias1 = bias1.contiguous() - weight2 = weight2.contiguous() - bias2 = bias2.contiguous() - batch_shape, n = x.shape[:-1], x.shape[-1] - batch_dim = batch_shape.numel() - if is_bf16: - act_input = fused_dense_cuda.linear_bias_forward( - x.reshape(batch_dim, n), weight1, bias1 - ) - output1 = sqrelu_fwd(act_input) - else: - save_act_input = checkpoint_lvl != 2 - result = triton_linear_act( - x.reshape(batch_dim, n), - weight1, - bias1, - activation="squared_relu", - save_act_input=save_act_input, - ) - if save_act_input: - output1, act_input = result - else: - output1 = result - output2 = fused_dense_cuda.linear_bias_forward(output1, weight2, bias2) - ctx.checkpoint_lvl = checkpoint_lvl - if checkpoint_lvl == 0: - ctx.save_for_backward(x, weight1, bias1, weight2, act_input, output1) - elif checkpoint_lvl == 1: - ctx.save_for_backward(x, weight1, bias1, weight2, act_input) - elif checkpoint_lvl == 2: - ctx.save_for_backward(x, weight1, bias1, weight2) - return output2.reshape(*batch_shape, output2.shape[-1]) - - @staticmethod - @custom_bwd - def backward(ctx, grad_output): - grad_output = grad_output.contiguous() - checkpoint_lvl = ctx.checkpoint_lvl - x, weight1, bias1, weight2, *rest = ctx.saved_tensors - batch_shape, n = x.shape[:-1], x.shape[-1] - batch_dim = batch_shape.numel() - is_bf16 = x.dtype == torch.bfloat16 - if checkpoint_lvl == 0: - act_input, output1 = rest - elif checkpoint_lvl == 1: - (act_input,) = rest - output1 = sqrelu_fwd(act_input) - elif checkpoint_lvl == 2: - if is_bf16: - act_input = fused_dense_cuda.linear_bias_forward( - x.reshape(batch_dim, n), weight1, bias1 - ) - output1 = sqrelu_fwd(act_input) - else: - output1, act_input = triton_linear_act( - x.reshape(batch_dim, n), - weight1, - bias1, - activation="squared_relu", - save_act_input=True, - ) - - if is_bf16: - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad(output1, grad_output) - grad_output1 = grad_output @ weight2 - grad_act_input = sqrelu_bwd(grad_output1, act_input) - grad_input, grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_backward( - x.reshape(batch_dim, n), weight1, grad_act_input - ) - else: - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad(output1, grad_output) - grad_act_input = triton_dgrad_act( - grad_output, weight2, activation="squared_relu", act_input=act_input - ) - grad_input, grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_backward( - x.reshape(batch_dim, n), weight1, grad_act_input - ) - return grad_input.reshape_as(x), grad_weight1, grad_bias1, grad_weight2, grad_bias2, None - - -fused_dense_sqrelu_dense_function = FusedDenseSqreluDenseFunc.apply - - -class FusedDenseSqreluDense(nn.Module): - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - bias1=True, - bias2=True, - checkpoint_lvl=0, - device=None, - dtype=None, - ): - """ - checkpoint_lvl (increasing lvl means slower but more memory saving): - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute gelu_in and gelu_out in the bwd - """ - assert checkpoint_lvl in [0, 1, 2] - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features * 4 - assert bias1 == True, "DenseSqreluDense module without bias is currently not supported" - assert bias2 == True, "DenseSqreluDense module without bias is currently not supported" - self.checkpoint_lvl = checkpoint_lvl - self.fc1 = nn.Linear(in_features, hidden_features, bias=bias1, **factory_kwargs) - self.fc2 = nn.Linear(hidden_features, out_features, bias=bias2, **factory_kwargs) - - def forward(self, x): - assert x.is_cuda - return fused_dense_sqrelu_dense_function( - x, self.fc1.weight, self.fc1.bias, self.fc2.weight, self.fc2.bias, self.checkpoint_lvl - ) diff --git a/build/torch211-cxx11-xpu20253-x86_64-linux/ops/triton/rotary.py b/build/torch211-cxx11-xpu20253-x86_64-linux/ops/triton/rotary.py deleted file mode 100644 index 6ad3ab412a084270120167ad32b67ad1720f2c37..0000000000000000000000000000000000000000 --- a/build/torch211-cxx11-xpu20253-x86_64-linux/ops/triton/rotary.py +++ /dev/null @@ -1,186 +0,0 @@ -# Copyright (c) 2025, Tri Dao. -# As of 2025-04-23, we require triton >= 3.0 - -from typing import Optional, Union - -import torch - -import triton -import triton.language as tl - - -@triton.jit -def rotary_kernel( - OUT, # Pointers to matrices - X, - COS, - SIN, - CU_SEQLENS, - SEQLEN_OFFSETS, # this could be int or a pointer - # Matrix dimensions - seqlen, - nheads, - seqlen_ro, - # strides - stride_out_batch, - stride_out_seqlen, - stride_out_nheads, - stride_out_headdim, - stride_x_batch, - stride_x_seqlen, - stride_x_nheads, - stride_x_headdim, - # Meta-parameters - # We want ROTARY_DIM to be constexpr, otherwise the triton compiler doesn't know that - # the mask is constant every 8 elements, and it will generate LDG.16 instead of LDG.128 - ROTARY_DIM: tl.constexpr, - IS_SEQLEN_OFFSETS_TENSOR: tl.constexpr, - IS_VARLEN: tl.constexpr, - INTERLEAVED: tl.constexpr, - CONJUGATE: tl.constexpr, - BLOCK_H: tl.constexpr, - BLOCK_M: tl.constexpr, -): - BLOCK_K: tl.constexpr = triton.next_power_of_2(ROTARY_DIM) - ROTARY_DIM_HALF = ROTARY_DIM // 2 - pid_head = tl.program_id(axis=0) - pid_m = tl.program_id(axis=1) - pid_batch = tl.program_id(axis=2) - - if not IS_VARLEN: - X = X + pid_batch * stride_x_batch - OUT = OUT + pid_batch * stride_out_batch - else: - start_idx = tl.load(CU_SEQLENS + pid_batch) - seqlen = tl.load(CU_SEQLENS + pid_batch + 1) - start_idx - X = X + start_idx * stride_x_seqlen - OUT = OUT + start_idx * stride_out_seqlen - - if pid_m * BLOCK_M >= seqlen: - return - - rh = pid_head * BLOCK_H + tl.arange(0, BLOCK_H) - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - if not IS_SEQLEN_OFFSETS_TENSOR: - rm_cs = rm + SEQLEN_OFFSETS - else: - rm_cs = rm + tl.load(SEQLEN_OFFSETS + pid_batch) - - rk_half = tl.arange(0, BLOCK_K // 2) - COS = COS + (rm_cs[:, None] * ROTARY_DIM_HALF + rk_half[None, :]) - SIN = SIN + (rm_cs[:, None] * ROTARY_DIM_HALF + rk_half[None, :]) - mask_cs = (rm_cs[:, None] < seqlen_ro) & (rk_half[None, :] < ROTARY_DIM_HALF) - cos = tl.load(COS, mask=mask_cs, other=1.0).to(tl.float32) - sin = tl.load(SIN, mask=mask_cs, other=0.0).to(tl.float32) - if CONJUGATE: - sin = -sin - - if not INTERLEAVED: - # Load the 1st and 2nd halves of X, do calculation, then store to 1st and 2nd halves of OUT - X = X + (rh[:, None, None] * stride_x_nheads + rm[None, :, None] * stride_x_seqlen + rk_half[None, None, :] * stride_x_headdim) - OUT = OUT + (rh[:, None, None] * stride_out_nheads + rm[None, :, None] * stride_out_seqlen + rk_half[None, None, :] * stride_out_headdim) - mask = (rh[:, None, None] < nheads) & (rm[None, :, None] < seqlen) & (rk_half[None, None, :] < ROTARY_DIM_HALF) - x0 = tl.load(X, mask=mask, other=0.0).to(tl.float32) - x1 = tl.load(X + ROTARY_DIM_HALF * stride_x_headdim, mask=mask, other=0.0,).to(tl.float32) - o0 = x0 * cos - x1 * sin - o1 = x0 * sin + x1 * cos - tl.store(OUT, o0, mask=mask) - tl.store(OUT + ROTARY_DIM_HALF * stride_out_headdim, o1, mask=mask) - else: - rk = tl.arange(0, BLOCK_K) - X = X + (rh[:, None, None] * stride_x_nheads + rm[None, :, None] * stride_x_seqlen + rk[None, None, :] * stride_x_headdim) - OUT = OUT + (rh[:, None, None] * stride_out_nheads + rm[None, :, None] * stride_out_seqlen + rk[None, None, :] * stride_out_headdim) - mask = (rh[:, None, None] < nheads) & (rm[None, :, None] < seqlen) & (rk[None, None, :] < ROTARY_DIM) - x = tl.load(X, mask=mask, other=0.0).to(tl.float32) - x0, x1 = tl.split(tl.reshape(x, [BLOCK_H, BLOCK_M, BLOCK_K // 2, 2])) - o0 = x0 * cos - x1 * sin - o1 = x0 * sin + x1 * cos - o = tl.reshape(tl.join(o0, o1), [BLOCK_H, BLOCK_M, BLOCK_K]) - tl.store(OUT, o, mask=mask) - - -def apply_rotary( - x: torch.Tensor, - cos: torch.Tensor, - sin: torch.Tensor, - seqlen_offsets: Union[int, torch.Tensor] = 0, - cu_seqlens: Optional[torch.Tensor] = None, - max_seqlen: Optional[int] = None, - interleaved=False, - inplace=False, - conjugate=False, -) -> torch.Tensor: - """ - Arguments: - x: (batch, seqlen, nheads, headdim) if cu_seqlens is None - else (total_seqlen, nheads, headdim). - cos: (seqlen_ro, rotary_dim / 2) - sin: (seqlen_ro, rotary_dim / 2) - seqlen_offsets: integer or integer tensor of size (batch,) - cu_seqlens: (batch + 1,) or None - max_seqlen: int - Returns: - y: (batch, seqlen, nheads, headdim) - """ - is_varlen = cu_seqlens is not None - if not is_varlen: - batch, seqlen, nheads, headdim = x.shape - else: - assert max_seqlen is not None, "If cu_seqlens is passed in, then max_seqlen must be passed" - total_seqlen, nheads, headdim = x.shape - batch_p_1 = cu_seqlens.shape[0] - batch = batch_p_1 - 1 - seqlen = max_seqlen - seqlen_ro, rotary_dim = cos.shape - assert sin.shape == cos.shape - rotary_dim *= 2 - assert rotary_dim <= headdim, "rotary_dim must be <= headdim" - assert headdim <= 256, "Only support headdim <= 256" - assert seqlen_ro >= seqlen, "seqlen_ro must be >= seqlen" - - cos, sin = cos.contiguous(), sin.contiguous() - if isinstance(seqlen_offsets, torch.Tensor): - assert seqlen_offsets.shape == (batch,) - assert seqlen_offsets.dtype in [torch.int32, torch.int64] - seqlen_offsets = seqlen_offsets.contiguous() - else: - assert seqlen_offsets + seqlen <= seqlen_ro - - output = torch.empty_like(x) if not inplace else x - if rotary_dim < headdim and not inplace: - output[..., rotary_dim:].copy_(x[..., rotary_dim:]) - - grid = lambda META: (triton.cdiv(nheads, META["BLOCK_H"]), triton.cdiv(seqlen, META["BLOCK_M"]), batch) # noqa - BLOCK_M = 8 if rotary_dim <= 128 else 4 - - # Need this, otherwise Triton tries to launch from cuda:0 and we get - # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?) - device_ctx = torch.cuda.device(x.device.index) if x.device.type == 'cuda' else torch.xpu.device(x.device.index) - with device_ctx: - torch.library.wrap_triton(rotary_kernel)[grid]( - output, # data ptrs - x, - cos, - sin, - cu_seqlens, - seqlen_offsets, - seqlen, # shapes - nheads, - seqlen_ro, - output.stride(0) if not is_varlen else 0, # batch_strides if not varlen else 0 - output.stride(-3), # seqlen_stride or total_seqlen_stride - output.stride(-2), # nheads_stride - output.stride(-1), # headdim_stride - x.stride(0) if not is_varlen else 0, # batch_strides if not varlen else 0 - x.stride(-3), # seqlen stride or total_seqlen_stride - x.stride(-2), # nheads stride - x.stride(-1), # headdim stride - rotary_dim, - isinstance(seqlen_offsets, torch.Tensor), - is_varlen, - interleaved, - conjugate, - BLOCK_M=BLOCK_M, - BLOCK_H=2, - ) - return output diff --git a/build/torch212-cxx11-cpu-x86_64-linux/__init__.py b/build/torch212-cxx11-cpu-x86_64-linux/__init__.py deleted file mode 100644 index ecc2f9d896b6c93f90b0a1499856dc0612177422..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cpu-x86_64-linux/__init__.py +++ /dev/null @@ -1,393 +0,0 @@ -from typing import Optional, List -import torch -from ._ops import ops as flash_attn_ops -from .flash_attn_interface import ( - flash_attn_func, - flash_attn_kvpacked_func, - flash_attn_qkvpacked_func, - flash_attn_varlen_func, - flash_attn_varlen_kvpacked_func, - flash_attn_varlen_qkvpacked_func, - flash_attn_with_kvcache, -) - - -def fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - return_softmax: bool = False, - gen: Optional[torch.Generator] = None, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention. - - Args: - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Optional output tensor, same shape as q - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - return_softmax: Whether to return softmax weights - gen: Optional random number generator - - Returns: - List of tensors: [output, softmax_lse, (softmax if return_softmax)] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.fwd( - q, - k, - v, - out, - alibi_slopes, - p_dropout, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - gen, - ) - - -def varlen_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - out: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - block_table: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - max_seqlen_q: int = 0, - max_seqlen_k: int = 0, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - zero_tensors: bool = False, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - return_softmax: bool = False, - gen: Optional[torch.Generator] = None, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention with variable sequence lengths. - - Args: - q: Query tensor of shape [total_q, num_heads, head_size] - k: Key tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - v: Value tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1] - cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1] - out: Optional output tensor of shape [total_q, num_heads, head_size] - seqused_k: Optional tensor specifying how many keys to use per batch element [batch_size] - leftpad_k: Optional left padding for keys of shape [batch_size] - block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq] - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - max_seqlen_q: Maximum sequence length for queries - max_seqlen_k: Maximum sequence length for keys - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - zero_tensors: Whether to zero tensors before computation - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - return_softmax: Whether to return softmax weights - gen: Optional random number generator - - Returns: - List of tensors: [output, softmax_lse, (softmax if return_softmax)] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.varlen_fwd( - q, - k, - v, - out, - cu_seqlens_q, - cu_seqlens_k, - seqused_k, - leftpad_k, - block_table, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - p_dropout, - softmax_scale, - zero_tensors, - is_causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - gen, - ) - - -def bwd( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - p_dropout: float = 0.0, - 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, - gen: Optional[torch.Generator] = None, - rng_state: Optional[torch.Tensor] = None, -) -> List[torch.Tensor]: - """ - Backward pass for multi-head attention. - - Args: - dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size] - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size] - softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q] - dq: Optional gradient tensor for queries, same shape as q - dk: Optional gradient tensor for keys, same shape as k - dv: Optional gradient tensor for values, same shape as v - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - deterministic: Whether to use deterministic algorithms - gen: Optional random number generator - rng_state: Optional RNG state from forward pass - - Returns: - List of tensors: [dq, dk, dv] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - alibi_slopes, - p_dropout, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - deterministic, - gen, - rng_state, - ) - - -def varlen_bwd( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - max_seqlen_q: int = 0, - max_seqlen_k: int = 0, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - zero_tensors: bool = False, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - deterministic: bool = False, - gen: Optional[torch.Generator] = None, - rng_state: Optional[torch.Tensor] = None, -) -> List[torch.Tensor]: - """ - Backward pass for multi-head attention with variable sequence lengths. - - Args: - dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size] - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size] - softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q] - cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1] - cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1] - dq: Optional gradient tensor for queries, same shape as q - dk: Optional gradient tensor for keys, same shape as k - dv: Optional gradient tensor for values, same shape as v - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - max_seqlen_q: Maximum sequence length for queries - max_seqlen_k: Maximum sequence length for keys - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - zero_tensors: Whether to zero tensors before computation - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - deterministic: Whether to use deterministic algorithms - gen: Optional random number generator - rng_state: Optional RNG state from forward pass - - Returns: - List of tensors: [dq, dk, dv] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.varlen_bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - p_dropout, - softmax_scale, - zero_tensors, - is_causal, - window_size_left, - window_size_right, - softcap, - deterministic, - gen, - rng_state, - ) - - -def fwd_kvcache( - q: torch.Tensor, - kcache: torch.Tensor, - vcache: torch.Tensor, - k: Optional[torch.Tensor] = None, - v: Optional[torch.Tensor] = None, - seqlens_k: Optional[torch.Tensor] = None, - rotary_cos: Optional[torch.Tensor] = None, - rotary_sin: Optional[torch.Tensor] = None, - cache_batch_idx: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - block_table: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - out: 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, - is_rotary_interleaved: bool = False, - num_splits: int = 1, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention with KV cache. - - Args: - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - kcache: Key cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - vcache: Value cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - k: Optional new keys tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size] - v: Optional new values tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size] - seqlens_k: Optional sequence lengths for keys of shape [batch_size] - rotary_cos: Optional rotary cosine tensor of shape [seqlen_ro, rotary_dim/2] - rotary_sin: Optional rotary sine tensor of shape [seqlen_ro, rotary_dim/2] - cache_batch_idx: Optional indices to index into the KV cache - leftpad_k: Optional left padding for keys of shape [batch_size] - block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq] - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - out: Optional output tensor, same shape as q - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - is_rotary_interleaved: Whether rotary embeddings are interleaved - num_splits: Number of splits for computation - - Returns: - List of tensors: [output, softmax_lse] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.fwd_kvcache( - q, - kcache, - vcache, - k, - v, - seqlens_k, - rotary_cos, - rotary_sin, - cache_batch_idx, - leftpad_k, - block_table, - alibi_slopes, - out, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - is_rotary_interleaved, - num_splits, - ) diff --git a/build/torch212-cxx11-cpu-x86_64-linux/_flash_attn2_cpu_f12afc9.abi3.so b/build/torch212-cxx11-cpu-x86_64-linux/_flash_attn2_cpu_f12afc9.abi3.so deleted file mode 100644 index 3396cd3fc067e3c8d36040356cdcb4f1fb590287..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cpu-x86_64-linux/_flash_attn2_cpu_f12afc9.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:218245911f36b3a573b9a4f47d8054b7784c7ad6cd11fa322c2426b75f1e2489 -size 206008 diff --git a/build/torch212-cxx11-cpu-x86_64-linux/_ops.py b/build/torch212-cxx11-cpu-x86_64-linux/_ops.py deleted file mode 100644 index 90538036e3208c65483d0140dc37ea54c295f477..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cpu-x86_64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _flash_attn2_cpu_f12afc9 -ops = torch.ops._flash_attn2_cpu_f12afc9 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_flash_attn2_cpu_f12afc9::{op_name}" diff --git a/build/torch212-cxx11-cpu-x86_64-linux/bert_padding.py b/build/torch212-cxx11-cpu-x86_64-linux/bert_padding.py deleted file mode 100644 index 3c2d35159a014a9d03aabead9e52e009168696ea..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cpu-x86_64-linux/bert_padding.py +++ /dev/null @@ -1,218 +0,0 @@ -# Adapted from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/padding.py - -import torch -import torch.nn.functional as F -from einops import rearrange, repeat - - -class IndexFirstAxis(torch.autograd.Function): - @staticmethod - def forward(ctx, input, indices): - ctx.save_for_backward(indices) - assert input.ndim >= 2 - ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] - second_dim = other_shape.numel() - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - # return input[indices] - return torch.gather( - rearrange(input, "b ... -> b (...)"), 0, repeat(indices, "z -> z d", d=second_dim) - ).reshape(-1, *other_shape) - - @staticmethod - def backward(ctx, grad_output): - (indices,) = ctx.saved_tensors - assert grad_output.ndim >= 2 - other_shape = grad_output.shape[1:] - grad_output = rearrange(grad_output, "b ... -> b (...)") - grad_input = torch.zeros( - [ctx.first_axis_dim, grad_output.shape[1]], - device=grad_output.device, - dtype=grad_output.dtype, - ) - # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. - # grad_input[indices] = grad_output - grad_input.scatter_(0, repeat(indices, "z -> z d", d=grad_output.shape[1]), grad_output) - return grad_input.reshape(ctx.first_axis_dim, *other_shape), None - - -index_first_axis = IndexFirstAxis.apply - - -class IndexPutFirstAxis(torch.autograd.Function): - @staticmethod - def forward(ctx, values, indices, first_axis_dim): - ctx.save_for_backward(indices) - assert indices.ndim == 1 - assert values.ndim >= 2 - output = torch.zeros( - first_axis_dim, *values.shape[1:], device=values.device, dtype=values.dtype - ) - # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. - output[indices] = values - # output.scatter_(0, repeat(indices, 'z -> z d', d=values.shape[1]), values) - return output - - @staticmethod - def backward(ctx, grad_output): - (indices,) = ctx.saved_tensors - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - grad_values = grad_output[indices] - # grad_values = torch.gather(grad_output, 0, repeat(indices, 'z -> z d', d=grad_output.shape[1])) - return grad_values, None, None - - -index_put_first_axis = IndexPutFirstAxis.apply - - -class IndexFirstAxisResidual(torch.autograd.Function): - @staticmethod - def forward(ctx, input, indices): - ctx.save_for_backward(indices) - assert input.ndim >= 2 - ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] - second_dim = other_shape.numel() - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - output = input[indices] - # We don't want to reshape input (b ... -> b (...)) since it could change the channel_last - # memory format to channel_first. In other words, input might not be contiguous. - # If we don't detach, Pytorch complains about output being a view and is being modified inplace - return output, input.detach() - - @staticmethod - def backward(ctx, grad_output, grad_residual): - (indices,) = ctx.saved_tensors - assert grad_output.ndim >= 2 - other_shape = grad_output.shape[1:] - assert grad_residual.shape[1:] == other_shape - grad_input = grad_residual - # grad_input[indices] += grad_output - indices = indices.reshape(indices.shape[0], *((1,) * (grad_output.ndim - 1))) - indices = indices.expand_as(grad_output) - grad_input.scatter_add_(0, indices, grad_output) - return grad_input.reshape(ctx.first_axis_dim, *other_shape), None - - -index_first_axis_residual = IndexFirstAxisResidual.apply - - -def unpad_input(hidden_states, attention_mask, unused_mask=None): - """ - Arguments: - hidden_states: (batch, seqlen, ...) - attention_mask: (batch, seqlen), bool / int, 1 means valid and 0 means not valid. - unused_mask: (batch, seqlen), bool / int, 1 means the element is allocated but unused. - Return: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens selected in attention_mask + unused_mask. - indices: (total_nnz), the indices of masked tokens from the flattened input sequence. - cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. - max_seqlen_in_batch: int - seqused: (batch), returns the number of tokens selected in attention_mask + unused_mask. - """ - all_masks = (attention_mask + unused_mask) if unused_mask is not None else attention_mask - seqlens_in_batch = all_masks.sum(dim=-1, dtype=torch.int32) - used_seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) - indices = torch.nonzero(all_masks.flatten(), as_tuple=False).flatten() - max_seqlen_in_batch = seqlens_in_batch.max().item() - cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) - # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the - # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim - # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to - # index with integer indices. Moreover, torch's index is a bit slower than it needs to be, - # so we write custom forward and backward to make it a bit faster. - return ( - index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices), - indices, - cu_seqlens, - max_seqlen_in_batch, - used_seqlens_in_batch, - ) - - -def unpad_input_for_concatenated_sequences(hidden_states, attention_mask_in_length): - """ - Supports concatenating short samples in one sequence. The attention_mask_in_length is utilized to mask other short samples. It helps efficient training of variant lengths-based samples (e.g., the supervised fine-tuning task in large language model). - The motivation for this function is explained [here](https://github.com/Dao-AILab/flash-attention/issues/432#issuecomment-1668822286). - - For example, if batch = 3 and seqlen = 6, the attention_mask_in_length is: - ``` - [ - [2, 3, 0, 0, 0, 0], - [3, 2, 0, 0, 0, 0], - [6, 0, 0, 0, 0, 0] - ] - ``` - , which refers to the 3D-attention mask: - ``` - [ - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [0, 0, 1, 0, 0, 0], - [0, 0, 1, 1, 0, 0], - [0, 0, 1, 1, 1, 0], - [0, 0, 0, 0, 0, 1] - ], - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [0, 0, 0, 1, 0, 0], - [0, 0, 0, 1, 1, 0], - [0, 0, 0, 0, 0, 1] - ], - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [1, 1, 1, 1, 0, 0], - [1, 1, 1, 1, 1, 0], - [1, 1, 1, 1, 1, 1] - ] - ] - ```. - - Arguments: - hidden_states: (batch, seqlen, ...) - attention_mask_in_length: (batch, seqlen), int, a nonzero number (e.g., 1, 2, 3, etc.) means length of concatenated sequence in b-th batch, and 0 means none. - Return: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. - indices: (total_nnz), the indices of non-masked tokens from the flattened input sequence. - cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. - max_seqlen_in_batch: int - """ - length = attention_mask_in_length.sum(dim=-1) - seqlen = attention_mask_in_length.size(-1) - attention_mask_2d = torch.arange(seqlen, device=length.device, dtype=length.dtype).expand(len(length), seqlen) < length.unsqueeze(1) - real_indices_idx = torch.nonzero(attention_mask_in_length.flatten(), as_tuple=False).flatten() - seqlens_in_batch = attention_mask_in_length.flatten()[real_indices_idx] - indices = torch.nonzero(attention_mask_2d.flatten(), as_tuple=False).flatten() - max_seqlen_in_batch = seqlens_in_batch.max().item() - cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) - # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the - # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim - # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to - # index with integer indices. Moreover, torch's index is a bit slower than it needs to be, - # so we write custom forward and backward to make it a bit faster. - return ( - index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices), - indices, - cu_seqlens, - max_seqlen_in_batch, - ) - - -def pad_input(hidden_states, indices, batch, seqlen): - """ - Arguments: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. - indices: (total_nnz), the indices that represent the non-masked tokens of the original padded input sequence. - batch: int, batch size for the padded sequence. - seqlen: int, maximum sequence length for the padded sequence. - Return: - hidden_states: (batch, seqlen, ...) - """ - dim = hidden_states.shape[-1] - # output = torch.zeros((batch * seqlen), dim, device=hidden_states.device, dtype=hidden_states.dtype) - # output[indices] = hidden_states - output = index_put_first_axis(hidden_states, indices, batch * seqlen) - return rearrange(output, "(b s) ... -> b s ...", b=batch) diff --git a/build/torch212-cxx11-cpu-x86_64-linux/flash_attn2/__init__.py b/build/torch212-cxx11-cpu-x86_64-linux/flash_attn2/__init__.py deleted file mode 100644 index a9b2672c1cd85b74c1b3ded0fc0b2100e1aeac23..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cpu-x86_64-linux/flash_attn2/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -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/torch212-cxx11-cpu-x86_64-linux/flash_attn_interface.py b/build/torch212-cxx11-cpu-x86_64-linux/flash_attn_interface.py deleted file mode 100644 index b014e569dc9812a9c840d1043f2cebe767b4d477..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cpu-x86_64-linux/flash_attn_interface.py +++ /dev/null @@ -1,1631 +0,0 @@ -# Copyright (c) 2023, Tri Dao. - -from typing import Optional, Sequence, Tuple, Union - -import torch -import torch.nn as nn -import os - -# # isort: off -# # We need to import the CUDA kernels after importing torch -# USE_TRITON_ROCM = os.getenv("FLASH_ATTENTION_TRITON_AMD_ENABLE", "FALSE") == "TRUE" -# if USE_TRITON_ROCM: -# from .flash_attn_triton_amd import interface_fa as flash_attn -# else: -# import flash_attn_2_cuda as flash_attn - - -from . import _ops -from ._ops import ops as flash_attn -from ._ops import add_op_namespace_prefix - -# # isort: on - -def maybe_contiguous(x): - return x.contiguous() if x is not None and x.stride(-1) != 1 else x - - -def _get_device(): - if torch.xpu.is_available(): - return "xpu" - elif torch.cuda.is_available(): - return "cuda" - else: - return "cpu" - - -def _get_block_size_n(device, head_dim, is_dropout, is_causal): - assert head_dim <= 256 - - if device.type == "xpu": - if head_dim <= 96: - return 64 - elif head_dim <= 128: - return 32 - elif head_dim <= 256: - return 64 - else: - return 32 - - # This should match the block sizes in the CUDA kernel - major, minor = torch.cuda.get_device_capability(device) - is_sm8x = major == 8 and minor > 0 # Only include sm86 and sm89, exclude sm80 (A100) - is_sm80 = major == 8 and minor == 0 - is_sm90 = major == 9 and minor == 0 - if head_dim <= 32: - return 128 - if head_dim <= 64: - return 128 if not is_dropout else 64 - elif head_dim <= 96: - return 64 - elif head_dim <= 128: - if is_sm8x: - return 64 if (not is_dropout and is_causal) else 32 - else: - return 64 if not is_dropout else 32 - elif head_dim <= 192: - return 64 - elif head_dim <= 224: - return 64 - elif head_dim <= 256: - return 64 - - -def round_multiple(x, m): - return (x + m - 1) // m * m - - -# torch.compile() support is only enabled for pytorch >= 2.4 -# The reason for this is that we are using the new custom_op and register_fake -# APIs, which support inplace modification of inputs in the function itself -if torch.__version__ >= "2.4.0": - _torch_custom_op_wrapper = torch.library.custom_op - _torch_register_fake_wrapper = torch.library.register_fake -else: - def noop_custom_op_wrapper(name, fn=None, /, *, mutates_args, device_types=None, schema=None): - def wrap(func): - return func - if fn is None: - return wrap - return fn - def noop_register_fake_wrapper(op, fn=None, /, *, lib=None, _stacklevel=1): - def wrap(func): - return func - if fn is None: - return wrap - return fn - _torch_custom_op_wrapper = noop_custom_op_wrapper - _torch_register_fake_wrapper = noop_register_fake_wrapper - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_forward"), mutates_args=(), device_types=_get_device()) -def _flash_attn_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - return_softmax: bool -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - out, softmax_lse, S_dmask, rng_state = flash_attn.fwd( - q, - k, - v, - None, - alibi_slopes, - dropout_p, - softmax_scale, - causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - None, - ) - return out, softmax_lse, S_dmask, rng_state - - -@_torch_register_fake_wrapper(add_op_namespace_prefix("_flash_attn_forward")) -def _flash_attn_forward_fake( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - return_softmax: bool -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - batch_size, seqlen_q, num_heads, head_size = q.shape - seqlen_k = k.shape[1] - out = torch.empty_like(q) - softmax_lse = torch.empty((batch_size, num_heads, seqlen_q), dtype=torch.float32, device=q.device, layout=q.layout) - p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) - if return_softmax: - p = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128), round_multiple(seqlen_k, 128)), dtype=q.dtype, device=q.device, layout=q.layout) - rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) - - return out, softmax_lse, p, rng_state - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_forward = _ops.ops._flash_attn_forward -else: - _wrapped_flash_attn_forward = _flash_attn_forward - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_varlen_forward"), mutates_args=(), device_types=_get_device()) -def _flash_attn_varlen_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - alibi_slopes: Optional[torch.Tensor] = None, - return_softmax: bool = False, - block_table: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - out, softmax_lse, S_dmask, rng_state = flash_attn.varlen_fwd( - q, - k, - v, - None, - cu_seqlens_q, - cu_seqlens_k, - seqused_k, - leftpad_k, - block_table, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - zero_tensors, - causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - None, - ) - # if out.isnan().any() or softmax_lse.isnan().any(): - # breakpoint() - return out, softmax_lse, S_dmask, rng_state - - -@_torch_register_fake_wrapper(add_op_namespace_prefix("_flash_attn_varlen_forward")) -def _flash_attn_varlen_forward_fake( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - alibi_slopes: Optional[torch.Tensor] = None, - return_softmax: bool = False, - block_table: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - paged_kv = block_table is not None - batch_size = cu_seqlens_q.numel() - 1 - total_q, num_heads, _ = q.shape - - out = torch.empty_like(q) - softmax_lse = torch.empty((num_heads, total_q), dtype=torch.float32, device=q.device, layout=q.layout) - p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) - seqlen_q_rounded = round_multiple(max_seqlen_q, 128) - seqlen_k_rounded = round_multiple(max_seqlen_k, 128) - if return_softmax: - p = torch.empty((batch_size, num_heads, seqlen_q_rounded, seqlen_k_rounded), dtype=q.dtype, device=q.device, layout=q.layout) - rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) - return out, softmax_lse, p, rng_state - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_varlen_forward = _ops.ops._flash_attn_varlen_forward -else: - _wrapped_flash_attn_varlen_forward = _flash_attn_varlen_forward - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_backward"), mutates_args=("dq", "dk", "dv"), device_types=_get_device()) -def _flash_attn_backward( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, -) -> 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)] - ( - dq, - dk, - dv, - softmax_d, - ) = flash_attn.bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - alibi_slopes, - dropout_p, - softmax_scale, - causal, - window_size_left, - window_size_right, - softcap, - deterministic, - None, - rng_state, - ) - return softmax_d - - -@_torch_register_fake_wrapper(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, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, -) -> torch.Tensor: - dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] - if dq is None: - dq = torch.empty_like(q) - if dk is None: - dk = torch.empty_like(k) - if dv is None: - dv = torch.empty_like(v) - batch_size, seqlen_q, num_heads, _ = q.shape - softmax_d = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128)), device=q.device, dtype=torch.float32) - - return softmax_d - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_backward = _ops.ops._flash_attn_backward -else: - _wrapped_flash_attn_backward = _flash_attn_backward - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_varlen_backward"), mutates_args=("dq", "dk", "dv"), device_types=_get_device()) -def _flash_attn_varlen_backward( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> 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)] - ( - dq, - dk, - dv, - softmax_d, - ) = flash_attn.varlen_bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - zero_tensors, - causal, - window_size_left, - window_size_right, - softcap, - deterministic, - None, - rng_state, - ) - # if dk.isnan().any() or dk.isnan().any() or dv.isnan().any() or softmax_d.isnan().any(): - # breakpoint() - return softmax_d - - -@_torch_register_fake_wrapper(add_op_namespace_prefix("_flash_attn_varlen_backward")) -def _flash_attn_varlen_backward_fake( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> torch.Tensor: - dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] - batch_size = cu_seqlens_q.numel() - 1 - total_q, num_heads, _ = q.shape - - if dq is None: - dq = torch.empty_like(q) - if dk is None: - dk = torch.empty_like(k) - if dv is None: - dv = torch.empty_like(v) - softmax_d = torch.empty((num_heads, total_q + 128 * batch_size), device=q.device, dtype=torch.float32) - - return softmax_d - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_varlen_backward = _ops.ops._flash_attn_varlen_backward -else: - _wrapped_flash_attn_varlen_backward = _flash_attn_varlen_backward - - -class FlashAttnQKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and qkv.requires_grad - if softmax_scale is None: - softmax_scale = qkv.shape[-1] ** (-0.5) - q, k, v = qkv[:, :, 0].detach(), qkv[:, :, 1].detach(), qkv[:, :, 2].detach() - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) - dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dqkv[:, :, 0], - dqkv[:, :, 1], - dqkv[:, :, 2], - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dqkv = dqkv[..., : dout.shape[-1]] # We could have padded the head dimension - return dqkv, None, None, None, None, None, None, None, None, None - - -class FlashAttnVarlenQKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - cu_seqlens, - max_seqlen, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and qkv.requires_grad - if softmax_scale is None: - softmax_scale = qkv.shape[-1] ** (-0.5) - q, k, v = qkv[:, 0].detach(), qkv[:, 1].detach(), qkv[:, 2].detach() - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens, - cu_seqlens, - max_seqlen, - max_seqlen, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=None, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens, rng_state) - ctx.dropout_p = dropout_p - ctx.max_seqlen = max_seqlen - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens, rng_state = ctx.saved_tensors - qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) - dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dqkv[:, 0], - dqkv[:, 1], - dqkv[:, 2], - cu_seqlens, - cu_seqlens, - ctx.max_seqlen, - ctx.max_seqlen, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - 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 FlashAttnKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - kv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, kv] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - k, v = kv[:, :, 0].detach(), kv[:, :, 1].detach() - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - dq = torch.empty_like(q) - kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) - dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dkv[:, :, 0], - dkv[:, :, 1], - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dkv = dkv[..., : dout.shape[-1]] - return dq, dkv, None, None, None, None, None, None, None, None, None - - -class FlashAttnVarlenKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, kv] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - k, v = kv[:, 0].detach(), kv[:, 1].detach() - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=None, - ) - if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state - ) - ctx.dropout_p = dropout_p - 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.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors - dq = torch.empty_like(q) - kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) - dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dkv[:, 0], - dkv[:, 1], - cu_seqlens_q, - cu_seqlens_k, - ctx.max_seqlen_q, - ctx.max_seqlen_k, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dkv = dkv[..., : dout.shape[-1]] - return dq, dkv, None, None, None, None, None, None, None, None, None, None, None, None, None - - -class FlashAttnFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - k, - v, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, k, v] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dk = dk[..., : dout.shape[-1]] - dv = dv[..., : dout.shape[-1]] - return dq, dk, dv, 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, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - block_table, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, k, v] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=block_table, - ) - if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state - ) - ctx.dropout_p = dropout_p - 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.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors - dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - ctx.max_seqlen_q, - ctx.max_seqlen_k, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dk = dk[..., : dout.shape[-1]] - dv = dv[..., : dout.shape[-1]] - return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None - - -def flash_attn_qkvpacked_func( - qkv, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # <=0.0 means deactivate - alibi_slopes=None, - deterministic=False, - 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, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_kvpacked_func( - q, - kv, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, -): - """dropout_p should be set to 0.0 during evaluation - If 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 K, V. - 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) - kv: (batch_size, seqlen, 2, 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. - softcap: float. Anything > 0 activates softcapping 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). - 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 FlashAttnKVPackedFunc.apply( - q, - kv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_func( - q, - k, - v, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - 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). - 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 FlashAttnFunc.apply( - q, - k, - v, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_qkvpacked_func( - qkv, - cu_seqlens, - max_seqlen, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - 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_varlen_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_varlen_kvpacked_func and flash_attn_varlen_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: (total, 3, nheads, headdim), where total = total number of tokens in the batch. - cu_seqlens: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into qkv. - max_seqlen: int. Maximum sequence length in the batch. - 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenQKVPackedFunc.apply( - qkv, - cu_seqlens, - max_seqlen, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_kvpacked_func( - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, -): - """dropout_p should be set to 0.0 during evaluation - If 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 K, V. - 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: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch. - kv: (total_k, 2, nheads_k, headdim), where total_k = total number of key tokens in the batch. - cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into q. - cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into kv. - max_seqlen_q: int. Maximum query sequence length in the batch. - max_seqlen_k: int. Maximum key sequence length in the batch. - 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 + 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenKVPackedFunc.apply( - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_func( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, - block_table=None, -): - """dropout_p should be set to 0.0 during evaluation - Supports multi-query and grouped-query attention (MQA/GQA) by passing in K, V 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: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch. - k: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. - v: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. - cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into q. - cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into kv. - max_seqlen_q: int. Maximum query sequence length in the batch. - max_seqlen_k: int. Maximum key sequence length in the batch. - 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 + 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenFunc.apply( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - block_table, - False if q.device.type == "cpu" else torch.is_grad_enabled(), - ) - - -def flash_attn_with_kvcache( - q, - k_cache, - v_cache, - k=None, - v=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, - block_table: Optional[torch.Tensor] = None, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - rotary_interleaved=True, - alibi_slopes=None, - num_splits=0, - 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 block_table, - or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_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) if there's no block_table, - or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_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). Similar to k. - 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. - block_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). - 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. - 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" - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-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) - cache_batch_idx = maybe_contiguous(cache_batch_idx) - block_table = maybe_contiguous(block_table) - out, softmax_lse = flash_attn.fwd_kvcache( - q, - k_cache, - v_cache, - k, - v, - cache_seqlens, - rotary_cos, - rotary_sin, - cache_batch_idx, - cache_leftpad, - block_table, - alibi_slopes, - None, - softmax_scale, - causal, - window_size[0], - window_size[1], - softcap, - rotary_interleaved, - num_splits, - ) - return (out, softmax_lse) if return_softmax_lse else out diff --git a/build/torch212-cxx11-cpu-x86_64-linux/layers/__init__.py b/build/torch212-cxx11-cpu-x86_64-linux/layers/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/build/torch212-cxx11-cpu-x86_64-linux/layers/patch_embed.py b/build/torch212-cxx11-cpu-x86_64-linux/layers/patch_embed.py deleted file mode 100644 index 05562f8e8bcdb58e947c6f402a49eacd2d031871..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cpu-x86_64-linux/layers/patch_embed.py +++ /dev/null @@ -1,67 +0,0 @@ -# We use the same API as https://github.com/rwightman/pytorch-image-models/blob/v0.6.11/timm/models/layers/patch_embed.py -# But we use nn.Linear instead of Conv2d and it's about 8x faster. - -from functools import partial - -import torch.nn as nn -from einops import rearrange -from torch import _assert -from torch.nn.modules.utils import _pair - -try: - from flash_attn.ops.fused_dense import FusedDense -except ImportError: - FusedDense = None - - -class PatchEmbed(nn.Module): - """2D Image to Patch Embedding""" - - def __init__( - self, - img_size=224, - patch_size=16, - in_chans=3, - embed_dim=768, - norm_layer=None, - flatten=True, - bias=True, - fused_bias_fc=False, - ): - super().__init__() - img_size = _pair(img_size) - patch_size = _pair(patch_size) - self.img_size = img_size - self.patch_size = patch_size - self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1]) - self.num_patches = self.grid_size[0] * self.grid_size[1] - self.flatten = flatten - if fused_bias_fc and FusedDense is None: - raise ImportError("fused_dense is not installed") - - linear_cls = nn.Linear if not fused_bias_fc or not bias else FusedDense - self.proj = linear_cls(in_chans * patch_size[0] * patch_size[1], embed_dim, bias=bias) - self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() - - def forward(self, x): - _, _, H, W = x.shape - _assert( - H == self.img_size[0], - f"Input image height ({H}) doesn't match model ({self.img_size[0]}).", - ) - _assert( - W == self.img_size[1], - f"Input image width ({W}) doesn't match model ({self.img_size[1]}).", - ) - x = self.proj( - rearrange( - x, - "b c (h p1) (w p2) -> b h w (c p1 p2)", - p1=self.patch_size[0], - p2=self.patch_size[1], - ) - ) - if self.flatten: - x = rearrange(x, "b h w c -> b (h w) c") - x = self.norm(x) - return x diff --git a/build/torch212-cxx11-cpu-x86_64-linux/layers/rotary.py b/build/torch212-cxx11-cpu-x86_64-linux/layers/rotary.py deleted file mode 100644 index d1bfc21fc7de1dd287e8f382847b194a48075981..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cpu-x86_64-linux/layers/rotary.py +++ /dev/null @@ -1,483 +0,0 @@ -# Copyright (c) 2025, Tri Dao - -import math -from functools import partial -from typing import Optional, Tuple, Union - -import torch -from torch import Tensor - -from einops import rearrange, repeat -# from flash_attn.ops.triton.rotary import apply_rotary -from ..ops.triton.rotary import apply_rotary - - -def rotate_half(x, interleaved=False): - if not interleaved: - x1, x2 = x.chunk(2, dim=-1) - return torch.cat((-x2, x1), dim=-1) - else: - x1, x2 = x[..., ::2], x[..., 1::2] - return rearrange(torch.stack((-x2, x1), dim=-1), "... d two -> ... (d two)", two=2) - - -def apply_rotary_emb_torch(x, cos, sin, interleaved=False): - """ - x: (batch_size, seqlen, nheads, headdim) - cos, sin: (seqlen, rotary_dim / 2) or (batch_size, seqlen, rotary_dim / 2) - """ - ro_dim = cos.shape[-1] * 2 - assert ro_dim <= x.shape[-1] - cos = repeat(cos, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)") - sin = repeat(sin, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)") - return torch.cat( - [x[..., :ro_dim] * cos + rotate_half(x[..., :ro_dim], interleaved) * sin, x[..., ro_dim:]], - dim=-1, - ) - - -class ApplyRotaryEmb(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x, - cos, - sin, - interleaved=False, - inplace=False, - seqlen_offsets: Union[int, Tensor] = 0, - cu_seqlens: Optional[Tensor] = None, - max_seqlen: Optional[int] = None, - ): - out = apply_rotary( - x, - cos, - sin, - seqlen_offsets=seqlen_offsets, - cu_seqlens=cu_seqlens, - max_seqlen=max_seqlen, - interleaved=interleaved, - inplace=inplace, - ) - if isinstance(seqlen_offsets, int): - ctx.save_for_backward(cos, sin, cu_seqlens) # Can't save int with save_for_backward - ctx.seqlen_offsets = seqlen_offsets - else: - ctx.save_for_backward(cos, sin, cu_seqlens, seqlen_offsets) - ctx.seqlen_offsets = None - ctx.interleaved = interleaved - ctx.inplace = inplace - ctx.max_seqlen = max_seqlen - return out if not inplace else x - - @staticmethod - def backward(ctx, do): - seqlen_offsets = ctx.seqlen_offsets - if seqlen_offsets is None: - cos, sin, cu_seqlens, seqlen_offsets = ctx.saved_tensors - else: - cos, sin, cu_seqlens = ctx.saved_tensors - dx = apply_rotary( - do, - cos, - sin, - seqlen_offsets=seqlen_offsets, - cu_seqlens=cu_seqlens, - max_seqlen=ctx.max_seqlen, - interleaved=ctx.interleaved, - inplace=ctx.inplace, - conjugate=True, - ) - return dx, None, None, None, None, None, None, None - - -def apply_rotary_emb( - x, - cos, - sin, - interleaved=False, - inplace=False, - seqlen_offsets: Union[int, Tensor] = 0, - cu_seqlens: Optional[Tensor] = None, - max_seqlen: Optional[int] = None, -): - """ - Arguments: - x: (batch_size, seqlen, nheads, headdim) if cu_seqlens is None - else (total_seqlen, nheads, headdim) - cos, sin: (seqlen_rotary, rotary_dim / 2) - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead - of 1st half and 2nd half (GPT-NeoX style). - inplace: if True, apply rotary embedding in-place. - seqlen_offsets: (batch_size,) or int. Each sequence in x is shifted by this amount. - Most commonly used in inference when we have KV cache. - cu_seqlens: (batch + 1,) or None - max_seqlen: int - Return: - out: (batch_size, seqlen, nheads, headdim) if cu_seqlens is None - else (total_seqlen, nheads, headdim) - rotary_dim must be <= headdim - Apply rotary embedding to the first rotary_dim of x. - """ - return ApplyRotaryEmb.apply( - x, cos, sin, interleaved, inplace, seqlen_offsets, cu_seqlens, max_seqlen - ) - - -# For backward compatibility -apply_rotary_emb_func = apply_rotary_emb - - -def _apply_rotary_emb_qkv( - qkv, - cos, - sin, - cos_k=None, - sin_k=None, - interleaved=False, - inplace=False, - conjugate=False, - seqlen_offsets: Union[int, Tensor] = 0, - num_heads_q: Optional[int] = None, -): - apply_rotary_fn = partial( - apply_rotary, - interleaved=interleaved, - inplace=inplace, - conjugate=conjugate, - seqlen_offsets=seqlen_offsets - ) - if cos_k is None and sin_k is None and qkv.is_contiguous(): - # Call 1 kernel instead of 2 kernels - # We need qkv to be contiguous so that when we reshape to combine (3, nheads) - # dimensions, we get the same tensor - if qkv.dim() == 5: - batch, seqlen, three, nheads, headdim = qkv.shape - assert three == 3 - # qk = rearrange(qkv[:, :, :2], "b s t h d -> b s (t h) d") - qk = qkv[:, :, :2].reshape(batch, seqlen, -1, headdim) - qk = apply_rotary_fn(qk, cos, sin) - else: - assert qkv.dim() == 4 - assert num_heads_q is not None - num_heads_k = (qkv.shape[2] - num_heads_q) // 2 - assert qkv.shape[2] == num_heads_q + 2 * num_heads_k - qk = qkv[:, :, :num_heads_q + num_heads_k] - qk = apply_rotary_fn(qk, cos, sin) - if not inplace: - if qkv.dim() == 5: - qkv = torch.cat([rearrange(qk, "b s (t h) d -> b s t h d", t=2), qkv[:, :, 2:]], dim=2) - else: - qkv = torch.cat([qk, qkv[:, :, num_heads_q + num_heads_k :]], dim=2) - else: - cos_k = cos if cos_k is None else cos_k - sin_k = sin if sin_k is None else sin_k - if qkv.dim() == 5: - batch, seqlen, three, nheads, headdim = qkv.shape - assert three == 3 - q, k = qkv[:, :, 0], qkv[:, :, 1] - else: - assert qkv.dim() == 4 - assert num_heads_q is not None - num_heads_k = (qkv.shape[2] - num_heads_q) // 2 - assert qkv.shape[2] == num_heads_q + 2 * num_heads_k - q, k = qkv[:, :, :num_heads_q], qkv[:, :, num_heads_q : num_heads_q + num_heads_k] - q = apply_rotary_fn(q, cos, sin) - k = apply_rotary_fn(k, cos_k, sin_k) - if not inplace: - if qkv.dim() == 5: - qkv = torch.stack([q, k, qkv[:, :, 2]], dim=2) - else: - qkv = torch.cat([q, k, qkv[:, :, num_heads_q + num_heads_k:]], dim=2) - return qkv - - -class ApplyRotaryEmbQKV_(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - cos, - sin, - cos_k=None, - sin_k=None, - interleaved=False, - seqlen_offsets: Union[int, torch.Tensor] = 0, - num_heads_q: Optional[int] = None, - ): - # apply_rotary_emb_qkv_inplace( - qkv = _apply_rotary_emb_qkv( - qkv, cos, sin, cos_k, sin_k, interleaved=interleaved, inplace=True, - seqlen_offsets=seqlen_offsets, num_heads_q=num_heads_q, - ) - if isinstance(seqlen_offsets, int): - ctx.save_for_backward(cos, sin, cos_k, sin_k) - ctx.seqlen_offsets = seqlen_offsets - else: - ctx.save_for_backward(cos, sin, cos_k, sin_k, seqlen_offsets) - ctx.seqlen_offsets = None - ctx.interleaved = interleaved - ctx.num_heads_q = num_heads_q - return qkv - - @staticmethod - def backward(ctx, dqkv): - seqlen_offsets = ctx.seqlen_offsets - if seqlen_offsets is None: - cos, sin, cos_k, sin_k, seqlen_offsets = ctx.saved_tensors - else: - cos, sin, cos_k, sin_k = ctx.saved_tensors - dqkv = _apply_rotary_emb_qkv( - dqkv, cos, sin, cos_k, sin_k, interleaved=ctx.interleaved, inplace=True, - seqlen_offsets=seqlen_offsets, num_heads_q=ctx.num_heads_q, conjugate=True, - ) - return dqkv, None, None, None, None, None, None, None - - -def apply_rotary_emb_qkv_( - qkv, - cos, - sin, - cos_k=None, - sin_k=None, - interleaved=False, - seqlen_offsets: Union[int, torch.Tensor] = 0, - num_heads_q: Optional[int] = None, -): - """ - Arguments: - qkv: (batch_size, seqlen, 3, nheads, headdim) or (batch_size, seqlen, num_heads_q + 2 * num_heads_k, headdim). - If qkv has shape (batch_size, seqlen, num_heads_q + 2 * num_heads_k, headdim) (e.g. MQA / GQA), - then num_heads_q must be provided. - cos, sin: (seqlen, rotary_dim / 2) - cos_k, sin_k: (seqlen, rotary_dim / 2), optional - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead of - 1st half and 2nd half (GPT-NeoX style). - seqlen_offsets: (batch_size,) or int. Each sequence in Q and K is shifted by this amount. - Most commonly used in inference when we have KV cache. - Return: - qkv: (batch_size, seqlen, 3, nheads, headdim) or (batch_size, seqlen, num_heads_q + 2 * num_heads_k, headdim) - rotary_dim must be <= headdim - Apply rotary embedding *inplace* to the first rotary_dim of Q and K. - """ - return ApplyRotaryEmbQKV_.apply( - qkv, cos, sin, cos_k, sin_k, interleaved, seqlen_offsets, num_heads_q - ) - - -class ApplyRotaryEmbKV_(torch.autograd.Function): - - @staticmethod - def forward(ctx, kv, cos, sin, interleaved=False, seqlen_offsets: Union[int, torch.Tensor] = 0): - batch, seqlen, two, nheads, headdim = kv.shape - assert two == 2 - k = kv[:, :, 0] - apply_rotary( - k, cos, sin, seqlen_offsets=seqlen_offsets, interleaved=interleaved, inplace=True - ) - if isinstance(seqlen_offsets, int): - ctx.save_for_backward(cos, sin) # Can't save int with save_for_backward - ctx.seqlen_offsets = seqlen_offsets - else: - ctx.save_for_backward(cos, sin, seqlen_offsets) - ctx.seqlen_offsets = None - ctx.interleaved = interleaved - return kv - - @staticmethod - def backward(ctx, dkv): - seqlen_offsets = ctx.seqlen_offsets - if seqlen_offsets is None: - cos, sin, seqlen_offsets = ctx.saved_tensors - else: - cos, sin = ctx.saved_tensors - apply_rotary( - dkv[:, :, 0], - cos, - sin, - seqlen_offsets=seqlen_offsets, - interleaved=ctx.interleaved, - inplace=True, - conjugate=True, - ) - return dkv, None, None, None, None - - -apply_rotary_emb_kv_ = ApplyRotaryEmbKV_.apply - - -def apply_rotary_emb_kv_( - kv, - cos, - sin, - interleaved=False, - seqlen_offsets: Union[int, torch.Tensor] = 0, -): - """ - Arguments: - kv: (batch_size, seqlen, 2, nheads, headdim) - cos, sin: (seqlen, rotary_dim / 2) - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead of - 1st half and 2nd half (GPT-NeoX style). - seqlen_offsets: (batch_size,) or int. Each sequence in Q and K is shifted by this amount. - Most commonly used in inference when we have KV cache. - Return: - kv: (batch_size, seqlen, 2, nheads, headdim) - rotary_dim must be <= headdim - Apply rotary embedding *inplace* to the first rotary_dim of K. - """ - return ApplyRotaryEmbKV_.apply(kv, cos, sin, interleaved, seqlen_offsets) - - -class RotaryEmbedding(torch.nn.Module): - """ - The rotary position embeddings from RoFormer_ (Su et. al). - A crucial insight from the method is that the query and keys are - transformed by rotation matrices which depend on the relative positions. - - Other implementations are available in the Rotary Transformer repo_ and in - GPT-NeoX_, GPT-NeoX was an inspiration - - .. _RoFormer: https://arxiv.org/abs/2104.09864 - .. _repo: https://github.com/ZhuiyiTechnology/roformer - .. _GPT-NeoX: https://github.com/EleutherAI/gpt-neox - - If scale_base is not None, this implements XPos (Sun et al., https://arxiv.org/abs/2212.10554). - A recommended value for scale_base is 512: https://github.com/HazyResearch/flash-attention/issues/96 - Reference: https://github.com/sunyt32/torchscale/blob/main/torchscale/component/xpos_relative_position.py - """ - - def __init__( - self, - dim: int, - base=10000.0, - interleaved=False, - scale_base=None, - device=None, - ): - """ - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead - of 1st half and 2nd half (GPT-NeoX style). - """ - super().__init__() - self.dim = dim - self.base = float(base) - # Generate and save the inverse frequency buffer (non trainable) - inv_freq = self._compute_inv_freq(device) - self.register_buffer("inv_freq", inv_freq, persistent=False) - self.interleaved = interleaved - self.scale_base = scale_base - scale = ( - (torch.arange(0, dim, 2, device=device, dtype=torch.float32) + 0.4 * dim) / (1.4 * dim) - if scale_base is not None - else None - ) - self.register_buffer("scale", scale, persistent=False) - - self._seq_len_cached = 0 - self._cos_cached = None - self._sin_cached = None - self._cos_k_cached = None - self._sin_k_cached = None - - def _compute_inv_freq(self, device=None): - return 1.0 / ( - self.base - ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim) - ) - - def _update_cos_sin_cache(self, seqlen, device=None, dtype=None): - # Reset the tables if the sequence length has changed, - # if we're on a new device (possibly due to tracing for instance), - # or if we're switching from inference mode to training - if ( - seqlen > self._seq_len_cached - or self._cos_cached is None - or self._cos_cached.device != device - or self._cos_cached.dtype != dtype - or (self.training and self._cos_cached.is_inference()) - ): - self._seq_len_cached = seqlen - # We want fp32 here, not self.inv_freq.dtype, since the model could be loaded in bf16 - # And the output of arange can be quite large, so bf16 would lose a lot of precision. - t = torch.arange(seqlen, device=device, dtype=torch.float32) - # We want fp32 here as well since inv_freq will be multiplied with t, and the output - # will be large. Having it in bf16 will lose a lot of precision and cause the - # cos & sin output to change significantly. - # We want to recompute self.inv_freq if it was not loaded in fp32 - if self.inv_freq.dtype != torch.float32: - inv_freq = self._compute_inv_freq(device=device) - else: - inv_freq = self.inv_freq - # Don't do einsum, it converts fp32 to bf16 under AMP - # freqs = torch.einsum("i,j->ij", t, self.inv_freq) - freqs = torch.outer(t, inv_freq) - if self.scale is None: - self._cos_cached = torch.cos(freqs).to(dtype) - self._sin_cached = torch.sin(freqs).to(dtype) - else: - power = ( - torch.arange(seqlen, dtype=self.scale.dtype, device=self.scale.device) - - seqlen // 2 - ) / self.scale_base - scale = self.scale.to(device=power.device) ** rearrange(power, "s -> s 1") - # We want the multiplication by scale to happen in fp32 - self._cos_cached = (torch.cos(freqs) * scale).to(dtype) - self._sin_cached = (torch.sin(freqs) * scale).to(dtype) - self._cos_k_cached = (torch.cos(freqs) / scale).to(dtype) - self._sin_k_cached = (torch.sin(freqs) / scale).to(dtype) - - def forward( - self, - qkv: torch.Tensor, - kv: Optional[torch.Tensor] = None, - seqlen_offset: Union[int, torch.Tensor] = 0, - max_seqlen: Optional[int] = None, - num_heads_q: Optional[int] = None, - ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: - """ - qkv: (batch, seqlen, 3, nheads, headdim) or (batch, seqlen, num_heads_q + 2 * num_heads_k, headdim) - if kv is none, else it's just q of shape (batch, seqlen, nheads, headdim). - If qkv has shape (batch, seqlen, num_heads_q + 2 * num_heads_k, headdim) (e.g. MQA / GQA), - then num_heads_q must be provided. - kv: (batch, seqlen, 2, nheads, headdim) - seqlen_offset: (batch_size,) or int. Each sequence in x is shifted by this amount. - Most commonly used in inference when we have KV cache. - If it's a tensor of shape (batch_size,), then to update the cos / sin cache, one - should pass in max_seqlen, which will update the cos / sin cache up to that length. - Apply rotary embedding *inplace* to qkv and / or kv. - """ - seqlen = qkv.shape[1] - if max_seqlen is not None: - self._update_cos_sin_cache(max_seqlen, device=qkv.device, dtype=qkv.dtype) - elif isinstance(seqlen_offset, int): - self._update_cos_sin_cache(seqlen + seqlen_offset, device=qkv.device, dtype=qkv.dtype) - if kv is None: - return apply_rotary_emb_qkv_( - qkv, - self._cos_cached, - self._sin_cached, - self._cos_k_cached if self.scale is not None else None, - self._sin_k_cached if self.scale is not None else None, - interleaved=self.interleaved, - seqlen_offsets=seqlen_offset, - num_heads_q=num_heads_q, - ) - else: - q = qkv - q = apply_rotary_emb_func( - q, - self._cos_cached, - self._sin_cached, - interleaved=self.interleaved, - inplace=True, - seqlen_offsets=seqlen_offset, - ) - kv = apply_rotary_emb_kv_( - kv, - self._cos_cached if self.scale is None else self._cos_k_cached, - self._sin_cached if self.scale is None else self._sin_k_cached, - interleaved=self.interleaved, - seqlen_offsets=seqlen_offset, - ) - return q, kv diff --git a/build/torch212-cxx11-cpu-x86_64-linux/metadata.json b/build/torch212-cxx11-cpu-x86_64-linux/metadata.json deleted file mode 100644 index fd8fd6de8e88fcd272a06c6632905dbd0cd11270..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cpu-x86_64-linux/metadata.json +++ /dev/null @@ -1,36 +0,0 @@ -{ - "name": "flash-attn2", - "id": "_flash_attn2_cpu_f12afc9", - "version": 1, - "license": "BSD-3-Clause", - "python-depends": [], - "backend": { - "type": "cpu" - }, - "digest": { - "algorithm": "sha256", - "files": { - "__init__.py": "Cd9MEdd/WUYnBmfiQpC4SyPqAxrTilTjvsvswNju4zc=", - "_flash_attn2_cpu_f12afc9.abi3.so": "IYJFkR82s6VzuaT0fYBUt3hMetbNEfoyLCQmt18eJIk=", - "_ops.py": "Ru9fdaV75iST1v660HbpH/kA76qidjALxUmNhGcmOz0=", - "bert_padding.py": "gF1EmsdJ+HpQ86MRQ4VxDw+Sb/RVISdQALdNnoByHlw=", - "flash_attn2/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY=", - "flash_attn_interface.py": "W7a9bHR2ly6QL3aKqsBqxWA7M7ZWYcZPmtsPJu6rAfY=", - "layers/__init__.py": "47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU=", - "layers/patch_embed.py": "H58CgME/qSOPTZLOG08wFgrQS1j34pvNwMPrkTj3Ek4=", - "layers/rotary.py": "rFIdNrZQFNawwLD6d7t5LE2m+M93wwDZ8EuExw+gfD8=", - "ops/__init__.py": "47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU=", - "ops/activations.py": "t5lzNg1In8LP6bKeTnyeMizwqjv27JGbJ6ylPdGvZYg=", - "ops/fused_dense.py": "BqT8VTSCGdLJAjh6WBMkCRFYrdVodPeoElWBvH6AIzA=", - "ops/layer_norm.py": "zr7NXIm+2mtEynTp1CS0fbFGI2Mqdp41dY4AfDWF6EQ=", - "ops/rms_norm.py": "XEnihcj0a4aSz4LO55m5iKGVn4HKTeKN8TIyHjuDgxI=", - "ops/triton/__init__.py": "AbpHGcgLb+kRsJGnwFEktk7uzpZOCcBY74+YBdrKVGs=", - "ops/triton/cross_entropy.py": "y56nmuOtpSNxM7sDZcUSgXKIsIEayfuvrU8bAVog/nw=", - "ops/triton/k_activations.py": "+Z3vIyO4JkqBMipKsPvhzmxljtBdIhJCsl/M+/ESqBo=", - "ops/triton/layer_norm.py": "BBQdJiqopbIS12viUqzx3SvOQCNbDSAYLJXhpyMGyC8=", - "ops/triton/linear.py": "OtRvKz8xdpl+7v3q/ZTaS9fdBt9XrzMyapgRr50uBbM=", - "ops/triton/mlp.py": "uWA+D7HyuRUsOT+4nhi3vvCDNZnNEXwYcHMQ02UREfQ=", - "ops/triton/rotary.py": "0YSrlbep0lJmzLpTiANs0slTFeGcQQqR+el6uzbSLvI=" - } - } -} \ No newline at end of file diff --git a/build/torch212-cxx11-cpu-x86_64-linux/metadata.json.sigstore b/build/torch212-cxx11-cpu-x86_64-linux/metadata.json.sigstore deleted file mode 100644 index d9eeab724f8bf5caa83a7b4e57f74d0ef2cce8e4..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cpu-x86_64-linux/metadata.json.sigstore +++ /dev/null @@ -1 +0,0 @@ -{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json","verificationMaterial":{"certificate":{"rawBytes":"MIIHSzCCBtKgAwIBAgIUf6xfRlv9WKH/fMDNtZIBrlVr7rowCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwNjIzMTU0OTEzWhcNMjYwNjIzMTU1OTEzWjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEBjW+po/EOJPOtAyyRpH4rXoy8xoBs9FRJvvUpZ6aNUOw9VXJvJEIw4EyXbsJyn+oajtmRa7Rw2mUXbTEtxWthKOCBfEwggXtMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUZedgEqLb2AXqrPCRPqjthAa3NwswHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoZjEyYWZjOTJjNDBiYmQ3NmQ5OTBkOTExZWJhYmYyOGU0OTBhN2JkZTATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoZjEyYWZjOTJjNDBiYmQ3NmQ5OTBkOTExZWJhYmYyOGU0OTBhN2JkZTAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoZjEyYWZjOTJjNDBiYmQ3NmQ5OTBkOTExZWJhYmYyOGU0OTBhN2JkZTAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKGYxMmFmYzkyYzQwYmJkNzZkOTkwZDkxMWViYWJmMjhlNDkwYTdiZGUwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMjgwMzYyNjg5MzUvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBiwYKKwYBBAHWeQIEAgR9BHsAeQB3AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABnvUr3WsAAAQDAEgwRgIhAIXK/LbSdScSeowDsn95l/EVjPjAvveRmvKjTkCBuFvwAiEAq1uZ6fVR/JG+LgBB1jQrJjnJocLSBHFLr921BF+Bm7AwCgYIKoZIzj0EAwMDZwAwZAIwZyx8hKDzRzlquc/WUBsgQz6j/3i/7k/W5NkqsZCKyttJNchzcw93OlnLN6J69cGMAjAWcaa938o3z0LhYt7PjMHCuCbKJ0mabihbjOh3IzC+BvSEq0Ai9QyNgcxXrbcUhVA="},"tlogEntries":[{"logIndex":"1928479721","logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="},"kindVersion":{"kind":"hashedrekord","version":"0.0.1"},"integratedTime":"1782229753","inclusionPromise":{"signedEntryTimestamp":"MEUCIEbl5WnryFHiB22NV8hB8o32BDtTOCVhrXwAr+kHZGC9AiEA3Tfld2ZruT+bVZp38qXjQmRK8GUD2z05z8md1s3OzPw="},"inclusionProof":{"logIndex":"1806575459","rootHash":"Nnw1/UMCylaiqNENyUbMBERkdFbrLKsI9RMNASjhnVA=","treeSize":"1806575504","hashes":["T0QeXuXMwFi8YvpdvLRj/nI+NnBbVxPNAOqpipe5hF4=","WYu80oKczFlg8Rna610u1qyWhIWRn6Dk9zmnQSug0cI=","A8sLSt5vKCuBnFIOxGBWvconxIYXM8eVdNBVDET7DQI=","iA43TuD75ddhmpTrxjHmyPQfjA1XZ8oJ1VlCuY9Hyxc=","7tIwuDm0arifI+EFnnin0JCJtaxZOqhFfoI7Y3WBJRU=","+BMUSUrkoJHTbmhXyk/7feuWpYv0Q7j+n7k1Km9FoOM=","6upHB/KVez3f+imFNj2HNvhBBOg7M/3kaeOyWHKTyt8=","YYXwJqO7ZXeZ/v7VvkX0gCEGZvh/TtevMtX7gorj1Z4=","oAtde2NUMdGJ/erAg4oMNMz6CPIhuVMayYtTUY1WeGg=","V5gig8R2JeCm7rMID/aKGPW9/XTTywm3LyvVmOXyKWU=","FzVnFZVWkh/v9DZ4ulW796qT3F1KfQgNLEMT9xdl+5A=","BeJEZvGhIqI3T3bd4SKopQeAtyUUubTj1394a27605U=","TSzTqGkRADCpuZ5PwHaPV9gpzMTICOfnAYVjz2hJjQY=","sYnu6tOuM31ewt4jLBjWIwrC5Sd9suVDOXASLgbkP4M=","Y4xfXQrKMXTYkmle88lw9ofERyxOD79P+2S/XJkZfJ8=","V9c1wGJ9L8DpcbTK6ljbIoNaakhzwNhY/80E6RI15pM=","OMfDh4k9YD0E+t4sjqJsh+vg8D1ueo1cC07exyxmr6s=","+vHKm4wNwGJ46ceyZ8mePD7Vn1SediStNHCStRoTSWQ=","5HwWjuGNDisetbiWM0cWIXM65TIM3EJlqk13ntCQlcU=","NDIRlqxO0IaGr94q86MIU7VCAN3yEEgOHS4loya5lu4=","daxmZaajRpZV+JxHiOYZhJBiSKN5ucqjh2WnGbHhirw=","DOCeoSMovIvLExkhIvisow9AuNXgeWs4ECkyR6EcqYU="],"checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n1806575504\nNnw1/UMCylaiqNENyUbMBERkdFbrLKsI9RMNASjhnVA=\n\n— rekor.sigstore.dev wNI9ajBEAiAKd41wCB57E+435UAaJ2vPmWjPh10UJ8pgSA34/O2d6QIgJWOtw0LF+jTVVo7qwyzwny11V+bYk5kRejX1KMw1uXM=\n"}},"canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiIzMTIzYTlkNTIxNDA4YjAyOGMwZWIwMGJiOWExOTlhNWZmYzkxZjZmZjhmMDJiNmU3MTgxNjYzNTVmYmY3Y2E2In19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FWUNJUUNGczg2blBjcS93WEpETkxoWWlsQ014bmlEeHVocXlsN1RCaUh0SlBHTTZRSWhBSUx0S2NQZ3ZMbGlLMmlYTEVDUjBydTUvaXlyNGlDc3NoNVZ5WUQ3YVlSRyIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRla05EUW5STFowRjNTVUpCWjBsVlpqWjRabEpzZGpsWFMwZ3ZaazFFVG5SYVNVSnliRlp5TjNKdmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDVxU1hwTlZGVXdUMVJGZWxkb1kwNU5hbGwzVG1wSmVrMVVWVEZQVkVWNlYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZDYWxjcmNHOHZSVTlLVUU5MFFYbDVVbkJJTkhKWWIzazRlRzlDY3psR1VrcDJkbFVLY0ZvMllVNVZUM2M1VmxoS2RrcEZTWGMwUlhsWVluTktlVzRyYjJGcWRHMVNZVGRTZHpKdFZWaGlWRVYwZUZkMGFFdFBRMEptUlhkbloxaDBUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlZhWldSbkNrVnhUR0l5UVZoeGNsQkRVbEJ4YW5Sb1FXRXpUbmR6ZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOWFha1Y1V1ZkYWFrOVVTbXBPUkVKcFdXMVJNMDV0VVRWUFZFSnJDazlVUlhoYVYwcG9XVzFaZVU5SFZUQlBWRUpvVGpKS2ExcFVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMXBxUlhsWlYxcHFUMVJLYWs1RVFtbFpiVkV6VG0xUk5VOVVRbXRQVkVWNFdsZEthRmx0V1hrS1QwZFZNRTlVUW1oT01rcHJXbFJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMXBxUlhsWlYxcHFUMVJLYWs1RVFtbFpiVkV6VG0xUk5VOVVRbXNLVDFSRmVGcFhTbWhaYlZsNVQwZFZNRTlVUW1oT01rcHJXbFJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwZFplRTF0Um0wS1dYcHJlVmw2VVhkWmJVcHJUbnBhYTA5VWEzZGFSR3Q0VFZkV2FWbFhTbTFOYW1oc1RrUnJkMWxVWkdsYVIxVjNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOYW1kM1RYcFplVTVxWnpWTmVsVjJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwZDFsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTT1VKSWMwRUtaVkZDTTBGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtNTJWWEl6VjNOQlFVRlJSQXBCUldkM1VtZEphRUZKV0VzdlRHSlRaRk5qVTJWdmQwUnpiamsxYkM5RlZtcFFha0YyZG1WU2JYWkxhbFJyUTBKMVJuWjNRV2xGUVhFeGRWbzJabFpTQ2k5S1J5dE1aMEpDTVdwUmNrcHFia3B2WTB4VFFraEdUSEk1TWpGQ1JpdENiVGRCZDBObldVbExiMXBKZW1vd1JVRjNUVVJhZDBGM1drRkpkMXA1ZURnS2FFdEVlbEo2YkhGMVl5OVhWVUp6WjFGNk5tb3ZNMmt2TjJzdlZ6Vk9hM0Z6V2tOTGVYUjBTazVqYUhwamR6a3pUMnh1VEU0MlNqWTVZMGROUVdwQlZ3cGpZV0U1TXpodk0zb3dUR2haZERkUWFrMUlRM1ZEWWt0S01HMWhZbWxvWW1wUGFETkpla01yUW5aVFJYRXdRV2s1VVhsT1oyTjRXSEppWTFWb1ZrRTlDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}],"timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyzADAgEAMIICwgYJKoZIhvcNAQcCoIICszCCAq8CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQg7pehXzDeh3DWJbjIKI6Ry9cfXVp/oeSBGlMOgp2E3BYCFQCXZQUrDCyVVls4LutwXYN0tHKeqhgPMjAyNjA2MjMxNTQ5MTNaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHcMIIB2AIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MDYyMzE1NDkxM1owLwYJKoZIhvcNAQkEMSIEIPALpze9xNrecyCd1AB5vNRNR5uklFRTssAgRuWWUeC8MIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRoMGYCMQCRIhoNh4SkXkTxx0Jpy1uBtlVKBzDKBZCVStFUiPctQdjbbzq/02Sb5aG7kYLLAHsCMQCW3ejDQeg7FsQIW1AdcqmwhCug8AZ3ZUauHnB6a49Mov9NNs6HQuDE8h9oYl7gufE="}]}},"messageSignature":{"messageDigest":{"algorithm":"SHA2_256","digest":"MSOp1SFAiwKMDrALuaGZpf/JH2/48CtucYFmNV+/fKY="},"signature":"MEYCIQCFs86nPcq/wXJDNLhYilCMxniDxuhqyl7TBiHtJPGM6QIhAILtKcPgvLliK2iXLECR0ru5/iyr4iCssh5VyYD7aYRG"}} \ No newline at end of file diff --git a/build/torch212-cxx11-cpu-x86_64-linux/ops/__init__.py b/build/torch212-cxx11-cpu-x86_64-linux/ops/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/build/torch212-cxx11-cpu-x86_64-linux/ops/activations.py b/build/torch212-cxx11-cpu-x86_64-linux/ops/activations.py deleted file mode 100644 index 7c09649fc41e12d5a360c5672825d8380bc7ec80..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cpu-x86_64-linux/ops/activations.py +++ /dev/null @@ -1,135 +0,0 @@ -# Copied from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/model/layers/activations.py -import math - -import torch -import torch.nn as nn -import torch.nn.functional as F - -# 1/sqrt(2*pi)-> 0.3989423 -# 1/sqrt(2) -> 0.70710678 -# sqrt(2/pi) -> 0.79788456 - -# this function is tanh approximation of gelu -# actual gelu is: -# x * 0.5 * (1.0 + torch.erf(x * 0.70710678)) -@torch.jit.script -def bias_gelu(y, bias): - x = bias + y - return (x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))).to(dtype=y.dtype) - - -# gradient of tanh approximation of gelu -# gradient of actual gelu is: -# 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x) -@torch.jit.script -def bias_gelu_back(g, y, bias): - """Assume that y has shape (B, D) and bias has shape (D)""" - x = bias + y - tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)) - # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243 - ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( - 1 + tanh_out - ) - grad_y = ff * g - return grad_y.to(dtype=y.dtype), grad_y.sum(dim=(0), dtype=bias.dtype) - - -class GeLUFunction(torch.autograd.Function): - @staticmethod - # bias is an optional argument - def forward(ctx, input, bias): - ctx.save_for_backward(input, bias) - return bias_gelu(input, bias) - - @staticmethod - def backward(ctx, grad_output): - input, bias = ctx.saved_tensors - tmp = bias_gelu_back(grad_output, input, bias) - return tmp, tmp - - -bias_gelu_impl = GeLUFunction.apply - -# this function is tanh approximation of gelu -# actual gelu is: -# x * 0.5 * (1.0 + torch.erf(x * 0.70710678)) -@torch.jit.script -def gelu_fwd(x): - return (x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))).to(dtype=x.dtype) - - -# gradient of tanh approximation of gelu -# gradient of actual gelu is: -# 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x) -@torch.jit.script -def gelu_bwd(g, x): - tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)) - # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243 - ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( - 1 + tanh_out - ) - return (ff * g).to(dtype=x.dtype) - - -class FastGeLUFunction(torch.autograd.Function): - @staticmethod - # bias is an optional argument - def forward(ctx, input): - ctx.save_for_backward(input) - return gelu_fwd(input) - - @staticmethod - def backward(ctx, grad_output): - (input,) = ctx.saved_tensors - tmp = gelu_bwd(grad_output, input) - return tmp - - -fast_gelu_impl = FastGeLUFunction.apply - - -@torch.jit.script -def relu_bwd(g, x): - return torch.where(x >= 0, g, 0.0).to(dtype=x.dtype) - - -@torch.jit.script -def sqrelu_fwd(x): - r = F.relu(x) - return (r * r).to(dtype=x.dtype) - - -@torch.jit.script -def sqrelu_bwd(g, x): - return (2.0 * g * F.relu(x)).to(dtype=x.dtype) - - -swiglu_fwd_codestring = """ -template T swiglu_fwd(T x, T y) { - return float(x) * float(y) / (1.0f + ::exp(-float(x))); -} -""" -swiglu_bwd_codestring = """ -template void swiglu_bwd(T x, T y, T g, T& dx, T& dy) { - float x_sigmoid = 1.0f / (1.0f + ::exp(-float(x))); - dx = x_sigmoid * (1 + float(x) * (1.0f - x_sigmoid)) * float(g) * float(y); - dy = float(x) * x_sigmoid * float(g); -} -""" -swiglu_fwd = torch.cuda.jiterator._create_jit_fn(swiglu_fwd_codestring) -swiglu_bwd = torch.cuda.jiterator._create_multi_output_jit_fn(swiglu_bwd_codestring, num_outputs=2) - - -class SwiGLUFunction(torch.autograd.Function): - - @staticmethod - def forward(ctx, x, y): - ctx.save_for_backward(x, y) - return swiglu_fwd(x, y) - - @staticmethod - def backward(ctx, dout): - x, y = ctx.saved_tensors - return swiglu_bwd(x, y, dout) - -swiglu = SwiGLUFunction.apply diff --git a/build/torch212-cxx11-cpu-x86_64-linux/ops/fused_dense.py b/build/torch212-cxx11-cpu-x86_64-linux/ops/fused_dense.py deleted file mode 100644 index 6b4033d134e4093fe278f7b3f8c7d3128ce9f36d..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cpu-x86_64-linux/ops/fused_dense.py +++ /dev/null @@ -1,688 +0,0 @@ -# Copyright (c) 2023, Tri Dao. -# Inspired by https://github.com/NVIDIA/apex/blob/master/apex/fused_dense/fused_dense.py -# We make it work with pytorch amp and with bfloat16. -# The TensorParallel linear modules are inspired by https://github.com/NVIDIA/apex/blob/master/apex/transformer/tensor_parallel/layers.py -from functools import partial -from typing import Optional - -# import fused_dense_cuda # from apex -import fused_dense_lib as fused_dense_cuda -import torch -import torch.nn as nn -import torch.nn.functional as F -from torch import Tensor -from torch.distributed import ProcessGroup - -from flash_attn.utils.torch import custom_fwd, custom_bwd -from flash_attn.ops.activations import gelu_bwd, relu_bwd, sqrelu_bwd, sqrelu_fwd -from flash_attn.utils.distributed import ( - all_gather_raw, - all_reduce, - all_reduce_raw, - reduce_scatter, - reduce_scatter_raw, -) - - -class FusedDenseFunc(torch.autograd.Function): - @staticmethod - @custom_fwd - def forward( - ctx, x, weight, bias, return_residual=False, process_group=None, sequence_parallel=True - ): - """ - If process_group is not None and sequence_parallel=True, we're doing Tensor Parallel - with sequence parallelism: we do an all_gather_raw of x before doing the matmul. - """ - ctx.compute_weight_gradient = weight.requires_grad - ctx.return_residual = return_residual - ctx.process_group = process_group - ctx.sequence_parallel = sequence_parallel - - if torch.is_autocast_enabled(): - x = x.to(dtype=torch.get_autocast_gpu_dtype()) - x = x.contiguous() - if process_group is not None and sequence_parallel: - # We want to kick off the all_gather early, before weight dtype conversion - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - else: - total_x = x - - if torch.is_autocast_enabled(): - weight = weight.to(dtype=torch.get_autocast_gpu_dtype()) - bias = bias.to(dtype=torch.get_autocast_gpu_dtype()) if bias is not None else None - weight = weight.contiguous() - if process_group is not None and sequence_parallel: - handle_x.wait() - batch_shape, n = total_x.shape[:-1], total_x.shape[-1] - batch_dim = batch_shape.numel() - # https://github.com/pytorch/pytorch/blob/5b51849b48a7dbccd297286cc0110def4706f9e7/aten/src/ATen/native/cuda/Blas.cpp#L174 - if min(batch_dim, n, *weight.shape) > 65535 * 32: - raise RuntimeError("fused_dense only supports matrix dims <= 2M") - output = F.linear(total_x, weight, bias) - if ctx.compute_weight_gradient: - ctx.save_for_backward(x, weight) - else: - ctx.save_for_backward(weight) - return output if not return_residual else (output, x) - - @staticmethod - @custom_bwd - def backward(ctx, grad_output, *args): - grad_output = grad_output.contiguous() - if ctx.return_residual: - (grad_input,) = args - grad_input = grad_input.contiguous() - process_group = ctx.process_group - sequence_parallel = ctx.sequence_parallel - if ctx.compute_weight_gradient: - x, weight = ctx.saved_tensors - if process_group is not None and sequence_parallel: - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - else: - total_x = x - else: - (weight,) = ctx.saved_tensors - total_x = None - batch_shape = grad_output.shape[:-1] - batch_dim = batch_shape.numel() - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - if ctx.needs_input_grad[0]: - if not ctx.return_residual: - grad_input = F.linear(grad_output, weight.t()) - else: - grad_input = torch.addmm( - grad_input.reshape(batch_dim, grad_input.shape[-1]), grad_output, weight - ) - grad_input = grad_input.reshape(*batch_shape, grad_input.shape[-1]) - if process_group is not None: - reduce_fn = reduce_scatter_raw if sequence_parallel else all_reduce_raw - grad_input, handle_grad_input = reduce_fn(grad_input, process_group, async_op=True) - else: - grad_input = None - if ctx.needs_input_grad[1]: - assert ctx.compute_weight_gradient - if process_group is not None and sequence_parallel: - handle_x.wait() - grad_weight, grad_bias = fused_dense_cuda.linear_bias_wgrad( - total_x.reshape(batch_dim, total_x.shape[-1]), grad_output, ctx.needs_input_grad[2] - ) - else: - grad_weight = None - grad_bias = grad_output if ctx.needs_input_grad[2] else None - if process_group is not None and ctx.needs_input_grad[0]: - handle_grad_input.wait() - return grad_input, grad_weight, grad_bias, None, None, None - - -def fused_dense_func( - x: Tensor, - weight: Tensor, - bias: Optional[Tensor] = None, - return_residual: bool = False, - process_group: Optional[ProcessGroup] = None, - sequence_parallel: bool = True, -): - dtype_eligible = x.dtype in [torch.float16, torch.bfloat16] or ( - x.dtype == torch.float32 and torch.is_autocast_enabled() - ) - if x.is_cuda and weight.is_cuda and (bias is None or bias.is_cuda) and dtype_eligible: - return FusedDenseFunc.apply( - x, weight, bias, return_residual, process_group, sequence_parallel - ) - else: - assert process_group is None - out = F.linear(x, weight, bias) - return out if not return_residual else (out, x) - - -class FusedDense(nn.Linear): - def __init__( - self, - in_features: int, - out_features: int, - bias: bool = True, - return_residual: bool = False, - device=None, - dtype=None, - ) -> None: - super().__init__(in_features, out_features, bias=bias, device=device, dtype=dtype) - self.return_residual = return_residual - - def forward(self, x, process_group=None): - """ - If process_group is not None, we're doing Tensor Parallel with sequence parallelism: - we do an all_gather of x before doing the matmul. - """ - return fused_dense_func( - x, - self.weight, - self.bias, - return_residual=self.return_residual, - process_group=process_group, - ) - - -class ColumnParallelLinear(nn.Linear): - def __init__( - self, - in_features: int, - out_features: int, - process_group: ProcessGroup, - bias: bool = True, - sequence_parallel=True, - multiple_of=1, - device=None, - dtype=None, - ) -> None: - world_size = torch.distributed.get_world_size(process_group) - if out_features % multiple_of: - raise ValueError(f"out_features ({out_features}) must be a multiple of {multiple_of}") - multiple = out_features // multiple_of - # We want to split @multiple across world_size, but it could be an uneven split - div = multiple // world_size - mod = multiple % world_size - # The first @mod ranks get @div + 1 copies, the rest get @div copies - local_multiple = div + int(torch.distributed.get_rank(process_group) < mod) - super().__init__( - in_features, local_multiple * multiple_of, bias=bias, device=device, dtype=dtype - ) - self.process_group = process_group - self.sequence_parallel = sequence_parallel - - def forward(self, x): - # If self.sequence_parallel is True, we're doing Tensor Parallel with sequence parallelism: - # we do an all_gather of x before doing the matmul. - # If not, then the input is already gathered. - return fused_dense_func( - x, - self.weight, - self.bias, - process_group=self.process_group, - sequence_parallel=self.sequence_parallel, - ) - - -class RowParallelLinear(nn.Linear): - def __init__( - self, - in_features: int, - out_features: int, - process_group: ProcessGroup, - bias: bool = True, - sequence_parallel=True, - multiple_of=1, - device=None, - dtype=None, - ) -> None: - world_size = torch.distributed.get_world_size(process_group) - rank = torch.distributed.get_rank(process_group) - if in_features % multiple_of: - raise ValueError(f"in_features ({in_features}) must be a multiple of {multiple_of}") - multiple = in_features // multiple_of - # We want to split @multiple across world_size, but it could be an uneven split - div = multiple // world_size - mod = multiple % world_size - # The first @mod ranks get @div + 1 copies, the rest get @div copies - local_multiple = div + int(torch.distributed.get_rank(process_group) < mod) - # Only rank 0 will have bias - super().__init__( - local_multiple * multiple_of, - out_features, - bias=bias and rank == 0, - device=device, - dtype=dtype, - ) - self.process_group = process_group - self.sequence_parallel = sequence_parallel - - def forward(self, x): - """ - We're doing Tensor Parallel with sequence parallelism: we do the matmul and then - a reduce_scatter of the result. - """ - out = fused_dense_func(x, self.weight, self.bias) - reduce_fn = reduce_scatter if self.sequence_parallel else all_reduce - return reduce_fn(out, self.process_group) - - -class FusedMLPFunc(torch.autograd.Function): - @staticmethod - @custom_fwd - def forward( - ctx, - x, - weight1, - bias1, - weight2, - bias2, - activation="gelu_approx", - save_pre_act=True, - return_residual=False, - checkpoint_lvl=0, - heuristic=0, - process_group=None, - sequence_parallel=True, - ): - """ - If process_group is not None and sequence_parallel=True, we're doing Tensor Parallel - with sequence parallelism: we do an all_gather of x before doing the matmul. - If sequence_parallel=False, then the input is already gathered. - - checkpoint_lvl: - 0: no recomputation in the bwd - 1: recompute gelu_out / relu_out in the bwd - 2: recompute pre_act and gelu_out / relu_out in the bwd - """ - assert -1 <= heuristic <= 4 - assert activation in ["gelu_approx", "relu", "sqrelu"] - if activation == "sqrelu": - assert heuristic == -1 - if not save_pre_act: - checkpoint_lvl = 2 - assert checkpoint_lvl in [0, 1, 2] - ctx.return_residual = return_residual - ctx.process_group = process_group - ctx.sequence_parallel = sequence_parallel - ctx.checkpoint_lvl = checkpoint_lvl - ctx.activation = activation - ctx.heuristic = heuristic - - if torch.is_autocast_enabled(): - x = x.to(dtype=torch.get_autocast_gpu_dtype()) - x = x.contiguous() - if process_group is not None and sequence_parallel: - # We want to kick off the all_gather early, before weight dtype conversion - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - else: - total_x = x - - if torch.is_autocast_enabled(): - dtype = torch.get_autocast_gpu_dtype() - weight1, weight2 = [a.to(dtype=dtype) for a in [weight1, weight2]] - bias1 = bias1.to(dtype=dtype) if bias1 is not None else None - bias2 = bias2.to(dtype=dtype) if bias2 is not None else None - weight1 = weight1.contiguous() - bias1 = bias1.contiguous() if bias1 is not None else None - weight2 = weight2.contiguous() - bias2 = bias2.contiguous() if bias2 is not None else None - if process_group is not None and sequence_parallel: - handle_x.wait() - batch_shape, n = total_x.shape[:-1], total_x.shape[-1] - batch_dim = batch_shape.numel() - # https://github.com/pytorch/pytorch/blob/5b51849b48a7dbccd297286cc0110def4706f9e7/aten/src/ATen/native/cuda/Blas.cpp#L174 - if min(batch_dim, n, *weight1.shape, *weight2.shape) > 65535 * 32: - raise RuntimeError("fused_dense only supports matrix dims <= 2M") - if heuristic == -1: - pre_act = F.linear(total_x, weight1, bias1) - activation_fn = ( - partial(F.gelu, approximate="tanh") - if activation == "gelu_approx" - else (sqrelu_fwd if activation == "sqrelu" else F.relu) - ) - with torch.jit.fuser("fuser2"): - output1 = activation_fn(pre_act) - # This is before adding bias1 - # pre_act = F.linear(total_x.reshape(batch_dim, n), weight1) - # with torch.jit.fuser('fuser2'): - # output1 = bias_gelu(pre_act, bias1) - else: - is_gelu = activation == "gelu_approx" - output1, *rest = fused_dense_cuda.linear_act_forward( - total_x.reshape(batch_dim, n), weight1, bias1, is_gelu, save_pre_act, heuristic - ) - if save_pre_act: - pre_act = rest[0] - output2 = F.linear(output1, weight2, bias2) - if checkpoint_lvl == 0 or (checkpoint_lvl == 1 and activation == "relu"): - # For RELU the pre_act is very small (just a bit-mask) so we just save it - ctx.save_for_backward(x, weight1, weight2, pre_act, output1) - elif checkpoint_lvl == 1: - ctx.save_for_backward(x, weight1, weight2, pre_act) - elif checkpoint_lvl == 2: - ctx.save_for_backward(x, weight1, weight2, bias1) - output2 = output2.reshape(*batch_shape, output2.shape[-1]) - return output2 if not return_residual else (output2, x) - - @staticmethod - @custom_bwd - def backward(ctx, grad_output, *args): - grad_output = grad_output.contiguous() - checkpoint_lvl = ctx.checkpoint_lvl - activation = ctx.activation - activation_fn = ( - partial(F.gelu, approximate="tanh") - if activation == "gelu_approx" - else (sqrelu_fwd if activation == "sqrelu" else F.relu) - ) - if ctx.return_residual: - (grad_input,) = args - grad_input = grad_input.contiguous() - process_group = ctx.process_group - sequence_parallel = ctx.sequence_parallel - x, weight1, weight2, *rest = ctx.saved_tensors - if process_group is None or not sequence_parallel: - total_x = x - batch_shape = grad_output.shape[:-1] - batch_dim = batch_shape.numel() - if checkpoint_lvl in [0, 1]: - if process_group is not None and sequence_parallel: - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - if checkpoint_lvl == 0 or (checkpoint_lvl == 1 and activation == "relu"): - pre_act, output1 = rest - elif checkpoint_lvl == 1: - (pre_act,) = rest - with torch.jit.fuser("fuser2"): - output1 = activation_fn(pre_act) - elif checkpoint_lvl == 2: - (bias1,) = rest - if process_group is not None and sequence_parallel: - total_x, _ = all_gather_raw(x, process_group) - if ctx.heuristic == -1: - pre_act = F.linear(total_x, weight1, bias1) - with torch.jit.fuser("fuser2"): - output1 = activation_fn(pre_act) - else: - output1, pre_act = fused_dense_cuda.linear_act_forward( - total_x.reshape(batch_dim, total_x.shape[-1]), - weight1, - bias1, - activation == "gelu_approx", - True, - ctx.heuristic, - ) - - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - output1 = output1.reshape(batch_dim, output1.shape[-1]) - pre_act = pre_act.reshape(batch_dim, pre_act.shape[-1]) - if ctx.needs_input_grad[3]: - grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad( - output1, grad_output, ctx.needs_input_grad[4] - ) - else: - grad_weight2 = None - grad_bias2 = grad_output if ctx.needs_input_grad[4] else None - if ctx.heuristic == -1: - # grad_pre_act = matmul_dgelu(grad_output, weight2, pre_act) - grad_output1 = F.linear(grad_output, weight2.t()) - activation_grad_fn = ( - gelu_bwd - if activation == "gelu_approx" - else (sqrelu_bwd if activation == "sqrelu" else relu_bwd) - ) - with torch.jit.fuser("fuser2"): - grad_pre_act = activation_grad_fn(grad_output1, pre_act) - else: - # The cublasLt epilogue has to compute both gelu/relu grad and bias grad, we can't - # just compute gelu/relu grad - grad_pre_act, grad_bias1 = fused_dense_cuda.bias_act_linear_dgrad_bgrad( - weight2, grad_output, pre_act, activation == "gelu_approx", ctx.heuristic - ) - if not ctx.needs_input_grad[2]: - grad_bias1 = None - if ctx.needs_input_grad[0]: - if not ctx.return_residual: - grad_input = F.linear(grad_pre_act, weight1.t()) - else: - grad_input = torch.addmm( - grad_input.reshape(batch_dim, grad_input.shape[-1]), grad_pre_act, weight1 - ) - grad_input = grad_input.reshape(*batch_shape, grad_input.shape[-1]) - if process_group is not None: - reduce_fn = reduce_scatter_raw if sequence_parallel else all_reduce_raw - grad_input, handle_grad_input = reduce_fn(grad_input, process_group, async_op=True) - else: - grad_input = None - if ctx.heuristic == -1: - if ctx.needs_input_grad[1]: - if process_group is not None and sequence_parallel and checkpoint_lvl != 2: - handle_x.wait() - grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_wgrad( - total_x.reshape(batch_dim, total_x.shape[-1]), - grad_pre_act, - ctx.needs_input_grad[2], - ) - else: - grad_weight1 = None - grad_bias1 = grad_pre_act if ctx.needs_input_grad[2] else None - else: - if ctx.needs_input_grad[1]: - if process_group is not None and sequence_parallel and checkpoint_lvl != 2: - handle_x.wait() - grad_weight1 = F.linear( - grad_pre_act.t(), total_x.reshape(batch_dim, total_x.shape[-1]).t() - ) - else: - grad_weight1 = None - if process_group is not None and ctx.needs_input_grad[0]: - handle_grad_input.wait() - return ( - grad_input, - grad_weight1, - grad_bias1, - grad_weight2, - grad_bias2, - None, - None, - None, - None, - None, - None, - None, - ) - - -def fused_mlp_func( - x: Tensor, - weight1: Tensor, - weight2: Tensor, - bias1: Optional[Tensor] = None, - bias2: Optional[Tensor] = None, - activation: str = "gelu_approx", - save_pre_act: bool = True, - return_residual: bool = False, - checkpoint_lvl: int = 0, - heuristic: int = 0, - process_group: Optional[ProcessGroup] = None, - sequence_parallel: bool = True, -): - assert activation in ["gelu_approx", "relu", "sqrelu"] - dtype_eligible = x.dtype in [torch.float16, torch.bfloat16] or ( - x.dtype == torch.float32 and torch.is_autocast_enabled() - ) - # If we save pre-activation, dimension must be divisible by 128 (relu) or 8 (gelu) - dim_eligible = not save_pre_act or (x.shape[-1] % (128 if activation == "relu" else 8) == 0) - if ( - x.is_cuda - and weight1.is_cuda - and weight2.is_cuda - and (bias1 is None or bias1.is_cuda) - and (bias2 is None or bias2.is_cuda) - and dtype_eligible - and dim_eligible - ): - return FusedMLPFunc.apply( - x, - weight1, - bias1, - weight2, - bias2, - activation, - save_pre_act, - return_residual, - checkpoint_lvl, - heuristic, - process_group, - sequence_parallel, - ) - else: - assert process_group is None - pre_act = F.linear(x, weight1, bias1) - activation_fn = ( - partial(F.gelu, approximate="tanh") - if activation == "gelu_approx" - else partial(F.relu, inplace=True) - ) - output1 = activation_fn(pre_act) - output2 = F.linear(output1, weight2, bias2) - return output2 if not return_residual else (output2, x) - - -class FusedMLP(nn.Module): - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - bias1=True, - bias2=True, - activation="gelu_approx", - return_residual=False, - checkpoint_lvl=0, - heuristic="auto", - device=None, - dtype=None, - ): - """ - If process_group is not None, we're doing Tensor Parallel with sequence parallelism: - we do an all_gather of x before doing the matmul, gelu, then matmul. - Finally we do a reduce_scatter of the output. - - checkpoint_lvl (increasing lvl means slower but more memory saving): - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute pre_act and gelu_out in the bwd - heuristic: - -1: don't fuse gemm + gelu (separate kernel) - 0..4: use this heuristic for the algo section in the fused gemm + gelu - 'auto': heuristic will be picked automatically: - For CUDA >= 11.8, we set heuristic=0 for both fp16 and bf16 for best perf. - For CUDA <= 11.7, we set heuristic=1 for fp16 and heuristic=-1 for bf16. - For H100, we set heuristic=-1 for both fp16 and bf16 as the fused cuBlasLt implementation - is slower than the unfused version. - return_residual: whether to return the input x along with the output. This is for - performance reason: for post-norm architecture, returning the input allows us - to fuse the backward of nn.Linear with the residual connection. - """ - assert checkpoint_lvl in [0, 1, 2] - assert activation in ["gelu_approx", "relu", "sqrelu"] - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features * 4 - self.activation = activation - self.return_residual = return_residual - self.checkpoint_lvl = checkpoint_lvl - self.heuristic = heuristic if activation != "sqrelu" else -1 - self.fc1 = nn.Linear(in_features, hidden_features, bias=bias1, **factory_kwargs) - self.fc2 = nn.Linear(hidden_features, out_features, bias=bias2, **factory_kwargs) - - def forward(self, x, process_group=None): - dtype = x.dtype if not torch.is_autocast_enabled() else torch.get_autocast_gpu_dtype() - if self.heuristic == "auto": - if self.activation == "gelu_approx": - if torch.cuda.get_device_capability("cuda") == (9, 0): - heuristic = -1 - else: - cuda_ver = tuple(map(int, torch.version.cuda.split("."))) - heuristic = 0 if cuda_ver >= (11, 8) else (1 if dtype == torch.float16 else -1) - else: - heuristic = 0 - else: - heuristic = self.heuristic - out = fused_mlp_func( - x, - self.fc1.weight, - self.fc2.weight, - self.fc1.bias, - self.fc2.bias, - activation=self.activation, - save_pre_act=self.training, - return_residual=self.return_residual, - checkpoint_lvl=self.checkpoint_lvl, - heuristic=heuristic, - process_group=process_group, - ) - if self.return_residual: - out, x = out - if process_group is not None: - out = reduce_scatter(out, process_group) - return out if not self.return_residual else (out, x) - - -class ParallelFusedMLP(nn.Module): - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - activation="gelu_approx", - process_group: ProcessGroup = None, - bias1=True, - bias2=True, - sequence_parallel=True, - checkpoint_lvl=0, - heuristic="auto", - device=None, - dtype=None, - ): - """ - process_group is required. We're doing Tensor Parallel with sequence parallelism: - we do an all_gather of x before doing the matmul, gelu, then matmul. - Finally we do a reduce_scatter of the output. - - checkpoint_lvl (increasing lvl means slower but more memory saving): - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute pre_act and gelu_out in the bwd - heuristic: - -1: don't fuse gemm + gelu (separate kernel) - 0..4: use this heuristic for the algo section in the fused gemm + gelu - 'auto': heuristic will be picked automatically: - For CUDA >= 11.8, we set heuristic=0 for both fp16 and bf16 for best perf. - For CUDA <= 11.7, we set heuristic=1 for fp16 and heuristic=-1 for bf16. - """ - assert checkpoint_lvl in [0, 1, 2] - assert activation in ["gelu_approx", "relu", "sqrelu"] - assert process_group is not None - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features * 4 - self.activation = activation - self.process_group = process_group - self.sequence_parallel = sequence_parallel - self.checkpoint_lvl = checkpoint_lvl - self.heuristic = heuristic if activation != "sqrelu" else -1 - self.fc1 = ColumnParallelLinear( - in_features, hidden_features, process_group, bias=bias1, **factory_kwargs - ) - self.fc2 = RowParallelLinear( - hidden_features, out_features, process_group, bias=bias2, **factory_kwargs - ) - - def forward(self, x): - dtype = x.dtype if not torch.is_autocast_enabled() else torch.get_autocast_gpu_dtype() - if self.heuristic == "auto": - if self.activation == "gelu_approx": - cuda_ver = tuple(map(int, torch.version.cuda.split("."))) - heuristic = 0 if cuda_ver >= (11, 8) else (1 if dtype == torch.float16 else -1) - else: - heuristic = 0 - else: - heuristic = self.heuristic - out = fused_mlp_func( - x, - self.fc1.weight, - self.fc2.weight, - self.fc1.bias, - self.fc2.bias, - activation=self.activation, - save_pre_act=self.training, - checkpoint_lvl=self.checkpoint_lvl, - heuristic=heuristic, - process_group=self.process_group, - sequence_parallel=self.sequence_parallel, - ) - reduce_fn = reduce_scatter if self.sequence_parallel else all_reduce - return reduce_fn(out, self.process_group) diff --git a/build/torch212-cxx11-cpu-x86_64-linux/ops/layer_norm.py b/build/torch212-cxx11-cpu-x86_64-linux/ops/layer_norm.py deleted file mode 100644 index 4b6cd798fd02844ef9cd3897f8ab95e490e638bf..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cpu-x86_64-linux/ops/layer_norm.py +++ /dev/null @@ -1,800 +0,0 @@ -# Copyright (c) 2022, Tri Dao. -# Adapted from https://github.com/NVIDIA/apex/blob/master/apex/contrib/layer_norm/layer_norm.py - -import dropout_layer_norm -import torch -from torch.nn import init - - -def maybe_align(x, alignment_in_bytes=16): - """Assume that x already has last dim divisible by alignment_in_bytes""" - # TD [2023-07-04] I'm not 100% sure that clone will align the memory - # https://discuss.pytorch.org/t/how-to-ensure-that-tensor-data-ptr-is-aligned-to-16-bytes/183440 - return x if x.data_ptr() % alignment_in_bytes == 0 else x.clone() - - -def _dropout_add_layer_norm_forward( - x0, - residual, - gamma, - beta, - rowscale, - colscale, - dropout_p, - epsilon, - residual_in_fp32=False, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes""" - hidden_size = gamma.numel() - x0mat = x0.view((-1, hidden_size)) - residualmat = residual.view((-1, hidden_size)) if residual is not None else None - rowscale = rowscale.view(-1) if rowscale is not None else None - zmat, xmat, dmask, mu, rsigma = dropout_layer_norm.dropout_add_ln_fwd( - x0mat, - residualmat, - gamma, - beta, - rowscale, - colscale, - None, - None, - dropout_p, - epsilon, - 1.0, - 0, - None, - residual_in_fp32, - is_rms_norm, - ) - # dmask is None if dropout_p == 0.0 - # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype - return zmat, xmat if xmat is not None else x0mat, dmask, mu, rsigma - - -def _dropout_add_layer_norm_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - rowscale, - colscale, - dropout_p, - has_residual, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes - dx == None means that it was a post-norm architecture - (x = drop(x0) + residual was not returned in the fwd). - x0 must not be None if we have colscale. - """ - hidden_size = gamma.numel() - xmat = x.view((-1, hidden_size)) - dzmat = dz.view(xmat.shape) - dxmat = dx.view(xmat.shape) if dx is not None else None - x0mat = x0.view((-1, hidden_size)) if x0 is not None else None - rowscale = rowscale.view(-1) if rowscale is not None else None - if colscale is not None: - assert x0 is not None, "x0 is required to compute the gradient of colscale" - dx0mat, dresidualmat, dgamma, dbeta, _, _, *rest = dropout_layer_norm.dropout_add_ln_bwd( - dzmat, - dxmat, - xmat, - x0mat, - dmask, - mu, - rsigma, - gamma, - rowscale, - colscale, - None, - None, - dropout_p, - 1.0, - 0, - has_residual, - is_rms_norm, - ) - # dresidualmat is None if not has_residual - if colscale is None: - return dx0mat, dresidualmat, dgamma, dbeta - else: - dcolscale = rest[0] - return dx0mat, dresidualmat, dgamma, dbeta, dcolscale - - -def _dropout_add_layer_norm_subset_forward( - x0, - residual, - gamma, - beta, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32=False, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes""" - hidden_size = gamma.numel() - x0mat = x0.view((-1, hidden_size)) - residualmat = residual.view((-1, hidden_size)) if residual is not None else None - x0_subset = x0_subset.view(-1) if x0_subset is not None else None - out_subset = out_subset.view(-1) if out_subset is not None else None - zmat, xmat, dmask, mu, rsigma = dropout_layer_norm.dropout_add_ln_fwd( - x0mat, - residualmat, - gamma, - beta, - None, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - None, - residual_in_fp32, - is_rms_norm, - ) - # dmask is None if dropout_p == 0.0 - # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype - return zmat, xmat if xmat is not None else x0mat, dmask, mu, rsigma - - -def _dropout_add_layer_norm_subset_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - colscale, - x0_subset, - out_subset, - dropout_p, - rowscale_const, - x0_numrows, - has_residual, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes - dx == None means that it was a post-norm architecture - (x = drop(x0) + residual was not returned in the fwd). - x0 must not be None if we have colscale. - """ - hidden_size = gamma.numel() - xmat = x.view((-1, hidden_size)) - dzmat = dz.view(-1, hidden_size) - dxmat = dx.view(xmat.shape) if dx is not None else None - x0mat = x0.view((-1, hidden_size)) if x0 is not None else None - x0_subset = x0_subset.view(-1) if x0_subset is not None else None - out_subset = out_subset.view(-1) if out_subset is not None else None - if colscale is not None: - assert x0 is not None, "x0 is required to compute the gradient of colscale" - dx0mat, dresidualmat, dgamma, dbeta, _, _, *rest = dropout_layer_norm.dropout_add_ln_bwd( - dzmat, - dxmat, - xmat, - x0mat, - dmask, - mu, - rsigma, - gamma, - None, - colscale, - x0_subset, - out_subset, - dropout_p, - rowscale_const, - x0_numrows, - has_residual, - is_rms_norm, - ) - # dresidualmat is None if not has_residual - if colscale is None: - return dx0mat, dresidualmat, dgamma, dbeta - else: - dcolscale = rest[0] - return dx0mat, dresidualmat, dgamma, dbeta, dcolscale - - -def _dropout_add_layer_norm_parallel_residual_forward( - x0, - x1, - residual, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - residual_in_fp32=False, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes""" - hidden_size = gamma0.numel() - x0mat = x0.view((-1, hidden_size)) - x1mat = x1.view((-1, hidden_size)) if x1 is not None else None - residualmat = residual.view((-1, hidden_size)) if residual is not None else None - ( - z0mat, - z1mat, - xmat, - dmask0, - dmask1, - mu, - rsigma, - ) = dropout_layer_norm.dropout_add_ln_parallel_residual_fwd( - x0mat, - x1mat, - residualmat, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - None, - residual_in_fp32, - is_rms_norm, - ) - # dmask0 and dmask1 are None if dropout_p == 0.0 - # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype - return z0mat, z1mat, xmat if xmat is not None else x0mat, dmask0, dmask1, mu, rsigma - - -def _dropout_add_layer_norm_parallel_residual_backward( - dz0, - dz1, - dx, - x, - dmask0, - dmask1, - mu, - rsigma, - gamma0, - gamma1, - dropout_p, - has_x1, - has_residual, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes - dx == None means that it was a post-norm architecture - (x = drop(x0) + residual was not returned in the fwd). - """ - hidden_size = gamma0.numel() - xmat = x.view((-1, hidden_size)) - dz0mat = dz0.view(xmat.shape) - dz1mat = dz1.view(xmat.shape) if dz1 is not None else None - dxmat = dx.view(xmat.shape) if dx is not None else None - ( - dx0mat, - dx1mat, - dresidualmat, - dgamma0, - dbeta0, - dgamma1, - dbeta1, - *rest, - ) = dropout_layer_norm.dropout_add_ln_parallel_residual_bwd( - dz0mat, - dz1mat, - dxmat, - xmat, - dmask0, - dmask1, - mu, - rsigma, - gamma0, - gamma1, - dropout_p, - has_x1, - has_residual, - is_rms_norm, - ) - # dresidualmat is None if not has_residual - return dx0mat, dx1mat, dresidualmat, dgamma0, dbeta0, dgamma1, dbeta1 - - -class DropoutAddLayerNormFn(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x0, - residual, - gamma, - beta, - rowscale, - colscale, - dropout_p, - epsilon, - residual_in_fp32=False, - prenorm=False, - is_rms_norm=False, - return_dmask=False, - ): - x0 = maybe_align(x0.contiguous(), 16) - residual = maybe_align(residual.contiguous(), 16) if residual is not None else None - gamma = maybe_align(gamma.contiguous(), 16) - beta = maybe_align(beta.contiguous(), 16) if beta is not None else None - rowscale = maybe_align(rowscale.contiguous(), 16) if rowscale is not None else None - colscale = maybe_align(colscale.contiguous(), 16) if colscale is not None else None - zmat, xmat, dmask, mu, rsigma = _dropout_add_layer_norm_forward( - x0, - residual, - gamma, - beta, - rowscale, - colscale, - dropout_p, - epsilon, - residual_in_fp32, - is_rms_norm, - ) - # Only need to save x0 if we need to compute gradient wrt colscale - x0_saved = x0 if colscale is not None else None - ctx.save_for_backward( - xmat.view(x0.shape), x0_saved, dmask, gamma, mu, rsigma, rowscale, colscale - ) - ctx.prenorm = prenorm - ctx.dropout_p = dropout_p - ctx.has_residual = residual is not None - ctx.is_rms_norm = is_rms_norm - ctx.has_beta = beta is not None - if not return_dmask: - return ( - zmat.view(x0.shape) if not prenorm else (zmat.view(x0.shape), xmat.view(x0.shape)) - ) - else: - dmask = ( - dmask.view(x0.shape) - if dropout_p > 0.0 - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - ctx.mark_non_differentiable(dmask) - return ( - (zmat.view(x0.shape), dmask) - if not prenorm - else (zmat.view(x0.shape), xmat.view(x0.shape), dmask) - ) - - @staticmethod - def backward(ctx, dz, *args): - # assert dz.is_contiguous() - dz = maybe_align(dz.contiguous(), 16) # this happens! - dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None - x, x0, dmask, gamma, mu, rsigma, rowscale, colscale = ctx.saved_tensors - # x0 is None if colscale is None - dropout_p = ctx.dropout_p - has_residual = ctx.has_residual - dx0mat, dresidualmat, dgamma, dbeta, *rest = _dropout_add_layer_norm_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - rowscale, - colscale, - dropout_p, - has_residual, - ctx.is_rms_norm, - ) - dx0 = dx0mat.view(x.shape) - dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None - dcolscale = rest[0] if colscale is not None else None - return ( - dx0, - dresidual, - dgamma, - dbeta if ctx.has_beta else None, - None, - dcolscale, - None, - None, - None, - None, - None, - None, - ) - - -class DropoutAddLayerNormSubsetFn(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x0, - residual, - gamma, - beta, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32=False, - prenorm=False, - is_rms_norm=False, - return_dmask=False, - ): - x0 = maybe_align(x0.contiguous(), 16) - residual = maybe_align(residual.contiguous(), 16) if residual is not None else None - gamma = maybe_align(gamma.contiguous(), 16) - beta = maybe_align(beta.contiguous(), 16) if beta is not None else None - colscale = maybe_align(colscale.contiguous(), 16) if colscale is not None else None - zmat, xmat, dmask, mu, rsigma = _dropout_add_layer_norm_subset_forward( - x0, - residual, - gamma, - beta, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32, - is_rms_norm, - ) - # Only need to save x0 if we need to compute gradient wrt colscale - x0_saved = x0 if colscale is not None else None - x_shape = (-1, *x0.shape[1:]) - ctx.save_for_backward( - xmat.view(x_shape), x0_saved, dmask, gamma, mu, rsigma, colscale, x0_subset, out_subset - ) - ctx.prenorm = prenorm - ctx.dropout_p = dropout_p - ctx.rowscale_const = rowscale_const - ctx.x0_numrows = x0.shape[:-1].numel() - ctx.has_residual = residual is not None - ctx.is_rms_norm = is_rms_norm - ctx.has_beta = beta is not None - z_shape = (-1, *x0.shape[1:]) - if not return_dmask: - return zmat.view(z_shape) if not prenorm else (zmat.view(z_shape), xmat.view(x0.shape)) - else: - z = zmat.view(z_shape) - dmask = ( - dmask.view(x0.shape) - if dropout_p > 0.0 - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - ctx.mark_non_differentiable(dmask) - return (z, dmask) if not prenorm else (z, xmat.view(x_shape), dmask) - - @staticmethod - def backward(ctx, dz, *args): - # assert dz.is_contiguous() - dz = maybe_align(dz.contiguous(), 16) # this happens! - dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None - x, x0, dmask, gamma, mu, rsigma, colscale, x0_subset, out_subset = ctx.saved_tensors - # x0 is None if colscale is None - dropout_p = ctx.dropout_p - has_residual = ctx.has_residual - dx0mat, dresidualmat, dgamma, dbeta, *rest = _dropout_add_layer_norm_subset_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - colscale, - x0_subset, - out_subset, - dropout_p, - ctx.rowscale_const, - ctx.x0_numrows, - has_residual, - ctx.is_rms_norm, - ) - dx0 = dx0mat.view(-1, *x.shape[1:]) - dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None - dcolscale = rest[0] if colscale is not None else None - return ( - dx0, - dresidual, - dgamma, - dbeta if ctx.has_beta else None, - dcolscale, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) - - -class DropoutAddLayerNormParallelResidualFn(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x0, - x1, - residual, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - residual_in_fp32=False, - prenorm=False, - is_rms_norm=False, - return_dmask=False, - ): - x0 = maybe_align(x0.contiguous(), 16) - x1 = maybe_align(x1.contiguous(), 16) if x1 is not None else None - residual = maybe_align(residual.contiguous(), 16) if residual is not None else None - gamma0 = maybe_align(gamma0.contiguous(), 16) - beta0 = maybe_align(beta0.contiguous(), 16) if beta0 is not None else None - gamma1 = maybe_align(gamma1.contiguous(), 16) if gamma1 is not None else None - beta1 = maybe_align(beta1.contiguous(), 16) if beta1 is not None else None - ( - z0mat, - z1mat, - xmat, - dmask0, - dmask1, - mu, - rsigma, - ) = _dropout_add_layer_norm_parallel_residual_forward( - x0, - x1, - residual, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - residual_in_fp32, - is_rms_norm, - ) - ctx.save_for_backward(xmat.view(x0.shape), dmask0, dmask1, gamma0, gamma1, mu, rsigma) - ctx.prenorm = prenorm - ctx.dropout_p = dropout_p - ctx.has_x1 = x1 is not None - ctx.has_residual = residual is not None - ctx.is_rms_norm = is_rms_norm - ctx.has_beta = beta0 is not None - z = (z0mat.view(x0.shape), z1mat.view(x0.shape) if z1mat is not None else None) - if not return_dmask: - return z if not prenorm else (*z, xmat.view(x0.shape)) - else: - dmask0 = ( - dmask0.view(x0.shape) - if dropout_p > 0.0 - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - dmask1 = ( - dmask1.view(x0.shape) - if dropout_p > 0.0 and x1 is not None - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - ctx.mark_non_differentiable(dmask0) - ctx.mark_non_differentiable(dmask1) - return ( - (*z, dmask0, dmask1) if not prenorm else (*z, xmat.view(x0.shape), dmask0, dmask1) - ) - - @staticmethod - def backward(ctx, dz0, dz1, *args): - dz0 = maybe_align(dz0.contiguous(), 16) # this happens! - dz1 = maybe_align(dz1.contiguous(), 16) if dz1 is not None else None - dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None - x, dmask0, dmask1, gamma0, gamma1, mu, rsigma = ctx.saved_tensors - dropout_p = ctx.dropout_p - has_x1 = ctx.has_x1 - has_residual = ctx.has_residual - ( - dx0mat, - dx1mat, - dresidualmat, - dgamma0, - dbeta0, - dgamma1, - dbeta1, - ) = _dropout_add_layer_norm_parallel_residual_backward( - dz0, - dz1, - dx, - x, - dmask0, - dmask1, - mu, - rsigma, - gamma0, - gamma1, - dropout_p, - has_x1, - has_residual, - ctx.is_rms_norm, - ) - dx0 = dx0mat.view(x.shape) - dx1 = dx1mat.view(x.shape) if dx1mat is not None else None - dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None - return ( - dx0, - dx1, - dresidual, - dgamma0, - dbeta0 if ctx.has_beta else None, - dgamma1, - dbeta1 if ctx.has_beta else None, - None, - None, - None, - None, - None, - None, - ) - - -def layer_norm(x, weight, bias, epsilon): - return DropoutAddLayerNormFn.apply(x, None, weight, bias, None, None, 0.0, epsilon, False) - - -def dropout_add_layer_norm( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - rowscale=None, - layerscale=None, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormFn.apply( - x0, - residual, - weight, - bias, - rowscale, - layerscale, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - False, - return_dropout_mask, - ) - - -def dropout_add_layer_norm_subset( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - layerscale=None, - x0_subset=None, - out_subset=None, - rowscale_const=1.0, - out_numrows=0, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormSubsetFn.apply( - x0, - residual, - weight, - bias, - layerscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32, - prenorm, - False, - return_dropout_mask, - ) - - -def dropout_add_layer_norm_parallel_residual( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormParallelResidualFn.apply( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - False, - return_dropout_mask, - ) - - -class DropoutAddLayerNorm(torch.nn.Module): - def __init__( - self, - hidden_size, - prenorm=False, - p=0.0, - eps=1e-5, - residual_in_fp32=False, - device=None, - dtype=None, - ): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.prenorm = prenorm - self.p = p - self.eps = eps - self.residual_in_fp32 = residual_in_fp32 - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.bias = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.reset_parameters() - - def reset_parameters(self): - init.ones_(self.weight) - init.zeros_(self.bias) - - def forward(self, x0, residual=None): - return dropout_add_layer_norm( - x0, - residual, - self.weight, - self.bias, - self.p if self.training else 0.0, - self.eps, - prenorm=self.prenorm, - residual_in_fp32=self.residual_in_fp32, - ) diff --git a/build/torch212-cxx11-cpu-x86_64-linux/ops/rms_norm.py b/build/torch212-cxx11-cpu-x86_64-linux/ops/rms_norm.py deleted file mode 100644 index 068348d61290e3839dd082b540d898578ba1e8e2..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cpu-x86_64-linux/ops/rms_norm.py +++ /dev/null @@ -1,174 +0,0 @@ -# Copyright (c) 2022, Tri Dao. -# Adapted from https://github.com/NVIDIA/apex/blob/master/apex/contrib/layer_norm/layer_norm.py - -import torch -from torch.nn import init - -from flash_attn.ops.layer_norm import ( - DropoutAddLayerNormFn, - DropoutAddLayerNormParallelResidualFn, - DropoutAddLayerNormSubsetFn, -) - - -def rms_norm(x, weight, epsilon): - return DropoutAddLayerNormFn.apply( - x, None, weight, None, None, None, 0.0, epsilon, False, False, True - ) - - -def dropout_add_rms_norm( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - rowscale=None, - layerscale=None, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormFn.apply( - x0, - residual, - weight, - bias, - rowscale, - layerscale, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - True, - return_dropout_mask, - ) - - -def dropout_add_rms_norm_subset( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - layerscale=None, - x0_subset=None, - out_subset=None, - rowscale_const=1.0, - out_numrows=0, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormSubsetFn.apply( - x0, - residual, - weight, - bias, - layerscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32, - prenorm, - True, - return_dropout_mask, - ) - - -def dropout_add_rms_norm_parallel_residual( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormParallelResidualFn.apply( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - True, - return_dropout_mask, - ) - - -class RMSNorm(torch.nn.Module): - def __init__(self, hidden_size, eps=1e-5, device=None, dtype=None): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.eps = eps - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.register_parameter("bias", None) - self.reset_parameters() - - def reset_parameters(self): - init.ones_(self.weight) - - def forward(self, x): - return rms_norm(x, self.weight, self.eps) - - -class DropoutAddRMSNorm(torch.nn.Module): - def __init__( - self, - hidden_size, - prenorm=False, - p=0.0, - eps=1e-5, - residual_in_fp32=False, - device=None, - dtype=None, - ): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.prenorm = prenorm - self.p = p - self.eps = eps - self.residual_in_fp32 = residual_in_fp32 - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.register_parameter("bias", None) - self.reset_parameters() - - def reset_parameters(self): - init.ones_(self.weight) - - def forward(self, x0, residual=None): - return dropout_add_rms_norm( - x0, - residual, - self.weight, - None, - self.p if self.training else 0.0, - self.eps, - prenorm=self.prenorm, - residual_in_fp32=self.residual_in_fp32, - ) diff --git a/build/torch212-cxx11-cpu-x86_64-linux/ops/triton/__init__.py b/build/torch212-cxx11-cpu-x86_64-linux/ops/triton/__init__.py deleted file mode 100644 index 8b137891791fe96927ad78e64b0aad7bded08bdc..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cpu-x86_64-linux/ops/triton/__init__.py +++ /dev/null @@ -1 +0,0 @@ - diff --git a/build/torch212-cxx11-cpu-x86_64-linux/ops/triton/cross_entropy.py b/build/torch212-cxx11-cpu-x86_64-linux/ops/triton/cross_entropy.py deleted file mode 100644 index 1b5a415b73f236f3e05fb14b9141959559e18526..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cpu-x86_64-linux/ops/triton/cross_entropy.py +++ /dev/null @@ -1,330 +0,0 @@ -# Copyright (c) 2023, Tri Dao. - -from typing import Tuple, Optional, Union - -import torch -import torch.nn.functional as F - -import triton -import triton.language as tl - -# `all_gather_into_tensor` and `reduce_scatter_tensor` are new placeholders for -# `_all_gather_base` and `_reduce_scatter_base`. They require the most recent -# version of PyTorch. The following 2 lines are for backward compatibility with -# older PyTorch. -if "all_gather_into_tensor" not in dir(torch.distributed): - torch.distributed.all_gather_into_tensor = torch.distributed._all_gather_base - - -@triton.heuristics( - { - "HAS_SMOOTHING": lambda args: args["smoothing"] > 0.0, - } -) -@triton.jit -def cross_entropy_fwd_kernel( - loss_ptr, # data ptrs - lse_ptr, - z_loss_ptr, - logits_ptr, - labels_ptr, - smoothing, - logit_scale, - lse_square_scale, - ignore_index, - total_classes, - class_start_idx, # Useful for tensor parallel when each rank only has a subset of classes - n_cols, # shapes - logits_row_stride, # strides - BLOCK_SIZE: tl.constexpr, - HAS_SMOOTHING: tl.constexpr, - # if SPLIT (e.g. tensor parallel), don't include the LSE in the loss since it's not the final LSE - SPLIT: tl.constexpr, - PRECOMPUTED_LSE: tl.constexpr, # If LSE is already computed (also no smoothing and logit_scale == 1.0) -): - row_idx = tl.program_id(0) - logits_ptr = logits_ptr + row_idx * logits_row_stride.to(tl.int64) - sum_logits = 0.0 # For smoothing - if not PRECOMPUTED_LSE: - # Statistics for online softmax - m_i = -float("inf") - l_i = 0.0 - for col_offset in range(0, n_cols, BLOCK_SIZE): - cols = col_offset + tl.arange(0, BLOCK_SIZE) - logits = tl.load(logits_ptr + cols, mask=cols < n_cols, other=-float("inf")).to( - tl.float32 - ) * logit_scale - if HAS_SMOOTHING: - sum_logits += tl.sum(tl.where(cols < n_cols, logits, 0.0)) - m_i_new = tl.maximum(m_i, tl.max(logits)) - l_i = tl.exp(m_i - m_i_new) * l_i + tl.sum(tl.exp(logits - m_i_new)) - m_i = m_i_new - lse = tl.log(l_i) + m_i - tl.store(lse_ptr + row_idx, lse) - else: - lse = tl.load(lse_ptr + row_idx) - label_idx = tl.load(labels_ptr + row_idx) - if label_idx == ignore_index: - loss = 0.0 - z_loss = 0.0 - else: - label_idx -= class_start_idx - if label_idx >= 0 and label_idx < n_cols: - logits_label = tl.load(logits_ptr + label_idx) * logit_scale - if HAS_SMOOTHING: - loss = ( - (lse if not SPLIT else 0.0) - - smoothing * sum_logits / total_classes - - (1 - smoothing) * logits_label - ) - else: - loss = (lse if not SPLIT else 0.0) - logits_label - else: - # If label is out of bounds, we set the CE loss to 0.0. But we still want the smoothing loss - if HAS_SMOOTHING: - loss = smoothing * ((lse if not SPLIT else 0.0) - sum_logits / total_classes) - else: - loss = 0.0 - if not SPLIT: - z_loss = lse_square_scale * lse * lse - loss += z_loss - else: - z_loss = 0.0 - tl.store(loss_ptr + row_idx, loss) - if not SPLIT: - tl.store(z_loss_ptr + row_idx, z_loss) - - -@triton.heuristics( - { - "HAS_SMOOTHING": lambda args: args["smoothing"] > 0.0, - } -) -@triton.jit -def cross_entropy_bwd_kernel( - dlogits_ptr, # data ptrs - dloss_ptr, - logits_ptr, - lse_ptr, - labels_ptr, - smoothing, - logit_scale, - lse_square_scale, - ignore_index, - total_classes, - class_start_idx, # Useful for tensor parallel when each rank only has a subset of classes - n_cols, # shapes - logits_row_stride, # strides - dlogits_row_stride, - dloss_row_stride, - BLOCK_SIZE: tl.constexpr, - HAS_SMOOTHING: tl.constexpr, -): - row_idx = tl.program_id(0) - col_block_idx = tl.program_id(1) - logits_ptr = logits_ptr + row_idx * logits_row_stride.to(tl.int64) - dlogits_ptr = dlogits_ptr + row_idx * dlogits_row_stride.to(tl.int64) - col_offsets = col_block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - label_idx = tl.load(labels_ptr + row_idx) - if label_idx != ignore_index: - dloss = tl.load(dloss_ptr + row_idx * dloss_row_stride) - else: - dloss = 0.0 - logits = tl.load(logits_ptr + col_offsets, mask=col_offsets < n_cols, other=-float("inf")).to( - tl.float32 - ) * logit_scale - lse = tl.load(lse_ptr + row_idx) - probs = tl.exp(logits - lse) - probs += 2.0 * lse_square_scale * lse * probs - label_idx -= class_start_idx - if HAS_SMOOTHING: - smooth_positive = 1.0 - smoothing - smooth_negative = smoothing / total_classes - probs = tl.where(col_offsets == label_idx, probs - smooth_positive, probs) - smooth_negative - else: - probs = tl.where(col_offsets == label_idx, probs - 1.0, probs) - tl.store(dlogits_ptr + col_offsets, (dloss * logit_scale) * probs, mask=col_offsets < n_cols) - - -class CrossEntropyLoss(torch.autograd.Function): - - @staticmethod - def forward( - ctx, - logits, - labels, - precomputed_lse=None, - smoothing=0.0, - logit_scale=1.0, - lse_square_scale=0.0, - ignore_index=-100, - inplace_backward=False, - process_group=None, - ): - # For some reason Triton generates wrong code when labels has dtype long and its address - # is not aligned to 16 bytes. The ld.global.b64 seems to load the wrong label index. - if labels.dtype == torch.long and labels.data_ptr() % 16 != 0: - labels = F.pad(labels, (0, 1))[..., :-1] - assert labels.data_ptr() % 16 == 0 - assert logit_scale > 0.0 - n_rows, n_cols = logits.shape - assert labels.shape == (n_rows,) - world_size = 1 if process_group is None else torch.distributed.get_world_size(process_group) - total_classes = world_size * n_cols - rank = 0 if process_group is None else torch.distributed.get_rank(process_group) - class_start_idx = rank * n_cols - use_precomputed_lse = precomputed_lse is not None and logit_scale == 1.0 and smoothing == 0.0 - - if logits.stride(-1) != 1: - logits = logits.contiguous() - MAX_BLOCK_SIZE = 16 * 1024 - BLOCK_SIZE = min(triton.next_power_of_2(n_cols), MAX_BLOCK_SIZE) - num_warps = ( - 4 - if BLOCK_SIZE < 2048 - else (8 if BLOCK_SIZE < 8192 else (16 if BLOCK_SIZE < 128 * 1024 else 32)) - ) - losses = torch.empty(n_rows, dtype=torch.float, device=logits.device) - if use_precomputed_lse: - assert precomputed_lse.shape == (n_rows,) - lse = precomputed_lse.contiguous() - else: - lse = torch.empty(n_rows, dtype=torch.float, device=logits.device) - z_losses = torch.empty(n_rows, dtype=torch.float, device=logits.device) - # Need this, otherwise Triton tries to launch from cuda:0 and we get - # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?) - with torch.cuda.device(logits.device.index): - cross_entropy_fwd_kernel[(n_rows,)]( - losses, # data ptrs - lse, - z_losses, - logits, - labels, - smoothing, - logit_scale, - lse_square_scale, - ignore_index, - total_classes, - class_start_idx, - n_cols, # shapes - logits.stride(0), # strides - BLOCK_SIZE=BLOCK_SIZE, # constants - SPLIT=world_size > 1, - PRECOMPUTED_LSE=use_precomputed_lse, - num_warps=num_warps, - ) - - if world_size > 1: - # If there's no smoothing, if labels are in the vocab of this partition, losses contains - # - predicted logit, and 0 otherwise. - # If there's smoothing=0.1, for labels in the vocab of this partition, losses contains - # -0.9 * predicted logit - 0.1 * sum logit / total_classes. - # For labels not in the vocab of this partition, losses contains - # -0.1 * sum logit / total_classes. - if world_size > 1: - lse_allgather = torch.empty(world_size, n_rows, dtype=lse.dtype, device=lse.device) - torch.distributed.all_gather_into_tensor(lse_allgather, lse, group=process_group) - handle_losses = torch.distributed.all_reduce( - losses, op=torch.distributed.ReduceOp.SUM, group=process_group, async_op=True - ) - lse = torch.logsumexp(lse_allgather, dim=0) - handle_losses.wait() - # After the allreduce, if there's no smoothing, the total losses are - predicted_logit, - # we just have to add the (global) lse. - # If there's smoothing=0.1, the total losses are - # -0.9 * predicted_logit - 0.1 * sum logit / total_classes. - # Again, we just have to add the (global) lse. - losses += lse - if lse_square_scale != 0.0: - z_losses = lse_square_scale * lse.square() - z_losses.masked_fill_(labels == ignore_index, 0.0) - losses += z_losses - else: - z_losses = torch.zeros_like(losses) - losses.masked_fill_(labels == ignore_index, 0.0) - - ctx.save_for_backward(logits, lse, labels) - ctx.mark_non_differentiable(z_losses) - ctx.smoothing = smoothing - ctx.logit_scale = logit_scale - ctx.lse_square_scale = lse_square_scale - ctx.ignore_index = ignore_index - ctx.total_classes = total_classes - ctx.class_start_idx = class_start_idx - ctx.inplace_backward = inplace_backward - return losses, z_losses - - @staticmethod - def backward(ctx, grad_losses, grad_z_losses): - del grad_z_losses # z_losses are only for logging. - - logits, lse, labels = ctx.saved_tensors - dlogits = logits if ctx.inplace_backward else torch.empty_like(logits) - n_rows, n_cols = logits.shape - BLOCK_SIZE = min(triton.next_power_of_2(n_cols), 4 * 1024) - num_warps = 4 if BLOCK_SIZE < 2048 else (8 if BLOCK_SIZE < 8192 else 16) - grid = lambda META: (n_rows, triton.cdiv(n_cols, META["BLOCK_SIZE"])) # noqa - # Need this, otherwise Triton tries to launch from cuda:0 and we get - # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?) - with torch.cuda.device(logits.device.index): - cross_entropy_bwd_kernel[grid]( - dlogits, # data ptrs - grad_losses, - logits, - lse, - labels, - ctx.smoothing, - ctx.logit_scale, - ctx.lse_square_scale, - ctx.ignore_index, - ctx.total_classes, - ctx.class_start_idx, - n_cols, # shapes - logits.stride(0), # strides - dlogits.stride(0), - grad_losses.stride(0), - BLOCK_SIZE=BLOCK_SIZE, # constants - num_warps=num_warps, - ) - return dlogits, None, None, None, None, None, None, None, None, None - - -def cross_entropy_loss( - logits: torch.Tensor, - labels: torch.Tensor, - precomputed_lse: Optional[torch.Tensor] = None, - label_smoothing: float = 0.0, - logit_scale: float = 1.0, - lse_square_scale: float = 0.0, - ignore_index=-100, - inplace_backward: bool = False, - process_group=None, -) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Arguments: - logits: (batch, vocab_size) - labels: (batch,) - label_smoothing: float - logit_scale: float. Multiply logits by this scale before calculating the loss. - lse_square_scale: float. If > 0, we add lse_square_scale * lse(logits) ^ 2 to the loss. - This is also referred to as "z-loss". - ignore_index: int. If labels == ignore_index, the loss is set to 0.0. - inplace_backward: bool. If True, we do the backward pass in-place by modifying the logits. - This saves memory. - process_group: if not None, we're doing Tensor Parallel: each process is responsible for - one part of the vocab. The loss will be aggregated across processes. - Returns: - losses: (batch,), float - z_losses: (batch,), float - """ - return CrossEntropyLoss.apply( - logits, - labels, - precomputed_lse, - label_smoothing, - logit_scale, - lse_square_scale, - ignore_index, - inplace_backward, - process_group, - ) diff --git a/build/torch212-cxx11-cpu-x86_64-linux/ops/triton/k_activations.py b/build/torch212-cxx11-cpu-x86_64-linux/ops/triton/k_activations.py deleted file mode 100644 index efb83c358eb4a85d069ee340a3c83f418f9a805b..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cpu-x86_64-linux/ops/triton/k_activations.py +++ /dev/null @@ -1,162 +0,0 @@ -# Adapted from https://github.com/facebookresearch/xformers/blob/main/xformers/triton/k_activations.py -# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved. -# -# This source code is licensed under the BSD license found in the -# LICENSE file in the root directory of this source tree. - -import math -from enum import Enum -from typing import Optional - -import triton -import triton.language as tl - -_sqrt2pi = math.sqrt(2.0 / math.pi) -_sqrt1_2 = math.sqrt(1.0 / 2) -_gaussian_pdf_normalization = 1.0 / math.sqrt(2 * math.pi) - - -class Activation(str, Enum): - SquaredReLU = "squared_relu" - GeLU = "gelu" - GeLUApprox = "gelu_approx" - LeakyReLU = "leaky_relu" - ReLU = "relu" - - -def get_triton_activation_kernel(activation: Optional[Activation]): - return ( - { - Activation.ReLU: relu, - Activation.LeakyReLU: leaky_relu, - Activation.GeLU: gelu, - Activation.GeLUApprox: gelu_approx, - Activation.SquaredReLU: squared_relu, - }[activation] - if activation - else None - ) - - -def get_triton_activation_bwd_kernel(activation: Optional[Activation]): - return ( - { - Activation.ReLU: relu_grad, - Activation.LeakyReLU: leaky_relu_grad, - Activation.GeLU: gelu_grad, - Activation.GeLUApprox: gelu_approx_grad, - Activation.SquaredReLU: squared_relu_grad, - }[activation] - if activation - else None - ) - - -@triton.jit -def tanh(x): - # Tanh is just a scaled sigmoid - return 2 * tl.sigmoid(2 * x) - 1 - - -@triton.jit -def cosh(x): - exp_x = tl.exp(x) - return (exp_x + 1.0 / exp_x) * 0.5 - - -# a Triton implementation of the most used activations -# See for instance http://arxiv.org/abs/1606.08415 for an overview - -# ReLU -@triton.jit -def relu(x): - """ - ReLU_ activation function - - .. _ReLU: https://pytorch.org/docs/stable/generated/torch.nn.ReLU.html - """ - zero = 0.0 - return tl.where(x >= 0, x, zero.to(x.dtype)) - - -@triton.jit -def relu_grad(x): - # ReLU is different from other activations - # in that it does not require the input to retrospectively compute its gradient - # here the input is the downstream gradient, and we return the upstream gradient directly - zero = 0.0 - one = 1.0 - return tl.where(x >= 0, one.to(x.dtype), zero.to(x.dtype)) - - -@triton.jit -def squared_relu(x): - """ - Squared ReLU activation, as proposed in the Primer_ paper. - - .. _Primer: https://arxiv.org/abs/2109.08668 - """ - x_ = relu(x) - return (x_ * x_).to(x.dtype) - - -@triton.jit -def squared_relu_grad(x): - return tl.where(x >= 0, 2.0 * x, 0.0) - - -# Leaky ReLU -@triton.jit -def leaky_relu(x): - """ - LeakyReLU_ activation - - .. _LeakyReLU: https://pytorch.org/docs/stable/generated/torch.nn.LeakyReLU.html - """ - scale = 0.01 + 0.0 - scale = scale.to(x.dtype) - return tl.where(x >= 0, x, scale * x) - - -@triton.jit -def leaky_relu_grad(x): - min_grad = 0.01 - max_grad = 1 - - min_grad = min_grad.to(x.dtype) - max_grad = max_grad.to(x.dtype) - - return tl.where(x >= 0, max_grad, min_grad) - - -@triton.jit -def gelu(x): - """Gaussian Error Linear Unit (GELU)""" - return x * 0.5 * (1.0 + tl.libdevice.erf(x * _sqrt1_2)) - - -@triton.jit -def gelu_grad(x): - cdf = 0.5 * (1.0 + tl.libdevice.erf(x * _sqrt1_2)) - pdf = tl.exp(-0.5 * x * x) * _gaussian_pdf_normalization - return cdf + x * pdf - - -@triton.jit -def gelu_approx(x): - """ - GeLU_ activation - Gaussian error linear unit, with tanh approximation - - .. _GeLU: https://arxiv.org/pdf/1606.08415.pdf - """ - return 0.5 * x * (1.0 + tanh(_sqrt2pi * x * (1.0 + 0.044715 * x * x))) - - -@triton.jit -def gelu_approx_grad(x): - # CREDITS: Fast implementation proposed in - # https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/fused_bias_gelu.py#L30 - tanh_out = tanh(0.79788456 * x * (1 + 0.044715 * x * x)) - return 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( - 1 + tanh_out - ) diff --git a/build/torch212-cxx11-cpu-x86_64-linux/ops/triton/layer_norm.py b/build/torch212-cxx11-cpu-x86_64-linux/ops/triton/layer_norm.py deleted file mode 100644 index 5e6f0b3697b612e6a19031bbc55ac76ad233dc85..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cpu-x86_64-linux/ops/triton/layer_norm.py +++ /dev/null @@ -1,1254 +0,0 @@ -# Copyright (c) 2024, Tri Dao. -# Implement dropout + residual + layer_norm / rms_norm. - -# Based on the Triton LayerNorm tutorial: https://triton-lang.org/main/getting-started/tutorials/05-layer-norm.html -# For the backward pass, we keep weight_grad and bias_grad in registers and accumulate. -# This is faster for dimensions up to 8k, but after that it's much slower due to register spilling. -# The models we train have hidden dim up to 8k anyway (e.g. Llama 70B), so this is fine. - -import math -from typing import Optional, List - -import torch -import torch.nn.functional as F -from torch import Tensor - -import triton -import triton.language as tl - -from flash_attn.utils.torch import custom_fwd, custom_bwd -from flash_attn.utils.library import triton_op - -from ..._ops import add_op_namespace_prefix - - -def maybe_contiguous_lastdim(x): - return x.contiguous() if x is not None and x.stride(-1) != 1 else x - - -def maybe_contiguous(x): - return x.contiguous() if x is not None else None - - -def triton_autotune_configs(): - # Return configs with a valid warp count for the current device - configs = [] - # Maximum threads per block is architecture-dependent in theory, but in reality all are 1024 - max_threads_per_block = 1024 - # Default to warp size 32 if not defined by device - warp_size = getattr(torch.cuda.get_device_properties(torch.cuda.current_device()), "warp_size", 32) - # Autotune for warp counts which are powers of 2 and do not exceed thread per block limit - return [triton.Config({}, num_warps=warp_count) for warp_count in [1, 2, 4, 8, 16, 32] - if warp_count * warp_size <= max_threads_per_block] - # return [triton.Config({}, num_warps=8)] - - -def layer_norm_ref( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - zero_centered_weight=False, - dropout_mask=None, - dropout_mask1=None, - upcast=False, -): - dtype = x.dtype - if upcast: - x = x.float() - weight = weight.float() - bias = bias.float() if bias is not None else None - residual = residual.float() if residual is not None else residual - x1 = x1.float() if x1 is not None else None - weight1 = weight1.float() if weight1 is not None else None - bias1 = bias1.float() if bias1 is not None else None - if zero_centered_weight: - weight = weight + 1.0 - if weight1 is not None: - weight1 = weight1 + 1.0 - if x1 is not None: - assert rowscale is None, "rowscale is not supported with parallel LayerNorm" - if rowscale is not None: - x = x * rowscale[..., None] - if dropout_p > 0.0: - if dropout_mask is not None: - x = x.masked_fill(~dropout_mask, 0.0) / (1.0 - dropout_p) - else: - x = F.dropout(x, p=dropout_p) - if x1 is not None: - if dropout_mask1 is not None: - x1 = x1.masked_fill(~dropout_mask1, 0.0) / (1.0 - dropout_p) - else: - x1 = F.dropout(x1, p=dropout_p) - if x1 is not None: - x = x + x1 - if residual is not None: - x = (x + residual).to(x.dtype) - out = F.layer_norm(x.to(weight.dtype), x.shape[-1:], weight=weight, bias=bias, eps=eps).to( - dtype - ) - if weight1 is None: - return out if not prenorm else (out, x) - else: - out1 = F.layer_norm( - x.to(weight1.dtype), x.shape[-1:], weight=weight1, bias=bias1, eps=eps - ).to(dtype) - return (out, out1) if not prenorm else (out, out1, x) - - -def rms_norm_ref( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - zero_centered_weight=False, - dropout_mask=None, - dropout_mask1=None, - upcast=False, -): - dtype = x.dtype - if upcast: - x = x.float() - weight = weight.float() - bias = bias.float() if bias is not None else None - residual = residual.float() if residual is not None else residual - x1 = x1.float() if x1 is not None else None - weight1 = weight1.float() if weight1 is not None else None - bias1 = bias1.float() if bias1 is not None else None - if zero_centered_weight: - weight = weight + 1.0 - if weight1 is not None: - weight1 = weight1 + 1.0 - if x1 is not None: - assert rowscale is None, "rowscale is not supported with parallel LayerNorm" - if rowscale is not None: - x = x * rowscale[..., None] - if dropout_p > 0.0: - if dropout_mask is not None: - x = x.masked_fill(~dropout_mask, 0.0) / (1.0 - dropout_p) - else: - x = F.dropout(x, p=dropout_p) - if x1 is not None: - if dropout_mask1 is not None: - x1 = x1.masked_fill(~dropout_mask1, 0.0) / (1.0 - dropout_p) - else: - x1 = F.dropout(x1, p=dropout_p) - if x1 is not None: - x = x + x1 - if residual is not None: - x = (x + residual).to(x.dtype) - rstd = 1 / torch.sqrt((x.square()).mean(dim=-1, keepdim=True) + eps) - out = ((x * rstd * weight) + bias if bias is not None else (x * rstd * weight)).to(dtype) - if weight1 is None: - return out if not prenorm else (out, x) - else: - out1 = ((x * rstd * weight1) + bias1 if bias1 is not None else (x * rstd * weight1)).to( - dtype - ) - return (out, out1) if not prenorm else (out, out1, x) - - -@triton.autotune( - configs=triton_autotune_configs(), - key=["N", "HAS_RESIDUAL", "STORE_RESIDUAL_OUT", "IS_RMS_NORM", "HAS_BIAS", "HAS_X1", "HAS_W1", "HAS_B1"], -) -# torch compile doesn't like triton.heuristics, so we set these manually when calling the kernel -# @triton.heuristics({"HAS_BIAS": lambda args: args["B"] is not None}) -# @triton.heuristics({"HAS_RESIDUAL": lambda args: args["RESIDUAL"] is not None}) -# @triton.heuristics({"HAS_X1": lambda args: args["X1"] is not None}) -# @triton.heuristics({"HAS_W1": lambda args: args["W1"] is not None}) -# @triton.heuristics({"HAS_B1": lambda args: args["B1"] is not None}) -@triton.jit -def _layer_norm_fwd_1pass_kernel( - X, # pointer to the input - Y, # pointer to the output - W, # pointer to the weights - B, # pointer to the biases - RESIDUAL, # pointer to the residual - X1, - W1, - B1, - Y1, - RESIDUAL_OUT, # pointer to the residual - ROWSCALE, - SEEDS, # Dropout seeds for each row - DROPOUT_MASK, - DROPOUT_MASK1, - Mean, # pointer to the mean - Rstd, # pointer to the 1/std - stride_x_row, # how much to increase the pointer when moving by 1 row - stride_y_row, - stride_res_row, - stride_res_out_row, - stride_x1_row, - stride_y1_row, - M, # number of rows in X - N, # number of columns in X - eps, # epsilon to avoid division by zero - dropout_p, # Dropout probability - zero_centered_weight, # If true, add 1.0 to the weight - IS_RMS_NORM: tl.constexpr, - BLOCK_N: tl.constexpr, - HAS_RESIDUAL: tl.constexpr, - STORE_RESIDUAL_OUT: tl.constexpr, - HAS_BIAS: tl.constexpr, - HAS_DROPOUT: tl.constexpr, - STORE_DROPOUT_MASK: tl.constexpr, - HAS_ROWSCALE: tl.constexpr, - HAS_X1: tl.constexpr, - HAS_W1: tl.constexpr, - HAS_B1: tl.constexpr, -): - # Map the program id to the row of X and Y it should compute. - row = tl.program_id(0) - X += row * stride_x_row - Y += row * stride_y_row - if HAS_RESIDUAL: - RESIDUAL += row * stride_res_row - if STORE_RESIDUAL_OUT: - RESIDUAL_OUT += row * stride_res_out_row - if HAS_X1: - X1 += row * stride_x1_row - if HAS_W1: - Y1 += row * stride_y1_row - # Compute mean and variance - cols = tl.arange(0, BLOCK_N) - x = tl.load(X + cols, mask=cols < N, other=0.0).to(tl.float32) - if HAS_ROWSCALE: - rowscale = tl.load(ROWSCALE + row).to(tl.float32) - x *= rowscale - if HAS_DROPOUT: - # Compute dropout mask - # 7 rounds is good enough, and reduces register pressure - keep_mask = tl.rand(tl.load(SEEDS + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - x = tl.where(keep_mask, x / (1.0 - dropout_p), 0.0) - if STORE_DROPOUT_MASK: - tl.store(DROPOUT_MASK + row * N + cols, keep_mask, mask=cols < N) - if HAS_X1: - x1 = tl.load(X1 + cols, mask=cols < N, other=0.0).to(tl.float32) - if HAS_ROWSCALE: - rowscale = tl.load(ROWSCALE + M + row).to(tl.float32) - x1 *= rowscale - if HAS_DROPOUT: - # Compute dropout mask - # 7 rounds is good enough, and reduces register pressure - keep_mask = ( - tl.rand(tl.load(SEEDS + M + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - ) - x1 = tl.where(keep_mask, x1 / (1.0 - dropout_p), 0.0) - if STORE_DROPOUT_MASK: - tl.store(DROPOUT_MASK1 + row * N + cols, keep_mask, mask=cols < N) - x += x1 - if HAS_RESIDUAL: - residual = tl.load(RESIDUAL + cols, mask=cols < N, other=0.0).to(tl.float32) - x += residual - if STORE_RESIDUAL_OUT: - tl.store(RESIDUAL_OUT + cols, x, mask=cols < N) - if not IS_RMS_NORM: - mean = tl.sum(x, axis=0) / N - tl.store(Mean + row, mean) - xbar = tl.where(cols < N, x - mean, 0.0) - var = tl.sum(xbar * xbar, axis=0) / N - else: - xbar = tl.where(cols < N, x, 0.0) - var = tl.sum(xbar * xbar, axis=0) / N - rstd = 1 / tl.sqrt(var + eps) - tl.store(Rstd + row, rstd) - # Normalize and apply linear transformation - mask = cols < N - w = tl.load(W + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w += 1.0 - if HAS_BIAS: - b = tl.load(B + cols, mask=mask).to(tl.float32) - x_hat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd - y = x_hat * w + b if HAS_BIAS else x_hat * w - # Write output - tl.store(Y + cols, y, mask=mask) - if HAS_W1: - w1 = tl.load(W1 + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w1 += 1.0 - if HAS_B1: - b1 = tl.load(B1 + cols, mask=mask).to(tl.float32) - y1 = x_hat * w1 + b1 if HAS_B1 else x_hat * w1 - tl.store(Y1 + cols, y1, mask=mask) - - -def _layer_norm_fwd( - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - residual: Optional[Tensor] = None, - x1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - out_dtype: Optional[torch.dtype] = None, - residual_dtype: Optional[torch.dtype] = None, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - return_dropout_mask: bool = False, - out: Optional[Tensor] = None, - residual_out: Optional[Tensor] = None -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - # Need to wrap to handle the case where residual_out is a alias of x, which makes torch.library - # and torch.compile unhappy. Also allocate memory for out and residual_out if they are None - # so that _layer_norm_fwd_impl doesn't have to return them. - if out is None: - out = torch.empty_like(x, dtype=x.dtype if out_dtype is None else out_dtype) - if residual is not None: - residual_dtype = residual.dtype - if residual_out is None and ( - residual is not None - or (residual_dtype is not None and residual_dtype != x.dtype) - or dropout_p > 0.0 - or rowscale is not None - or x1 is not None - ): - residual_out = torch.empty_like( - x, dtype=residual_dtype if residual_dtype is not None else x.dtype - ) - else: - residual_out = None - y1, mean, rstd, seeds, dropout_mask, dropout_mask1 = _layer_norm_fwd_impl( - x, - weight, - bias, - eps, - out, - residual=residual, - x1=x1, - weight1=weight1, - bias1=bias1, - dropout_p=dropout_p, - rowscale=rowscale, - zero_centered_weight=zero_centered_weight, - is_rms_norm=is_rms_norm, - return_dropout_mask=return_dropout_mask, - residual_out=residual_out, - ) - # residual_out is None if residual is None and residual_dtype == input_dtype and dropout_p == 0.0 - if residual_out is None: - residual_out = x - return out, y1, mean, rstd, residual_out, seeds, dropout_mask, dropout_mask1 - - -# [2025-04-28] torch.library.triton_op ignores the schema argument, but here we need the schema -# since we're returning a tuple of tensors -@triton_op(add_op_namespace_prefix("layer_norm_fwd_impl"), mutates_args={"out", "residual_out"}, - schema="(Tensor x, Tensor weight, Tensor bias, float eps, Tensor(a!) out, Tensor? residual, Tensor? x1, Tensor? weight1, Tensor? bias1, float dropout_p, Tensor? rowscale, bool zero_centered_weight, bool is_rms_norm, bool return_dropout_mask, Tensor(a!)? residual_out) -> (Tensor y1, Tensor mean, Tensor rstd, Tensor seeds, Tensor dropout_mask, Tensor dropout_mask1)") -def _layer_norm_fwd_impl( - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - out: Tensor, - residual: Optional[Tensor] = None, - x1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - return_dropout_mask: bool = False, - residual_out: Optional[Tensor] = None -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - M, N = x.shape - assert x.stride(-1) == 1 - if residual is not None: - assert residual.stride(-1) == 1 - assert residual.shape == (M, N) - assert weight.shape == (N,) - assert weight.stride(-1) == 1 - if bias is not None: - assert bias.stride(-1) == 1 - assert bias.shape == (N,) - if x1 is not None: - assert x1.shape == x.shape - assert rowscale is None - assert x1.stride(-1) == 1 - if weight1 is not None: - assert weight1.shape == (N,) - assert weight1.stride(-1) == 1 - if bias1 is not None: - assert bias1.shape == (N,) - assert bias1.stride(-1) == 1 - if rowscale is not None: - assert rowscale.is_contiguous() - assert rowscale.shape == (M,) - assert out.shape == x.shape - assert out.stride(-1) == 1 - if residual_out is not None: - assert residual_out.shape == x.shape - assert residual_out.stride(-1) == 1 - if weight1 is not None: - y1 = torch.empty_like(out) - assert y1.stride(-1) == 1 - else: - y1 = None - mean = torch.empty((M,), dtype=torch.float32, device=x.device) if not is_rms_norm else None - rstd = torch.empty((M,), dtype=torch.float32, device=x.device) - if dropout_p > 0.0: - seeds = torch.randint( - 2**32, (M if x1 is None else 2 * M,), device=x.device, dtype=torch.int64 - ) - else: - seeds = None - if return_dropout_mask and dropout_p > 0.0: - dropout_mask = torch.empty(M, N, device=x.device, dtype=torch.bool) - if x1 is not None: - dropout_mask1 = torch.empty(M, N, device=x.device, dtype=torch.bool) - else: - dropout_mask1 = None - else: - dropout_mask, dropout_mask1 = None, None - # Less than 64KB per feature: enqueue fused kernel - MAX_FUSED_SIZE = 65536 // x.element_size() - BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) - if N > BLOCK_N: - raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") - with torch.cuda.device(x.device.index): - torch.library.wrap_triton(_layer_norm_fwd_1pass_kernel)[(M,)]( - x, - out, - weight, - bias, - residual, - x1, - weight1, - bias1, - y1, - residual_out, - rowscale, - seeds, - dropout_mask, - dropout_mask1, - mean, - rstd, - x.stride(0), - out.stride(0), - residual.stride(0) if residual is not None else 0, - residual_out.stride(0) if residual_out is not None else 0, - x1.stride(0) if x1 is not None else 0, - y1.stride(0) if y1 is not None else 0, - M, - N, - eps, - dropout_p, - # Passing bool make torch inductor very unhappy since it then tries to compare to int_max - int(zero_centered_weight), - is_rms_norm, - BLOCK_N, - residual is not None, - residual_out is not None, - bias is not None, - dropout_p > 0.0, - dropout_mask is not None, - rowscale is not None, - HAS_X1=x1 is not None, - HAS_W1=weight1 is not None, - HAS_B1=bias1 is not None, - ) - return y1, mean, rstd, seeds, dropout_mask, dropout_mask1 - - -@triton.autotune( - configs=triton_autotune_configs(), - key=["N", "HAS_DRESIDUAL", "STORE_DRESIDUAL", "IS_RMS_NORM", "HAS_BIAS", "HAS_DROPOUT"], -) -# torch compile doesn't like triton.heuristics, so we set these manually when calling the kernel -# @triton.heuristics({"HAS_BIAS": lambda args: args["B"] is not None}) -# @triton.heuristics({"HAS_DRESIDUAL": lambda args: args["DRESIDUAL"] is not None}) -# @triton.heuristics({"STORE_DRESIDUAL": lambda args: args["DRESIDUAL_IN"] is not None}) -# @triton.heuristics({"HAS_ROWSCALE": lambda args: args["ROWSCALE"] is not None}) -# @triton.heuristics({"HAS_DY1": lambda args: args["DY1"] is not None}) -# @triton.heuristics({"HAS_DX1": lambda args: args["DX1"] is not None}) -# @triton.heuristics({"HAS_B1": lambda args: args["DB1"] is not None}) -# @triton.heuristics({"RECOMPUTE_OUTPUT": lambda args: args["Y"] is not None}) -@triton.jit -def _layer_norm_bwd_kernel( - X, # pointer to the input - W, # pointer to the weights - B, # pointer to the biases - Y, # pointer to the output to be recomputed - DY, # pointer to the output gradient - DX, # pointer to the input gradient - DW, # pointer to the partial sum of weights gradient - DB, # pointer to the partial sum of biases gradient - DRESIDUAL, - W1, - DY1, - DX1, - DW1, - DB1, - DRESIDUAL_IN, - ROWSCALE, - SEEDS, - Mean, # pointer to the mean - Rstd, # pointer to the 1/std - stride_x_row, # how much to increase the pointer when moving by 1 row - stride_y_row, - stride_dy_row, - stride_dx_row, - stride_dres_row, - stride_dy1_row, - stride_dx1_row, - stride_dres_in_row, - M, # number of rows in X - N, # number of columns in X - eps, # epsilon to avoid division by zero - dropout_p, - zero_centered_weight, - rows_per_program, - IS_RMS_NORM: tl.constexpr, - BLOCK_N: tl.constexpr, - HAS_DRESIDUAL: tl.constexpr, - STORE_DRESIDUAL: tl.constexpr, - HAS_BIAS: tl.constexpr, - HAS_DROPOUT: tl.constexpr, - HAS_ROWSCALE: tl.constexpr, - HAS_DY1: tl.constexpr, - HAS_DX1: tl.constexpr, - HAS_B1: tl.constexpr, - RECOMPUTE_OUTPUT: tl.constexpr, -): - # Map the program id to the elements of X, DX, and DY it should compute. - row_block_id = tl.program_id(0) - row_start = row_block_id * rows_per_program - # Do not early exit if row_start >= M, because we need to write DW and DB - cols = tl.arange(0, BLOCK_N) - mask = cols < N - X += row_start * stride_x_row - if HAS_DRESIDUAL: - DRESIDUAL += row_start * stride_dres_row - if STORE_DRESIDUAL: - DRESIDUAL_IN += row_start * stride_dres_in_row - DY += row_start * stride_dy_row - DX += row_start * stride_dx_row - if HAS_DY1: - DY1 += row_start * stride_dy1_row - if HAS_DX1: - DX1 += row_start * stride_dx1_row - if RECOMPUTE_OUTPUT: - Y += row_start * stride_y_row - w = tl.load(W + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w += 1.0 - if RECOMPUTE_OUTPUT and HAS_BIAS: - b = tl.load(B + cols, mask=mask, other=0.0).to(tl.float32) - if HAS_DY1: - w1 = tl.load(W1 + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w1 += 1.0 - dw = tl.zeros((BLOCK_N,), dtype=tl.float32) - if HAS_BIAS: - db = tl.zeros((BLOCK_N,), dtype=tl.float32) - if HAS_DY1: - dw1 = tl.zeros((BLOCK_N,), dtype=tl.float32) - if HAS_B1: - db1 = tl.zeros((BLOCK_N,), dtype=tl.float32) - row_end = min((row_block_id + 1) * rows_per_program, M) - for row in range(row_start, row_end): - # Load data to SRAM - x = tl.load(X + cols, mask=mask, other=0).to(tl.float32) - dy = tl.load(DY + cols, mask=mask, other=0).to(tl.float32) - if HAS_DY1: - dy1 = tl.load(DY1 + cols, mask=mask, other=0).to(tl.float32) - if not IS_RMS_NORM: - mean = tl.load(Mean + row) - rstd = tl.load(Rstd + row) - # Compute dx - xhat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd - xhat = tl.where(mask, xhat, 0.0) - if RECOMPUTE_OUTPUT: - y = xhat * w + b if HAS_BIAS else xhat * w - tl.store(Y + cols, y, mask=mask) - wdy = w * dy - dw += dy * xhat - if HAS_BIAS: - db += dy - if HAS_DY1: - wdy += w1 * dy1 - dw1 += dy1 * xhat - if HAS_B1: - db1 += dy1 - if not IS_RMS_NORM: - c1 = tl.sum(xhat * wdy, axis=0) / N - c2 = tl.sum(wdy, axis=0) / N - dx = (wdy - (xhat * c1 + c2)) * rstd - else: - c1 = tl.sum(xhat * wdy, axis=0) / N - dx = (wdy - xhat * c1) * rstd - if HAS_DRESIDUAL: - dres = tl.load(DRESIDUAL + cols, mask=mask, other=0).to(tl.float32) - dx += dres - # Write dx - if STORE_DRESIDUAL: - tl.store(DRESIDUAL_IN + cols, dx, mask=mask) - if HAS_DX1: - if HAS_DROPOUT: - keep_mask = ( - tl.rand(tl.load(SEEDS + M + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - ) - dx1 = tl.where(keep_mask, dx / (1.0 - dropout_p), 0.0) - else: - dx1 = dx - tl.store(DX1 + cols, dx1, mask=mask) - if HAS_DROPOUT: - keep_mask = tl.rand(tl.load(SEEDS + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - dx = tl.where(keep_mask, dx / (1.0 - dropout_p), 0.0) - if HAS_ROWSCALE: - rowscale = tl.load(ROWSCALE + row).to(tl.float32) - dx *= rowscale - tl.store(DX + cols, dx, mask=mask) - - X += stride_x_row - if HAS_DRESIDUAL: - DRESIDUAL += stride_dres_row - if STORE_DRESIDUAL: - DRESIDUAL_IN += stride_dres_in_row - if RECOMPUTE_OUTPUT: - Y += stride_y_row - DY += stride_dy_row - DX += stride_dx_row - if HAS_DY1: - DY1 += stride_dy1_row - if HAS_DX1: - DX1 += stride_dx1_row - tl.store(DW + row_block_id * N + cols, dw, mask=mask) - if HAS_BIAS: - tl.store(DB + row_block_id * N + cols, db, mask=mask) - if HAS_DY1: - tl.store(DW1 + row_block_id * N + cols, dw1, mask=mask) - if HAS_B1: - tl.store(DB1 + row_block_id * N + cols, db1, mask=mask) - - -def _layer_norm_bwd( - dy: Tensor, - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - mean: Tensor, - rstd: Tensor, - dresidual: Optional[Tensor] = None, - dy1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - seeds: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - has_residual: bool = False, - has_x1: bool = False, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - x_dtype: Optional[torch.dtype] = None, - recompute_output: bool = False, -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - # Need to wrap to handle the case where dresidual_in or dx1 are aliases of x, - # which makes torch.library unhappy - dx, dw, db, dresidual_in, dx1, dw1, db1, y = _layer_norm_bwd_impl( - dy, - x, - weight, - bias, - eps, - mean, - rstd, - dresidual, - dy1, - weight1, - bias1, - seeds, - dropout_p, - rowscale, - has_residual, - has_x1, - zero_centered_weight, - is_rms_norm, - x_dtype=x_dtype, - recompute_output=recompute_output, - ) - # Don't need to compute dresidual_in separately in this case - if has_residual and dx.dtype == x.dtype and dropout_p == 0.0 and rowscale is None: - dresidual_in = dx - if has_x1 and dropout_p == 0.0: - dx1 = dx - return dx, dw, db, dresidual_in, dx1, dw1, db1, y - - - -@triton_op(add_op_namespace_prefix("layer_norm_bwd_impl"), mutates_args={}, - schema="(Tensor dy, Tensor x, Tensor weight, Tensor bias, float eps, Tensor mean, Tensor rstd, Tensor? dresidual, Tensor? dy1, Tensor? weight1, Tensor? bias1, Tensor? seeds, float dropout_p, Tensor? rowscale, bool has_residual, bool has_x1, bool zero_centered_weight, bool is_rms_norm, ScalarType? x_dtype, bool recompute_output) -> (Tensor dx, Tensor dw, Tensor db, Tensor dresidual_in, Tensor dx1, Tensor dw1, Tensor db1, Tensor y)", - allow_decomposition=False, # Don't let torch.compile trace inside - ) -def _layer_norm_bwd_impl( - dy: Tensor, - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - mean: Tensor, - rstd: Tensor, - dresidual: Optional[Tensor] = None, - dy1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - seeds: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - has_residual: bool = False, - has_x1: bool = False, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - x_dtype: Optional[torch.dtype] = None, - recompute_output: bool = False, -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - M, N = x.shape - assert x.stride(-1) == 1 - dy = maybe_contiguous_lastdim(dy) - assert dy.stride(-1) == 1 - assert dy.shape == (M, N) - if dresidual is not None: - dresidual = maybe_contiguous_lastdim(dresidual) - assert dresidual.stride(-1) == 1 - assert dresidual.shape == (M, N) - assert weight.shape == (N,) - assert weight.stride(-1) == 1 - if bias is not None: - assert bias.stride(-1) == 1 - assert bias.shape == (N,) - if dy1 is not None: - dy1 = maybe_contiguous_lastdim(dy1) - assert weight1 is not None - assert dy1.shape == dy.shape - assert dy1.stride(-1) == 1 - if weight1 is not None: - assert weight1.shape == (N,) - assert weight1.stride(-1) == 1 - if bias1 is not None: - assert bias1.shape == (N,) - assert bias1.stride(-1) == 1 - if seeds is not None: - assert seeds.is_contiguous() - assert seeds.shape == (M if not has_x1 else M * 2,) - if rowscale is not None: - assert rowscale.is_contiguous() - assert rowscale.shape == (M,) - # allocate output - dx = ( - torch.empty_like(x) - if x_dtype is None - else torch.empty(M, N, dtype=x_dtype, device=x.device) - ) - dresidual_in = ( - torch.empty_like(x) - if has_residual - and (dx.dtype != x.dtype or dropout_p > 0.0 or rowscale is not None or has_x1) - else None - ) - dx1 = torch.empty_like(dx) if (has_x1 and dropout_p > 0.0) else None - y = torch.empty(M, N, dtype=dy.dtype, device=dy.device) if recompute_output else None - if recompute_output: - assert weight1 is None, "recompute_output is not supported with parallel LayerNorm" - - # Less than 64KB per feature: enqueue fused kernel - MAX_FUSED_SIZE = 65536 // x.element_size() - BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) - if N > BLOCK_N: - raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") - # Increasing the multiple (e.g. 8) will allow more thread blocks to be launched and hide the - # latency of the gmem reads/writes, but will increase the time of summing up dw / db. - sm_count = torch.cuda.get_device_properties(x.device).multi_processor_count * 8 - _dw = torch.empty((sm_count, N), dtype=torch.float32, device=weight.device) - _db = ( - torch.empty((sm_count, N), dtype=torch.float32, device=bias.device) - if bias is not None - else None - ) - _dw1 = torch.empty_like(_dw) if weight1 is not None else None - _db1 = torch.empty_like(_db) if bias1 is not None else None - rows_per_program = math.ceil(M / sm_count) - grid = (sm_count,) - with torch.cuda.device(x.device.index): - torch.library.wrap_triton(_layer_norm_bwd_kernel)[grid]( - x, - weight, - bias, - y, - dy, - dx, - _dw, - _db, - dresidual, - weight1, - dy1, - dx1, - _dw1, - _db1, - dresidual_in, - rowscale, - seeds, - mean, - rstd, - x.stride(0), - 0 if not recompute_output else y.stride(0), - dy.stride(0), - dx.stride(0), - dresidual.stride(0) if dresidual is not None else 0, - dy1.stride(0) if dy1 is not None else 0, - dx1.stride(0) if dx1 is not None else 0, - dresidual_in.stride(0) if dresidual_in is not None else 0, - M, - N, - eps, - dropout_p, - # Passing bool make torch inductor very unhappy since it then tries to compare to int_max - int(zero_centered_weight), - rows_per_program, - is_rms_norm, - BLOCK_N, - dresidual is not None, - dresidual_in is not None, - bias is not None, - dropout_p > 0.0, - HAS_ROWSCALE=rowscale is not None, - HAS_DY1=dy1 is not None, - HAS_DX1=dx1 is not None, - HAS_B1=bias1 is not None, - RECOMPUTE_OUTPUT=y is not None, - ) - dw = _dw.sum(0).to(weight.dtype) - db = _db.sum(0).to(bias.dtype) if bias is not None else None - dw1 = _dw1.sum(0).to(weight1.dtype) if weight1 is not None else None - db1 = _db1.sum(0).to(bias1.dtype) if bias1 is not None else None - # dresidual_in and dx1 could be None, the wrapper will handle assigning them from dx - return dx, dw, db, dresidual_in, dx1, dw1, db1, y - - -class LayerNormFn(torch.autograd.Function): - - @staticmethod - def forward( - ctx, - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - residual_in_fp32=False, - zero_centered_weight=False, - is_rms_norm=False, - return_dropout_mask=False, - out_dtype=None, - out=None, - residual_out=None - ): - x_shape_og = x.shape - # reshape input data into 2D tensor - x = maybe_contiguous_lastdim(x.reshape(-1, x.shape[-1])) - if residual is not None: - assert residual.shape == x_shape_og - residual = maybe_contiguous_lastdim(residual.reshape(-1, residual.shape[-1])) - if x1 is not None: - assert x1.shape == x_shape_og - assert rowscale is None, "rowscale is not supported with parallel LayerNorm" - x1 = maybe_contiguous_lastdim(x1.reshape(-1, x1.shape[-1])) - weight = weight.contiguous() - bias = maybe_contiguous(bias) - weight1 = maybe_contiguous(weight1) - bias1 = maybe_contiguous(bias1) - if rowscale is not None: - rowscale = rowscale.reshape(-1).contiguous() - residual_dtype = ( - residual.dtype - if residual is not None - else (torch.float32 if residual_in_fp32 else None) - ) - if out is not None: - out = out.reshape(-1, out.shape[-1]) - if residual_out is not None: - residual_out = residual_out.reshape(-1, residual_out.shape[-1]) - y, y1, mean, rstd, residual_out, seeds, dropout_mask, dropout_mask1 = _layer_norm_fwd( - x, - weight, - bias, - eps, - residual, - x1, - weight1, - bias1, - dropout_p=dropout_p, - rowscale=rowscale, - out_dtype=out_dtype, - residual_dtype=residual_dtype, - zero_centered_weight=zero_centered_weight, - is_rms_norm=is_rms_norm, - return_dropout_mask=return_dropout_mask, - out=out, - residual_out=residual_out, - ) - ctx.save_for_backward( - residual_out, weight, bias, weight1, bias1, rowscale, seeds, mean, rstd - ) - ctx.x_shape_og = x_shape_og - ctx.eps = eps - ctx.dropout_p = dropout_p - ctx.is_rms_norm = is_rms_norm - ctx.has_residual = residual is not None - ctx.has_x1 = x1 is not None - ctx.prenorm = prenorm - ctx.x_dtype = x.dtype - ctx.zero_centered_weight = zero_centered_weight - y = y.reshape(x_shape_og) - y1 = y1.reshape(x_shape_og) if y1 is not None else None - residual_out = residual_out.reshape(x_shape_og) if residual_out is not None else None - dropout_mask = dropout_mask.reshape(x_shape_og) if dropout_mask is not None else None - dropout_mask1 = dropout_mask1.reshape(x_shape_og) if dropout_mask1 is not None else None - if not return_dropout_mask: - if weight1 is None: - return y if not prenorm else (y, residual_out) - else: - return (y, y1) if not prenorm else (y, y1, residual_out) - else: - if weight1 is None: - return ( - (y, dropout_mask, dropout_mask1) - if not prenorm - else (y, residual_out, dropout_mask, dropout_mask1) - ) - else: - return ( - (y, y1, dropout_mask, dropout_mask1) - if not prenorm - else (y, y1, residual_out, dropout_mask, dropout_mask1) - ) - - @staticmethod - def backward(ctx, dy, *args): - x, weight, bias, weight1, bias1, rowscale, seeds, mean, rstd = ctx.saved_tensors - dy = dy.reshape(-1, dy.shape[-1]) - if weight1 is not None: - dy1, args = args[0], args[1:] - dy1 = dy1.reshape(-1, dy1.shape[-1]) - assert dy1.shape == x.shape - else: - dy1 = None - if ctx.prenorm: - dresidual = args[0] - dresidual = dresidual.reshape(-1, dresidual.shape[-1]) - assert dresidual.shape == x.shape - else: - dresidual = None - dx, dw, db, dresidual_in, dx1, dw1, db1, _ = _layer_norm_bwd( - dy, - x, - weight, - bias, - ctx.eps, - mean, - rstd, - dresidual, - dy1, - weight1, - bias1, - seeds, - ctx.dropout_p, - rowscale, - ctx.has_residual, - ctx.has_x1, - ctx.zero_centered_weight, - ctx.is_rms_norm, - x_dtype=ctx.x_dtype, - recompute_output=False, - ) - return ( - dx.reshape(ctx.x_shape_og), - dw, - db, - dresidual_in.reshape(ctx.x_shape_og) if ctx.has_residual else None, - dx1.reshape(ctx.x_shape_og) if dx1 is not None else None, - dw1, - db1, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) - - -def layer_norm_fn( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - residual_in_fp32=False, - zero_centered_weight=False, - is_rms_norm=False, - return_dropout_mask=False, - out_dtype=None, - out=None, - residual_out=None -): - return LayerNormFn.apply( - x, - weight, - bias, - residual, - x1, - weight1, - bias1, - eps, - dropout_p, - rowscale, - prenorm, - residual_in_fp32, - zero_centered_weight, - is_rms_norm, - return_dropout_mask, - out_dtype, - out, - residual_out - ) - - -def rms_norm_fn( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - residual_in_fp32=False, - zero_centered_weight=False, - return_dropout_mask=False, - out_dtype=None, - out=None, - residual_out=None -): - return LayerNormFn.apply( - x, - weight, - bias, - residual, - x1, - weight1, - bias1, - eps, - dropout_p, - rowscale, - prenorm, - residual_in_fp32, - zero_centered_weight, - True, - return_dropout_mask, - out_dtype, - out, - residual_out - ) - - -class RMSNorm(torch.nn.Module): - - def __init__(self, hidden_size, eps=1e-5, dropout_p=0.0, zero_centered_weight=False, - device=None, dtype=None): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.eps = eps - if dropout_p > 0.0: - self.drop = torch.nn.Dropout(dropout_p) - else: - self.drop = None - self.zero_centered_weight = zero_centered_weight - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.register_parameter("bias", None) - self.reset_parameters() - - def reset_parameters(self): - if not self.zero_centered_weight: - torch.nn.init.ones_(self.weight) - else: - torch.nn.init.zeros_(self.weight) - - def forward(self, x, residual=None, prenorm=False, residual_in_fp32=False): - return rms_norm_fn( - x, - self.weight, - self.bias, - residual=residual, - eps=self.eps, - dropout_p=self.drop.p if self.drop is not None and self.training else 0.0, - prenorm=prenorm, - residual_in_fp32=residual_in_fp32, - zero_centered_weight=self.zero_centered_weight, - ) - - -class LayerNormLinearFn(torch.autograd.Function): - - @staticmethod - @custom_fwd - def forward( - ctx, - x, - norm_weight, - norm_bias, - linear_weight, - linear_bias, - residual=None, - eps=1e-6, - prenorm=False, - residual_in_fp32=False, - is_rms_norm=False, - ): - x_shape_og = x.shape - # reshape input data into 2D tensor - x = maybe_contiguous_lastdim(x.reshape(-1, x.shape[-1])) - if residual is not None: - assert residual.shape == x_shape_og - residual = maybe_contiguous_lastdim(residual.reshape(-1, residual.shape[-1])) - norm_weight = norm_weight.contiguous() - norm_bias = maybe_contiguous(norm_bias) - residual_dtype = ( - residual.dtype - if residual is not None - else (torch.float32 if residual_in_fp32 else None) - ) - y, _, mean, rstd, residual_out, *rest = _layer_norm_fwd( - x, - norm_weight, - norm_bias, - eps, - residual, - out_dtype=None if not torch.is_autocast_enabled() else torch.get_autocast_dtype("cuda"), - residual_dtype=residual_dtype, - is_rms_norm=is_rms_norm, - ) - y = y.reshape(x_shape_og) - dtype = torch.get_autocast_dtype("cuda") if torch.is_autocast_enabled() else y.dtype - linear_weight = linear_weight.to(dtype) - linear_bias = linear_bias.to(dtype) if linear_bias is not None else None - out = F.linear(y.to(linear_weight.dtype), linear_weight, linear_bias) - # We don't store y, will be recomputed in the backward pass to save memory - ctx.save_for_backward(residual_out, norm_weight, norm_bias, linear_weight, mean, rstd) - ctx.x_shape_og = x_shape_og - ctx.eps = eps - ctx.is_rms_norm = is_rms_norm - ctx.has_residual = residual is not None - ctx.prenorm = prenorm - ctx.x_dtype = x.dtype - ctx.linear_bias_is_none = linear_bias is None - return out if not prenorm else (out, residual_out.reshape(x_shape_og)) - - @staticmethod - @custom_bwd - def backward(ctx, dout, *args): - x, norm_weight, norm_bias, linear_weight, mean, rstd = ctx.saved_tensors - dout = dout.reshape(-1, dout.shape[-1]) - dy = F.linear(dout, linear_weight.t()) - dlinear_bias = None if ctx.linear_bias_is_none else dout.sum(0) - dy = maybe_contiguous_lastdim(dy) - assert dy.shape == x.shape - if ctx.prenorm: - dresidual = args[0] - dresidual = maybe_contiguous_lastdim(dresidual.reshape(-1, dresidual.shape[-1])) - assert dresidual.shape == x.shape - else: - dresidual = None - dx, dnorm_weight, dnorm_bias, dresidual_in, _, _, _, y = _layer_norm_bwd( - dy, - x, - norm_weight, - norm_bias, - ctx.eps, - mean, - rstd, - dresidual=dresidual, - has_residual=ctx.has_residual, - is_rms_norm=ctx.is_rms_norm, - x_dtype=ctx.x_dtype, - recompute_output=True, - ) - dlinear_weight = torch.einsum("bo,bi->oi", dout, y) - return ( - dx.reshape(ctx.x_shape_og), - dnorm_weight, - dnorm_bias, - dlinear_weight, - dlinear_bias, - dresidual_in.reshape(ctx.x_shape_og) if ctx.has_residual else None, - None, - None, - None, - None, - ) - - -def layer_norm_linear_fn( - x, - norm_weight, - norm_bias, - linear_weight, - linear_bias, - residual=None, - eps=1e-6, - prenorm=False, - residual_in_fp32=False, - is_rms_norm=False, -): - return LayerNormLinearFn.apply( - x, - norm_weight, - norm_bias, - linear_weight, - linear_bias, - residual, - eps, - prenorm, - residual_in_fp32, - is_rms_norm, - ) diff --git a/build/torch212-cxx11-cpu-x86_64-linux/ops/triton/linear.py b/build/torch212-cxx11-cpu-x86_64-linux/ops/triton/linear.py deleted file mode 100644 index a8966dbc345ab0e593df0124451ee7be3dae131a..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cpu-x86_64-linux/ops/triton/linear.py +++ /dev/null @@ -1,594 +0,0 @@ -# Adapted from https://github.com/ELS-RD/kernl/blob/main/src/kernl/implementations/linear_layer.py -# and https://github.com/openai/triton/blob/master/python/triton/ops/matmul.py -from typing import Optional - -import torch -import triton -import triton.language as tl -from triton.ops.matmul_perf_model import early_config_prune, estimate_matmul_time - -from flash_attn.ops.triton.k_activations import ( - gelu, - gelu_approx, - gelu_approx_grad, - gelu_grad, - squared_relu, - squared_relu_grad, -) - -# CREDITS: Initially inspired by the Triton tutorial on matrix multiplications - - -def init_to_zero(name): - return lambda nargs: nargs[name].zero_() - - -def get_configs_io_bound(): - configs = [] - for num_stages in [2, 3, 4, 5, 6]: - for block_m in [16, 32]: - for block_k in [32, 64]: - for block_n in [32, 64, 128, 256]: - num_warps = 2 if block_n <= 64 else 4 - configs.append( - triton.Config( - { - "BLOCK_M": block_m, - "BLOCK_N": block_n, - "BLOCK_K": block_k, - "SPLIT_K": 1, - }, - num_stages=num_stages, - num_warps=num_warps, - ) - ) - # split_k not used - # for split_k in [2, 4, 8, 16]: - # configs.append(triton.Config( - # {'BLOCK_M': block_m, 'BLOCK_N': block_n, 'BLOCK_K': block_k, 'SPLIT_K': split_k}, - # num_stages=num_stages, num_warps=num_warps, pre_hook=init_to_zero('C'))) - return configs - - -@triton.autotune( - configs=[ - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - # good for int8 - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=4, - num_warps=4, - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - ] - + get_configs_io_bound(), - key=["CACHE_KEY_M", "CACHE_KEY_N", "CACHE_KEY_K"], - prune_configs_by={ - "early_config_prune": early_config_prune, - "perf_model": estimate_matmul_time, - "top_k": 10, - }, -) -@triton.heuristics( - { - "EVEN_K": lambda args: args["K"] % (args["BLOCK_K"] * args["SPLIT_K"]) == 0, - } -) -@triton.jit -def kernel_fwd( - C, # Pointers to matrices - ACT_INPUT, - A, - B, - bias, - # Matrix dimensions - M, - N, - K, - CACHE_KEY_M, - CACHE_KEY_N, - CACHE_KEY_K, - # The stride variables represent how much to increase the ptr by when moving by 1 - # element in a particular dimension. E.g. stride_am is how much to increase a_ptr - # by to get the element one row down (A has M rows) - stride_cm, - # stride_cn, # Assume that stride_cn == 1 - stride_am, - stride_ak, - stride_bn, - stride_bk, - # Meta-parameters - BLOCK_M: tl.constexpr, - GROUP_M: tl.constexpr, - BLOCK_N: tl.constexpr, - BLOCK_K: tl.constexpr, - # split k not used, not performant with activation, kept because early_config_prune is expecting it - SPLIT_K: tl.constexpr, - EVEN_K: tl.constexpr, - A_ROWMAJOR: tl.constexpr, - B_COLMAJOR: tl.constexpr, - BIAS: tl.constexpr, - SAVE_ACT_INPUT: tl.constexpr, - ACTIVATION: tl.constexpr, -): - - """ - Kernel for computing Out = activation(A x W + C) - - Input has shape (M, K) - - Weight has shape (K, N) - - Bias has shape (N,) - - Output has shape (M, N) - - ActInputs (optional) has shape (M, N) - 'ActInputs' optionally saves the A x W + C intermediate for backward computations - This kernel will consolidate over K - """ - - pid = tl.program_id(axis=0) - - grid_m = (M + BLOCK_M - 1) // BLOCK_M - grid_n = (N + BLOCK_N - 1) // BLOCK_N - # re-order program ID for better L2 performance - width = GROUP_M * grid_n - group_id = pid // width - group_size = min(grid_m - group_id * GROUP_M, GROUP_M) - pid_m = group_id * GROUP_M + (pid % group_size) - pid_n = (pid % width) // (group_size) - - # now compute the block that each program will go through - # rm (resp. rn) denotes a range of indices - # for rows (resp. col) of C - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - # trick to avoid masking on M and N axis - ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M) - rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N) - rk = tl.arange(0, BLOCK_K) - - if A_ROWMAJOR: - A = A + (ram[:, None] * stride_am + rk[None, :]) - else: - A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak) - if B_COLMAJOR: - B = B + (rk[:, None] + rbn[None, :] * stride_bn) - else: - B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn) - - acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) - - for k in range(K, 0, -BLOCK_K): - if EVEN_K: - a = tl.load(A) - b = tl.load(B) - else: - a = tl.load(A, mask=rk[None, :] < k, other=0.0) - b = tl.load(B, mask=rk[:, None] < k, other=0.0) - acc += tl.dot(a, b) - - if A_ROWMAJOR: - A += BLOCK_K - else: - A += BLOCK_K * stride_ak - if B_COLMAJOR: - B += BLOCK_K - else: - B += BLOCK_K * stride_bk - - # Putting bias after the matmul (instead of before) is faster, idk why - if BIAS: - bias = tl.load(bias + rn, mask=rn < N, other=0.0).to(tl.float32) - acc += bias[None, :] - - # optional: save the activation inputs - if SAVE_ACT_INPUT: - # act_in_ptrs = ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :] * stride_cn - act_in_ptrs = ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :] - tl.store(act_in_ptrs, acc) - - # optional: fused activation (while the data is in shared memory) - if ACTIVATION == "gelu": - acc = gelu(acc) - elif ACTIVATION == "gelu_approx": - acc = gelu_approx(acc) - elif ACTIVATION == "squared_relu": - acc = squared_relu(acc) - # rematerialize rm and rn to save registers - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - - # write back result - # C = C + rm[:, None] * stride_cm + rn[None, :] * stride_cn - C = C + rm[:, None] * stride_cm + rn[None, :] - mask = (rm < M)[:, None] & (rn < N)[None, :] - tl.store(C, acc) - - -def triton_linear_act( - x: torch.Tensor, - weight: torch.Tensor, - bias: Optional[torch.Tensor] = None, - activation: str = "id", - save_act_input: bool = False, -) -> torch.Tensor: - """ - Compute e = activation(x @ weight.T + bias). - This wrapper kicks the `kernel_fwd` Triton kernel - :param x: input tensor - :param weight: weight matrix - :param bias: an optional bias tensor - :param activation: Activation name. Needs to be a Triton kernel. - :param act_input: an optional tensor to save the activation inputs (for backward) - :return: result tensor - """ - # if torch.is_autocast_enabled(): - # dtype = torch.get_autocast_gpu_dtype() - # x, weight, bias = [a.to(dtype=dtype) for a in [x, weight, bias]] - - assert activation in ["id", "gelu", "gelu_approx", "squared_relu"] - - batch_shape, n = x.shape[:-1], x.shape[-1] - batch_dim = batch_shape.numel() - x_reshaped = x.reshape(batch_dim, n) - - if x_reshaped.stride(0) > 1 and x_reshaped.stride(1) > 1: - x_reshaped = x_reshaped.contiguous() - if weight.stride(0) > 1 and weight.stride(1) > 1: - weight = weight.contiguous() - bias = bias.contiguous() if bias is not None else None - - assert ( - x.dtype == weight.dtype - ), f"Input and weight must have the same dtype, got {x.dtype} and {weight.dtype}" - if bias is not None: - assert ( - x.dtype == bias.dtype - ), f"Input and bias must have the same dtype, got {x.dtype} and {bias.dtype}" - assert ( - x_reshaped.shape[1] == weight.shape[1] - ), f"Incompatible dimensions: {x_reshaped.shape} - {weight.shape}" - - assert ( - bias is None or bias.shape[0] == weight.shape[0] - ), "Incompatible dimensions in between weight and bias" - - M, K = x_reshaped.shape - N, K = weight.shape - - output = torch.empty((M, N), device=x.device, dtype=x.dtype) - act_input = torch.empty_like(output) if save_act_input else None - - # 1D launch kernel where each block gets its own program. - grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),) # noqa - - kernel_fwd[grid]( - output, - act_input, - x_reshaped, - weight, # data ptrs - bias if bias is not None else x, # auto skip bias if not present - M, # shapes - N, - K, - M // 32, # key for triton cache (limit number of compilations) - N // 32, - K // 32, - stride_cm=output.stride(0), # strides - # stride_cn=output.stride(1), - stride_am=x_reshaped.stride(0), - stride_ak=x_reshaped.stride(1), - stride_bk=weight.stride(1), - stride_bn=weight.stride(0), - BIAS=bias is not None, # optional fused bias - SAVE_ACT_INPUT=save_act_input, # optional save activation inputs - ACTIVATION=activation, # optional fused activation - A_ROWMAJOR=x_reshaped.stride(1) == 1, - B_COLMAJOR=weight.stride(1) == 1, - GROUP_M=8, # speed optimization: group the programs - ) - - if not save_act_input: - return output.reshape(*batch_shape, output.shape[-1]) - else: - return ( - output.reshape(*batch_shape, output.shape[-1]), - act_input.reshape(*batch_shape, act_input.shape[-1]), - ) - - -@triton.autotune( - configs=[ - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - # good for int8 - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=4, - num_warps=4, - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - ] - + get_configs_io_bound(), - key=["CACHE_KEY_M", "CACHE_KEY_N", "CACHE_KEY_K"], - prune_configs_by={ - "early_config_prune": early_config_prune, - "perf_model": estimate_matmul_time, - "top_k": 10, - }, -) -@triton.heuristics( - { - "EVEN_K": lambda args: args["K"] % (args["BLOCK_K"] * args["SPLIT_K"]) == 0, - } -) -@triton.jit -def kernel_bwd( - C, # Pointers to matrices - ACT_INPUT, - A, - B, - # Matrix dimensions - M, - N, - K, - CACHE_KEY_M, - CACHE_KEY_N, - CACHE_KEY_K, - # The stride variables represent how much to increase the ptr by when moving by 1 - # element in a particular dimension. E.g. stride_am is how much to increase a_ptr - # by to get the element one row down (A has M rows) - stride_cm, - # stride_cn, # Assume that stride_cn == 1 - stride_am, - stride_ak, - stride_bk, - stride_bn, - # Meta-parameters - BLOCK_M: tl.constexpr, - GROUP_M: tl.constexpr, - BLOCK_N: tl.constexpr, - BLOCK_K: tl.constexpr, - # split k not used, not performant with activation, kept because early_config_prune is expecting it - SPLIT_K: tl.constexpr, - EVEN_K: tl.constexpr, - ACTIVATION: tl.constexpr, -): - - """ - Kernel for computing Out = activation(A x W + C) - - Input has shape (M, K) - - Weight has shape (K, N) - - Output has shape (M, N) - - ActInputs (optional) has shape (M, N) - 'ActInputs' optionally saves the A x W + C intermediate for backward computations - This kernel will consolidate over K - """ - - pid = tl.program_id(axis=0) - - grid_m = (M + BLOCK_M - 1) // BLOCK_M - grid_n = (N + BLOCK_N - 1) // BLOCK_N - # re-order program ID for better L2 performance - width = GROUP_M * grid_n - group_id = pid // width - group_size = min(grid_m - group_id * GROUP_M, GROUP_M) - pid_m = group_id * GROUP_M + (pid % group_size) - pid_n = (pid % width) // (group_size) - - # now compute the block that each program will go through - # rm (resp. rn) denotes a range of indices - # for rows (resp. col) of C - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - # trick to avoid masking on M and N axis - ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M) - rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N) - rk = tl.arange(0, BLOCK_K) - - A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak) - B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn) - - acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) - - for k in range(K, 0, -BLOCK_K): - if EVEN_K: - a = tl.load(A) - b = tl.load(B) - else: - a = tl.load(A, mask=rk[None, :] < k, other=0.0) - b = tl.load(B, mask=rk[:, None] < k, other=0.0) - acc += tl.dot(a, b) - - A += BLOCK_K * stride_ak - B += BLOCK_K * stride_bk - - # optional: fused activation (while the data is in shared memory) - if ACTIVATION != "id": - act_in_ptrs = ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :] - act_input = tl.load(act_in_ptrs).to(acc.dtype) - if ACTIVATION == "gelu": - acc *= gelu_grad(act_input) - elif ACTIVATION == "gelu_approx": - acc *= gelu_approx_grad(act_input) - elif ACTIVATION == "squared_relu": - acc *= squared_relu_grad(act_input) - - # rematerialize rm and rn to save registers - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - - # write back result - C = C + rm[:, None] * stride_cm + rn[None, :] - mask = (rm < M)[:, None] & (rn < N)[None, :] - tl.store(C, acc, mask=mask) - - -def triton_dgrad_act( - grad_output: torch.Tensor, - weight: torch.Tensor, - activation: str = "id", - act_input: Optional[torch.Tensor] = None, -) -> torch.Tensor: - """ - Compute e = activation(grad_output @ weight + bias). - This wrapper kicks the `kernel_fwd` Triton kernel - :param grad_output: input tensor - :param weight: weight matrix - :param activation: Activation name. Needs to be a Triton kernel. - :param act_input: an optional tensor to save the activation inputs (for backward) - :return: result tensor - """ - assert activation in ["id", "gelu", "gelu_approx", "squared_relu"] - - batch_shape, n = grad_output.shape[:-1], grad_output.shape[-1] - batch_dim = batch_shape.numel() - grad_output_reshaped = grad_output.reshape(batch_dim, n) - - if grad_output_reshaped.stride(0) > 1 and grad_output_reshaped.stride(1) > 1: - grad_output_reshaped = grad_output_reshaped.contiguous() - if weight.stride(0) > 1 and weight.stride(1) > 1: - weight = weight.contiguous() - - assert ( - grad_output.dtype == weight.dtype - ), f"grad_output and weight must have the same dtype, got {grad_output.dtype} and {weight.dtype}" - assert ( - grad_output_reshaped.shape[1] == weight.shape[0] - ), f"Incompatible dimensions: {grad_output_reshaped.shape} - {weight.shape}" - if activation != "id": - assert act_input is not None, f"act_input is required for activation {activation}" - - # M, N, K in bwd are different from M, N, K in fwd - M, K = grad_output_reshaped.shape - K, N = weight.shape - - grad_input = torch.empty((M, N), device=grad_output.device, dtype=grad_output.dtype) - - # 1D launch kernel where each block gets its own program. - grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),) # noqa - - kernel_bwd[grid]( - grad_input, - act_input, - grad_output_reshaped, - weight, # data ptrs - M, # shapes - N, - K, - M // 32, # key for triton cache (limit number of compilations) - N // 32, - K // 32, - stride_cm=grad_input.stride(0), # strides - # stride_cn=grad_input.stride(1), - stride_am=grad_output_reshaped.stride(0), - stride_ak=grad_output_reshaped.stride(1), - stride_bk=weight.stride(0), - stride_bn=weight.stride(1), - ACTIVATION=activation, # optional fused activation - GROUP_M=8, # speed optimization: group the programs - ) - - return grad_input.reshape(*batch_shape, grad_input.shape[-1]) diff --git a/build/torch212-cxx11-cpu-x86_64-linux/ops/triton/mlp.py b/build/torch212-cxx11-cpu-x86_64-linux/ops/triton/mlp.py deleted file mode 100644 index 059f4f8a5e174c1f4824e43d313fca18eaa799b8..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cpu-x86_64-linux/ops/triton/mlp.py +++ /dev/null @@ -1,149 +0,0 @@ -# The triton fused matmul + sqrelu is faster for fp16 but slower for bf16, compared -# to naive implementation. -import fused_dense_lib as fused_dense_cuda -import torch -import torch.nn as nn -import torch.nn.functional as F - -from flash_attn.utils.torch import custom_fwd, custom_bwd -from flash_attn.ops.activations import sqrelu_bwd, sqrelu_fwd -from flash_attn.ops.triton.linear import triton_dgrad_act, triton_linear_act - - -class FusedDenseSqreluDenseFunc(torch.autograd.Function): - @staticmethod - @custom_fwd - def forward(ctx, x, weight1, bias1, weight2, bias2, checkpoint_lvl=0): - """checkpoint_lvl: - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute act_input and gelu_out in the bwd - """ - if torch.is_autocast_enabled(): - dtype = torch.get_autocast_gpu_dtype() - x, weight1, bias1, weight2, bias2 = [ - a.to(dtype=dtype) for a in [x, weight1, bias1, weight2, bias2] - ] - is_bf16 = x.dtype == torch.bfloat16 - assert checkpoint_lvl in [0, 1, 2] - x = x.contiguous() - weight1 = weight1.contiguous() - bias1 = bias1.contiguous() - weight2 = weight2.contiguous() - bias2 = bias2.contiguous() - batch_shape, n = x.shape[:-1], x.shape[-1] - batch_dim = batch_shape.numel() - if is_bf16: - act_input = fused_dense_cuda.linear_bias_forward( - x.reshape(batch_dim, n), weight1, bias1 - ) - output1 = sqrelu_fwd(act_input) - else: - save_act_input = checkpoint_lvl != 2 - result = triton_linear_act( - x.reshape(batch_dim, n), - weight1, - bias1, - activation="squared_relu", - save_act_input=save_act_input, - ) - if save_act_input: - output1, act_input = result - else: - output1 = result - output2 = fused_dense_cuda.linear_bias_forward(output1, weight2, bias2) - ctx.checkpoint_lvl = checkpoint_lvl - if checkpoint_lvl == 0: - ctx.save_for_backward(x, weight1, bias1, weight2, act_input, output1) - elif checkpoint_lvl == 1: - ctx.save_for_backward(x, weight1, bias1, weight2, act_input) - elif checkpoint_lvl == 2: - ctx.save_for_backward(x, weight1, bias1, weight2) - return output2.reshape(*batch_shape, output2.shape[-1]) - - @staticmethod - @custom_bwd - def backward(ctx, grad_output): - grad_output = grad_output.contiguous() - checkpoint_lvl = ctx.checkpoint_lvl - x, weight1, bias1, weight2, *rest = ctx.saved_tensors - batch_shape, n = x.shape[:-1], x.shape[-1] - batch_dim = batch_shape.numel() - is_bf16 = x.dtype == torch.bfloat16 - if checkpoint_lvl == 0: - act_input, output1 = rest - elif checkpoint_lvl == 1: - (act_input,) = rest - output1 = sqrelu_fwd(act_input) - elif checkpoint_lvl == 2: - if is_bf16: - act_input = fused_dense_cuda.linear_bias_forward( - x.reshape(batch_dim, n), weight1, bias1 - ) - output1 = sqrelu_fwd(act_input) - else: - output1, act_input = triton_linear_act( - x.reshape(batch_dim, n), - weight1, - bias1, - activation="squared_relu", - save_act_input=True, - ) - - if is_bf16: - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad(output1, grad_output) - grad_output1 = grad_output @ weight2 - grad_act_input = sqrelu_bwd(grad_output1, act_input) - grad_input, grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_backward( - x.reshape(batch_dim, n), weight1, grad_act_input - ) - else: - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad(output1, grad_output) - grad_act_input = triton_dgrad_act( - grad_output, weight2, activation="squared_relu", act_input=act_input - ) - grad_input, grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_backward( - x.reshape(batch_dim, n), weight1, grad_act_input - ) - return grad_input.reshape_as(x), grad_weight1, grad_bias1, grad_weight2, grad_bias2, None - - -fused_dense_sqrelu_dense_function = FusedDenseSqreluDenseFunc.apply - - -class FusedDenseSqreluDense(nn.Module): - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - bias1=True, - bias2=True, - checkpoint_lvl=0, - device=None, - dtype=None, - ): - """ - checkpoint_lvl (increasing lvl means slower but more memory saving): - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute gelu_in and gelu_out in the bwd - """ - assert checkpoint_lvl in [0, 1, 2] - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features * 4 - assert bias1 == True, "DenseSqreluDense module without bias is currently not supported" - assert bias2 == True, "DenseSqreluDense module without bias is currently not supported" - self.checkpoint_lvl = checkpoint_lvl - self.fc1 = nn.Linear(in_features, hidden_features, bias=bias1, **factory_kwargs) - self.fc2 = nn.Linear(hidden_features, out_features, bias=bias2, **factory_kwargs) - - def forward(self, x): - assert x.is_cuda - return fused_dense_sqrelu_dense_function( - x, self.fc1.weight, self.fc1.bias, self.fc2.weight, self.fc2.bias, self.checkpoint_lvl - ) diff --git a/build/torch212-cxx11-cpu-x86_64-linux/ops/triton/rotary.py b/build/torch212-cxx11-cpu-x86_64-linux/ops/triton/rotary.py deleted file mode 100644 index 6ad3ab412a084270120167ad32b67ad1720f2c37..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cpu-x86_64-linux/ops/triton/rotary.py +++ /dev/null @@ -1,186 +0,0 @@ -# Copyright (c) 2025, Tri Dao. -# As of 2025-04-23, we require triton >= 3.0 - -from typing import Optional, Union - -import torch - -import triton -import triton.language as tl - - -@triton.jit -def rotary_kernel( - OUT, # Pointers to matrices - X, - COS, - SIN, - CU_SEQLENS, - SEQLEN_OFFSETS, # this could be int or a pointer - # Matrix dimensions - seqlen, - nheads, - seqlen_ro, - # strides - stride_out_batch, - stride_out_seqlen, - stride_out_nheads, - stride_out_headdim, - stride_x_batch, - stride_x_seqlen, - stride_x_nheads, - stride_x_headdim, - # Meta-parameters - # We want ROTARY_DIM to be constexpr, otherwise the triton compiler doesn't know that - # the mask is constant every 8 elements, and it will generate LDG.16 instead of LDG.128 - ROTARY_DIM: tl.constexpr, - IS_SEQLEN_OFFSETS_TENSOR: tl.constexpr, - IS_VARLEN: tl.constexpr, - INTERLEAVED: tl.constexpr, - CONJUGATE: tl.constexpr, - BLOCK_H: tl.constexpr, - BLOCK_M: tl.constexpr, -): - BLOCK_K: tl.constexpr = triton.next_power_of_2(ROTARY_DIM) - ROTARY_DIM_HALF = ROTARY_DIM // 2 - pid_head = tl.program_id(axis=0) - pid_m = tl.program_id(axis=1) - pid_batch = tl.program_id(axis=2) - - if not IS_VARLEN: - X = X + pid_batch * stride_x_batch - OUT = OUT + pid_batch * stride_out_batch - else: - start_idx = tl.load(CU_SEQLENS + pid_batch) - seqlen = tl.load(CU_SEQLENS + pid_batch + 1) - start_idx - X = X + start_idx * stride_x_seqlen - OUT = OUT + start_idx * stride_out_seqlen - - if pid_m * BLOCK_M >= seqlen: - return - - rh = pid_head * BLOCK_H + tl.arange(0, BLOCK_H) - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - if not IS_SEQLEN_OFFSETS_TENSOR: - rm_cs = rm + SEQLEN_OFFSETS - else: - rm_cs = rm + tl.load(SEQLEN_OFFSETS + pid_batch) - - rk_half = tl.arange(0, BLOCK_K // 2) - COS = COS + (rm_cs[:, None] * ROTARY_DIM_HALF + rk_half[None, :]) - SIN = SIN + (rm_cs[:, None] * ROTARY_DIM_HALF + rk_half[None, :]) - mask_cs = (rm_cs[:, None] < seqlen_ro) & (rk_half[None, :] < ROTARY_DIM_HALF) - cos = tl.load(COS, mask=mask_cs, other=1.0).to(tl.float32) - sin = tl.load(SIN, mask=mask_cs, other=0.0).to(tl.float32) - if CONJUGATE: - sin = -sin - - if not INTERLEAVED: - # Load the 1st and 2nd halves of X, do calculation, then store to 1st and 2nd halves of OUT - X = X + (rh[:, None, None] * stride_x_nheads + rm[None, :, None] * stride_x_seqlen + rk_half[None, None, :] * stride_x_headdim) - OUT = OUT + (rh[:, None, None] * stride_out_nheads + rm[None, :, None] * stride_out_seqlen + rk_half[None, None, :] * stride_out_headdim) - mask = (rh[:, None, None] < nheads) & (rm[None, :, None] < seqlen) & (rk_half[None, None, :] < ROTARY_DIM_HALF) - x0 = tl.load(X, mask=mask, other=0.0).to(tl.float32) - x1 = tl.load(X + ROTARY_DIM_HALF * stride_x_headdim, mask=mask, other=0.0,).to(tl.float32) - o0 = x0 * cos - x1 * sin - o1 = x0 * sin + x1 * cos - tl.store(OUT, o0, mask=mask) - tl.store(OUT + ROTARY_DIM_HALF * stride_out_headdim, o1, mask=mask) - else: - rk = tl.arange(0, BLOCK_K) - X = X + (rh[:, None, None] * stride_x_nheads + rm[None, :, None] * stride_x_seqlen + rk[None, None, :] * stride_x_headdim) - OUT = OUT + (rh[:, None, None] * stride_out_nheads + rm[None, :, None] * stride_out_seqlen + rk[None, None, :] * stride_out_headdim) - mask = (rh[:, None, None] < nheads) & (rm[None, :, None] < seqlen) & (rk[None, None, :] < ROTARY_DIM) - x = tl.load(X, mask=mask, other=0.0).to(tl.float32) - x0, x1 = tl.split(tl.reshape(x, [BLOCK_H, BLOCK_M, BLOCK_K // 2, 2])) - o0 = x0 * cos - x1 * sin - o1 = x0 * sin + x1 * cos - o = tl.reshape(tl.join(o0, o1), [BLOCK_H, BLOCK_M, BLOCK_K]) - tl.store(OUT, o, mask=mask) - - -def apply_rotary( - x: torch.Tensor, - cos: torch.Tensor, - sin: torch.Tensor, - seqlen_offsets: Union[int, torch.Tensor] = 0, - cu_seqlens: Optional[torch.Tensor] = None, - max_seqlen: Optional[int] = None, - interleaved=False, - inplace=False, - conjugate=False, -) -> torch.Tensor: - """ - Arguments: - x: (batch, seqlen, nheads, headdim) if cu_seqlens is None - else (total_seqlen, nheads, headdim). - cos: (seqlen_ro, rotary_dim / 2) - sin: (seqlen_ro, rotary_dim / 2) - seqlen_offsets: integer or integer tensor of size (batch,) - cu_seqlens: (batch + 1,) or None - max_seqlen: int - Returns: - y: (batch, seqlen, nheads, headdim) - """ - is_varlen = cu_seqlens is not None - if not is_varlen: - batch, seqlen, nheads, headdim = x.shape - else: - assert max_seqlen is not None, "If cu_seqlens is passed in, then max_seqlen must be passed" - total_seqlen, nheads, headdim = x.shape - batch_p_1 = cu_seqlens.shape[0] - batch = batch_p_1 - 1 - seqlen = max_seqlen - seqlen_ro, rotary_dim = cos.shape - assert sin.shape == cos.shape - rotary_dim *= 2 - assert rotary_dim <= headdim, "rotary_dim must be <= headdim" - assert headdim <= 256, "Only support headdim <= 256" - assert seqlen_ro >= seqlen, "seqlen_ro must be >= seqlen" - - cos, sin = cos.contiguous(), sin.contiguous() - if isinstance(seqlen_offsets, torch.Tensor): - assert seqlen_offsets.shape == (batch,) - assert seqlen_offsets.dtype in [torch.int32, torch.int64] - seqlen_offsets = seqlen_offsets.contiguous() - else: - assert seqlen_offsets + seqlen <= seqlen_ro - - output = torch.empty_like(x) if not inplace else x - if rotary_dim < headdim and not inplace: - output[..., rotary_dim:].copy_(x[..., rotary_dim:]) - - grid = lambda META: (triton.cdiv(nheads, META["BLOCK_H"]), triton.cdiv(seqlen, META["BLOCK_M"]), batch) # noqa - BLOCK_M = 8 if rotary_dim <= 128 else 4 - - # Need this, otherwise Triton tries to launch from cuda:0 and we get - # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?) - device_ctx = torch.cuda.device(x.device.index) if x.device.type == 'cuda' else torch.xpu.device(x.device.index) - with device_ctx: - torch.library.wrap_triton(rotary_kernel)[grid]( - output, # data ptrs - x, - cos, - sin, - cu_seqlens, - seqlen_offsets, - seqlen, # shapes - nheads, - seqlen_ro, - output.stride(0) if not is_varlen else 0, # batch_strides if not varlen else 0 - output.stride(-3), # seqlen_stride or total_seqlen_stride - output.stride(-2), # nheads_stride - output.stride(-1), # headdim_stride - x.stride(0) if not is_varlen else 0, # batch_strides if not varlen else 0 - x.stride(-3), # seqlen stride or total_seqlen_stride - x.stride(-2), # nheads stride - x.stride(-1), # headdim stride - rotary_dim, - isinstance(seqlen_offsets, torch.Tensor), - is_varlen, - interleaved, - conjugate, - BLOCK_M=BLOCK_M, - BLOCK_H=2, - ) - return output diff --git a/build/torch212-cxx11-cu126-aarch64-linux/__init__.py b/build/torch212-cxx11-cu126-aarch64-linux/__init__.py deleted file mode 100644 index ecc2f9d896b6c93f90b0a1499856dc0612177422..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu126-aarch64-linux/__init__.py +++ /dev/null @@ -1,393 +0,0 @@ -from typing import Optional, List -import torch -from ._ops import ops as flash_attn_ops -from .flash_attn_interface import ( - flash_attn_func, - flash_attn_kvpacked_func, - flash_attn_qkvpacked_func, - flash_attn_varlen_func, - flash_attn_varlen_kvpacked_func, - flash_attn_varlen_qkvpacked_func, - flash_attn_with_kvcache, -) - - -def fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - return_softmax: bool = False, - gen: Optional[torch.Generator] = None, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention. - - Args: - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Optional output tensor, same shape as q - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - return_softmax: Whether to return softmax weights - gen: Optional random number generator - - Returns: - List of tensors: [output, softmax_lse, (softmax if return_softmax)] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.fwd( - q, - k, - v, - out, - alibi_slopes, - p_dropout, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - gen, - ) - - -def varlen_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - out: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - block_table: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - max_seqlen_q: int = 0, - max_seqlen_k: int = 0, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - zero_tensors: bool = False, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - return_softmax: bool = False, - gen: Optional[torch.Generator] = None, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention with variable sequence lengths. - - Args: - q: Query tensor of shape [total_q, num_heads, head_size] - k: Key tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - v: Value tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1] - cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1] - out: Optional output tensor of shape [total_q, num_heads, head_size] - seqused_k: Optional tensor specifying how many keys to use per batch element [batch_size] - leftpad_k: Optional left padding for keys of shape [batch_size] - block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq] - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - max_seqlen_q: Maximum sequence length for queries - max_seqlen_k: Maximum sequence length for keys - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - zero_tensors: Whether to zero tensors before computation - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - return_softmax: Whether to return softmax weights - gen: Optional random number generator - - Returns: - List of tensors: [output, softmax_lse, (softmax if return_softmax)] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.varlen_fwd( - q, - k, - v, - out, - cu_seqlens_q, - cu_seqlens_k, - seqused_k, - leftpad_k, - block_table, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - p_dropout, - softmax_scale, - zero_tensors, - is_causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - gen, - ) - - -def bwd( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - p_dropout: float = 0.0, - 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, - gen: Optional[torch.Generator] = None, - rng_state: Optional[torch.Tensor] = None, -) -> List[torch.Tensor]: - """ - Backward pass for multi-head attention. - - Args: - dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size] - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size] - softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q] - dq: Optional gradient tensor for queries, same shape as q - dk: Optional gradient tensor for keys, same shape as k - dv: Optional gradient tensor for values, same shape as v - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - deterministic: Whether to use deterministic algorithms - gen: Optional random number generator - rng_state: Optional RNG state from forward pass - - Returns: - List of tensors: [dq, dk, dv] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - alibi_slopes, - p_dropout, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - deterministic, - gen, - rng_state, - ) - - -def varlen_bwd( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - max_seqlen_q: int = 0, - max_seqlen_k: int = 0, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - zero_tensors: bool = False, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - deterministic: bool = False, - gen: Optional[torch.Generator] = None, - rng_state: Optional[torch.Tensor] = None, -) -> List[torch.Tensor]: - """ - Backward pass for multi-head attention with variable sequence lengths. - - Args: - dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size] - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size] - softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q] - cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1] - cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1] - dq: Optional gradient tensor for queries, same shape as q - dk: Optional gradient tensor for keys, same shape as k - dv: Optional gradient tensor for values, same shape as v - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - max_seqlen_q: Maximum sequence length for queries - max_seqlen_k: Maximum sequence length for keys - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - zero_tensors: Whether to zero tensors before computation - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - deterministic: Whether to use deterministic algorithms - gen: Optional random number generator - rng_state: Optional RNG state from forward pass - - Returns: - List of tensors: [dq, dk, dv] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.varlen_bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - p_dropout, - softmax_scale, - zero_tensors, - is_causal, - window_size_left, - window_size_right, - softcap, - deterministic, - gen, - rng_state, - ) - - -def fwd_kvcache( - q: torch.Tensor, - kcache: torch.Tensor, - vcache: torch.Tensor, - k: Optional[torch.Tensor] = None, - v: Optional[torch.Tensor] = None, - seqlens_k: Optional[torch.Tensor] = None, - rotary_cos: Optional[torch.Tensor] = None, - rotary_sin: Optional[torch.Tensor] = None, - cache_batch_idx: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - block_table: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - out: 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, - is_rotary_interleaved: bool = False, - num_splits: int = 1, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention with KV cache. - - Args: - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - kcache: Key cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - vcache: Value cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - k: Optional new keys tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size] - v: Optional new values tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size] - seqlens_k: Optional sequence lengths for keys of shape [batch_size] - rotary_cos: Optional rotary cosine tensor of shape [seqlen_ro, rotary_dim/2] - rotary_sin: Optional rotary sine tensor of shape [seqlen_ro, rotary_dim/2] - cache_batch_idx: Optional indices to index into the KV cache - leftpad_k: Optional left padding for keys of shape [batch_size] - block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq] - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - out: Optional output tensor, same shape as q - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - is_rotary_interleaved: Whether rotary embeddings are interleaved - num_splits: Number of splits for computation - - Returns: - List of tensors: [output, softmax_lse] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.fwd_kvcache( - q, - kcache, - vcache, - k, - v, - seqlens_k, - rotary_cos, - rotary_sin, - cache_batch_idx, - leftpad_k, - block_table, - alibi_slopes, - out, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - is_rotary_interleaved, - num_splits, - ) diff --git a/build/torch212-cxx11-cu126-aarch64-linux/_flash_attn2_cuda_f12afc9.abi3.so b/build/torch212-cxx11-cu126-aarch64-linux/_flash_attn2_cuda_f12afc9.abi3.so deleted file mode 100644 index aac7e7ff5e9f0f027d8cf51951144aed6f4d9aa5..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu126-aarch64-linux/_flash_attn2_cuda_f12afc9.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:37f77edec6546c4cdcd0c566a6b095b84e70354ae6c0e0f6f6374b4e6845738c -size 446586496 diff --git a/build/torch212-cxx11-cu126-aarch64-linux/_ops.py b/build/torch212-cxx11-cu126-aarch64-linux/_ops.py deleted file mode 100644 index af5b731cb04b5e6ce54e862db0501f3683c3163f..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu126-aarch64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _flash_attn2_cuda_f12afc9 -ops = torch.ops._flash_attn2_cuda_f12afc9 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_flash_attn2_cuda_f12afc9::{op_name}" diff --git a/build/torch212-cxx11-cu126-aarch64-linux/bert_padding.py b/build/torch212-cxx11-cu126-aarch64-linux/bert_padding.py deleted file mode 100644 index 3c2d35159a014a9d03aabead9e52e009168696ea..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu126-aarch64-linux/bert_padding.py +++ /dev/null @@ -1,218 +0,0 @@ -# Adapted from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/padding.py - -import torch -import torch.nn.functional as F -from einops import rearrange, repeat - - -class IndexFirstAxis(torch.autograd.Function): - @staticmethod - def forward(ctx, input, indices): - ctx.save_for_backward(indices) - assert input.ndim >= 2 - ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] - second_dim = other_shape.numel() - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - # return input[indices] - return torch.gather( - rearrange(input, "b ... -> b (...)"), 0, repeat(indices, "z -> z d", d=second_dim) - ).reshape(-1, *other_shape) - - @staticmethod - def backward(ctx, grad_output): - (indices,) = ctx.saved_tensors - assert grad_output.ndim >= 2 - other_shape = grad_output.shape[1:] - grad_output = rearrange(grad_output, "b ... -> b (...)") - grad_input = torch.zeros( - [ctx.first_axis_dim, grad_output.shape[1]], - device=grad_output.device, - dtype=grad_output.dtype, - ) - # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. - # grad_input[indices] = grad_output - grad_input.scatter_(0, repeat(indices, "z -> z d", d=grad_output.shape[1]), grad_output) - return grad_input.reshape(ctx.first_axis_dim, *other_shape), None - - -index_first_axis = IndexFirstAxis.apply - - -class IndexPutFirstAxis(torch.autograd.Function): - @staticmethod - def forward(ctx, values, indices, first_axis_dim): - ctx.save_for_backward(indices) - assert indices.ndim == 1 - assert values.ndim >= 2 - output = torch.zeros( - first_axis_dim, *values.shape[1:], device=values.device, dtype=values.dtype - ) - # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. - output[indices] = values - # output.scatter_(0, repeat(indices, 'z -> z d', d=values.shape[1]), values) - return output - - @staticmethod - def backward(ctx, grad_output): - (indices,) = ctx.saved_tensors - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - grad_values = grad_output[indices] - # grad_values = torch.gather(grad_output, 0, repeat(indices, 'z -> z d', d=grad_output.shape[1])) - return grad_values, None, None - - -index_put_first_axis = IndexPutFirstAxis.apply - - -class IndexFirstAxisResidual(torch.autograd.Function): - @staticmethod - def forward(ctx, input, indices): - ctx.save_for_backward(indices) - assert input.ndim >= 2 - ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] - second_dim = other_shape.numel() - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - output = input[indices] - # We don't want to reshape input (b ... -> b (...)) since it could change the channel_last - # memory format to channel_first. In other words, input might not be contiguous. - # If we don't detach, Pytorch complains about output being a view and is being modified inplace - return output, input.detach() - - @staticmethod - def backward(ctx, grad_output, grad_residual): - (indices,) = ctx.saved_tensors - assert grad_output.ndim >= 2 - other_shape = grad_output.shape[1:] - assert grad_residual.shape[1:] == other_shape - grad_input = grad_residual - # grad_input[indices] += grad_output - indices = indices.reshape(indices.shape[0], *((1,) * (grad_output.ndim - 1))) - indices = indices.expand_as(grad_output) - grad_input.scatter_add_(0, indices, grad_output) - return grad_input.reshape(ctx.first_axis_dim, *other_shape), None - - -index_first_axis_residual = IndexFirstAxisResidual.apply - - -def unpad_input(hidden_states, attention_mask, unused_mask=None): - """ - Arguments: - hidden_states: (batch, seqlen, ...) - attention_mask: (batch, seqlen), bool / int, 1 means valid and 0 means not valid. - unused_mask: (batch, seqlen), bool / int, 1 means the element is allocated but unused. - Return: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens selected in attention_mask + unused_mask. - indices: (total_nnz), the indices of masked tokens from the flattened input sequence. - cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. - max_seqlen_in_batch: int - seqused: (batch), returns the number of tokens selected in attention_mask + unused_mask. - """ - all_masks = (attention_mask + unused_mask) if unused_mask is not None else attention_mask - seqlens_in_batch = all_masks.sum(dim=-1, dtype=torch.int32) - used_seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) - indices = torch.nonzero(all_masks.flatten(), as_tuple=False).flatten() - max_seqlen_in_batch = seqlens_in_batch.max().item() - cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) - # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the - # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim - # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to - # index with integer indices. Moreover, torch's index is a bit slower than it needs to be, - # so we write custom forward and backward to make it a bit faster. - return ( - index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices), - indices, - cu_seqlens, - max_seqlen_in_batch, - used_seqlens_in_batch, - ) - - -def unpad_input_for_concatenated_sequences(hidden_states, attention_mask_in_length): - """ - Supports concatenating short samples in one sequence. The attention_mask_in_length is utilized to mask other short samples. It helps efficient training of variant lengths-based samples (e.g., the supervised fine-tuning task in large language model). - The motivation for this function is explained [here](https://github.com/Dao-AILab/flash-attention/issues/432#issuecomment-1668822286). - - For example, if batch = 3 and seqlen = 6, the attention_mask_in_length is: - ``` - [ - [2, 3, 0, 0, 0, 0], - [3, 2, 0, 0, 0, 0], - [6, 0, 0, 0, 0, 0] - ] - ``` - , which refers to the 3D-attention mask: - ``` - [ - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [0, 0, 1, 0, 0, 0], - [0, 0, 1, 1, 0, 0], - [0, 0, 1, 1, 1, 0], - [0, 0, 0, 0, 0, 1] - ], - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [0, 0, 0, 1, 0, 0], - [0, 0, 0, 1, 1, 0], - [0, 0, 0, 0, 0, 1] - ], - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [1, 1, 1, 1, 0, 0], - [1, 1, 1, 1, 1, 0], - [1, 1, 1, 1, 1, 1] - ] - ] - ```. - - Arguments: - hidden_states: (batch, seqlen, ...) - attention_mask_in_length: (batch, seqlen), int, a nonzero number (e.g., 1, 2, 3, etc.) means length of concatenated sequence in b-th batch, and 0 means none. - Return: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. - indices: (total_nnz), the indices of non-masked tokens from the flattened input sequence. - cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. - max_seqlen_in_batch: int - """ - length = attention_mask_in_length.sum(dim=-1) - seqlen = attention_mask_in_length.size(-1) - attention_mask_2d = torch.arange(seqlen, device=length.device, dtype=length.dtype).expand(len(length), seqlen) < length.unsqueeze(1) - real_indices_idx = torch.nonzero(attention_mask_in_length.flatten(), as_tuple=False).flatten() - seqlens_in_batch = attention_mask_in_length.flatten()[real_indices_idx] - indices = torch.nonzero(attention_mask_2d.flatten(), as_tuple=False).flatten() - max_seqlen_in_batch = seqlens_in_batch.max().item() - cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) - # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the - # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim - # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to - # index with integer indices. Moreover, torch's index is a bit slower than it needs to be, - # so we write custom forward and backward to make it a bit faster. - return ( - index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices), - indices, - cu_seqlens, - max_seqlen_in_batch, - ) - - -def pad_input(hidden_states, indices, batch, seqlen): - """ - Arguments: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. - indices: (total_nnz), the indices that represent the non-masked tokens of the original padded input sequence. - batch: int, batch size for the padded sequence. - seqlen: int, maximum sequence length for the padded sequence. - Return: - hidden_states: (batch, seqlen, ...) - """ - dim = hidden_states.shape[-1] - # output = torch.zeros((batch * seqlen), dim, device=hidden_states.device, dtype=hidden_states.dtype) - # output[indices] = hidden_states - output = index_put_first_axis(hidden_states, indices, batch * seqlen) - return rearrange(output, "(b s) ... -> b s ...", b=batch) diff --git a/build/torch212-cxx11-cu126-aarch64-linux/flash_attn2/__init__.py b/build/torch212-cxx11-cu126-aarch64-linux/flash_attn2/__init__.py deleted file mode 100644 index a9b2672c1cd85b74c1b3ded0fc0b2100e1aeac23..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu126-aarch64-linux/flash_attn2/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -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/torch212-cxx11-cu126-aarch64-linux/flash_attn_interface.py b/build/torch212-cxx11-cu126-aarch64-linux/flash_attn_interface.py deleted file mode 100644 index b014e569dc9812a9c840d1043f2cebe767b4d477..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu126-aarch64-linux/flash_attn_interface.py +++ /dev/null @@ -1,1631 +0,0 @@ -# Copyright (c) 2023, Tri Dao. - -from typing import Optional, Sequence, Tuple, Union - -import torch -import torch.nn as nn -import os - -# # isort: off -# # We need to import the CUDA kernels after importing torch -# USE_TRITON_ROCM = os.getenv("FLASH_ATTENTION_TRITON_AMD_ENABLE", "FALSE") == "TRUE" -# if USE_TRITON_ROCM: -# from .flash_attn_triton_amd import interface_fa as flash_attn -# else: -# import flash_attn_2_cuda as flash_attn - - -from . import _ops -from ._ops import ops as flash_attn -from ._ops import add_op_namespace_prefix - -# # isort: on - -def maybe_contiguous(x): - return x.contiguous() if x is not None and x.stride(-1) != 1 else x - - -def _get_device(): - if torch.xpu.is_available(): - return "xpu" - elif torch.cuda.is_available(): - return "cuda" - else: - return "cpu" - - -def _get_block_size_n(device, head_dim, is_dropout, is_causal): - assert head_dim <= 256 - - if device.type == "xpu": - if head_dim <= 96: - return 64 - elif head_dim <= 128: - return 32 - elif head_dim <= 256: - return 64 - else: - return 32 - - # This should match the block sizes in the CUDA kernel - major, minor = torch.cuda.get_device_capability(device) - is_sm8x = major == 8 and minor > 0 # Only include sm86 and sm89, exclude sm80 (A100) - is_sm80 = major == 8 and minor == 0 - is_sm90 = major == 9 and minor == 0 - if head_dim <= 32: - return 128 - if head_dim <= 64: - return 128 if not is_dropout else 64 - elif head_dim <= 96: - return 64 - elif head_dim <= 128: - if is_sm8x: - return 64 if (not is_dropout and is_causal) else 32 - else: - return 64 if not is_dropout else 32 - elif head_dim <= 192: - return 64 - elif head_dim <= 224: - return 64 - elif head_dim <= 256: - return 64 - - -def round_multiple(x, m): - return (x + m - 1) // m * m - - -# torch.compile() support is only enabled for pytorch >= 2.4 -# The reason for this is that we are using the new custom_op and register_fake -# APIs, which support inplace modification of inputs in the function itself -if torch.__version__ >= "2.4.0": - _torch_custom_op_wrapper = torch.library.custom_op - _torch_register_fake_wrapper = torch.library.register_fake -else: - def noop_custom_op_wrapper(name, fn=None, /, *, mutates_args, device_types=None, schema=None): - def wrap(func): - return func - if fn is None: - return wrap - return fn - def noop_register_fake_wrapper(op, fn=None, /, *, lib=None, _stacklevel=1): - def wrap(func): - return func - if fn is None: - return wrap - return fn - _torch_custom_op_wrapper = noop_custom_op_wrapper - _torch_register_fake_wrapper = noop_register_fake_wrapper - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_forward"), mutates_args=(), device_types=_get_device()) -def _flash_attn_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - return_softmax: bool -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - out, softmax_lse, S_dmask, rng_state = flash_attn.fwd( - q, - k, - v, - None, - alibi_slopes, - dropout_p, - softmax_scale, - causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - None, - ) - return out, softmax_lse, S_dmask, rng_state - - -@_torch_register_fake_wrapper(add_op_namespace_prefix("_flash_attn_forward")) -def _flash_attn_forward_fake( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - return_softmax: bool -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - batch_size, seqlen_q, num_heads, head_size = q.shape - seqlen_k = k.shape[1] - out = torch.empty_like(q) - softmax_lse = torch.empty((batch_size, num_heads, seqlen_q), dtype=torch.float32, device=q.device, layout=q.layout) - p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) - if return_softmax: - p = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128), round_multiple(seqlen_k, 128)), dtype=q.dtype, device=q.device, layout=q.layout) - rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) - - return out, softmax_lse, p, rng_state - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_forward = _ops.ops._flash_attn_forward -else: - _wrapped_flash_attn_forward = _flash_attn_forward - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_varlen_forward"), mutates_args=(), device_types=_get_device()) -def _flash_attn_varlen_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - alibi_slopes: Optional[torch.Tensor] = None, - return_softmax: bool = False, - block_table: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - out, softmax_lse, S_dmask, rng_state = flash_attn.varlen_fwd( - q, - k, - v, - None, - cu_seqlens_q, - cu_seqlens_k, - seqused_k, - leftpad_k, - block_table, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - zero_tensors, - causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - None, - ) - # if out.isnan().any() or softmax_lse.isnan().any(): - # breakpoint() - return out, softmax_lse, S_dmask, rng_state - - -@_torch_register_fake_wrapper(add_op_namespace_prefix("_flash_attn_varlen_forward")) -def _flash_attn_varlen_forward_fake( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - alibi_slopes: Optional[torch.Tensor] = None, - return_softmax: bool = False, - block_table: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - paged_kv = block_table is not None - batch_size = cu_seqlens_q.numel() - 1 - total_q, num_heads, _ = q.shape - - out = torch.empty_like(q) - softmax_lse = torch.empty((num_heads, total_q), dtype=torch.float32, device=q.device, layout=q.layout) - p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) - seqlen_q_rounded = round_multiple(max_seqlen_q, 128) - seqlen_k_rounded = round_multiple(max_seqlen_k, 128) - if return_softmax: - p = torch.empty((batch_size, num_heads, seqlen_q_rounded, seqlen_k_rounded), dtype=q.dtype, device=q.device, layout=q.layout) - rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) - return out, softmax_lse, p, rng_state - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_varlen_forward = _ops.ops._flash_attn_varlen_forward -else: - _wrapped_flash_attn_varlen_forward = _flash_attn_varlen_forward - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_backward"), mutates_args=("dq", "dk", "dv"), device_types=_get_device()) -def _flash_attn_backward( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, -) -> 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)] - ( - dq, - dk, - dv, - softmax_d, - ) = flash_attn.bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - alibi_slopes, - dropout_p, - softmax_scale, - causal, - window_size_left, - window_size_right, - softcap, - deterministic, - None, - rng_state, - ) - return softmax_d - - -@_torch_register_fake_wrapper(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, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, -) -> torch.Tensor: - dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] - if dq is None: - dq = torch.empty_like(q) - if dk is None: - dk = torch.empty_like(k) - if dv is None: - dv = torch.empty_like(v) - batch_size, seqlen_q, num_heads, _ = q.shape - softmax_d = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128)), device=q.device, dtype=torch.float32) - - return softmax_d - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_backward = _ops.ops._flash_attn_backward -else: - _wrapped_flash_attn_backward = _flash_attn_backward - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_varlen_backward"), mutates_args=("dq", "dk", "dv"), device_types=_get_device()) -def _flash_attn_varlen_backward( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> 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)] - ( - dq, - dk, - dv, - softmax_d, - ) = flash_attn.varlen_bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - zero_tensors, - causal, - window_size_left, - window_size_right, - softcap, - deterministic, - None, - rng_state, - ) - # if dk.isnan().any() or dk.isnan().any() or dv.isnan().any() or softmax_d.isnan().any(): - # breakpoint() - return softmax_d - - -@_torch_register_fake_wrapper(add_op_namespace_prefix("_flash_attn_varlen_backward")) -def _flash_attn_varlen_backward_fake( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> torch.Tensor: - dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] - batch_size = cu_seqlens_q.numel() - 1 - total_q, num_heads, _ = q.shape - - if dq is None: - dq = torch.empty_like(q) - if dk is None: - dk = torch.empty_like(k) - if dv is None: - dv = torch.empty_like(v) - softmax_d = torch.empty((num_heads, total_q + 128 * batch_size), device=q.device, dtype=torch.float32) - - return softmax_d - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_varlen_backward = _ops.ops._flash_attn_varlen_backward -else: - _wrapped_flash_attn_varlen_backward = _flash_attn_varlen_backward - - -class FlashAttnQKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and qkv.requires_grad - if softmax_scale is None: - softmax_scale = qkv.shape[-1] ** (-0.5) - q, k, v = qkv[:, :, 0].detach(), qkv[:, :, 1].detach(), qkv[:, :, 2].detach() - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) - dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dqkv[:, :, 0], - dqkv[:, :, 1], - dqkv[:, :, 2], - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dqkv = dqkv[..., : dout.shape[-1]] # We could have padded the head dimension - return dqkv, None, None, None, None, None, None, None, None, None - - -class FlashAttnVarlenQKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - cu_seqlens, - max_seqlen, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and qkv.requires_grad - if softmax_scale is None: - softmax_scale = qkv.shape[-1] ** (-0.5) - q, k, v = qkv[:, 0].detach(), qkv[:, 1].detach(), qkv[:, 2].detach() - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens, - cu_seqlens, - max_seqlen, - max_seqlen, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=None, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens, rng_state) - ctx.dropout_p = dropout_p - ctx.max_seqlen = max_seqlen - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens, rng_state = ctx.saved_tensors - qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) - dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dqkv[:, 0], - dqkv[:, 1], - dqkv[:, 2], - cu_seqlens, - cu_seqlens, - ctx.max_seqlen, - ctx.max_seqlen, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - 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 FlashAttnKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - kv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, kv] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - k, v = kv[:, :, 0].detach(), kv[:, :, 1].detach() - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - dq = torch.empty_like(q) - kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) - dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dkv[:, :, 0], - dkv[:, :, 1], - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dkv = dkv[..., : dout.shape[-1]] - return dq, dkv, None, None, None, None, None, None, None, None, None - - -class FlashAttnVarlenKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, kv] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - k, v = kv[:, 0].detach(), kv[:, 1].detach() - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=None, - ) - if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state - ) - ctx.dropout_p = dropout_p - 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.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors - dq = torch.empty_like(q) - kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) - dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dkv[:, 0], - dkv[:, 1], - cu_seqlens_q, - cu_seqlens_k, - ctx.max_seqlen_q, - ctx.max_seqlen_k, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dkv = dkv[..., : dout.shape[-1]] - return dq, dkv, None, None, None, None, None, None, None, None, None, None, None, None, None - - -class FlashAttnFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - k, - v, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, k, v] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dk = dk[..., : dout.shape[-1]] - dv = dv[..., : dout.shape[-1]] - return dq, dk, dv, 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, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - block_table, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, k, v] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=block_table, - ) - if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state - ) - ctx.dropout_p = dropout_p - 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.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors - dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - ctx.max_seqlen_q, - ctx.max_seqlen_k, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dk = dk[..., : dout.shape[-1]] - dv = dv[..., : dout.shape[-1]] - return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None - - -def flash_attn_qkvpacked_func( - qkv, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # <=0.0 means deactivate - alibi_slopes=None, - deterministic=False, - 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, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_kvpacked_func( - q, - kv, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, -): - """dropout_p should be set to 0.0 during evaluation - If 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 K, V. - 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) - kv: (batch_size, seqlen, 2, 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. - softcap: float. Anything > 0 activates softcapping 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). - 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 FlashAttnKVPackedFunc.apply( - q, - kv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_func( - q, - k, - v, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - 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). - 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 FlashAttnFunc.apply( - q, - k, - v, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_qkvpacked_func( - qkv, - cu_seqlens, - max_seqlen, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - 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_varlen_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_varlen_kvpacked_func and flash_attn_varlen_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: (total, 3, nheads, headdim), where total = total number of tokens in the batch. - cu_seqlens: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into qkv. - max_seqlen: int. Maximum sequence length in the batch. - 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenQKVPackedFunc.apply( - qkv, - cu_seqlens, - max_seqlen, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_kvpacked_func( - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, -): - """dropout_p should be set to 0.0 during evaluation - If 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 K, V. - 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: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch. - kv: (total_k, 2, nheads_k, headdim), where total_k = total number of key tokens in the batch. - cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into q. - cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into kv. - max_seqlen_q: int. Maximum query sequence length in the batch. - max_seqlen_k: int. Maximum key sequence length in the batch. - 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 + 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenKVPackedFunc.apply( - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_func( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, - block_table=None, -): - """dropout_p should be set to 0.0 during evaluation - Supports multi-query and grouped-query attention (MQA/GQA) by passing in K, V 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: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch. - k: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. - v: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. - cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into q. - cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into kv. - max_seqlen_q: int. Maximum query sequence length in the batch. - max_seqlen_k: int. Maximum key sequence length in the batch. - 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 + 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenFunc.apply( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - block_table, - False if q.device.type == "cpu" else torch.is_grad_enabled(), - ) - - -def flash_attn_with_kvcache( - q, - k_cache, - v_cache, - k=None, - v=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, - block_table: Optional[torch.Tensor] = None, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - rotary_interleaved=True, - alibi_slopes=None, - num_splits=0, - 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 block_table, - or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_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) if there's no block_table, - or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_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). Similar to k. - 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. - block_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). - 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. - 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" - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-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) - cache_batch_idx = maybe_contiguous(cache_batch_idx) - block_table = maybe_contiguous(block_table) - out, softmax_lse = flash_attn.fwd_kvcache( - q, - k_cache, - v_cache, - k, - v, - cache_seqlens, - rotary_cos, - rotary_sin, - cache_batch_idx, - cache_leftpad, - block_table, - alibi_slopes, - None, - softmax_scale, - causal, - window_size[0], - window_size[1], - softcap, - rotary_interleaved, - num_splits, - ) - return (out, softmax_lse) if return_softmax_lse else out diff --git a/build/torch212-cxx11-cu126-aarch64-linux/layers/__init__.py b/build/torch212-cxx11-cu126-aarch64-linux/layers/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/build/torch212-cxx11-cu126-aarch64-linux/layers/patch_embed.py b/build/torch212-cxx11-cu126-aarch64-linux/layers/patch_embed.py deleted file mode 100644 index 05562f8e8bcdb58e947c6f402a49eacd2d031871..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu126-aarch64-linux/layers/patch_embed.py +++ /dev/null @@ -1,67 +0,0 @@ -# We use the same API as https://github.com/rwightman/pytorch-image-models/blob/v0.6.11/timm/models/layers/patch_embed.py -# But we use nn.Linear instead of Conv2d and it's about 8x faster. - -from functools import partial - -import torch.nn as nn -from einops import rearrange -from torch import _assert -from torch.nn.modules.utils import _pair - -try: - from flash_attn.ops.fused_dense import FusedDense -except ImportError: - FusedDense = None - - -class PatchEmbed(nn.Module): - """2D Image to Patch Embedding""" - - def __init__( - self, - img_size=224, - patch_size=16, - in_chans=3, - embed_dim=768, - norm_layer=None, - flatten=True, - bias=True, - fused_bias_fc=False, - ): - super().__init__() - img_size = _pair(img_size) - patch_size = _pair(patch_size) - self.img_size = img_size - self.patch_size = patch_size - self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1]) - self.num_patches = self.grid_size[0] * self.grid_size[1] - self.flatten = flatten - if fused_bias_fc and FusedDense is None: - raise ImportError("fused_dense is not installed") - - linear_cls = nn.Linear if not fused_bias_fc or not bias else FusedDense - self.proj = linear_cls(in_chans * patch_size[0] * patch_size[1], embed_dim, bias=bias) - self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() - - def forward(self, x): - _, _, H, W = x.shape - _assert( - H == self.img_size[0], - f"Input image height ({H}) doesn't match model ({self.img_size[0]}).", - ) - _assert( - W == self.img_size[1], - f"Input image width ({W}) doesn't match model ({self.img_size[1]}).", - ) - x = self.proj( - rearrange( - x, - "b c (h p1) (w p2) -> b h w (c p1 p2)", - p1=self.patch_size[0], - p2=self.patch_size[1], - ) - ) - if self.flatten: - x = rearrange(x, "b h w c -> b (h w) c") - x = self.norm(x) - return x diff --git a/build/torch212-cxx11-cu126-aarch64-linux/layers/rotary.py b/build/torch212-cxx11-cu126-aarch64-linux/layers/rotary.py deleted file mode 100644 index d1bfc21fc7de1dd287e8f382847b194a48075981..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu126-aarch64-linux/layers/rotary.py +++ /dev/null @@ -1,483 +0,0 @@ -# Copyright (c) 2025, Tri Dao - -import math -from functools import partial -from typing import Optional, Tuple, Union - -import torch -from torch import Tensor - -from einops import rearrange, repeat -# from flash_attn.ops.triton.rotary import apply_rotary -from ..ops.triton.rotary import apply_rotary - - -def rotate_half(x, interleaved=False): - if not interleaved: - x1, x2 = x.chunk(2, dim=-1) - return torch.cat((-x2, x1), dim=-1) - else: - x1, x2 = x[..., ::2], x[..., 1::2] - return rearrange(torch.stack((-x2, x1), dim=-1), "... d two -> ... (d two)", two=2) - - -def apply_rotary_emb_torch(x, cos, sin, interleaved=False): - """ - x: (batch_size, seqlen, nheads, headdim) - cos, sin: (seqlen, rotary_dim / 2) or (batch_size, seqlen, rotary_dim / 2) - """ - ro_dim = cos.shape[-1] * 2 - assert ro_dim <= x.shape[-1] - cos = repeat(cos, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)") - sin = repeat(sin, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)") - return torch.cat( - [x[..., :ro_dim] * cos + rotate_half(x[..., :ro_dim], interleaved) * sin, x[..., ro_dim:]], - dim=-1, - ) - - -class ApplyRotaryEmb(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x, - cos, - sin, - interleaved=False, - inplace=False, - seqlen_offsets: Union[int, Tensor] = 0, - cu_seqlens: Optional[Tensor] = None, - max_seqlen: Optional[int] = None, - ): - out = apply_rotary( - x, - cos, - sin, - seqlen_offsets=seqlen_offsets, - cu_seqlens=cu_seqlens, - max_seqlen=max_seqlen, - interleaved=interleaved, - inplace=inplace, - ) - if isinstance(seqlen_offsets, int): - ctx.save_for_backward(cos, sin, cu_seqlens) # Can't save int with save_for_backward - ctx.seqlen_offsets = seqlen_offsets - else: - ctx.save_for_backward(cos, sin, cu_seqlens, seqlen_offsets) - ctx.seqlen_offsets = None - ctx.interleaved = interleaved - ctx.inplace = inplace - ctx.max_seqlen = max_seqlen - return out if not inplace else x - - @staticmethod - def backward(ctx, do): - seqlen_offsets = ctx.seqlen_offsets - if seqlen_offsets is None: - cos, sin, cu_seqlens, seqlen_offsets = ctx.saved_tensors - else: - cos, sin, cu_seqlens = ctx.saved_tensors - dx = apply_rotary( - do, - cos, - sin, - seqlen_offsets=seqlen_offsets, - cu_seqlens=cu_seqlens, - max_seqlen=ctx.max_seqlen, - interleaved=ctx.interleaved, - inplace=ctx.inplace, - conjugate=True, - ) - return dx, None, None, None, None, None, None, None - - -def apply_rotary_emb( - x, - cos, - sin, - interleaved=False, - inplace=False, - seqlen_offsets: Union[int, Tensor] = 0, - cu_seqlens: Optional[Tensor] = None, - max_seqlen: Optional[int] = None, -): - """ - Arguments: - x: (batch_size, seqlen, nheads, headdim) if cu_seqlens is None - else (total_seqlen, nheads, headdim) - cos, sin: (seqlen_rotary, rotary_dim / 2) - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead - of 1st half and 2nd half (GPT-NeoX style). - inplace: if True, apply rotary embedding in-place. - seqlen_offsets: (batch_size,) or int. Each sequence in x is shifted by this amount. - Most commonly used in inference when we have KV cache. - cu_seqlens: (batch + 1,) or None - max_seqlen: int - Return: - out: (batch_size, seqlen, nheads, headdim) if cu_seqlens is None - else (total_seqlen, nheads, headdim) - rotary_dim must be <= headdim - Apply rotary embedding to the first rotary_dim of x. - """ - return ApplyRotaryEmb.apply( - x, cos, sin, interleaved, inplace, seqlen_offsets, cu_seqlens, max_seqlen - ) - - -# For backward compatibility -apply_rotary_emb_func = apply_rotary_emb - - -def _apply_rotary_emb_qkv( - qkv, - cos, - sin, - cos_k=None, - sin_k=None, - interleaved=False, - inplace=False, - conjugate=False, - seqlen_offsets: Union[int, Tensor] = 0, - num_heads_q: Optional[int] = None, -): - apply_rotary_fn = partial( - apply_rotary, - interleaved=interleaved, - inplace=inplace, - conjugate=conjugate, - seqlen_offsets=seqlen_offsets - ) - if cos_k is None and sin_k is None and qkv.is_contiguous(): - # Call 1 kernel instead of 2 kernels - # We need qkv to be contiguous so that when we reshape to combine (3, nheads) - # dimensions, we get the same tensor - if qkv.dim() == 5: - batch, seqlen, three, nheads, headdim = qkv.shape - assert three == 3 - # qk = rearrange(qkv[:, :, :2], "b s t h d -> b s (t h) d") - qk = qkv[:, :, :2].reshape(batch, seqlen, -1, headdim) - qk = apply_rotary_fn(qk, cos, sin) - else: - assert qkv.dim() == 4 - assert num_heads_q is not None - num_heads_k = (qkv.shape[2] - num_heads_q) // 2 - assert qkv.shape[2] == num_heads_q + 2 * num_heads_k - qk = qkv[:, :, :num_heads_q + num_heads_k] - qk = apply_rotary_fn(qk, cos, sin) - if not inplace: - if qkv.dim() == 5: - qkv = torch.cat([rearrange(qk, "b s (t h) d -> b s t h d", t=2), qkv[:, :, 2:]], dim=2) - else: - qkv = torch.cat([qk, qkv[:, :, num_heads_q + num_heads_k :]], dim=2) - else: - cos_k = cos if cos_k is None else cos_k - sin_k = sin if sin_k is None else sin_k - if qkv.dim() == 5: - batch, seqlen, three, nheads, headdim = qkv.shape - assert three == 3 - q, k = qkv[:, :, 0], qkv[:, :, 1] - else: - assert qkv.dim() == 4 - assert num_heads_q is not None - num_heads_k = (qkv.shape[2] - num_heads_q) // 2 - assert qkv.shape[2] == num_heads_q + 2 * num_heads_k - q, k = qkv[:, :, :num_heads_q], qkv[:, :, num_heads_q : num_heads_q + num_heads_k] - q = apply_rotary_fn(q, cos, sin) - k = apply_rotary_fn(k, cos_k, sin_k) - if not inplace: - if qkv.dim() == 5: - qkv = torch.stack([q, k, qkv[:, :, 2]], dim=2) - else: - qkv = torch.cat([q, k, qkv[:, :, num_heads_q + num_heads_k:]], dim=2) - return qkv - - -class ApplyRotaryEmbQKV_(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - cos, - sin, - cos_k=None, - sin_k=None, - interleaved=False, - seqlen_offsets: Union[int, torch.Tensor] = 0, - num_heads_q: Optional[int] = None, - ): - # apply_rotary_emb_qkv_inplace( - qkv = _apply_rotary_emb_qkv( - qkv, cos, sin, cos_k, sin_k, interleaved=interleaved, inplace=True, - seqlen_offsets=seqlen_offsets, num_heads_q=num_heads_q, - ) - if isinstance(seqlen_offsets, int): - ctx.save_for_backward(cos, sin, cos_k, sin_k) - ctx.seqlen_offsets = seqlen_offsets - else: - ctx.save_for_backward(cos, sin, cos_k, sin_k, seqlen_offsets) - ctx.seqlen_offsets = None - ctx.interleaved = interleaved - ctx.num_heads_q = num_heads_q - return qkv - - @staticmethod - def backward(ctx, dqkv): - seqlen_offsets = ctx.seqlen_offsets - if seqlen_offsets is None: - cos, sin, cos_k, sin_k, seqlen_offsets = ctx.saved_tensors - else: - cos, sin, cos_k, sin_k = ctx.saved_tensors - dqkv = _apply_rotary_emb_qkv( - dqkv, cos, sin, cos_k, sin_k, interleaved=ctx.interleaved, inplace=True, - seqlen_offsets=seqlen_offsets, num_heads_q=ctx.num_heads_q, conjugate=True, - ) - return dqkv, None, None, None, None, None, None, None - - -def apply_rotary_emb_qkv_( - qkv, - cos, - sin, - cos_k=None, - sin_k=None, - interleaved=False, - seqlen_offsets: Union[int, torch.Tensor] = 0, - num_heads_q: Optional[int] = None, -): - """ - Arguments: - qkv: (batch_size, seqlen, 3, nheads, headdim) or (batch_size, seqlen, num_heads_q + 2 * num_heads_k, headdim). - If qkv has shape (batch_size, seqlen, num_heads_q + 2 * num_heads_k, headdim) (e.g. MQA / GQA), - then num_heads_q must be provided. - cos, sin: (seqlen, rotary_dim / 2) - cos_k, sin_k: (seqlen, rotary_dim / 2), optional - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead of - 1st half and 2nd half (GPT-NeoX style). - seqlen_offsets: (batch_size,) or int. Each sequence in Q and K is shifted by this amount. - Most commonly used in inference when we have KV cache. - Return: - qkv: (batch_size, seqlen, 3, nheads, headdim) or (batch_size, seqlen, num_heads_q + 2 * num_heads_k, headdim) - rotary_dim must be <= headdim - Apply rotary embedding *inplace* to the first rotary_dim of Q and K. - """ - return ApplyRotaryEmbQKV_.apply( - qkv, cos, sin, cos_k, sin_k, interleaved, seqlen_offsets, num_heads_q - ) - - -class ApplyRotaryEmbKV_(torch.autograd.Function): - - @staticmethod - def forward(ctx, kv, cos, sin, interleaved=False, seqlen_offsets: Union[int, torch.Tensor] = 0): - batch, seqlen, two, nheads, headdim = kv.shape - assert two == 2 - k = kv[:, :, 0] - apply_rotary( - k, cos, sin, seqlen_offsets=seqlen_offsets, interleaved=interleaved, inplace=True - ) - if isinstance(seqlen_offsets, int): - ctx.save_for_backward(cos, sin) # Can't save int with save_for_backward - ctx.seqlen_offsets = seqlen_offsets - else: - ctx.save_for_backward(cos, sin, seqlen_offsets) - ctx.seqlen_offsets = None - ctx.interleaved = interleaved - return kv - - @staticmethod - def backward(ctx, dkv): - seqlen_offsets = ctx.seqlen_offsets - if seqlen_offsets is None: - cos, sin, seqlen_offsets = ctx.saved_tensors - else: - cos, sin = ctx.saved_tensors - apply_rotary( - dkv[:, :, 0], - cos, - sin, - seqlen_offsets=seqlen_offsets, - interleaved=ctx.interleaved, - inplace=True, - conjugate=True, - ) - return dkv, None, None, None, None - - -apply_rotary_emb_kv_ = ApplyRotaryEmbKV_.apply - - -def apply_rotary_emb_kv_( - kv, - cos, - sin, - interleaved=False, - seqlen_offsets: Union[int, torch.Tensor] = 0, -): - """ - Arguments: - kv: (batch_size, seqlen, 2, nheads, headdim) - cos, sin: (seqlen, rotary_dim / 2) - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead of - 1st half and 2nd half (GPT-NeoX style). - seqlen_offsets: (batch_size,) or int. Each sequence in Q and K is shifted by this amount. - Most commonly used in inference when we have KV cache. - Return: - kv: (batch_size, seqlen, 2, nheads, headdim) - rotary_dim must be <= headdim - Apply rotary embedding *inplace* to the first rotary_dim of K. - """ - return ApplyRotaryEmbKV_.apply(kv, cos, sin, interleaved, seqlen_offsets) - - -class RotaryEmbedding(torch.nn.Module): - """ - The rotary position embeddings from RoFormer_ (Su et. al). - A crucial insight from the method is that the query and keys are - transformed by rotation matrices which depend on the relative positions. - - Other implementations are available in the Rotary Transformer repo_ and in - GPT-NeoX_, GPT-NeoX was an inspiration - - .. _RoFormer: https://arxiv.org/abs/2104.09864 - .. _repo: https://github.com/ZhuiyiTechnology/roformer - .. _GPT-NeoX: https://github.com/EleutherAI/gpt-neox - - If scale_base is not None, this implements XPos (Sun et al., https://arxiv.org/abs/2212.10554). - A recommended value for scale_base is 512: https://github.com/HazyResearch/flash-attention/issues/96 - Reference: https://github.com/sunyt32/torchscale/blob/main/torchscale/component/xpos_relative_position.py - """ - - def __init__( - self, - dim: int, - base=10000.0, - interleaved=False, - scale_base=None, - device=None, - ): - """ - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead - of 1st half and 2nd half (GPT-NeoX style). - """ - super().__init__() - self.dim = dim - self.base = float(base) - # Generate and save the inverse frequency buffer (non trainable) - inv_freq = self._compute_inv_freq(device) - self.register_buffer("inv_freq", inv_freq, persistent=False) - self.interleaved = interleaved - self.scale_base = scale_base - scale = ( - (torch.arange(0, dim, 2, device=device, dtype=torch.float32) + 0.4 * dim) / (1.4 * dim) - if scale_base is not None - else None - ) - self.register_buffer("scale", scale, persistent=False) - - self._seq_len_cached = 0 - self._cos_cached = None - self._sin_cached = None - self._cos_k_cached = None - self._sin_k_cached = None - - def _compute_inv_freq(self, device=None): - return 1.0 / ( - self.base - ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim) - ) - - def _update_cos_sin_cache(self, seqlen, device=None, dtype=None): - # Reset the tables if the sequence length has changed, - # if we're on a new device (possibly due to tracing for instance), - # or if we're switching from inference mode to training - if ( - seqlen > self._seq_len_cached - or self._cos_cached is None - or self._cos_cached.device != device - or self._cos_cached.dtype != dtype - or (self.training and self._cos_cached.is_inference()) - ): - self._seq_len_cached = seqlen - # We want fp32 here, not self.inv_freq.dtype, since the model could be loaded in bf16 - # And the output of arange can be quite large, so bf16 would lose a lot of precision. - t = torch.arange(seqlen, device=device, dtype=torch.float32) - # We want fp32 here as well since inv_freq will be multiplied with t, and the output - # will be large. Having it in bf16 will lose a lot of precision and cause the - # cos & sin output to change significantly. - # We want to recompute self.inv_freq if it was not loaded in fp32 - if self.inv_freq.dtype != torch.float32: - inv_freq = self._compute_inv_freq(device=device) - else: - inv_freq = self.inv_freq - # Don't do einsum, it converts fp32 to bf16 under AMP - # freqs = torch.einsum("i,j->ij", t, self.inv_freq) - freqs = torch.outer(t, inv_freq) - if self.scale is None: - self._cos_cached = torch.cos(freqs).to(dtype) - self._sin_cached = torch.sin(freqs).to(dtype) - else: - power = ( - torch.arange(seqlen, dtype=self.scale.dtype, device=self.scale.device) - - seqlen // 2 - ) / self.scale_base - scale = self.scale.to(device=power.device) ** rearrange(power, "s -> s 1") - # We want the multiplication by scale to happen in fp32 - self._cos_cached = (torch.cos(freqs) * scale).to(dtype) - self._sin_cached = (torch.sin(freqs) * scale).to(dtype) - self._cos_k_cached = (torch.cos(freqs) / scale).to(dtype) - self._sin_k_cached = (torch.sin(freqs) / scale).to(dtype) - - def forward( - self, - qkv: torch.Tensor, - kv: Optional[torch.Tensor] = None, - seqlen_offset: Union[int, torch.Tensor] = 0, - max_seqlen: Optional[int] = None, - num_heads_q: Optional[int] = None, - ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: - """ - qkv: (batch, seqlen, 3, nheads, headdim) or (batch, seqlen, num_heads_q + 2 * num_heads_k, headdim) - if kv is none, else it's just q of shape (batch, seqlen, nheads, headdim). - If qkv has shape (batch, seqlen, num_heads_q + 2 * num_heads_k, headdim) (e.g. MQA / GQA), - then num_heads_q must be provided. - kv: (batch, seqlen, 2, nheads, headdim) - seqlen_offset: (batch_size,) or int. Each sequence in x is shifted by this amount. - Most commonly used in inference when we have KV cache. - If it's a tensor of shape (batch_size,), then to update the cos / sin cache, one - should pass in max_seqlen, which will update the cos / sin cache up to that length. - Apply rotary embedding *inplace* to qkv and / or kv. - """ - seqlen = qkv.shape[1] - if max_seqlen is not None: - self._update_cos_sin_cache(max_seqlen, device=qkv.device, dtype=qkv.dtype) - elif isinstance(seqlen_offset, int): - self._update_cos_sin_cache(seqlen + seqlen_offset, device=qkv.device, dtype=qkv.dtype) - if kv is None: - return apply_rotary_emb_qkv_( - qkv, - self._cos_cached, - self._sin_cached, - self._cos_k_cached if self.scale is not None else None, - self._sin_k_cached if self.scale is not None else None, - interleaved=self.interleaved, - seqlen_offsets=seqlen_offset, - num_heads_q=num_heads_q, - ) - else: - q = qkv - q = apply_rotary_emb_func( - q, - self._cos_cached, - self._sin_cached, - interleaved=self.interleaved, - inplace=True, - seqlen_offsets=seqlen_offset, - ) - kv = apply_rotary_emb_kv_( - kv, - self._cos_cached if self.scale is None else self._cos_k_cached, - self._sin_cached if self.scale is None else self._sin_k_cached, - interleaved=self.interleaved, - seqlen_offsets=seqlen_offset, - ) - return q, kv diff --git a/build/torch212-cxx11-cu126-aarch64-linux/metadata.json b/build/torch212-cxx11-cu126-aarch64-linux/metadata.json deleted file mode 100644 index 0f17638a47f51e0cc72e4d5eb692f32203b90a90..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu126-aarch64-linux/metadata.json +++ /dev/null @@ -1,40 +0,0 @@ -{ - "name": "flash-attn2", - "id": "_flash_attn2_cuda_f12afc9", - "version": 1, - "license": "BSD-3-Clause", - "python-depends": [], - "backend": { - "type": "cuda", - "archs": [ - "8.0", - "9.0" - ] - }, - "digest": { - "algorithm": "sha256", - "files": { - "__init__.py": "Cd9MEdd/WUYnBmfiQpC4SyPqAxrTilTjvsvswNju4zc=", - "_flash_attn2_cuda_f12afc9.abi3.so": "N/d+3sZUbEzc0MVmprCVuE5wNUrmwOD29jdLTmhFc4w=", - "_ops.py": "GQbmTTeSOU+7Q2nI68Hz+1RvEogyfEFccx1ubUYw+2s=", - "bert_padding.py": "gF1EmsdJ+HpQ86MRQ4VxDw+Sb/RVISdQALdNnoByHlw=", - "flash_attn2/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY=", - "flash_attn_interface.py": "W7a9bHR2ly6QL3aKqsBqxWA7M7ZWYcZPmtsPJu6rAfY=", - "layers/__init__.py": "47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU=", - "layers/patch_embed.py": "H58CgME/qSOPTZLOG08wFgrQS1j34pvNwMPrkTj3Ek4=", - "layers/rotary.py": "rFIdNrZQFNawwLD6d7t5LE2m+M93wwDZ8EuExw+gfD8=", - "ops/__init__.py": "47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU=", - "ops/activations.py": "t5lzNg1In8LP6bKeTnyeMizwqjv27JGbJ6ylPdGvZYg=", - "ops/fused_dense.py": "BqT8VTSCGdLJAjh6WBMkCRFYrdVodPeoElWBvH6AIzA=", - "ops/layer_norm.py": "zr7NXIm+2mtEynTp1CS0fbFGI2Mqdp41dY4AfDWF6EQ=", - "ops/rms_norm.py": "XEnihcj0a4aSz4LO55m5iKGVn4HKTeKN8TIyHjuDgxI=", - "ops/triton/__init__.py": "AbpHGcgLb+kRsJGnwFEktk7uzpZOCcBY74+YBdrKVGs=", - "ops/triton/cross_entropy.py": "y56nmuOtpSNxM7sDZcUSgXKIsIEayfuvrU8bAVog/nw=", - "ops/triton/k_activations.py": "+Z3vIyO4JkqBMipKsPvhzmxljtBdIhJCsl/M+/ESqBo=", - "ops/triton/layer_norm.py": "BBQdJiqopbIS12viUqzx3SvOQCNbDSAYLJXhpyMGyC8=", - "ops/triton/linear.py": "OtRvKz8xdpl+7v3q/ZTaS9fdBt9XrzMyapgRr50uBbM=", - "ops/triton/mlp.py": "uWA+D7HyuRUsOT+4nhi3vvCDNZnNEXwYcHMQ02UREfQ=", - "ops/triton/rotary.py": "0YSrlbep0lJmzLpTiANs0slTFeGcQQqR+el6uzbSLvI=" - } - } -} \ No newline at end of file diff --git a/build/torch212-cxx11-cu126-aarch64-linux/metadata.json.sigstore b/build/torch212-cxx11-cu126-aarch64-linux/metadata.json.sigstore deleted file mode 100644 index d34dfa318943a6c34d85285c73d6006b4df69906..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu126-aarch64-linux/metadata.json.sigstore +++ /dev/null @@ -1 +0,0 @@ -{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHSzCCBtGgAwIBAgIUBMebfCFOM2GB20bwuYBkI+LQpJUwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwNjIzMjM0MDQ1WhcNMjYwNjIzMjM1MDQ1WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEgkuwCx59j5dlBuLLdb/nIf65drO7vB+0ue62bKxTqOn8mhyA3NHJ7pklkfRdPO5JWe+rQlBuqFbS5ztd4z4rjKOCBfAwggXsMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQU2pNxr2HwqhkfK+bz9KW0PH/ANgIwHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoZjEyYWZjOTJjNDBiYmQ3NmQ5OTBkOTExZWJhYmYyOGU0OTBhN2JkZTATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoZjEyYWZjOTJjNDBiYmQ3NmQ5OTBkOTExZWJhYmYyOGU0OTBhN2JkZTAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoZjEyYWZjOTJjNDBiYmQ3NmQ5OTBkOTExZWJhYmYyOGU0OTBhN2JkZTAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKGYxMmFmYzkyYzQwYmJkNzZkOTkwZDkxMWViYWJmMjhlNDkwYTdiZGUwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMjgwMzYyNjg5MzUvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBigYKKwYBBAHWeQIEAgR8BHoAeAB2AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABnvbbkY8AAAQDAEcwRQIhAI+OFULso99WAFEdYQJgdBZb+OZfesZ+27TRwDJlxsQCAiAkTqr2ktlHJoMW9Y7IvzAzi9/qF2/qO0TbnqsCYGj1njAKBggqhkjOPQQDAwNoADBlAjEA7SlmFUMVJLOauEgdPVgZDGu4+wqqNQYtCBqy46wZodZzDb8Xt+3i401DnH5/dBP3AjBzbDByJun1RlRktmOqCybGA+PiQ1m3mBDof1o7t5AYgj23FUe6xqx0MS2Fe2MxnDE="}, "tlogEntries":[{"logIndex":"1933039914", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1782258045", "inclusionPromise":{"signedEntryTimestamp":"MEUCIGLdmW73AMREKQzvRFWnhOZuUL1wNYtz/KgffX2mDaI5AiEAso0TAvQtODFwVNbTEviE/mI6v0ckgszYQMUxsgz/RzA="}, "inclusionProof":{"logIndex":"1811135652", "rootHash":"SN9x0a6PpjO7ieNnFIPGt+wRBBJEdnAdF5QvxFuksvE=", "treeSize":"1811135681", "hashes":["9LaqUsUGVbNVBo4NPYmNXY5lJjDUZHT3lUGV1pyGswU=", "MFAmhISOW3bTid4ibj4B0i7Mc8mDtNGmMJnF3RQ06q8=", "UaSy8xktF92SpV/3a01WGLUzqkYrvl7m7f5fDMOYxP0=", "/lbeQbDsrnsjtnwoZVFtbAewzaHVfQ75978YN/8KRo4=", "eOY5TyIYigI8fY5a7RHoId1pb+HQJEitbIC+d4VT8S0=", "HGcMIgE7Svu7RWSeOAHMon0GNFd9CuruITKsok7LW5c=", "y+hX8RhilaWIXRB3iYv1BtoL/dYGb37aDPC4LGRE9Dw=", "gQd98df/QmXvStstYQmpU/ij6N/YlujaCzVrj/QB+D0=", "LuDs4vs9YQGnXlArOp7boDhzWABrNJ1pOn55UTTP5B8=", "WD8KwmQq5R8K69koYj5k4K7EKo2VF3LqJcx2xoXcd7o=", "99JTcUz0aJeZZocenAppVOLl8v91qMn0v7CmEjq9lws=", "OSDrt6jTKV6LYAkaG63p2BFz4SmbNmGb3uerhF5pOpQ=", "yehxDZkBWOGIzNDb93qxmQa6VBVo7kacVSoHyXMm3ZI=", "oUh+N71k3jWcTmlDDB6Esa5uqlCnbmzxTmDdut2A16M=", "2BsqkdK2omLvu+CpsFjpwqGi6m22H/7zcwCVkPwslM8=", "S8Oe9ZicoJqbVuL1I0dIdTgVpLc82Cm2CMOcV83sKVs=", "E1O2ilWlpFRLWGdbmJUMDZ2qbGu/eDZ12m6YjQ9/JPA=", "nyv2TDlfkz0XfRwjY98zsbUbOkBZ+Tgz7PML0c2uN14=", "OMfDh4k9YD0E+t4sjqJsh+vg8D1ueo1cC07exyxmr6s=", "+vHKm4wNwGJ46ceyZ8mePD7Vn1SediStNHCStRoTSWQ=", "5HwWjuGNDisetbiWM0cWIXM65TIM3EJlqk13ntCQlcU=", "NDIRlqxO0IaGr94q86MIU7VCAN3yEEgOHS4loya5lu4=", "daxmZaajRpZV+JxHiOYZhJBiSKN5ucqjh2WnGbHhirw=", "DOCeoSMovIvLExkhIvisow9AuNXgeWs4ECkyR6EcqYU="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n1811135681\nSN9x0a6PpjO7ieNnFIPGt+wRBBJEdnAdF5QvxFuksvE=\n\n— rekor.sigstore.dev wNI9ajBFAiAOpd2UcnuBCo9eQCC/d2es/iLCz0hS60RPlZvNDASJxwIhAMGdTG3xSyYX3wIZe0C+oiikY/N0Nx9wcDoj3TG+1r1y\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiJjY2RkZDhhZWJkZWUzYmJlZDc2ZmJmZmYzNDM4NzZlOGU5OTE1ZTFhYWJlNzJlYTI2MTFkZTBlNjg4MmYzNzc1In19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJUURjRmZCREhydTZYMkx0UmNhRDJ3OWhmQnhncHZENm13UDZBNFBVai9vditnSWdFTGE0dk1zdWdlNHcyY0M4UTlTU0FVRW43bW91RnNWcVkxNmRXZEhMZmIwPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRla05EUW5SSFowRjNTVUpCWjBsVlFrMWxZbVpEUms5Tk1rZENNakJpZDNWWlFtdEpLMHhSY0VwVmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDVxU1hwTmFrMHdUVVJSTVZkb1kwNU5hbGwzVG1wSmVrMXFUVEZOUkZFeFYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZuYTNWM1EzZzFPV28xWkd4Q2RVeE1aR0l2YmtsbU5qVmtjazgzZGtJck1IVmxOaklLWWt0NFZIRlBiamh0YUhsQk0wNUlTamR3YTJ4clpsSmtVRTgxU2xkbEszSlJiRUoxY1VaaVV6VjZkR1EwZWpSeWFrdFBRMEptUVhkbloxaHpUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlV5Y0U1NENuSXlTSGR4YUd0bVN5dGllamxMVnpCUVNDOUJUbWRKZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOWFha1Y1V1ZkYWFrOVVTbXBPUkVKcFdXMVJNMDV0VVRWUFZFSnJDazlVUlhoYVYwcG9XVzFaZVU5SFZUQlBWRUpvVGpKS2ExcFVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMXBxUlhsWlYxcHFUMVJLYWs1RVFtbFpiVkV6VG0xUk5VOVVRbXRQVkVWNFdsZEthRmx0V1hrS1QwZFZNRTlVUW1oT01rcHJXbFJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMXBxUlhsWlYxcHFUMVJLYWs1RVFtbFpiVkV6VG0xUk5VOVVRbXNLVDFSRmVGcFhTbWhaYlZsNVQwZFZNRTlVUW1oT01rcHJXbFJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwZFplRTF0Um0wS1dYcHJlVmw2VVhkWmJVcHJUbnBhYTA5VWEzZGFSR3Q0VFZkV2FWbFhTbTFOYW1oc1RrUnJkMWxVWkdsYVIxVjNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOYW1kM1RYcFplVTVxWnpWTmVsVjJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwWjFsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTT0VKSWIwRUtaVUZDTWtGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtNTJZbUpyV1RoQlFVRlJSQXBCUldOM1VsRkphRUZKSzA5R1ZVeHpiems1VjBGR1JXUlpVVXBuWkVKYVlpdFBXbVpsYzFvck1qZFVVbmRFU214NGMxRkRRV2xCYTFSeGNqSnJkR3hJQ2twdlRWYzVXVGRKZG5wQmVtazVMM0ZHTWk5eFR6QlVZbTV4YzBOWlIyb3hibXBCUzBKblozRm9hMnBQVUZGUlJFRjNUbTlCUkVKc1FXcEZRVGRUYkcwS1JsVk5Wa3BNVDJGMVJXZGtVRlpuV2tSSGRUUXJkM0Z4VGxGWmRFTkNjWGswTm5kYWIyUmFla1JpT0ZoMEt6TnBOREF4Ukc1SU5TOWtRbEF6UVdwQ2VncGlSRUo1U25WdU1WSnNVbXQwYlU5eFEzbGlSMEVyVUdsUk1XMHpiVUpFYjJZeGJ6ZDBOVUZaWjJveU0wWlZaVFo0Y1hnd1RWTXlSbVV5VFhodVJFVTlDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyjADAgEAMIICwQYJKoZIhvcNAQcCoIICsjCCAq4CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgMOht2YeO+iu8Qfd6Or00IWnhxOB5O/xoEk7Opz4ysPACFQDEPUXRguFi9YmtV+aNC5i7XC0wxBgPMjAyNjA2MjMyMzQwNDVaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHbMIIB1wIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MDYyMzIzNDA0NVowLwYJKoZIhvcNAQkEMSIEIDC7vEKTGQNUPnFs7oLKuiE7nTTUtHkzqCbg8gz89edoMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRnMGUCMHMM/74C5tsjyKQngQ4NBA8QBCt3+smskK8KWFZy/aqUhCdLQ06szpEnijh6hiJIWwIxAOOfMdLnocgwBxxpVpOCPPn+toyIUGKKjT2TGaag89LBKb6msfP4OFA0L2+kQ3pn/w=="}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"zN3Yrr3uO77Xb7//NDh26OmRXhqr5y6iYR3g5ogvN3U="}, "signature":"MEUCIQDcFfBDHru6X2LtRcaD2w9hfBxgpvD6mwP6A4PUj/ov+gIgELa4vMsuge4w2cC8Q9SSAUEn7mouFsVqY16dWdHLfb0="}} \ No newline at end of file diff --git a/build/torch212-cxx11-cu126-aarch64-linux/ops/__init__.py b/build/torch212-cxx11-cu126-aarch64-linux/ops/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/build/torch212-cxx11-cu126-aarch64-linux/ops/activations.py b/build/torch212-cxx11-cu126-aarch64-linux/ops/activations.py deleted file mode 100644 index 7c09649fc41e12d5a360c5672825d8380bc7ec80..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu126-aarch64-linux/ops/activations.py +++ /dev/null @@ -1,135 +0,0 @@ -# Copied from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/model/layers/activations.py -import math - -import torch -import torch.nn as nn -import torch.nn.functional as F - -# 1/sqrt(2*pi)-> 0.3989423 -# 1/sqrt(2) -> 0.70710678 -# sqrt(2/pi) -> 0.79788456 - -# this function is tanh approximation of gelu -# actual gelu is: -# x * 0.5 * (1.0 + torch.erf(x * 0.70710678)) -@torch.jit.script -def bias_gelu(y, bias): - x = bias + y - return (x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))).to(dtype=y.dtype) - - -# gradient of tanh approximation of gelu -# gradient of actual gelu is: -# 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x) -@torch.jit.script -def bias_gelu_back(g, y, bias): - """Assume that y has shape (B, D) and bias has shape (D)""" - x = bias + y - tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)) - # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243 - ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( - 1 + tanh_out - ) - grad_y = ff * g - return grad_y.to(dtype=y.dtype), grad_y.sum(dim=(0), dtype=bias.dtype) - - -class GeLUFunction(torch.autograd.Function): - @staticmethod - # bias is an optional argument - def forward(ctx, input, bias): - ctx.save_for_backward(input, bias) - return bias_gelu(input, bias) - - @staticmethod - def backward(ctx, grad_output): - input, bias = ctx.saved_tensors - tmp = bias_gelu_back(grad_output, input, bias) - return tmp, tmp - - -bias_gelu_impl = GeLUFunction.apply - -# this function is tanh approximation of gelu -# actual gelu is: -# x * 0.5 * (1.0 + torch.erf(x * 0.70710678)) -@torch.jit.script -def gelu_fwd(x): - return (x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))).to(dtype=x.dtype) - - -# gradient of tanh approximation of gelu -# gradient of actual gelu is: -# 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x) -@torch.jit.script -def gelu_bwd(g, x): - tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)) - # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243 - ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( - 1 + tanh_out - ) - return (ff * g).to(dtype=x.dtype) - - -class FastGeLUFunction(torch.autograd.Function): - @staticmethod - # bias is an optional argument - def forward(ctx, input): - ctx.save_for_backward(input) - return gelu_fwd(input) - - @staticmethod - def backward(ctx, grad_output): - (input,) = ctx.saved_tensors - tmp = gelu_bwd(grad_output, input) - return tmp - - -fast_gelu_impl = FastGeLUFunction.apply - - -@torch.jit.script -def relu_bwd(g, x): - return torch.where(x >= 0, g, 0.0).to(dtype=x.dtype) - - -@torch.jit.script -def sqrelu_fwd(x): - r = F.relu(x) - return (r * r).to(dtype=x.dtype) - - -@torch.jit.script -def sqrelu_bwd(g, x): - return (2.0 * g * F.relu(x)).to(dtype=x.dtype) - - -swiglu_fwd_codestring = """ -template T swiglu_fwd(T x, T y) { - return float(x) * float(y) / (1.0f + ::exp(-float(x))); -} -""" -swiglu_bwd_codestring = """ -template void swiglu_bwd(T x, T y, T g, T& dx, T& dy) { - float x_sigmoid = 1.0f / (1.0f + ::exp(-float(x))); - dx = x_sigmoid * (1 + float(x) * (1.0f - x_sigmoid)) * float(g) * float(y); - dy = float(x) * x_sigmoid * float(g); -} -""" -swiglu_fwd = torch.cuda.jiterator._create_jit_fn(swiglu_fwd_codestring) -swiglu_bwd = torch.cuda.jiterator._create_multi_output_jit_fn(swiglu_bwd_codestring, num_outputs=2) - - -class SwiGLUFunction(torch.autograd.Function): - - @staticmethod - def forward(ctx, x, y): - ctx.save_for_backward(x, y) - return swiglu_fwd(x, y) - - @staticmethod - def backward(ctx, dout): - x, y = ctx.saved_tensors - return swiglu_bwd(x, y, dout) - -swiglu = SwiGLUFunction.apply diff --git a/build/torch212-cxx11-cu126-aarch64-linux/ops/fused_dense.py b/build/torch212-cxx11-cu126-aarch64-linux/ops/fused_dense.py deleted file mode 100644 index 6b4033d134e4093fe278f7b3f8c7d3128ce9f36d..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu126-aarch64-linux/ops/fused_dense.py +++ /dev/null @@ -1,688 +0,0 @@ -# Copyright (c) 2023, Tri Dao. -# Inspired by https://github.com/NVIDIA/apex/blob/master/apex/fused_dense/fused_dense.py -# We make it work with pytorch amp and with bfloat16. -# The TensorParallel linear modules are inspired by https://github.com/NVIDIA/apex/blob/master/apex/transformer/tensor_parallel/layers.py -from functools import partial -from typing import Optional - -# import fused_dense_cuda # from apex -import fused_dense_lib as fused_dense_cuda -import torch -import torch.nn as nn -import torch.nn.functional as F -from torch import Tensor -from torch.distributed import ProcessGroup - -from flash_attn.utils.torch import custom_fwd, custom_bwd -from flash_attn.ops.activations import gelu_bwd, relu_bwd, sqrelu_bwd, sqrelu_fwd -from flash_attn.utils.distributed import ( - all_gather_raw, - all_reduce, - all_reduce_raw, - reduce_scatter, - reduce_scatter_raw, -) - - -class FusedDenseFunc(torch.autograd.Function): - @staticmethod - @custom_fwd - def forward( - ctx, x, weight, bias, return_residual=False, process_group=None, sequence_parallel=True - ): - """ - If process_group is not None and sequence_parallel=True, we're doing Tensor Parallel - with sequence parallelism: we do an all_gather_raw of x before doing the matmul. - """ - ctx.compute_weight_gradient = weight.requires_grad - ctx.return_residual = return_residual - ctx.process_group = process_group - ctx.sequence_parallel = sequence_parallel - - if torch.is_autocast_enabled(): - x = x.to(dtype=torch.get_autocast_gpu_dtype()) - x = x.contiguous() - if process_group is not None and sequence_parallel: - # We want to kick off the all_gather early, before weight dtype conversion - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - else: - total_x = x - - if torch.is_autocast_enabled(): - weight = weight.to(dtype=torch.get_autocast_gpu_dtype()) - bias = bias.to(dtype=torch.get_autocast_gpu_dtype()) if bias is not None else None - weight = weight.contiguous() - if process_group is not None and sequence_parallel: - handle_x.wait() - batch_shape, n = total_x.shape[:-1], total_x.shape[-1] - batch_dim = batch_shape.numel() - # https://github.com/pytorch/pytorch/blob/5b51849b48a7dbccd297286cc0110def4706f9e7/aten/src/ATen/native/cuda/Blas.cpp#L174 - if min(batch_dim, n, *weight.shape) > 65535 * 32: - raise RuntimeError("fused_dense only supports matrix dims <= 2M") - output = F.linear(total_x, weight, bias) - if ctx.compute_weight_gradient: - ctx.save_for_backward(x, weight) - else: - ctx.save_for_backward(weight) - return output if not return_residual else (output, x) - - @staticmethod - @custom_bwd - def backward(ctx, grad_output, *args): - grad_output = grad_output.contiguous() - if ctx.return_residual: - (grad_input,) = args - grad_input = grad_input.contiguous() - process_group = ctx.process_group - sequence_parallel = ctx.sequence_parallel - if ctx.compute_weight_gradient: - x, weight = ctx.saved_tensors - if process_group is not None and sequence_parallel: - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - else: - total_x = x - else: - (weight,) = ctx.saved_tensors - total_x = None - batch_shape = grad_output.shape[:-1] - batch_dim = batch_shape.numel() - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - if ctx.needs_input_grad[0]: - if not ctx.return_residual: - grad_input = F.linear(grad_output, weight.t()) - else: - grad_input = torch.addmm( - grad_input.reshape(batch_dim, grad_input.shape[-1]), grad_output, weight - ) - grad_input = grad_input.reshape(*batch_shape, grad_input.shape[-1]) - if process_group is not None: - reduce_fn = reduce_scatter_raw if sequence_parallel else all_reduce_raw - grad_input, handle_grad_input = reduce_fn(grad_input, process_group, async_op=True) - else: - grad_input = None - if ctx.needs_input_grad[1]: - assert ctx.compute_weight_gradient - if process_group is not None and sequence_parallel: - handle_x.wait() - grad_weight, grad_bias = fused_dense_cuda.linear_bias_wgrad( - total_x.reshape(batch_dim, total_x.shape[-1]), grad_output, ctx.needs_input_grad[2] - ) - else: - grad_weight = None - grad_bias = grad_output if ctx.needs_input_grad[2] else None - if process_group is not None and ctx.needs_input_grad[0]: - handle_grad_input.wait() - return grad_input, grad_weight, grad_bias, None, None, None - - -def fused_dense_func( - x: Tensor, - weight: Tensor, - bias: Optional[Tensor] = None, - return_residual: bool = False, - process_group: Optional[ProcessGroup] = None, - sequence_parallel: bool = True, -): - dtype_eligible = x.dtype in [torch.float16, torch.bfloat16] or ( - x.dtype == torch.float32 and torch.is_autocast_enabled() - ) - if x.is_cuda and weight.is_cuda and (bias is None or bias.is_cuda) and dtype_eligible: - return FusedDenseFunc.apply( - x, weight, bias, return_residual, process_group, sequence_parallel - ) - else: - assert process_group is None - out = F.linear(x, weight, bias) - return out if not return_residual else (out, x) - - -class FusedDense(nn.Linear): - def __init__( - self, - in_features: int, - out_features: int, - bias: bool = True, - return_residual: bool = False, - device=None, - dtype=None, - ) -> None: - super().__init__(in_features, out_features, bias=bias, device=device, dtype=dtype) - self.return_residual = return_residual - - def forward(self, x, process_group=None): - """ - If process_group is not None, we're doing Tensor Parallel with sequence parallelism: - we do an all_gather of x before doing the matmul. - """ - return fused_dense_func( - x, - self.weight, - self.bias, - return_residual=self.return_residual, - process_group=process_group, - ) - - -class ColumnParallelLinear(nn.Linear): - def __init__( - self, - in_features: int, - out_features: int, - process_group: ProcessGroup, - bias: bool = True, - sequence_parallel=True, - multiple_of=1, - device=None, - dtype=None, - ) -> None: - world_size = torch.distributed.get_world_size(process_group) - if out_features % multiple_of: - raise ValueError(f"out_features ({out_features}) must be a multiple of {multiple_of}") - multiple = out_features // multiple_of - # We want to split @multiple across world_size, but it could be an uneven split - div = multiple // world_size - mod = multiple % world_size - # The first @mod ranks get @div + 1 copies, the rest get @div copies - local_multiple = div + int(torch.distributed.get_rank(process_group) < mod) - super().__init__( - in_features, local_multiple * multiple_of, bias=bias, device=device, dtype=dtype - ) - self.process_group = process_group - self.sequence_parallel = sequence_parallel - - def forward(self, x): - # If self.sequence_parallel is True, we're doing Tensor Parallel with sequence parallelism: - # we do an all_gather of x before doing the matmul. - # If not, then the input is already gathered. - return fused_dense_func( - x, - self.weight, - self.bias, - process_group=self.process_group, - sequence_parallel=self.sequence_parallel, - ) - - -class RowParallelLinear(nn.Linear): - def __init__( - self, - in_features: int, - out_features: int, - process_group: ProcessGroup, - bias: bool = True, - sequence_parallel=True, - multiple_of=1, - device=None, - dtype=None, - ) -> None: - world_size = torch.distributed.get_world_size(process_group) - rank = torch.distributed.get_rank(process_group) - if in_features % multiple_of: - raise ValueError(f"in_features ({in_features}) must be a multiple of {multiple_of}") - multiple = in_features // multiple_of - # We want to split @multiple across world_size, but it could be an uneven split - div = multiple // world_size - mod = multiple % world_size - # The first @mod ranks get @div + 1 copies, the rest get @div copies - local_multiple = div + int(torch.distributed.get_rank(process_group) < mod) - # Only rank 0 will have bias - super().__init__( - local_multiple * multiple_of, - out_features, - bias=bias and rank == 0, - device=device, - dtype=dtype, - ) - self.process_group = process_group - self.sequence_parallel = sequence_parallel - - def forward(self, x): - """ - We're doing Tensor Parallel with sequence parallelism: we do the matmul and then - a reduce_scatter of the result. - """ - out = fused_dense_func(x, self.weight, self.bias) - reduce_fn = reduce_scatter if self.sequence_parallel else all_reduce - return reduce_fn(out, self.process_group) - - -class FusedMLPFunc(torch.autograd.Function): - @staticmethod - @custom_fwd - def forward( - ctx, - x, - weight1, - bias1, - weight2, - bias2, - activation="gelu_approx", - save_pre_act=True, - return_residual=False, - checkpoint_lvl=0, - heuristic=0, - process_group=None, - sequence_parallel=True, - ): - """ - If process_group is not None and sequence_parallel=True, we're doing Tensor Parallel - with sequence parallelism: we do an all_gather of x before doing the matmul. - If sequence_parallel=False, then the input is already gathered. - - checkpoint_lvl: - 0: no recomputation in the bwd - 1: recompute gelu_out / relu_out in the bwd - 2: recompute pre_act and gelu_out / relu_out in the bwd - """ - assert -1 <= heuristic <= 4 - assert activation in ["gelu_approx", "relu", "sqrelu"] - if activation == "sqrelu": - assert heuristic == -1 - if not save_pre_act: - checkpoint_lvl = 2 - assert checkpoint_lvl in [0, 1, 2] - ctx.return_residual = return_residual - ctx.process_group = process_group - ctx.sequence_parallel = sequence_parallel - ctx.checkpoint_lvl = checkpoint_lvl - ctx.activation = activation - ctx.heuristic = heuristic - - if torch.is_autocast_enabled(): - x = x.to(dtype=torch.get_autocast_gpu_dtype()) - x = x.contiguous() - if process_group is not None and sequence_parallel: - # We want to kick off the all_gather early, before weight dtype conversion - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - else: - total_x = x - - if torch.is_autocast_enabled(): - dtype = torch.get_autocast_gpu_dtype() - weight1, weight2 = [a.to(dtype=dtype) for a in [weight1, weight2]] - bias1 = bias1.to(dtype=dtype) if bias1 is not None else None - bias2 = bias2.to(dtype=dtype) if bias2 is not None else None - weight1 = weight1.contiguous() - bias1 = bias1.contiguous() if bias1 is not None else None - weight2 = weight2.contiguous() - bias2 = bias2.contiguous() if bias2 is not None else None - if process_group is not None and sequence_parallel: - handle_x.wait() - batch_shape, n = total_x.shape[:-1], total_x.shape[-1] - batch_dim = batch_shape.numel() - # https://github.com/pytorch/pytorch/blob/5b51849b48a7dbccd297286cc0110def4706f9e7/aten/src/ATen/native/cuda/Blas.cpp#L174 - if min(batch_dim, n, *weight1.shape, *weight2.shape) > 65535 * 32: - raise RuntimeError("fused_dense only supports matrix dims <= 2M") - if heuristic == -1: - pre_act = F.linear(total_x, weight1, bias1) - activation_fn = ( - partial(F.gelu, approximate="tanh") - if activation == "gelu_approx" - else (sqrelu_fwd if activation == "sqrelu" else F.relu) - ) - with torch.jit.fuser("fuser2"): - output1 = activation_fn(pre_act) - # This is before adding bias1 - # pre_act = F.linear(total_x.reshape(batch_dim, n), weight1) - # with torch.jit.fuser('fuser2'): - # output1 = bias_gelu(pre_act, bias1) - else: - is_gelu = activation == "gelu_approx" - output1, *rest = fused_dense_cuda.linear_act_forward( - total_x.reshape(batch_dim, n), weight1, bias1, is_gelu, save_pre_act, heuristic - ) - if save_pre_act: - pre_act = rest[0] - output2 = F.linear(output1, weight2, bias2) - if checkpoint_lvl == 0 or (checkpoint_lvl == 1 and activation == "relu"): - # For RELU the pre_act is very small (just a bit-mask) so we just save it - ctx.save_for_backward(x, weight1, weight2, pre_act, output1) - elif checkpoint_lvl == 1: - ctx.save_for_backward(x, weight1, weight2, pre_act) - elif checkpoint_lvl == 2: - ctx.save_for_backward(x, weight1, weight2, bias1) - output2 = output2.reshape(*batch_shape, output2.shape[-1]) - return output2 if not return_residual else (output2, x) - - @staticmethod - @custom_bwd - def backward(ctx, grad_output, *args): - grad_output = grad_output.contiguous() - checkpoint_lvl = ctx.checkpoint_lvl - activation = ctx.activation - activation_fn = ( - partial(F.gelu, approximate="tanh") - if activation == "gelu_approx" - else (sqrelu_fwd if activation == "sqrelu" else F.relu) - ) - if ctx.return_residual: - (grad_input,) = args - grad_input = grad_input.contiguous() - process_group = ctx.process_group - sequence_parallel = ctx.sequence_parallel - x, weight1, weight2, *rest = ctx.saved_tensors - if process_group is None or not sequence_parallel: - total_x = x - batch_shape = grad_output.shape[:-1] - batch_dim = batch_shape.numel() - if checkpoint_lvl in [0, 1]: - if process_group is not None and sequence_parallel: - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - if checkpoint_lvl == 0 or (checkpoint_lvl == 1 and activation == "relu"): - pre_act, output1 = rest - elif checkpoint_lvl == 1: - (pre_act,) = rest - with torch.jit.fuser("fuser2"): - output1 = activation_fn(pre_act) - elif checkpoint_lvl == 2: - (bias1,) = rest - if process_group is not None and sequence_parallel: - total_x, _ = all_gather_raw(x, process_group) - if ctx.heuristic == -1: - pre_act = F.linear(total_x, weight1, bias1) - with torch.jit.fuser("fuser2"): - output1 = activation_fn(pre_act) - else: - output1, pre_act = fused_dense_cuda.linear_act_forward( - total_x.reshape(batch_dim, total_x.shape[-1]), - weight1, - bias1, - activation == "gelu_approx", - True, - ctx.heuristic, - ) - - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - output1 = output1.reshape(batch_dim, output1.shape[-1]) - pre_act = pre_act.reshape(batch_dim, pre_act.shape[-1]) - if ctx.needs_input_grad[3]: - grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad( - output1, grad_output, ctx.needs_input_grad[4] - ) - else: - grad_weight2 = None - grad_bias2 = grad_output if ctx.needs_input_grad[4] else None - if ctx.heuristic == -1: - # grad_pre_act = matmul_dgelu(grad_output, weight2, pre_act) - grad_output1 = F.linear(grad_output, weight2.t()) - activation_grad_fn = ( - gelu_bwd - if activation == "gelu_approx" - else (sqrelu_bwd if activation == "sqrelu" else relu_bwd) - ) - with torch.jit.fuser("fuser2"): - grad_pre_act = activation_grad_fn(grad_output1, pre_act) - else: - # The cublasLt epilogue has to compute both gelu/relu grad and bias grad, we can't - # just compute gelu/relu grad - grad_pre_act, grad_bias1 = fused_dense_cuda.bias_act_linear_dgrad_bgrad( - weight2, grad_output, pre_act, activation == "gelu_approx", ctx.heuristic - ) - if not ctx.needs_input_grad[2]: - grad_bias1 = None - if ctx.needs_input_grad[0]: - if not ctx.return_residual: - grad_input = F.linear(grad_pre_act, weight1.t()) - else: - grad_input = torch.addmm( - grad_input.reshape(batch_dim, grad_input.shape[-1]), grad_pre_act, weight1 - ) - grad_input = grad_input.reshape(*batch_shape, grad_input.shape[-1]) - if process_group is not None: - reduce_fn = reduce_scatter_raw if sequence_parallel else all_reduce_raw - grad_input, handle_grad_input = reduce_fn(grad_input, process_group, async_op=True) - else: - grad_input = None - if ctx.heuristic == -1: - if ctx.needs_input_grad[1]: - if process_group is not None and sequence_parallel and checkpoint_lvl != 2: - handle_x.wait() - grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_wgrad( - total_x.reshape(batch_dim, total_x.shape[-1]), - grad_pre_act, - ctx.needs_input_grad[2], - ) - else: - grad_weight1 = None - grad_bias1 = grad_pre_act if ctx.needs_input_grad[2] else None - else: - if ctx.needs_input_grad[1]: - if process_group is not None and sequence_parallel and checkpoint_lvl != 2: - handle_x.wait() - grad_weight1 = F.linear( - grad_pre_act.t(), total_x.reshape(batch_dim, total_x.shape[-1]).t() - ) - else: - grad_weight1 = None - if process_group is not None and ctx.needs_input_grad[0]: - handle_grad_input.wait() - return ( - grad_input, - grad_weight1, - grad_bias1, - grad_weight2, - grad_bias2, - None, - None, - None, - None, - None, - None, - None, - ) - - -def fused_mlp_func( - x: Tensor, - weight1: Tensor, - weight2: Tensor, - bias1: Optional[Tensor] = None, - bias2: Optional[Tensor] = None, - activation: str = "gelu_approx", - save_pre_act: bool = True, - return_residual: bool = False, - checkpoint_lvl: int = 0, - heuristic: int = 0, - process_group: Optional[ProcessGroup] = None, - sequence_parallel: bool = True, -): - assert activation in ["gelu_approx", "relu", "sqrelu"] - dtype_eligible = x.dtype in [torch.float16, torch.bfloat16] or ( - x.dtype == torch.float32 and torch.is_autocast_enabled() - ) - # If we save pre-activation, dimension must be divisible by 128 (relu) or 8 (gelu) - dim_eligible = not save_pre_act or (x.shape[-1] % (128 if activation == "relu" else 8) == 0) - if ( - x.is_cuda - and weight1.is_cuda - and weight2.is_cuda - and (bias1 is None or bias1.is_cuda) - and (bias2 is None or bias2.is_cuda) - and dtype_eligible - and dim_eligible - ): - return FusedMLPFunc.apply( - x, - weight1, - bias1, - weight2, - bias2, - activation, - save_pre_act, - return_residual, - checkpoint_lvl, - heuristic, - process_group, - sequence_parallel, - ) - else: - assert process_group is None - pre_act = F.linear(x, weight1, bias1) - activation_fn = ( - partial(F.gelu, approximate="tanh") - if activation == "gelu_approx" - else partial(F.relu, inplace=True) - ) - output1 = activation_fn(pre_act) - output2 = F.linear(output1, weight2, bias2) - return output2 if not return_residual else (output2, x) - - -class FusedMLP(nn.Module): - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - bias1=True, - bias2=True, - activation="gelu_approx", - return_residual=False, - checkpoint_lvl=0, - heuristic="auto", - device=None, - dtype=None, - ): - """ - If process_group is not None, we're doing Tensor Parallel with sequence parallelism: - we do an all_gather of x before doing the matmul, gelu, then matmul. - Finally we do a reduce_scatter of the output. - - checkpoint_lvl (increasing lvl means slower but more memory saving): - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute pre_act and gelu_out in the bwd - heuristic: - -1: don't fuse gemm + gelu (separate kernel) - 0..4: use this heuristic for the algo section in the fused gemm + gelu - 'auto': heuristic will be picked automatically: - For CUDA >= 11.8, we set heuristic=0 for both fp16 and bf16 for best perf. - For CUDA <= 11.7, we set heuristic=1 for fp16 and heuristic=-1 for bf16. - For H100, we set heuristic=-1 for both fp16 and bf16 as the fused cuBlasLt implementation - is slower than the unfused version. - return_residual: whether to return the input x along with the output. This is for - performance reason: for post-norm architecture, returning the input allows us - to fuse the backward of nn.Linear with the residual connection. - """ - assert checkpoint_lvl in [0, 1, 2] - assert activation in ["gelu_approx", "relu", "sqrelu"] - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features * 4 - self.activation = activation - self.return_residual = return_residual - self.checkpoint_lvl = checkpoint_lvl - self.heuristic = heuristic if activation != "sqrelu" else -1 - self.fc1 = nn.Linear(in_features, hidden_features, bias=bias1, **factory_kwargs) - self.fc2 = nn.Linear(hidden_features, out_features, bias=bias2, **factory_kwargs) - - def forward(self, x, process_group=None): - dtype = x.dtype if not torch.is_autocast_enabled() else torch.get_autocast_gpu_dtype() - if self.heuristic == "auto": - if self.activation == "gelu_approx": - if torch.cuda.get_device_capability("cuda") == (9, 0): - heuristic = -1 - else: - cuda_ver = tuple(map(int, torch.version.cuda.split("."))) - heuristic = 0 if cuda_ver >= (11, 8) else (1 if dtype == torch.float16 else -1) - else: - heuristic = 0 - else: - heuristic = self.heuristic - out = fused_mlp_func( - x, - self.fc1.weight, - self.fc2.weight, - self.fc1.bias, - self.fc2.bias, - activation=self.activation, - save_pre_act=self.training, - return_residual=self.return_residual, - checkpoint_lvl=self.checkpoint_lvl, - heuristic=heuristic, - process_group=process_group, - ) - if self.return_residual: - out, x = out - if process_group is not None: - out = reduce_scatter(out, process_group) - return out if not self.return_residual else (out, x) - - -class ParallelFusedMLP(nn.Module): - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - activation="gelu_approx", - process_group: ProcessGroup = None, - bias1=True, - bias2=True, - sequence_parallel=True, - checkpoint_lvl=0, - heuristic="auto", - device=None, - dtype=None, - ): - """ - process_group is required. We're doing Tensor Parallel with sequence parallelism: - we do an all_gather of x before doing the matmul, gelu, then matmul. - Finally we do a reduce_scatter of the output. - - checkpoint_lvl (increasing lvl means slower but more memory saving): - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute pre_act and gelu_out in the bwd - heuristic: - -1: don't fuse gemm + gelu (separate kernel) - 0..4: use this heuristic for the algo section in the fused gemm + gelu - 'auto': heuristic will be picked automatically: - For CUDA >= 11.8, we set heuristic=0 for both fp16 and bf16 for best perf. - For CUDA <= 11.7, we set heuristic=1 for fp16 and heuristic=-1 for bf16. - """ - assert checkpoint_lvl in [0, 1, 2] - assert activation in ["gelu_approx", "relu", "sqrelu"] - assert process_group is not None - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features * 4 - self.activation = activation - self.process_group = process_group - self.sequence_parallel = sequence_parallel - self.checkpoint_lvl = checkpoint_lvl - self.heuristic = heuristic if activation != "sqrelu" else -1 - self.fc1 = ColumnParallelLinear( - in_features, hidden_features, process_group, bias=bias1, **factory_kwargs - ) - self.fc2 = RowParallelLinear( - hidden_features, out_features, process_group, bias=bias2, **factory_kwargs - ) - - def forward(self, x): - dtype = x.dtype if not torch.is_autocast_enabled() else torch.get_autocast_gpu_dtype() - if self.heuristic == "auto": - if self.activation == "gelu_approx": - cuda_ver = tuple(map(int, torch.version.cuda.split("."))) - heuristic = 0 if cuda_ver >= (11, 8) else (1 if dtype == torch.float16 else -1) - else: - heuristic = 0 - else: - heuristic = self.heuristic - out = fused_mlp_func( - x, - self.fc1.weight, - self.fc2.weight, - self.fc1.bias, - self.fc2.bias, - activation=self.activation, - save_pre_act=self.training, - checkpoint_lvl=self.checkpoint_lvl, - heuristic=heuristic, - process_group=self.process_group, - sequence_parallel=self.sequence_parallel, - ) - reduce_fn = reduce_scatter if self.sequence_parallel else all_reduce - return reduce_fn(out, self.process_group) diff --git a/build/torch212-cxx11-cu126-aarch64-linux/ops/layer_norm.py b/build/torch212-cxx11-cu126-aarch64-linux/ops/layer_norm.py deleted file mode 100644 index 4b6cd798fd02844ef9cd3897f8ab95e490e638bf..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu126-aarch64-linux/ops/layer_norm.py +++ /dev/null @@ -1,800 +0,0 @@ -# Copyright (c) 2022, Tri Dao. -# Adapted from https://github.com/NVIDIA/apex/blob/master/apex/contrib/layer_norm/layer_norm.py - -import dropout_layer_norm -import torch -from torch.nn import init - - -def maybe_align(x, alignment_in_bytes=16): - """Assume that x already has last dim divisible by alignment_in_bytes""" - # TD [2023-07-04] I'm not 100% sure that clone will align the memory - # https://discuss.pytorch.org/t/how-to-ensure-that-tensor-data-ptr-is-aligned-to-16-bytes/183440 - return x if x.data_ptr() % alignment_in_bytes == 0 else x.clone() - - -def _dropout_add_layer_norm_forward( - x0, - residual, - gamma, - beta, - rowscale, - colscale, - dropout_p, - epsilon, - residual_in_fp32=False, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes""" - hidden_size = gamma.numel() - x0mat = x0.view((-1, hidden_size)) - residualmat = residual.view((-1, hidden_size)) if residual is not None else None - rowscale = rowscale.view(-1) if rowscale is not None else None - zmat, xmat, dmask, mu, rsigma = dropout_layer_norm.dropout_add_ln_fwd( - x0mat, - residualmat, - gamma, - beta, - rowscale, - colscale, - None, - None, - dropout_p, - epsilon, - 1.0, - 0, - None, - residual_in_fp32, - is_rms_norm, - ) - # dmask is None if dropout_p == 0.0 - # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype - return zmat, xmat if xmat is not None else x0mat, dmask, mu, rsigma - - -def _dropout_add_layer_norm_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - rowscale, - colscale, - dropout_p, - has_residual, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes - dx == None means that it was a post-norm architecture - (x = drop(x0) + residual was not returned in the fwd). - x0 must not be None if we have colscale. - """ - hidden_size = gamma.numel() - xmat = x.view((-1, hidden_size)) - dzmat = dz.view(xmat.shape) - dxmat = dx.view(xmat.shape) if dx is not None else None - x0mat = x0.view((-1, hidden_size)) if x0 is not None else None - rowscale = rowscale.view(-1) if rowscale is not None else None - if colscale is not None: - assert x0 is not None, "x0 is required to compute the gradient of colscale" - dx0mat, dresidualmat, dgamma, dbeta, _, _, *rest = dropout_layer_norm.dropout_add_ln_bwd( - dzmat, - dxmat, - xmat, - x0mat, - dmask, - mu, - rsigma, - gamma, - rowscale, - colscale, - None, - None, - dropout_p, - 1.0, - 0, - has_residual, - is_rms_norm, - ) - # dresidualmat is None if not has_residual - if colscale is None: - return dx0mat, dresidualmat, dgamma, dbeta - else: - dcolscale = rest[0] - return dx0mat, dresidualmat, dgamma, dbeta, dcolscale - - -def _dropout_add_layer_norm_subset_forward( - x0, - residual, - gamma, - beta, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32=False, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes""" - hidden_size = gamma.numel() - x0mat = x0.view((-1, hidden_size)) - residualmat = residual.view((-1, hidden_size)) if residual is not None else None - x0_subset = x0_subset.view(-1) if x0_subset is not None else None - out_subset = out_subset.view(-1) if out_subset is not None else None - zmat, xmat, dmask, mu, rsigma = dropout_layer_norm.dropout_add_ln_fwd( - x0mat, - residualmat, - gamma, - beta, - None, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - None, - residual_in_fp32, - is_rms_norm, - ) - # dmask is None if dropout_p == 0.0 - # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype - return zmat, xmat if xmat is not None else x0mat, dmask, mu, rsigma - - -def _dropout_add_layer_norm_subset_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - colscale, - x0_subset, - out_subset, - dropout_p, - rowscale_const, - x0_numrows, - has_residual, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes - dx == None means that it was a post-norm architecture - (x = drop(x0) + residual was not returned in the fwd). - x0 must not be None if we have colscale. - """ - hidden_size = gamma.numel() - xmat = x.view((-1, hidden_size)) - dzmat = dz.view(-1, hidden_size) - dxmat = dx.view(xmat.shape) if dx is not None else None - x0mat = x0.view((-1, hidden_size)) if x0 is not None else None - x0_subset = x0_subset.view(-1) if x0_subset is not None else None - out_subset = out_subset.view(-1) if out_subset is not None else None - if colscale is not None: - assert x0 is not None, "x0 is required to compute the gradient of colscale" - dx0mat, dresidualmat, dgamma, dbeta, _, _, *rest = dropout_layer_norm.dropout_add_ln_bwd( - dzmat, - dxmat, - xmat, - x0mat, - dmask, - mu, - rsigma, - gamma, - None, - colscale, - x0_subset, - out_subset, - dropout_p, - rowscale_const, - x0_numrows, - has_residual, - is_rms_norm, - ) - # dresidualmat is None if not has_residual - if colscale is None: - return dx0mat, dresidualmat, dgamma, dbeta - else: - dcolscale = rest[0] - return dx0mat, dresidualmat, dgamma, dbeta, dcolscale - - -def _dropout_add_layer_norm_parallel_residual_forward( - x0, - x1, - residual, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - residual_in_fp32=False, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes""" - hidden_size = gamma0.numel() - x0mat = x0.view((-1, hidden_size)) - x1mat = x1.view((-1, hidden_size)) if x1 is not None else None - residualmat = residual.view((-1, hidden_size)) if residual is not None else None - ( - z0mat, - z1mat, - xmat, - dmask0, - dmask1, - mu, - rsigma, - ) = dropout_layer_norm.dropout_add_ln_parallel_residual_fwd( - x0mat, - x1mat, - residualmat, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - None, - residual_in_fp32, - is_rms_norm, - ) - # dmask0 and dmask1 are None if dropout_p == 0.0 - # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype - return z0mat, z1mat, xmat if xmat is not None else x0mat, dmask0, dmask1, mu, rsigma - - -def _dropout_add_layer_norm_parallel_residual_backward( - dz0, - dz1, - dx, - x, - dmask0, - dmask1, - mu, - rsigma, - gamma0, - gamma1, - dropout_p, - has_x1, - has_residual, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes - dx == None means that it was a post-norm architecture - (x = drop(x0) + residual was not returned in the fwd). - """ - hidden_size = gamma0.numel() - xmat = x.view((-1, hidden_size)) - dz0mat = dz0.view(xmat.shape) - dz1mat = dz1.view(xmat.shape) if dz1 is not None else None - dxmat = dx.view(xmat.shape) if dx is not None else None - ( - dx0mat, - dx1mat, - dresidualmat, - dgamma0, - dbeta0, - dgamma1, - dbeta1, - *rest, - ) = dropout_layer_norm.dropout_add_ln_parallel_residual_bwd( - dz0mat, - dz1mat, - dxmat, - xmat, - dmask0, - dmask1, - mu, - rsigma, - gamma0, - gamma1, - dropout_p, - has_x1, - has_residual, - is_rms_norm, - ) - # dresidualmat is None if not has_residual - return dx0mat, dx1mat, dresidualmat, dgamma0, dbeta0, dgamma1, dbeta1 - - -class DropoutAddLayerNormFn(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x0, - residual, - gamma, - beta, - rowscale, - colscale, - dropout_p, - epsilon, - residual_in_fp32=False, - prenorm=False, - is_rms_norm=False, - return_dmask=False, - ): - x0 = maybe_align(x0.contiguous(), 16) - residual = maybe_align(residual.contiguous(), 16) if residual is not None else None - gamma = maybe_align(gamma.contiguous(), 16) - beta = maybe_align(beta.contiguous(), 16) if beta is not None else None - rowscale = maybe_align(rowscale.contiguous(), 16) if rowscale is not None else None - colscale = maybe_align(colscale.contiguous(), 16) if colscale is not None else None - zmat, xmat, dmask, mu, rsigma = _dropout_add_layer_norm_forward( - x0, - residual, - gamma, - beta, - rowscale, - colscale, - dropout_p, - epsilon, - residual_in_fp32, - is_rms_norm, - ) - # Only need to save x0 if we need to compute gradient wrt colscale - x0_saved = x0 if colscale is not None else None - ctx.save_for_backward( - xmat.view(x0.shape), x0_saved, dmask, gamma, mu, rsigma, rowscale, colscale - ) - ctx.prenorm = prenorm - ctx.dropout_p = dropout_p - ctx.has_residual = residual is not None - ctx.is_rms_norm = is_rms_norm - ctx.has_beta = beta is not None - if not return_dmask: - return ( - zmat.view(x0.shape) if not prenorm else (zmat.view(x0.shape), xmat.view(x0.shape)) - ) - else: - dmask = ( - dmask.view(x0.shape) - if dropout_p > 0.0 - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - ctx.mark_non_differentiable(dmask) - return ( - (zmat.view(x0.shape), dmask) - if not prenorm - else (zmat.view(x0.shape), xmat.view(x0.shape), dmask) - ) - - @staticmethod - def backward(ctx, dz, *args): - # assert dz.is_contiguous() - dz = maybe_align(dz.contiguous(), 16) # this happens! - dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None - x, x0, dmask, gamma, mu, rsigma, rowscale, colscale = ctx.saved_tensors - # x0 is None if colscale is None - dropout_p = ctx.dropout_p - has_residual = ctx.has_residual - dx0mat, dresidualmat, dgamma, dbeta, *rest = _dropout_add_layer_norm_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - rowscale, - colscale, - dropout_p, - has_residual, - ctx.is_rms_norm, - ) - dx0 = dx0mat.view(x.shape) - dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None - dcolscale = rest[0] if colscale is not None else None - return ( - dx0, - dresidual, - dgamma, - dbeta if ctx.has_beta else None, - None, - dcolscale, - None, - None, - None, - None, - None, - None, - ) - - -class DropoutAddLayerNormSubsetFn(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x0, - residual, - gamma, - beta, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32=False, - prenorm=False, - is_rms_norm=False, - return_dmask=False, - ): - x0 = maybe_align(x0.contiguous(), 16) - residual = maybe_align(residual.contiguous(), 16) if residual is not None else None - gamma = maybe_align(gamma.contiguous(), 16) - beta = maybe_align(beta.contiguous(), 16) if beta is not None else None - colscale = maybe_align(colscale.contiguous(), 16) if colscale is not None else None - zmat, xmat, dmask, mu, rsigma = _dropout_add_layer_norm_subset_forward( - x0, - residual, - gamma, - beta, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32, - is_rms_norm, - ) - # Only need to save x0 if we need to compute gradient wrt colscale - x0_saved = x0 if colscale is not None else None - x_shape = (-1, *x0.shape[1:]) - ctx.save_for_backward( - xmat.view(x_shape), x0_saved, dmask, gamma, mu, rsigma, colscale, x0_subset, out_subset - ) - ctx.prenorm = prenorm - ctx.dropout_p = dropout_p - ctx.rowscale_const = rowscale_const - ctx.x0_numrows = x0.shape[:-1].numel() - ctx.has_residual = residual is not None - ctx.is_rms_norm = is_rms_norm - ctx.has_beta = beta is not None - z_shape = (-1, *x0.shape[1:]) - if not return_dmask: - return zmat.view(z_shape) if not prenorm else (zmat.view(z_shape), xmat.view(x0.shape)) - else: - z = zmat.view(z_shape) - dmask = ( - dmask.view(x0.shape) - if dropout_p > 0.0 - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - ctx.mark_non_differentiable(dmask) - return (z, dmask) if not prenorm else (z, xmat.view(x_shape), dmask) - - @staticmethod - def backward(ctx, dz, *args): - # assert dz.is_contiguous() - dz = maybe_align(dz.contiguous(), 16) # this happens! - dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None - x, x0, dmask, gamma, mu, rsigma, colscale, x0_subset, out_subset = ctx.saved_tensors - # x0 is None if colscale is None - dropout_p = ctx.dropout_p - has_residual = ctx.has_residual - dx0mat, dresidualmat, dgamma, dbeta, *rest = _dropout_add_layer_norm_subset_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - colscale, - x0_subset, - out_subset, - dropout_p, - ctx.rowscale_const, - ctx.x0_numrows, - has_residual, - ctx.is_rms_norm, - ) - dx0 = dx0mat.view(-1, *x.shape[1:]) - dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None - dcolscale = rest[0] if colscale is not None else None - return ( - dx0, - dresidual, - dgamma, - dbeta if ctx.has_beta else None, - dcolscale, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) - - -class DropoutAddLayerNormParallelResidualFn(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x0, - x1, - residual, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - residual_in_fp32=False, - prenorm=False, - is_rms_norm=False, - return_dmask=False, - ): - x0 = maybe_align(x0.contiguous(), 16) - x1 = maybe_align(x1.contiguous(), 16) if x1 is not None else None - residual = maybe_align(residual.contiguous(), 16) if residual is not None else None - gamma0 = maybe_align(gamma0.contiguous(), 16) - beta0 = maybe_align(beta0.contiguous(), 16) if beta0 is not None else None - gamma1 = maybe_align(gamma1.contiguous(), 16) if gamma1 is not None else None - beta1 = maybe_align(beta1.contiguous(), 16) if beta1 is not None else None - ( - z0mat, - z1mat, - xmat, - dmask0, - dmask1, - mu, - rsigma, - ) = _dropout_add_layer_norm_parallel_residual_forward( - x0, - x1, - residual, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - residual_in_fp32, - is_rms_norm, - ) - ctx.save_for_backward(xmat.view(x0.shape), dmask0, dmask1, gamma0, gamma1, mu, rsigma) - ctx.prenorm = prenorm - ctx.dropout_p = dropout_p - ctx.has_x1 = x1 is not None - ctx.has_residual = residual is not None - ctx.is_rms_norm = is_rms_norm - ctx.has_beta = beta0 is not None - z = (z0mat.view(x0.shape), z1mat.view(x0.shape) if z1mat is not None else None) - if not return_dmask: - return z if not prenorm else (*z, xmat.view(x0.shape)) - else: - dmask0 = ( - dmask0.view(x0.shape) - if dropout_p > 0.0 - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - dmask1 = ( - dmask1.view(x0.shape) - if dropout_p > 0.0 and x1 is not None - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - ctx.mark_non_differentiable(dmask0) - ctx.mark_non_differentiable(dmask1) - return ( - (*z, dmask0, dmask1) if not prenorm else (*z, xmat.view(x0.shape), dmask0, dmask1) - ) - - @staticmethod - def backward(ctx, dz0, dz1, *args): - dz0 = maybe_align(dz0.contiguous(), 16) # this happens! - dz1 = maybe_align(dz1.contiguous(), 16) if dz1 is not None else None - dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None - x, dmask0, dmask1, gamma0, gamma1, mu, rsigma = ctx.saved_tensors - dropout_p = ctx.dropout_p - has_x1 = ctx.has_x1 - has_residual = ctx.has_residual - ( - dx0mat, - dx1mat, - dresidualmat, - dgamma0, - dbeta0, - dgamma1, - dbeta1, - ) = _dropout_add_layer_norm_parallel_residual_backward( - dz0, - dz1, - dx, - x, - dmask0, - dmask1, - mu, - rsigma, - gamma0, - gamma1, - dropout_p, - has_x1, - has_residual, - ctx.is_rms_norm, - ) - dx0 = dx0mat.view(x.shape) - dx1 = dx1mat.view(x.shape) if dx1mat is not None else None - dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None - return ( - dx0, - dx1, - dresidual, - dgamma0, - dbeta0 if ctx.has_beta else None, - dgamma1, - dbeta1 if ctx.has_beta else None, - None, - None, - None, - None, - None, - None, - ) - - -def layer_norm(x, weight, bias, epsilon): - return DropoutAddLayerNormFn.apply(x, None, weight, bias, None, None, 0.0, epsilon, False) - - -def dropout_add_layer_norm( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - rowscale=None, - layerscale=None, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormFn.apply( - x0, - residual, - weight, - bias, - rowscale, - layerscale, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - False, - return_dropout_mask, - ) - - -def dropout_add_layer_norm_subset( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - layerscale=None, - x0_subset=None, - out_subset=None, - rowscale_const=1.0, - out_numrows=0, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormSubsetFn.apply( - x0, - residual, - weight, - bias, - layerscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32, - prenorm, - False, - return_dropout_mask, - ) - - -def dropout_add_layer_norm_parallel_residual( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormParallelResidualFn.apply( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - False, - return_dropout_mask, - ) - - -class DropoutAddLayerNorm(torch.nn.Module): - def __init__( - self, - hidden_size, - prenorm=False, - p=0.0, - eps=1e-5, - residual_in_fp32=False, - device=None, - dtype=None, - ): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.prenorm = prenorm - self.p = p - self.eps = eps - self.residual_in_fp32 = residual_in_fp32 - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.bias = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.reset_parameters() - - def reset_parameters(self): - init.ones_(self.weight) - init.zeros_(self.bias) - - def forward(self, x0, residual=None): - return dropout_add_layer_norm( - x0, - residual, - self.weight, - self.bias, - self.p if self.training else 0.0, - self.eps, - prenorm=self.prenorm, - residual_in_fp32=self.residual_in_fp32, - ) diff --git a/build/torch212-cxx11-cu126-aarch64-linux/ops/rms_norm.py b/build/torch212-cxx11-cu126-aarch64-linux/ops/rms_norm.py deleted file mode 100644 index 068348d61290e3839dd082b540d898578ba1e8e2..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu126-aarch64-linux/ops/rms_norm.py +++ /dev/null @@ -1,174 +0,0 @@ -# Copyright (c) 2022, Tri Dao. -# Adapted from https://github.com/NVIDIA/apex/blob/master/apex/contrib/layer_norm/layer_norm.py - -import torch -from torch.nn import init - -from flash_attn.ops.layer_norm import ( - DropoutAddLayerNormFn, - DropoutAddLayerNormParallelResidualFn, - DropoutAddLayerNormSubsetFn, -) - - -def rms_norm(x, weight, epsilon): - return DropoutAddLayerNormFn.apply( - x, None, weight, None, None, None, 0.0, epsilon, False, False, True - ) - - -def dropout_add_rms_norm( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - rowscale=None, - layerscale=None, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormFn.apply( - x0, - residual, - weight, - bias, - rowscale, - layerscale, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - True, - return_dropout_mask, - ) - - -def dropout_add_rms_norm_subset( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - layerscale=None, - x0_subset=None, - out_subset=None, - rowscale_const=1.0, - out_numrows=0, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormSubsetFn.apply( - x0, - residual, - weight, - bias, - layerscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32, - prenorm, - True, - return_dropout_mask, - ) - - -def dropout_add_rms_norm_parallel_residual( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormParallelResidualFn.apply( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - True, - return_dropout_mask, - ) - - -class RMSNorm(torch.nn.Module): - def __init__(self, hidden_size, eps=1e-5, device=None, dtype=None): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.eps = eps - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.register_parameter("bias", None) - self.reset_parameters() - - def reset_parameters(self): - init.ones_(self.weight) - - def forward(self, x): - return rms_norm(x, self.weight, self.eps) - - -class DropoutAddRMSNorm(torch.nn.Module): - def __init__( - self, - hidden_size, - prenorm=False, - p=0.0, - eps=1e-5, - residual_in_fp32=False, - device=None, - dtype=None, - ): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.prenorm = prenorm - self.p = p - self.eps = eps - self.residual_in_fp32 = residual_in_fp32 - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.register_parameter("bias", None) - self.reset_parameters() - - def reset_parameters(self): - init.ones_(self.weight) - - def forward(self, x0, residual=None): - return dropout_add_rms_norm( - x0, - residual, - self.weight, - None, - self.p if self.training else 0.0, - self.eps, - prenorm=self.prenorm, - residual_in_fp32=self.residual_in_fp32, - ) diff --git a/build/torch212-cxx11-cu126-aarch64-linux/ops/triton/__init__.py b/build/torch212-cxx11-cu126-aarch64-linux/ops/triton/__init__.py deleted file mode 100644 index 8b137891791fe96927ad78e64b0aad7bded08bdc..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu126-aarch64-linux/ops/triton/__init__.py +++ /dev/null @@ -1 +0,0 @@ - diff --git a/build/torch212-cxx11-cu126-aarch64-linux/ops/triton/cross_entropy.py b/build/torch212-cxx11-cu126-aarch64-linux/ops/triton/cross_entropy.py deleted file mode 100644 index 1b5a415b73f236f3e05fb14b9141959559e18526..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu126-aarch64-linux/ops/triton/cross_entropy.py +++ /dev/null @@ -1,330 +0,0 @@ -# Copyright (c) 2023, Tri Dao. - -from typing import Tuple, Optional, Union - -import torch -import torch.nn.functional as F - -import triton -import triton.language as tl - -# `all_gather_into_tensor` and `reduce_scatter_tensor` are new placeholders for -# `_all_gather_base` and `_reduce_scatter_base`. They require the most recent -# version of PyTorch. The following 2 lines are for backward compatibility with -# older PyTorch. -if "all_gather_into_tensor" not in dir(torch.distributed): - torch.distributed.all_gather_into_tensor = torch.distributed._all_gather_base - - -@triton.heuristics( - { - "HAS_SMOOTHING": lambda args: args["smoothing"] > 0.0, - } -) -@triton.jit -def cross_entropy_fwd_kernel( - loss_ptr, # data ptrs - lse_ptr, - z_loss_ptr, - logits_ptr, - labels_ptr, - smoothing, - logit_scale, - lse_square_scale, - ignore_index, - total_classes, - class_start_idx, # Useful for tensor parallel when each rank only has a subset of classes - n_cols, # shapes - logits_row_stride, # strides - BLOCK_SIZE: tl.constexpr, - HAS_SMOOTHING: tl.constexpr, - # if SPLIT (e.g. tensor parallel), don't include the LSE in the loss since it's not the final LSE - SPLIT: tl.constexpr, - PRECOMPUTED_LSE: tl.constexpr, # If LSE is already computed (also no smoothing and logit_scale == 1.0) -): - row_idx = tl.program_id(0) - logits_ptr = logits_ptr + row_idx * logits_row_stride.to(tl.int64) - sum_logits = 0.0 # For smoothing - if not PRECOMPUTED_LSE: - # Statistics for online softmax - m_i = -float("inf") - l_i = 0.0 - for col_offset in range(0, n_cols, BLOCK_SIZE): - cols = col_offset + tl.arange(0, BLOCK_SIZE) - logits = tl.load(logits_ptr + cols, mask=cols < n_cols, other=-float("inf")).to( - tl.float32 - ) * logit_scale - if HAS_SMOOTHING: - sum_logits += tl.sum(tl.where(cols < n_cols, logits, 0.0)) - m_i_new = tl.maximum(m_i, tl.max(logits)) - l_i = tl.exp(m_i - m_i_new) * l_i + tl.sum(tl.exp(logits - m_i_new)) - m_i = m_i_new - lse = tl.log(l_i) + m_i - tl.store(lse_ptr + row_idx, lse) - else: - lse = tl.load(lse_ptr + row_idx) - label_idx = tl.load(labels_ptr + row_idx) - if label_idx == ignore_index: - loss = 0.0 - z_loss = 0.0 - else: - label_idx -= class_start_idx - if label_idx >= 0 and label_idx < n_cols: - logits_label = tl.load(logits_ptr + label_idx) * logit_scale - if HAS_SMOOTHING: - loss = ( - (lse if not SPLIT else 0.0) - - smoothing * sum_logits / total_classes - - (1 - smoothing) * logits_label - ) - else: - loss = (lse if not SPLIT else 0.0) - logits_label - else: - # If label is out of bounds, we set the CE loss to 0.0. But we still want the smoothing loss - if HAS_SMOOTHING: - loss = smoothing * ((lse if not SPLIT else 0.0) - sum_logits / total_classes) - else: - loss = 0.0 - if not SPLIT: - z_loss = lse_square_scale * lse * lse - loss += z_loss - else: - z_loss = 0.0 - tl.store(loss_ptr + row_idx, loss) - if not SPLIT: - tl.store(z_loss_ptr + row_idx, z_loss) - - -@triton.heuristics( - { - "HAS_SMOOTHING": lambda args: args["smoothing"] > 0.0, - } -) -@triton.jit -def cross_entropy_bwd_kernel( - dlogits_ptr, # data ptrs - dloss_ptr, - logits_ptr, - lse_ptr, - labels_ptr, - smoothing, - logit_scale, - lse_square_scale, - ignore_index, - total_classes, - class_start_idx, # Useful for tensor parallel when each rank only has a subset of classes - n_cols, # shapes - logits_row_stride, # strides - dlogits_row_stride, - dloss_row_stride, - BLOCK_SIZE: tl.constexpr, - HAS_SMOOTHING: tl.constexpr, -): - row_idx = tl.program_id(0) - col_block_idx = tl.program_id(1) - logits_ptr = logits_ptr + row_idx * logits_row_stride.to(tl.int64) - dlogits_ptr = dlogits_ptr + row_idx * dlogits_row_stride.to(tl.int64) - col_offsets = col_block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - label_idx = tl.load(labels_ptr + row_idx) - if label_idx != ignore_index: - dloss = tl.load(dloss_ptr + row_idx * dloss_row_stride) - else: - dloss = 0.0 - logits = tl.load(logits_ptr + col_offsets, mask=col_offsets < n_cols, other=-float("inf")).to( - tl.float32 - ) * logit_scale - lse = tl.load(lse_ptr + row_idx) - probs = tl.exp(logits - lse) - probs += 2.0 * lse_square_scale * lse * probs - label_idx -= class_start_idx - if HAS_SMOOTHING: - smooth_positive = 1.0 - smoothing - smooth_negative = smoothing / total_classes - probs = tl.where(col_offsets == label_idx, probs - smooth_positive, probs) - smooth_negative - else: - probs = tl.where(col_offsets == label_idx, probs - 1.0, probs) - tl.store(dlogits_ptr + col_offsets, (dloss * logit_scale) * probs, mask=col_offsets < n_cols) - - -class CrossEntropyLoss(torch.autograd.Function): - - @staticmethod - def forward( - ctx, - logits, - labels, - precomputed_lse=None, - smoothing=0.0, - logit_scale=1.0, - lse_square_scale=0.0, - ignore_index=-100, - inplace_backward=False, - process_group=None, - ): - # For some reason Triton generates wrong code when labels has dtype long and its address - # is not aligned to 16 bytes. The ld.global.b64 seems to load the wrong label index. - if labels.dtype == torch.long and labels.data_ptr() % 16 != 0: - labels = F.pad(labels, (0, 1))[..., :-1] - assert labels.data_ptr() % 16 == 0 - assert logit_scale > 0.0 - n_rows, n_cols = logits.shape - assert labels.shape == (n_rows,) - world_size = 1 if process_group is None else torch.distributed.get_world_size(process_group) - total_classes = world_size * n_cols - rank = 0 if process_group is None else torch.distributed.get_rank(process_group) - class_start_idx = rank * n_cols - use_precomputed_lse = precomputed_lse is not None and logit_scale == 1.0 and smoothing == 0.0 - - if logits.stride(-1) != 1: - logits = logits.contiguous() - MAX_BLOCK_SIZE = 16 * 1024 - BLOCK_SIZE = min(triton.next_power_of_2(n_cols), MAX_BLOCK_SIZE) - num_warps = ( - 4 - if BLOCK_SIZE < 2048 - else (8 if BLOCK_SIZE < 8192 else (16 if BLOCK_SIZE < 128 * 1024 else 32)) - ) - losses = torch.empty(n_rows, dtype=torch.float, device=logits.device) - if use_precomputed_lse: - assert precomputed_lse.shape == (n_rows,) - lse = precomputed_lse.contiguous() - else: - lse = torch.empty(n_rows, dtype=torch.float, device=logits.device) - z_losses = torch.empty(n_rows, dtype=torch.float, device=logits.device) - # Need this, otherwise Triton tries to launch from cuda:0 and we get - # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?) - with torch.cuda.device(logits.device.index): - cross_entropy_fwd_kernel[(n_rows,)]( - losses, # data ptrs - lse, - z_losses, - logits, - labels, - smoothing, - logit_scale, - lse_square_scale, - ignore_index, - total_classes, - class_start_idx, - n_cols, # shapes - logits.stride(0), # strides - BLOCK_SIZE=BLOCK_SIZE, # constants - SPLIT=world_size > 1, - PRECOMPUTED_LSE=use_precomputed_lse, - num_warps=num_warps, - ) - - if world_size > 1: - # If there's no smoothing, if labels are in the vocab of this partition, losses contains - # - predicted logit, and 0 otherwise. - # If there's smoothing=0.1, for labels in the vocab of this partition, losses contains - # -0.9 * predicted logit - 0.1 * sum logit / total_classes. - # For labels not in the vocab of this partition, losses contains - # -0.1 * sum logit / total_classes. - if world_size > 1: - lse_allgather = torch.empty(world_size, n_rows, dtype=lse.dtype, device=lse.device) - torch.distributed.all_gather_into_tensor(lse_allgather, lse, group=process_group) - handle_losses = torch.distributed.all_reduce( - losses, op=torch.distributed.ReduceOp.SUM, group=process_group, async_op=True - ) - lse = torch.logsumexp(lse_allgather, dim=0) - handle_losses.wait() - # After the allreduce, if there's no smoothing, the total losses are - predicted_logit, - # we just have to add the (global) lse. - # If there's smoothing=0.1, the total losses are - # -0.9 * predicted_logit - 0.1 * sum logit / total_classes. - # Again, we just have to add the (global) lse. - losses += lse - if lse_square_scale != 0.0: - z_losses = lse_square_scale * lse.square() - z_losses.masked_fill_(labels == ignore_index, 0.0) - losses += z_losses - else: - z_losses = torch.zeros_like(losses) - losses.masked_fill_(labels == ignore_index, 0.0) - - ctx.save_for_backward(logits, lse, labels) - ctx.mark_non_differentiable(z_losses) - ctx.smoothing = smoothing - ctx.logit_scale = logit_scale - ctx.lse_square_scale = lse_square_scale - ctx.ignore_index = ignore_index - ctx.total_classes = total_classes - ctx.class_start_idx = class_start_idx - ctx.inplace_backward = inplace_backward - return losses, z_losses - - @staticmethod - def backward(ctx, grad_losses, grad_z_losses): - del grad_z_losses # z_losses are only for logging. - - logits, lse, labels = ctx.saved_tensors - dlogits = logits if ctx.inplace_backward else torch.empty_like(logits) - n_rows, n_cols = logits.shape - BLOCK_SIZE = min(triton.next_power_of_2(n_cols), 4 * 1024) - num_warps = 4 if BLOCK_SIZE < 2048 else (8 if BLOCK_SIZE < 8192 else 16) - grid = lambda META: (n_rows, triton.cdiv(n_cols, META["BLOCK_SIZE"])) # noqa - # Need this, otherwise Triton tries to launch from cuda:0 and we get - # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?) - with torch.cuda.device(logits.device.index): - cross_entropy_bwd_kernel[grid]( - dlogits, # data ptrs - grad_losses, - logits, - lse, - labels, - ctx.smoothing, - ctx.logit_scale, - ctx.lse_square_scale, - ctx.ignore_index, - ctx.total_classes, - ctx.class_start_idx, - n_cols, # shapes - logits.stride(0), # strides - dlogits.stride(0), - grad_losses.stride(0), - BLOCK_SIZE=BLOCK_SIZE, # constants - num_warps=num_warps, - ) - return dlogits, None, None, None, None, None, None, None, None, None - - -def cross_entropy_loss( - logits: torch.Tensor, - labels: torch.Tensor, - precomputed_lse: Optional[torch.Tensor] = None, - label_smoothing: float = 0.0, - logit_scale: float = 1.0, - lse_square_scale: float = 0.0, - ignore_index=-100, - inplace_backward: bool = False, - process_group=None, -) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Arguments: - logits: (batch, vocab_size) - labels: (batch,) - label_smoothing: float - logit_scale: float. Multiply logits by this scale before calculating the loss. - lse_square_scale: float. If > 0, we add lse_square_scale * lse(logits) ^ 2 to the loss. - This is also referred to as "z-loss". - ignore_index: int. If labels == ignore_index, the loss is set to 0.0. - inplace_backward: bool. If True, we do the backward pass in-place by modifying the logits. - This saves memory. - process_group: if not None, we're doing Tensor Parallel: each process is responsible for - one part of the vocab. The loss will be aggregated across processes. - Returns: - losses: (batch,), float - z_losses: (batch,), float - """ - return CrossEntropyLoss.apply( - logits, - labels, - precomputed_lse, - label_smoothing, - logit_scale, - lse_square_scale, - ignore_index, - inplace_backward, - process_group, - ) diff --git a/build/torch212-cxx11-cu126-aarch64-linux/ops/triton/k_activations.py b/build/torch212-cxx11-cu126-aarch64-linux/ops/triton/k_activations.py deleted file mode 100644 index efb83c358eb4a85d069ee340a3c83f418f9a805b..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu126-aarch64-linux/ops/triton/k_activations.py +++ /dev/null @@ -1,162 +0,0 @@ -# Adapted from https://github.com/facebookresearch/xformers/blob/main/xformers/triton/k_activations.py -# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved. -# -# This source code is licensed under the BSD license found in the -# LICENSE file in the root directory of this source tree. - -import math -from enum import Enum -from typing import Optional - -import triton -import triton.language as tl - -_sqrt2pi = math.sqrt(2.0 / math.pi) -_sqrt1_2 = math.sqrt(1.0 / 2) -_gaussian_pdf_normalization = 1.0 / math.sqrt(2 * math.pi) - - -class Activation(str, Enum): - SquaredReLU = "squared_relu" - GeLU = "gelu" - GeLUApprox = "gelu_approx" - LeakyReLU = "leaky_relu" - ReLU = "relu" - - -def get_triton_activation_kernel(activation: Optional[Activation]): - return ( - { - Activation.ReLU: relu, - Activation.LeakyReLU: leaky_relu, - Activation.GeLU: gelu, - Activation.GeLUApprox: gelu_approx, - Activation.SquaredReLU: squared_relu, - }[activation] - if activation - else None - ) - - -def get_triton_activation_bwd_kernel(activation: Optional[Activation]): - return ( - { - Activation.ReLU: relu_grad, - Activation.LeakyReLU: leaky_relu_grad, - Activation.GeLU: gelu_grad, - Activation.GeLUApprox: gelu_approx_grad, - Activation.SquaredReLU: squared_relu_grad, - }[activation] - if activation - else None - ) - - -@triton.jit -def tanh(x): - # Tanh is just a scaled sigmoid - return 2 * tl.sigmoid(2 * x) - 1 - - -@triton.jit -def cosh(x): - exp_x = tl.exp(x) - return (exp_x + 1.0 / exp_x) * 0.5 - - -# a Triton implementation of the most used activations -# See for instance http://arxiv.org/abs/1606.08415 for an overview - -# ReLU -@triton.jit -def relu(x): - """ - ReLU_ activation function - - .. _ReLU: https://pytorch.org/docs/stable/generated/torch.nn.ReLU.html - """ - zero = 0.0 - return tl.where(x >= 0, x, zero.to(x.dtype)) - - -@triton.jit -def relu_grad(x): - # ReLU is different from other activations - # in that it does not require the input to retrospectively compute its gradient - # here the input is the downstream gradient, and we return the upstream gradient directly - zero = 0.0 - one = 1.0 - return tl.where(x >= 0, one.to(x.dtype), zero.to(x.dtype)) - - -@triton.jit -def squared_relu(x): - """ - Squared ReLU activation, as proposed in the Primer_ paper. - - .. _Primer: https://arxiv.org/abs/2109.08668 - """ - x_ = relu(x) - return (x_ * x_).to(x.dtype) - - -@triton.jit -def squared_relu_grad(x): - return tl.where(x >= 0, 2.0 * x, 0.0) - - -# Leaky ReLU -@triton.jit -def leaky_relu(x): - """ - LeakyReLU_ activation - - .. _LeakyReLU: https://pytorch.org/docs/stable/generated/torch.nn.LeakyReLU.html - """ - scale = 0.01 + 0.0 - scale = scale.to(x.dtype) - return tl.where(x >= 0, x, scale * x) - - -@triton.jit -def leaky_relu_grad(x): - min_grad = 0.01 - max_grad = 1 - - min_grad = min_grad.to(x.dtype) - max_grad = max_grad.to(x.dtype) - - return tl.where(x >= 0, max_grad, min_grad) - - -@triton.jit -def gelu(x): - """Gaussian Error Linear Unit (GELU)""" - return x * 0.5 * (1.0 + tl.libdevice.erf(x * _sqrt1_2)) - - -@triton.jit -def gelu_grad(x): - cdf = 0.5 * (1.0 + tl.libdevice.erf(x * _sqrt1_2)) - pdf = tl.exp(-0.5 * x * x) * _gaussian_pdf_normalization - return cdf + x * pdf - - -@triton.jit -def gelu_approx(x): - """ - GeLU_ activation - Gaussian error linear unit, with tanh approximation - - .. _GeLU: https://arxiv.org/pdf/1606.08415.pdf - """ - return 0.5 * x * (1.0 + tanh(_sqrt2pi * x * (1.0 + 0.044715 * x * x))) - - -@triton.jit -def gelu_approx_grad(x): - # CREDITS: Fast implementation proposed in - # https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/fused_bias_gelu.py#L30 - tanh_out = tanh(0.79788456 * x * (1 + 0.044715 * x * x)) - return 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( - 1 + tanh_out - ) diff --git a/build/torch212-cxx11-cu126-aarch64-linux/ops/triton/layer_norm.py b/build/torch212-cxx11-cu126-aarch64-linux/ops/triton/layer_norm.py deleted file mode 100644 index 5e6f0b3697b612e6a19031bbc55ac76ad233dc85..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu126-aarch64-linux/ops/triton/layer_norm.py +++ /dev/null @@ -1,1254 +0,0 @@ -# Copyright (c) 2024, Tri Dao. -# Implement dropout + residual + layer_norm / rms_norm. - -# Based on the Triton LayerNorm tutorial: https://triton-lang.org/main/getting-started/tutorials/05-layer-norm.html -# For the backward pass, we keep weight_grad and bias_grad in registers and accumulate. -# This is faster for dimensions up to 8k, but after that it's much slower due to register spilling. -# The models we train have hidden dim up to 8k anyway (e.g. Llama 70B), so this is fine. - -import math -from typing import Optional, List - -import torch -import torch.nn.functional as F -from torch import Tensor - -import triton -import triton.language as tl - -from flash_attn.utils.torch import custom_fwd, custom_bwd -from flash_attn.utils.library import triton_op - -from ..._ops import add_op_namespace_prefix - - -def maybe_contiguous_lastdim(x): - return x.contiguous() if x is not None and x.stride(-1) != 1 else x - - -def maybe_contiguous(x): - return x.contiguous() if x is not None else None - - -def triton_autotune_configs(): - # Return configs with a valid warp count for the current device - configs = [] - # Maximum threads per block is architecture-dependent in theory, but in reality all are 1024 - max_threads_per_block = 1024 - # Default to warp size 32 if not defined by device - warp_size = getattr(torch.cuda.get_device_properties(torch.cuda.current_device()), "warp_size", 32) - # Autotune for warp counts which are powers of 2 and do not exceed thread per block limit - return [triton.Config({}, num_warps=warp_count) for warp_count in [1, 2, 4, 8, 16, 32] - if warp_count * warp_size <= max_threads_per_block] - # return [triton.Config({}, num_warps=8)] - - -def layer_norm_ref( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - zero_centered_weight=False, - dropout_mask=None, - dropout_mask1=None, - upcast=False, -): - dtype = x.dtype - if upcast: - x = x.float() - weight = weight.float() - bias = bias.float() if bias is not None else None - residual = residual.float() if residual is not None else residual - x1 = x1.float() if x1 is not None else None - weight1 = weight1.float() if weight1 is not None else None - bias1 = bias1.float() if bias1 is not None else None - if zero_centered_weight: - weight = weight + 1.0 - if weight1 is not None: - weight1 = weight1 + 1.0 - if x1 is not None: - assert rowscale is None, "rowscale is not supported with parallel LayerNorm" - if rowscale is not None: - x = x * rowscale[..., None] - if dropout_p > 0.0: - if dropout_mask is not None: - x = x.masked_fill(~dropout_mask, 0.0) / (1.0 - dropout_p) - else: - x = F.dropout(x, p=dropout_p) - if x1 is not None: - if dropout_mask1 is not None: - x1 = x1.masked_fill(~dropout_mask1, 0.0) / (1.0 - dropout_p) - else: - x1 = F.dropout(x1, p=dropout_p) - if x1 is not None: - x = x + x1 - if residual is not None: - x = (x + residual).to(x.dtype) - out = F.layer_norm(x.to(weight.dtype), x.shape[-1:], weight=weight, bias=bias, eps=eps).to( - dtype - ) - if weight1 is None: - return out if not prenorm else (out, x) - else: - out1 = F.layer_norm( - x.to(weight1.dtype), x.shape[-1:], weight=weight1, bias=bias1, eps=eps - ).to(dtype) - return (out, out1) if not prenorm else (out, out1, x) - - -def rms_norm_ref( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - zero_centered_weight=False, - dropout_mask=None, - dropout_mask1=None, - upcast=False, -): - dtype = x.dtype - if upcast: - x = x.float() - weight = weight.float() - bias = bias.float() if bias is not None else None - residual = residual.float() if residual is not None else residual - x1 = x1.float() if x1 is not None else None - weight1 = weight1.float() if weight1 is not None else None - bias1 = bias1.float() if bias1 is not None else None - if zero_centered_weight: - weight = weight + 1.0 - if weight1 is not None: - weight1 = weight1 + 1.0 - if x1 is not None: - assert rowscale is None, "rowscale is not supported with parallel LayerNorm" - if rowscale is not None: - x = x * rowscale[..., None] - if dropout_p > 0.0: - if dropout_mask is not None: - x = x.masked_fill(~dropout_mask, 0.0) / (1.0 - dropout_p) - else: - x = F.dropout(x, p=dropout_p) - if x1 is not None: - if dropout_mask1 is not None: - x1 = x1.masked_fill(~dropout_mask1, 0.0) / (1.0 - dropout_p) - else: - x1 = F.dropout(x1, p=dropout_p) - if x1 is not None: - x = x + x1 - if residual is not None: - x = (x + residual).to(x.dtype) - rstd = 1 / torch.sqrt((x.square()).mean(dim=-1, keepdim=True) + eps) - out = ((x * rstd * weight) + bias if bias is not None else (x * rstd * weight)).to(dtype) - if weight1 is None: - return out if not prenorm else (out, x) - else: - out1 = ((x * rstd * weight1) + bias1 if bias1 is not None else (x * rstd * weight1)).to( - dtype - ) - return (out, out1) if not prenorm else (out, out1, x) - - -@triton.autotune( - configs=triton_autotune_configs(), - key=["N", "HAS_RESIDUAL", "STORE_RESIDUAL_OUT", "IS_RMS_NORM", "HAS_BIAS", "HAS_X1", "HAS_W1", "HAS_B1"], -) -# torch compile doesn't like triton.heuristics, so we set these manually when calling the kernel -# @triton.heuristics({"HAS_BIAS": lambda args: args["B"] is not None}) -# @triton.heuristics({"HAS_RESIDUAL": lambda args: args["RESIDUAL"] is not None}) -# @triton.heuristics({"HAS_X1": lambda args: args["X1"] is not None}) -# @triton.heuristics({"HAS_W1": lambda args: args["W1"] is not None}) -# @triton.heuristics({"HAS_B1": lambda args: args["B1"] is not None}) -@triton.jit -def _layer_norm_fwd_1pass_kernel( - X, # pointer to the input - Y, # pointer to the output - W, # pointer to the weights - B, # pointer to the biases - RESIDUAL, # pointer to the residual - X1, - W1, - B1, - Y1, - RESIDUAL_OUT, # pointer to the residual - ROWSCALE, - SEEDS, # Dropout seeds for each row - DROPOUT_MASK, - DROPOUT_MASK1, - Mean, # pointer to the mean - Rstd, # pointer to the 1/std - stride_x_row, # how much to increase the pointer when moving by 1 row - stride_y_row, - stride_res_row, - stride_res_out_row, - stride_x1_row, - stride_y1_row, - M, # number of rows in X - N, # number of columns in X - eps, # epsilon to avoid division by zero - dropout_p, # Dropout probability - zero_centered_weight, # If true, add 1.0 to the weight - IS_RMS_NORM: tl.constexpr, - BLOCK_N: tl.constexpr, - HAS_RESIDUAL: tl.constexpr, - STORE_RESIDUAL_OUT: tl.constexpr, - HAS_BIAS: tl.constexpr, - HAS_DROPOUT: tl.constexpr, - STORE_DROPOUT_MASK: tl.constexpr, - HAS_ROWSCALE: tl.constexpr, - HAS_X1: tl.constexpr, - HAS_W1: tl.constexpr, - HAS_B1: tl.constexpr, -): - # Map the program id to the row of X and Y it should compute. - row = tl.program_id(0) - X += row * stride_x_row - Y += row * stride_y_row - if HAS_RESIDUAL: - RESIDUAL += row * stride_res_row - if STORE_RESIDUAL_OUT: - RESIDUAL_OUT += row * stride_res_out_row - if HAS_X1: - X1 += row * stride_x1_row - if HAS_W1: - Y1 += row * stride_y1_row - # Compute mean and variance - cols = tl.arange(0, BLOCK_N) - x = tl.load(X + cols, mask=cols < N, other=0.0).to(tl.float32) - if HAS_ROWSCALE: - rowscale = tl.load(ROWSCALE + row).to(tl.float32) - x *= rowscale - if HAS_DROPOUT: - # Compute dropout mask - # 7 rounds is good enough, and reduces register pressure - keep_mask = tl.rand(tl.load(SEEDS + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - x = tl.where(keep_mask, x / (1.0 - dropout_p), 0.0) - if STORE_DROPOUT_MASK: - tl.store(DROPOUT_MASK + row * N + cols, keep_mask, mask=cols < N) - if HAS_X1: - x1 = tl.load(X1 + cols, mask=cols < N, other=0.0).to(tl.float32) - if HAS_ROWSCALE: - rowscale = tl.load(ROWSCALE + M + row).to(tl.float32) - x1 *= rowscale - if HAS_DROPOUT: - # Compute dropout mask - # 7 rounds is good enough, and reduces register pressure - keep_mask = ( - tl.rand(tl.load(SEEDS + M + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - ) - x1 = tl.where(keep_mask, x1 / (1.0 - dropout_p), 0.0) - if STORE_DROPOUT_MASK: - tl.store(DROPOUT_MASK1 + row * N + cols, keep_mask, mask=cols < N) - x += x1 - if HAS_RESIDUAL: - residual = tl.load(RESIDUAL + cols, mask=cols < N, other=0.0).to(tl.float32) - x += residual - if STORE_RESIDUAL_OUT: - tl.store(RESIDUAL_OUT + cols, x, mask=cols < N) - if not IS_RMS_NORM: - mean = tl.sum(x, axis=0) / N - tl.store(Mean + row, mean) - xbar = tl.where(cols < N, x - mean, 0.0) - var = tl.sum(xbar * xbar, axis=0) / N - else: - xbar = tl.where(cols < N, x, 0.0) - var = tl.sum(xbar * xbar, axis=0) / N - rstd = 1 / tl.sqrt(var + eps) - tl.store(Rstd + row, rstd) - # Normalize and apply linear transformation - mask = cols < N - w = tl.load(W + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w += 1.0 - if HAS_BIAS: - b = tl.load(B + cols, mask=mask).to(tl.float32) - x_hat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd - y = x_hat * w + b if HAS_BIAS else x_hat * w - # Write output - tl.store(Y + cols, y, mask=mask) - if HAS_W1: - w1 = tl.load(W1 + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w1 += 1.0 - if HAS_B1: - b1 = tl.load(B1 + cols, mask=mask).to(tl.float32) - y1 = x_hat * w1 + b1 if HAS_B1 else x_hat * w1 - tl.store(Y1 + cols, y1, mask=mask) - - -def _layer_norm_fwd( - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - residual: Optional[Tensor] = None, - x1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - out_dtype: Optional[torch.dtype] = None, - residual_dtype: Optional[torch.dtype] = None, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - return_dropout_mask: bool = False, - out: Optional[Tensor] = None, - residual_out: Optional[Tensor] = None -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - # Need to wrap to handle the case where residual_out is a alias of x, which makes torch.library - # and torch.compile unhappy. Also allocate memory for out and residual_out if they are None - # so that _layer_norm_fwd_impl doesn't have to return them. - if out is None: - out = torch.empty_like(x, dtype=x.dtype if out_dtype is None else out_dtype) - if residual is not None: - residual_dtype = residual.dtype - if residual_out is None and ( - residual is not None - or (residual_dtype is not None and residual_dtype != x.dtype) - or dropout_p > 0.0 - or rowscale is not None - or x1 is not None - ): - residual_out = torch.empty_like( - x, dtype=residual_dtype if residual_dtype is not None else x.dtype - ) - else: - residual_out = None - y1, mean, rstd, seeds, dropout_mask, dropout_mask1 = _layer_norm_fwd_impl( - x, - weight, - bias, - eps, - out, - residual=residual, - x1=x1, - weight1=weight1, - bias1=bias1, - dropout_p=dropout_p, - rowscale=rowscale, - zero_centered_weight=zero_centered_weight, - is_rms_norm=is_rms_norm, - return_dropout_mask=return_dropout_mask, - residual_out=residual_out, - ) - # residual_out is None if residual is None and residual_dtype == input_dtype and dropout_p == 0.0 - if residual_out is None: - residual_out = x - return out, y1, mean, rstd, residual_out, seeds, dropout_mask, dropout_mask1 - - -# [2025-04-28] torch.library.triton_op ignores the schema argument, but here we need the schema -# since we're returning a tuple of tensors -@triton_op(add_op_namespace_prefix("layer_norm_fwd_impl"), mutates_args={"out", "residual_out"}, - schema="(Tensor x, Tensor weight, Tensor bias, float eps, Tensor(a!) out, Tensor? residual, Tensor? x1, Tensor? weight1, Tensor? bias1, float dropout_p, Tensor? rowscale, bool zero_centered_weight, bool is_rms_norm, bool return_dropout_mask, Tensor(a!)? residual_out) -> (Tensor y1, Tensor mean, Tensor rstd, Tensor seeds, Tensor dropout_mask, Tensor dropout_mask1)") -def _layer_norm_fwd_impl( - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - out: Tensor, - residual: Optional[Tensor] = None, - x1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - return_dropout_mask: bool = False, - residual_out: Optional[Tensor] = None -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - M, N = x.shape - assert x.stride(-1) == 1 - if residual is not None: - assert residual.stride(-1) == 1 - assert residual.shape == (M, N) - assert weight.shape == (N,) - assert weight.stride(-1) == 1 - if bias is not None: - assert bias.stride(-1) == 1 - assert bias.shape == (N,) - if x1 is not None: - assert x1.shape == x.shape - assert rowscale is None - assert x1.stride(-1) == 1 - if weight1 is not None: - assert weight1.shape == (N,) - assert weight1.stride(-1) == 1 - if bias1 is not None: - assert bias1.shape == (N,) - assert bias1.stride(-1) == 1 - if rowscale is not None: - assert rowscale.is_contiguous() - assert rowscale.shape == (M,) - assert out.shape == x.shape - assert out.stride(-1) == 1 - if residual_out is not None: - assert residual_out.shape == x.shape - assert residual_out.stride(-1) == 1 - if weight1 is not None: - y1 = torch.empty_like(out) - assert y1.stride(-1) == 1 - else: - y1 = None - mean = torch.empty((M,), dtype=torch.float32, device=x.device) if not is_rms_norm else None - rstd = torch.empty((M,), dtype=torch.float32, device=x.device) - if dropout_p > 0.0: - seeds = torch.randint( - 2**32, (M if x1 is None else 2 * M,), device=x.device, dtype=torch.int64 - ) - else: - seeds = None - if return_dropout_mask and dropout_p > 0.0: - dropout_mask = torch.empty(M, N, device=x.device, dtype=torch.bool) - if x1 is not None: - dropout_mask1 = torch.empty(M, N, device=x.device, dtype=torch.bool) - else: - dropout_mask1 = None - else: - dropout_mask, dropout_mask1 = None, None - # Less than 64KB per feature: enqueue fused kernel - MAX_FUSED_SIZE = 65536 // x.element_size() - BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) - if N > BLOCK_N: - raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") - with torch.cuda.device(x.device.index): - torch.library.wrap_triton(_layer_norm_fwd_1pass_kernel)[(M,)]( - x, - out, - weight, - bias, - residual, - x1, - weight1, - bias1, - y1, - residual_out, - rowscale, - seeds, - dropout_mask, - dropout_mask1, - mean, - rstd, - x.stride(0), - out.stride(0), - residual.stride(0) if residual is not None else 0, - residual_out.stride(0) if residual_out is not None else 0, - x1.stride(0) if x1 is not None else 0, - y1.stride(0) if y1 is not None else 0, - M, - N, - eps, - dropout_p, - # Passing bool make torch inductor very unhappy since it then tries to compare to int_max - int(zero_centered_weight), - is_rms_norm, - BLOCK_N, - residual is not None, - residual_out is not None, - bias is not None, - dropout_p > 0.0, - dropout_mask is not None, - rowscale is not None, - HAS_X1=x1 is not None, - HAS_W1=weight1 is not None, - HAS_B1=bias1 is not None, - ) - return y1, mean, rstd, seeds, dropout_mask, dropout_mask1 - - -@triton.autotune( - configs=triton_autotune_configs(), - key=["N", "HAS_DRESIDUAL", "STORE_DRESIDUAL", "IS_RMS_NORM", "HAS_BIAS", "HAS_DROPOUT"], -) -# torch compile doesn't like triton.heuristics, so we set these manually when calling the kernel -# @triton.heuristics({"HAS_BIAS": lambda args: args["B"] is not None}) -# @triton.heuristics({"HAS_DRESIDUAL": lambda args: args["DRESIDUAL"] is not None}) -# @triton.heuristics({"STORE_DRESIDUAL": lambda args: args["DRESIDUAL_IN"] is not None}) -# @triton.heuristics({"HAS_ROWSCALE": lambda args: args["ROWSCALE"] is not None}) -# @triton.heuristics({"HAS_DY1": lambda args: args["DY1"] is not None}) -# @triton.heuristics({"HAS_DX1": lambda args: args["DX1"] is not None}) -# @triton.heuristics({"HAS_B1": lambda args: args["DB1"] is not None}) -# @triton.heuristics({"RECOMPUTE_OUTPUT": lambda args: args["Y"] is not None}) -@triton.jit -def _layer_norm_bwd_kernel( - X, # pointer to the input - W, # pointer to the weights - B, # pointer to the biases - Y, # pointer to the output to be recomputed - DY, # pointer to the output gradient - DX, # pointer to the input gradient - DW, # pointer to the partial sum of weights gradient - DB, # pointer to the partial sum of biases gradient - DRESIDUAL, - W1, - DY1, - DX1, - DW1, - DB1, - DRESIDUAL_IN, - ROWSCALE, - SEEDS, - Mean, # pointer to the mean - Rstd, # pointer to the 1/std - stride_x_row, # how much to increase the pointer when moving by 1 row - stride_y_row, - stride_dy_row, - stride_dx_row, - stride_dres_row, - stride_dy1_row, - stride_dx1_row, - stride_dres_in_row, - M, # number of rows in X - N, # number of columns in X - eps, # epsilon to avoid division by zero - dropout_p, - zero_centered_weight, - rows_per_program, - IS_RMS_NORM: tl.constexpr, - BLOCK_N: tl.constexpr, - HAS_DRESIDUAL: tl.constexpr, - STORE_DRESIDUAL: tl.constexpr, - HAS_BIAS: tl.constexpr, - HAS_DROPOUT: tl.constexpr, - HAS_ROWSCALE: tl.constexpr, - HAS_DY1: tl.constexpr, - HAS_DX1: tl.constexpr, - HAS_B1: tl.constexpr, - RECOMPUTE_OUTPUT: tl.constexpr, -): - # Map the program id to the elements of X, DX, and DY it should compute. - row_block_id = tl.program_id(0) - row_start = row_block_id * rows_per_program - # Do not early exit if row_start >= M, because we need to write DW and DB - cols = tl.arange(0, BLOCK_N) - mask = cols < N - X += row_start * stride_x_row - if HAS_DRESIDUAL: - DRESIDUAL += row_start * stride_dres_row - if STORE_DRESIDUAL: - DRESIDUAL_IN += row_start * stride_dres_in_row - DY += row_start * stride_dy_row - DX += row_start * stride_dx_row - if HAS_DY1: - DY1 += row_start * stride_dy1_row - if HAS_DX1: - DX1 += row_start * stride_dx1_row - if RECOMPUTE_OUTPUT: - Y += row_start * stride_y_row - w = tl.load(W + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w += 1.0 - if RECOMPUTE_OUTPUT and HAS_BIAS: - b = tl.load(B + cols, mask=mask, other=0.0).to(tl.float32) - if HAS_DY1: - w1 = tl.load(W1 + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w1 += 1.0 - dw = tl.zeros((BLOCK_N,), dtype=tl.float32) - if HAS_BIAS: - db = tl.zeros((BLOCK_N,), dtype=tl.float32) - if HAS_DY1: - dw1 = tl.zeros((BLOCK_N,), dtype=tl.float32) - if HAS_B1: - db1 = tl.zeros((BLOCK_N,), dtype=tl.float32) - row_end = min((row_block_id + 1) * rows_per_program, M) - for row in range(row_start, row_end): - # Load data to SRAM - x = tl.load(X + cols, mask=mask, other=0).to(tl.float32) - dy = tl.load(DY + cols, mask=mask, other=0).to(tl.float32) - if HAS_DY1: - dy1 = tl.load(DY1 + cols, mask=mask, other=0).to(tl.float32) - if not IS_RMS_NORM: - mean = tl.load(Mean + row) - rstd = tl.load(Rstd + row) - # Compute dx - xhat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd - xhat = tl.where(mask, xhat, 0.0) - if RECOMPUTE_OUTPUT: - y = xhat * w + b if HAS_BIAS else xhat * w - tl.store(Y + cols, y, mask=mask) - wdy = w * dy - dw += dy * xhat - if HAS_BIAS: - db += dy - if HAS_DY1: - wdy += w1 * dy1 - dw1 += dy1 * xhat - if HAS_B1: - db1 += dy1 - if not IS_RMS_NORM: - c1 = tl.sum(xhat * wdy, axis=0) / N - c2 = tl.sum(wdy, axis=0) / N - dx = (wdy - (xhat * c1 + c2)) * rstd - else: - c1 = tl.sum(xhat * wdy, axis=0) / N - dx = (wdy - xhat * c1) * rstd - if HAS_DRESIDUAL: - dres = tl.load(DRESIDUAL + cols, mask=mask, other=0).to(tl.float32) - dx += dres - # Write dx - if STORE_DRESIDUAL: - tl.store(DRESIDUAL_IN + cols, dx, mask=mask) - if HAS_DX1: - if HAS_DROPOUT: - keep_mask = ( - tl.rand(tl.load(SEEDS + M + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - ) - dx1 = tl.where(keep_mask, dx / (1.0 - dropout_p), 0.0) - else: - dx1 = dx - tl.store(DX1 + cols, dx1, mask=mask) - if HAS_DROPOUT: - keep_mask = tl.rand(tl.load(SEEDS + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - dx = tl.where(keep_mask, dx / (1.0 - dropout_p), 0.0) - if HAS_ROWSCALE: - rowscale = tl.load(ROWSCALE + row).to(tl.float32) - dx *= rowscale - tl.store(DX + cols, dx, mask=mask) - - X += stride_x_row - if HAS_DRESIDUAL: - DRESIDUAL += stride_dres_row - if STORE_DRESIDUAL: - DRESIDUAL_IN += stride_dres_in_row - if RECOMPUTE_OUTPUT: - Y += stride_y_row - DY += stride_dy_row - DX += stride_dx_row - if HAS_DY1: - DY1 += stride_dy1_row - if HAS_DX1: - DX1 += stride_dx1_row - tl.store(DW + row_block_id * N + cols, dw, mask=mask) - if HAS_BIAS: - tl.store(DB + row_block_id * N + cols, db, mask=mask) - if HAS_DY1: - tl.store(DW1 + row_block_id * N + cols, dw1, mask=mask) - if HAS_B1: - tl.store(DB1 + row_block_id * N + cols, db1, mask=mask) - - -def _layer_norm_bwd( - dy: Tensor, - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - mean: Tensor, - rstd: Tensor, - dresidual: Optional[Tensor] = None, - dy1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - seeds: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - has_residual: bool = False, - has_x1: bool = False, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - x_dtype: Optional[torch.dtype] = None, - recompute_output: bool = False, -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - # Need to wrap to handle the case where dresidual_in or dx1 are aliases of x, - # which makes torch.library unhappy - dx, dw, db, dresidual_in, dx1, dw1, db1, y = _layer_norm_bwd_impl( - dy, - x, - weight, - bias, - eps, - mean, - rstd, - dresidual, - dy1, - weight1, - bias1, - seeds, - dropout_p, - rowscale, - has_residual, - has_x1, - zero_centered_weight, - is_rms_norm, - x_dtype=x_dtype, - recompute_output=recompute_output, - ) - # Don't need to compute dresidual_in separately in this case - if has_residual and dx.dtype == x.dtype and dropout_p == 0.0 and rowscale is None: - dresidual_in = dx - if has_x1 and dropout_p == 0.0: - dx1 = dx - return dx, dw, db, dresidual_in, dx1, dw1, db1, y - - - -@triton_op(add_op_namespace_prefix("layer_norm_bwd_impl"), mutates_args={}, - schema="(Tensor dy, Tensor x, Tensor weight, Tensor bias, float eps, Tensor mean, Tensor rstd, Tensor? dresidual, Tensor? dy1, Tensor? weight1, Tensor? bias1, Tensor? seeds, float dropout_p, Tensor? rowscale, bool has_residual, bool has_x1, bool zero_centered_weight, bool is_rms_norm, ScalarType? x_dtype, bool recompute_output) -> (Tensor dx, Tensor dw, Tensor db, Tensor dresidual_in, Tensor dx1, Tensor dw1, Tensor db1, Tensor y)", - allow_decomposition=False, # Don't let torch.compile trace inside - ) -def _layer_norm_bwd_impl( - dy: Tensor, - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - mean: Tensor, - rstd: Tensor, - dresidual: Optional[Tensor] = None, - dy1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - seeds: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - has_residual: bool = False, - has_x1: bool = False, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - x_dtype: Optional[torch.dtype] = None, - recompute_output: bool = False, -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - M, N = x.shape - assert x.stride(-1) == 1 - dy = maybe_contiguous_lastdim(dy) - assert dy.stride(-1) == 1 - assert dy.shape == (M, N) - if dresidual is not None: - dresidual = maybe_contiguous_lastdim(dresidual) - assert dresidual.stride(-1) == 1 - assert dresidual.shape == (M, N) - assert weight.shape == (N,) - assert weight.stride(-1) == 1 - if bias is not None: - assert bias.stride(-1) == 1 - assert bias.shape == (N,) - if dy1 is not None: - dy1 = maybe_contiguous_lastdim(dy1) - assert weight1 is not None - assert dy1.shape == dy.shape - assert dy1.stride(-1) == 1 - if weight1 is not None: - assert weight1.shape == (N,) - assert weight1.stride(-1) == 1 - if bias1 is not None: - assert bias1.shape == (N,) - assert bias1.stride(-1) == 1 - if seeds is not None: - assert seeds.is_contiguous() - assert seeds.shape == (M if not has_x1 else M * 2,) - if rowscale is not None: - assert rowscale.is_contiguous() - assert rowscale.shape == (M,) - # allocate output - dx = ( - torch.empty_like(x) - if x_dtype is None - else torch.empty(M, N, dtype=x_dtype, device=x.device) - ) - dresidual_in = ( - torch.empty_like(x) - if has_residual - and (dx.dtype != x.dtype or dropout_p > 0.0 or rowscale is not None or has_x1) - else None - ) - dx1 = torch.empty_like(dx) if (has_x1 and dropout_p > 0.0) else None - y = torch.empty(M, N, dtype=dy.dtype, device=dy.device) if recompute_output else None - if recompute_output: - assert weight1 is None, "recompute_output is not supported with parallel LayerNorm" - - # Less than 64KB per feature: enqueue fused kernel - MAX_FUSED_SIZE = 65536 // x.element_size() - BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) - if N > BLOCK_N: - raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") - # Increasing the multiple (e.g. 8) will allow more thread blocks to be launched and hide the - # latency of the gmem reads/writes, but will increase the time of summing up dw / db. - sm_count = torch.cuda.get_device_properties(x.device).multi_processor_count * 8 - _dw = torch.empty((sm_count, N), dtype=torch.float32, device=weight.device) - _db = ( - torch.empty((sm_count, N), dtype=torch.float32, device=bias.device) - if bias is not None - else None - ) - _dw1 = torch.empty_like(_dw) if weight1 is not None else None - _db1 = torch.empty_like(_db) if bias1 is not None else None - rows_per_program = math.ceil(M / sm_count) - grid = (sm_count,) - with torch.cuda.device(x.device.index): - torch.library.wrap_triton(_layer_norm_bwd_kernel)[grid]( - x, - weight, - bias, - y, - dy, - dx, - _dw, - _db, - dresidual, - weight1, - dy1, - dx1, - _dw1, - _db1, - dresidual_in, - rowscale, - seeds, - mean, - rstd, - x.stride(0), - 0 if not recompute_output else y.stride(0), - dy.stride(0), - dx.stride(0), - dresidual.stride(0) if dresidual is not None else 0, - dy1.stride(0) if dy1 is not None else 0, - dx1.stride(0) if dx1 is not None else 0, - dresidual_in.stride(0) if dresidual_in is not None else 0, - M, - N, - eps, - dropout_p, - # Passing bool make torch inductor very unhappy since it then tries to compare to int_max - int(zero_centered_weight), - rows_per_program, - is_rms_norm, - BLOCK_N, - dresidual is not None, - dresidual_in is not None, - bias is not None, - dropout_p > 0.0, - HAS_ROWSCALE=rowscale is not None, - HAS_DY1=dy1 is not None, - HAS_DX1=dx1 is not None, - HAS_B1=bias1 is not None, - RECOMPUTE_OUTPUT=y is not None, - ) - dw = _dw.sum(0).to(weight.dtype) - db = _db.sum(0).to(bias.dtype) if bias is not None else None - dw1 = _dw1.sum(0).to(weight1.dtype) if weight1 is not None else None - db1 = _db1.sum(0).to(bias1.dtype) if bias1 is not None else None - # dresidual_in and dx1 could be None, the wrapper will handle assigning them from dx - return dx, dw, db, dresidual_in, dx1, dw1, db1, y - - -class LayerNormFn(torch.autograd.Function): - - @staticmethod - def forward( - ctx, - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - residual_in_fp32=False, - zero_centered_weight=False, - is_rms_norm=False, - return_dropout_mask=False, - out_dtype=None, - out=None, - residual_out=None - ): - x_shape_og = x.shape - # reshape input data into 2D tensor - x = maybe_contiguous_lastdim(x.reshape(-1, x.shape[-1])) - if residual is not None: - assert residual.shape == x_shape_og - residual = maybe_contiguous_lastdim(residual.reshape(-1, residual.shape[-1])) - if x1 is not None: - assert x1.shape == x_shape_og - assert rowscale is None, "rowscale is not supported with parallel LayerNorm" - x1 = maybe_contiguous_lastdim(x1.reshape(-1, x1.shape[-1])) - weight = weight.contiguous() - bias = maybe_contiguous(bias) - weight1 = maybe_contiguous(weight1) - bias1 = maybe_contiguous(bias1) - if rowscale is not None: - rowscale = rowscale.reshape(-1).contiguous() - residual_dtype = ( - residual.dtype - if residual is not None - else (torch.float32 if residual_in_fp32 else None) - ) - if out is not None: - out = out.reshape(-1, out.shape[-1]) - if residual_out is not None: - residual_out = residual_out.reshape(-1, residual_out.shape[-1]) - y, y1, mean, rstd, residual_out, seeds, dropout_mask, dropout_mask1 = _layer_norm_fwd( - x, - weight, - bias, - eps, - residual, - x1, - weight1, - bias1, - dropout_p=dropout_p, - rowscale=rowscale, - out_dtype=out_dtype, - residual_dtype=residual_dtype, - zero_centered_weight=zero_centered_weight, - is_rms_norm=is_rms_norm, - return_dropout_mask=return_dropout_mask, - out=out, - residual_out=residual_out, - ) - ctx.save_for_backward( - residual_out, weight, bias, weight1, bias1, rowscale, seeds, mean, rstd - ) - ctx.x_shape_og = x_shape_og - ctx.eps = eps - ctx.dropout_p = dropout_p - ctx.is_rms_norm = is_rms_norm - ctx.has_residual = residual is not None - ctx.has_x1 = x1 is not None - ctx.prenorm = prenorm - ctx.x_dtype = x.dtype - ctx.zero_centered_weight = zero_centered_weight - y = y.reshape(x_shape_og) - y1 = y1.reshape(x_shape_og) if y1 is not None else None - residual_out = residual_out.reshape(x_shape_og) if residual_out is not None else None - dropout_mask = dropout_mask.reshape(x_shape_og) if dropout_mask is not None else None - dropout_mask1 = dropout_mask1.reshape(x_shape_og) if dropout_mask1 is not None else None - if not return_dropout_mask: - if weight1 is None: - return y if not prenorm else (y, residual_out) - else: - return (y, y1) if not prenorm else (y, y1, residual_out) - else: - if weight1 is None: - return ( - (y, dropout_mask, dropout_mask1) - if not prenorm - else (y, residual_out, dropout_mask, dropout_mask1) - ) - else: - return ( - (y, y1, dropout_mask, dropout_mask1) - if not prenorm - else (y, y1, residual_out, dropout_mask, dropout_mask1) - ) - - @staticmethod - def backward(ctx, dy, *args): - x, weight, bias, weight1, bias1, rowscale, seeds, mean, rstd = ctx.saved_tensors - dy = dy.reshape(-1, dy.shape[-1]) - if weight1 is not None: - dy1, args = args[0], args[1:] - dy1 = dy1.reshape(-1, dy1.shape[-1]) - assert dy1.shape == x.shape - else: - dy1 = None - if ctx.prenorm: - dresidual = args[0] - dresidual = dresidual.reshape(-1, dresidual.shape[-1]) - assert dresidual.shape == x.shape - else: - dresidual = None - dx, dw, db, dresidual_in, dx1, dw1, db1, _ = _layer_norm_bwd( - dy, - x, - weight, - bias, - ctx.eps, - mean, - rstd, - dresidual, - dy1, - weight1, - bias1, - seeds, - ctx.dropout_p, - rowscale, - ctx.has_residual, - ctx.has_x1, - ctx.zero_centered_weight, - ctx.is_rms_norm, - x_dtype=ctx.x_dtype, - recompute_output=False, - ) - return ( - dx.reshape(ctx.x_shape_og), - dw, - db, - dresidual_in.reshape(ctx.x_shape_og) if ctx.has_residual else None, - dx1.reshape(ctx.x_shape_og) if dx1 is not None else None, - dw1, - db1, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) - - -def layer_norm_fn( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - residual_in_fp32=False, - zero_centered_weight=False, - is_rms_norm=False, - return_dropout_mask=False, - out_dtype=None, - out=None, - residual_out=None -): - return LayerNormFn.apply( - x, - weight, - bias, - residual, - x1, - weight1, - bias1, - eps, - dropout_p, - rowscale, - prenorm, - residual_in_fp32, - zero_centered_weight, - is_rms_norm, - return_dropout_mask, - out_dtype, - out, - residual_out - ) - - -def rms_norm_fn( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - residual_in_fp32=False, - zero_centered_weight=False, - return_dropout_mask=False, - out_dtype=None, - out=None, - residual_out=None -): - return LayerNormFn.apply( - x, - weight, - bias, - residual, - x1, - weight1, - bias1, - eps, - dropout_p, - rowscale, - prenorm, - residual_in_fp32, - zero_centered_weight, - True, - return_dropout_mask, - out_dtype, - out, - residual_out - ) - - -class RMSNorm(torch.nn.Module): - - def __init__(self, hidden_size, eps=1e-5, dropout_p=0.0, zero_centered_weight=False, - device=None, dtype=None): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.eps = eps - if dropout_p > 0.0: - self.drop = torch.nn.Dropout(dropout_p) - else: - self.drop = None - self.zero_centered_weight = zero_centered_weight - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.register_parameter("bias", None) - self.reset_parameters() - - def reset_parameters(self): - if not self.zero_centered_weight: - torch.nn.init.ones_(self.weight) - else: - torch.nn.init.zeros_(self.weight) - - def forward(self, x, residual=None, prenorm=False, residual_in_fp32=False): - return rms_norm_fn( - x, - self.weight, - self.bias, - residual=residual, - eps=self.eps, - dropout_p=self.drop.p if self.drop is not None and self.training else 0.0, - prenorm=prenorm, - residual_in_fp32=residual_in_fp32, - zero_centered_weight=self.zero_centered_weight, - ) - - -class LayerNormLinearFn(torch.autograd.Function): - - @staticmethod - @custom_fwd - def forward( - ctx, - x, - norm_weight, - norm_bias, - linear_weight, - linear_bias, - residual=None, - eps=1e-6, - prenorm=False, - residual_in_fp32=False, - is_rms_norm=False, - ): - x_shape_og = x.shape - # reshape input data into 2D tensor - x = maybe_contiguous_lastdim(x.reshape(-1, x.shape[-1])) - if residual is not None: - assert residual.shape == x_shape_og - residual = maybe_contiguous_lastdim(residual.reshape(-1, residual.shape[-1])) - norm_weight = norm_weight.contiguous() - norm_bias = maybe_contiguous(norm_bias) - residual_dtype = ( - residual.dtype - if residual is not None - else (torch.float32 if residual_in_fp32 else None) - ) - y, _, mean, rstd, residual_out, *rest = _layer_norm_fwd( - x, - norm_weight, - norm_bias, - eps, - residual, - out_dtype=None if not torch.is_autocast_enabled() else torch.get_autocast_dtype("cuda"), - residual_dtype=residual_dtype, - is_rms_norm=is_rms_norm, - ) - y = y.reshape(x_shape_og) - dtype = torch.get_autocast_dtype("cuda") if torch.is_autocast_enabled() else y.dtype - linear_weight = linear_weight.to(dtype) - linear_bias = linear_bias.to(dtype) if linear_bias is not None else None - out = F.linear(y.to(linear_weight.dtype), linear_weight, linear_bias) - # We don't store y, will be recomputed in the backward pass to save memory - ctx.save_for_backward(residual_out, norm_weight, norm_bias, linear_weight, mean, rstd) - ctx.x_shape_og = x_shape_og - ctx.eps = eps - ctx.is_rms_norm = is_rms_norm - ctx.has_residual = residual is not None - ctx.prenorm = prenorm - ctx.x_dtype = x.dtype - ctx.linear_bias_is_none = linear_bias is None - return out if not prenorm else (out, residual_out.reshape(x_shape_og)) - - @staticmethod - @custom_bwd - def backward(ctx, dout, *args): - x, norm_weight, norm_bias, linear_weight, mean, rstd = ctx.saved_tensors - dout = dout.reshape(-1, dout.shape[-1]) - dy = F.linear(dout, linear_weight.t()) - dlinear_bias = None if ctx.linear_bias_is_none else dout.sum(0) - dy = maybe_contiguous_lastdim(dy) - assert dy.shape == x.shape - if ctx.prenorm: - dresidual = args[0] - dresidual = maybe_contiguous_lastdim(dresidual.reshape(-1, dresidual.shape[-1])) - assert dresidual.shape == x.shape - else: - dresidual = None - dx, dnorm_weight, dnorm_bias, dresidual_in, _, _, _, y = _layer_norm_bwd( - dy, - x, - norm_weight, - norm_bias, - ctx.eps, - mean, - rstd, - dresidual=dresidual, - has_residual=ctx.has_residual, - is_rms_norm=ctx.is_rms_norm, - x_dtype=ctx.x_dtype, - recompute_output=True, - ) - dlinear_weight = torch.einsum("bo,bi->oi", dout, y) - return ( - dx.reshape(ctx.x_shape_og), - dnorm_weight, - dnorm_bias, - dlinear_weight, - dlinear_bias, - dresidual_in.reshape(ctx.x_shape_og) if ctx.has_residual else None, - None, - None, - None, - None, - ) - - -def layer_norm_linear_fn( - x, - norm_weight, - norm_bias, - linear_weight, - linear_bias, - residual=None, - eps=1e-6, - prenorm=False, - residual_in_fp32=False, - is_rms_norm=False, -): - return LayerNormLinearFn.apply( - x, - norm_weight, - norm_bias, - linear_weight, - linear_bias, - residual, - eps, - prenorm, - residual_in_fp32, - is_rms_norm, - ) diff --git a/build/torch212-cxx11-cu126-aarch64-linux/ops/triton/linear.py b/build/torch212-cxx11-cu126-aarch64-linux/ops/triton/linear.py deleted file mode 100644 index a8966dbc345ab0e593df0124451ee7be3dae131a..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu126-aarch64-linux/ops/triton/linear.py +++ /dev/null @@ -1,594 +0,0 @@ -# Adapted from https://github.com/ELS-RD/kernl/blob/main/src/kernl/implementations/linear_layer.py -# and https://github.com/openai/triton/blob/master/python/triton/ops/matmul.py -from typing import Optional - -import torch -import triton -import triton.language as tl -from triton.ops.matmul_perf_model import early_config_prune, estimate_matmul_time - -from flash_attn.ops.triton.k_activations import ( - gelu, - gelu_approx, - gelu_approx_grad, - gelu_grad, - squared_relu, - squared_relu_grad, -) - -# CREDITS: Initially inspired by the Triton tutorial on matrix multiplications - - -def init_to_zero(name): - return lambda nargs: nargs[name].zero_() - - -def get_configs_io_bound(): - configs = [] - for num_stages in [2, 3, 4, 5, 6]: - for block_m in [16, 32]: - for block_k in [32, 64]: - for block_n in [32, 64, 128, 256]: - num_warps = 2 if block_n <= 64 else 4 - configs.append( - triton.Config( - { - "BLOCK_M": block_m, - "BLOCK_N": block_n, - "BLOCK_K": block_k, - "SPLIT_K": 1, - }, - num_stages=num_stages, - num_warps=num_warps, - ) - ) - # split_k not used - # for split_k in [2, 4, 8, 16]: - # configs.append(triton.Config( - # {'BLOCK_M': block_m, 'BLOCK_N': block_n, 'BLOCK_K': block_k, 'SPLIT_K': split_k}, - # num_stages=num_stages, num_warps=num_warps, pre_hook=init_to_zero('C'))) - return configs - - -@triton.autotune( - configs=[ - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - # good for int8 - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=4, - num_warps=4, - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - ] - + get_configs_io_bound(), - key=["CACHE_KEY_M", "CACHE_KEY_N", "CACHE_KEY_K"], - prune_configs_by={ - "early_config_prune": early_config_prune, - "perf_model": estimate_matmul_time, - "top_k": 10, - }, -) -@triton.heuristics( - { - "EVEN_K": lambda args: args["K"] % (args["BLOCK_K"] * args["SPLIT_K"]) == 0, - } -) -@triton.jit -def kernel_fwd( - C, # Pointers to matrices - ACT_INPUT, - A, - B, - bias, - # Matrix dimensions - M, - N, - K, - CACHE_KEY_M, - CACHE_KEY_N, - CACHE_KEY_K, - # The stride variables represent how much to increase the ptr by when moving by 1 - # element in a particular dimension. E.g. stride_am is how much to increase a_ptr - # by to get the element one row down (A has M rows) - stride_cm, - # stride_cn, # Assume that stride_cn == 1 - stride_am, - stride_ak, - stride_bn, - stride_bk, - # Meta-parameters - BLOCK_M: tl.constexpr, - GROUP_M: tl.constexpr, - BLOCK_N: tl.constexpr, - BLOCK_K: tl.constexpr, - # split k not used, not performant with activation, kept because early_config_prune is expecting it - SPLIT_K: tl.constexpr, - EVEN_K: tl.constexpr, - A_ROWMAJOR: tl.constexpr, - B_COLMAJOR: tl.constexpr, - BIAS: tl.constexpr, - SAVE_ACT_INPUT: tl.constexpr, - ACTIVATION: tl.constexpr, -): - - """ - Kernel for computing Out = activation(A x W + C) - - Input has shape (M, K) - - Weight has shape (K, N) - - Bias has shape (N,) - - Output has shape (M, N) - - ActInputs (optional) has shape (M, N) - 'ActInputs' optionally saves the A x W + C intermediate for backward computations - This kernel will consolidate over K - """ - - pid = tl.program_id(axis=0) - - grid_m = (M + BLOCK_M - 1) // BLOCK_M - grid_n = (N + BLOCK_N - 1) // BLOCK_N - # re-order program ID for better L2 performance - width = GROUP_M * grid_n - group_id = pid // width - group_size = min(grid_m - group_id * GROUP_M, GROUP_M) - pid_m = group_id * GROUP_M + (pid % group_size) - pid_n = (pid % width) // (group_size) - - # now compute the block that each program will go through - # rm (resp. rn) denotes a range of indices - # for rows (resp. col) of C - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - # trick to avoid masking on M and N axis - ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M) - rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N) - rk = tl.arange(0, BLOCK_K) - - if A_ROWMAJOR: - A = A + (ram[:, None] * stride_am + rk[None, :]) - else: - A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak) - if B_COLMAJOR: - B = B + (rk[:, None] + rbn[None, :] * stride_bn) - else: - B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn) - - acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) - - for k in range(K, 0, -BLOCK_K): - if EVEN_K: - a = tl.load(A) - b = tl.load(B) - else: - a = tl.load(A, mask=rk[None, :] < k, other=0.0) - b = tl.load(B, mask=rk[:, None] < k, other=0.0) - acc += tl.dot(a, b) - - if A_ROWMAJOR: - A += BLOCK_K - else: - A += BLOCK_K * stride_ak - if B_COLMAJOR: - B += BLOCK_K - else: - B += BLOCK_K * stride_bk - - # Putting bias after the matmul (instead of before) is faster, idk why - if BIAS: - bias = tl.load(bias + rn, mask=rn < N, other=0.0).to(tl.float32) - acc += bias[None, :] - - # optional: save the activation inputs - if SAVE_ACT_INPUT: - # act_in_ptrs = ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :] * stride_cn - act_in_ptrs = ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :] - tl.store(act_in_ptrs, acc) - - # optional: fused activation (while the data is in shared memory) - if ACTIVATION == "gelu": - acc = gelu(acc) - elif ACTIVATION == "gelu_approx": - acc = gelu_approx(acc) - elif ACTIVATION == "squared_relu": - acc = squared_relu(acc) - # rematerialize rm and rn to save registers - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - - # write back result - # C = C + rm[:, None] * stride_cm + rn[None, :] * stride_cn - C = C + rm[:, None] * stride_cm + rn[None, :] - mask = (rm < M)[:, None] & (rn < N)[None, :] - tl.store(C, acc) - - -def triton_linear_act( - x: torch.Tensor, - weight: torch.Tensor, - bias: Optional[torch.Tensor] = None, - activation: str = "id", - save_act_input: bool = False, -) -> torch.Tensor: - """ - Compute e = activation(x @ weight.T + bias). - This wrapper kicks the `kernel_fwd` Triton kernel - :param x: input tensor - :param weight: weight matrix - :param bias: an optional bias tensor - :param activation: Activation name. Needs to be a Triton kernel. - :param act_input: an optional tensor to save the activation inputs (for backward) - :return: result tensor - """ - # if torch.is_autocast_enabled(): - # dtype = torch.get_autocast_gpu_dtype() - # x, weight, bias = [a.to(dtype=dtype) for a in [x, weight, bias]] - - assert activation in ["id", "gelu", "gelu_approx", "squared_relu"] - - batch_shape, n = x.shape[:-1], x.shape[-1] - batch_dim = batch_shape.numel() - x_reshaped = x.reshape(batch_dim, n) - - if x_reshaped.stride(0) > 1 and x_reshaped.stride(1) > 1: - x_reshaped = x_reshaped.contiguous() - if weight.stride(0) > 1 and weight.stride(1) > 1: - weight = weight.contiguous() - bias = bias.contiguous() if bias is not None else None - - assert ( - x.dtype == weight.dtype - ), f"Input and weight must have the same dtype, got {x.dtype} and {weight.dtype}" - if bias is not None: - assert ( - x.dtype == bias.dtype - ), f"Input and bias must have the same dtype, got {x.dtype} and {bias.dtype}" - assert ( - x_reshaped.shape[1] == weight.shape[1] - ), f"Incompatible dimensions: {x_reshaped.shape} - {weight.shape}" - - assert ( - bias is None or bias.shape[0] == weight.shape[0] - ), "Incompatible dimensions in between weight and bias" - - M, K = x_reshaped.shape - N, K = weight.shape - - output = torch.empty((M, N), device=x.device, dtype=x.dtype) - act_input = torch.empty_like(output) if save_act_input else None - - # 1D launch kernel where each block gets its own program. - grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),) # noqa - - kernel_fwd[grid]( - output, - act_input, - x_reshaped, - weight, # data ptrs - bias if bias is not None else x, # auto skip bias if not present - M, # shapes - N, - K, - M // 32, # key for triton cache (limit number of compilations) - N // 32, - K // 32, - stride_cm=output.stride(0), # strides - # stride_cn=output.stride(1), - stride_am=x_reshaped.stride(0), - stride_ak=x_reshaped.stride(1), - stride_bk=weight.stride(1), - stride_bn=weight.stride(0), - BIAS=bias is not None, # optional fused bias - SAVE_ACT_INPUT=save_act_input, # optional save activation inputs - ACTIVATION=activation, # optional fused activation - A_ROWMAJOR=x_reshaped.stride(1) == 1, - B_COLMAJOR=weight.stride(1) == 1, - GROUP_M=8, # speed optimization: group the programs - ) - - if not save_act_input: - return output.reshape(*batch_shape, output.shape[-1]) - else: - return ( - output.reshape(*batch_shape, output.shape[-1]), - act_input.reshape(*batch_shape, act_input.shape[-1]), - ) - - -@triton.autotune( - configs=[ - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - # good for int8 - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=4, - num_warps=4, - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - ] - + get_configs_io_bound(), - key=["CACHE_KEY_M", "CACHE_KEY_N", "CACHE_KEY_K"], - prune_configs_by={ - "early_config_prune": early_config_prune, - "perf_model": estimate_matmul_time, - "top_k": 10, - }, -) -@triton.heuristics( - { - "EVEN_K": lambda args: args["K"] % (args["BLOCK_K"] * args["SPLIT_K"]) == 0, - } -) -@triton.jit -def kernel_bwd( - C, # Pointers to matrices - ACT_INPUT, - A, - B, - # Matrix dimensions - M, - N, - K, - CACHE_KEY_M, - CACHE_KEY_N, - CACHE_KEY_K, - # The stride variables represent how much to increase the ptr by when moving by 1 - # element in a particular dimension. E.g. stride_am is how much to increase a_ptr - # by to get the element one row down (A has M rows) - stride_cm, - # stride_cn, # Assume that stride_cn == 1 - stride_am, - stride_ak, - stride_bk, - stride_bn, - # Meta-parameters - BLOCK_M: tl.constexpr, - GROUP_M: tl.constexpr, - BLOCK_N: tl.constexpr, - BLOCK_K: tl.constexpr, - # split k not used, not performant with activation, kept because early_config_prune is expecting it - SPLIT_K: tl.constexpr, - EVEN_K: tl.constexpr, - ACTIVATION: tl.constexpr, -): - - """ - Kernel for computing Out = activation(A x W + C) - - Input has shape (M, K) - - Weight has shape (K, N) - - Output has shape (M, N) - - ActInputs (optional) has shape (M, N) - 'ActInputs' optionally saves the A x W + C intermediate for backward computations - This kernel will consolidate over K - """ - - pid = tl.program_id(axis=0) - - grid_m = (M + BLOCK_M - 1) // BLOCK_M - grid_n = (N + BLOCK_N - 1) // BLOCK_N - # re-order program ID for better L2 performance - width = GROUP_M * grid_n - group_id = pid // width - group_size = min(grid_m - group_id * GROUP_M, GROUP_M) - pid_m = group_id * GROUP_M + (pid % group_size) - pid_n = (pid % width) // (group_size) - - # now compute the block that each program will go through - # rm (resp. rn) denotes a range of indices - # for rows (resp. col) of C - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - # trick to avoid masking on M and N axis - ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M) - rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N) - rk = tl.arange(0, BLOCK_K) - - A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak) - B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn) - - acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) - - for k in range(K, 0, -BLOCK_K): - if EVEN_K: - a = tl.load(A) - b = tl.load(B) - else: - a = tl.load(A, mask=rk[None, :] < k, other=0.0) - b = tl.load(B, mask=rk[:, None] < k, other=0.0) - acc += tl.dot(a, b) - - A += BLOCK_K * stride_ak - B += BLOCK_K * stride_bk - - # optional: fused activation (while the data is in shared memory) - if ACTIVATION != "id": - act_in_ptrs = ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :] - act_input = tl.load(act_in_ptrs).to(acc.dtype) - if ACTIVATION == "gelu": - acc *= gelu_grad(act_input) - elif ACTIVATION == "gelu_approx": - acc *= gelu_approx_grad(act_input) - elif ACTIVATION == "squared_relu": - acc *= squared_relu_grad(act_input) - - # rematerialize rm and rn to save registers - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - - # write back result - C = C + rm[:, None] * stride_cm + rn[None, :] - mask = (rm < M)[:, None] & (rn < N)[None, :] - tl.store(C, acc, mask=mask) - - -def triton_dgrad_act( - grad_output: torch.Tensor, - weight: torch.Tensor, - activation: str = "id", - act_input: Optional[torch.Tensor] = None, -) -> torch.Tensor: - """ - Compute e = activation(grad_output @ weight + bias). - This wrapper kicks the `kernel_fwd` Triton kernel - :param grad_output: input tensor - :param weight: weight matrix - :param activation: Activation name. Needs to be a Triton kernel. - :param act_input: an optional tensor to save the activation inputs (for backward) - :return: result tensor - """ - assert activation in ["id", "gelu", "gelu_approx", "squared_relu"] - - batch_shape, n = grad_output.shape[:-1], grad_output.shape[-1] - batch_dim = batch_shape.numel() - grad_output_reshaped = grad_output.reshape(batch_dim, n) - - if grad_output_reshaped.stride(0) > 1 and grad_output_reshaped.stride(1) > 1: - grad_output_reshaped = grad_output_reshaped.contiguous() - if weight.stride(0) > 1 and weight.stride(1) > 1: - weight = weight.contiguous() - - assert ( - grad_output.dtype == weight.dtype - ), f"grad_output and weight must have the same dtype, got {grad_output.dtype} and {weight.dtype}" - assert ( - grad_output_reshaped.shape[1] == weight.shape[0] - ), f"Incompatible dimensions: {grad_output_reshaped.shape} - {weight.shape}" - if activation != "id": - assert act_input is not None, f"act_input is required for activation {activation}" - - # M, N, K in bwd are different from M, N, K in fwd - M, K = grad_output_reshaped.shape - K, N = weight.shape - - grad_input = torch.empty((M, N), device=grad_output.device, dtype=grad_output.dtype) - - # 1D launch kernel where each block gets its own program. - grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),) # noqa - - kernel_bwd[grid]( - grad_input, - act_input, - grad_output_reshaped, - weight, # data ptrs - M, # shapes - N, - K, - M // 32, # key for triton cache (limit number of compilations) - N // 32, - K // 32, - stride_cm=grad_input.stride(0), # strides - # stride_cn=grad_input.stride(1), - stride_am=grad_output_reshaped.stride(0), - stride_ak=grad_output_reshaped.stride(1), - stride_bk=weight.stride(0), - stride_bn=weight.stride(1), - ACTIVATION=activation, # optional fused activation - GROUP_M=8, # speed optimization: group the programs - ) - - return grad_input.reshape(*batch_shape, grad_input.shape[-1]) diff --git a/build/torch212-cxx11-cu126-aarch64-linux/ops/triton/mlp.py b/build/torch212-cxx11-cu126-aarch64-linux/ops/triton/mlp.py deleted file mode 100644 index 059f4f8a5e174c1f4824e43d313fca18eaa799b8..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu126-aarch64-linux/ops/triton/mlp.py +++ /dev/null @@ -1,149 +0,0 @@ -# The triton fused matmul + sqrelu is faster for fp16 but slower for bf16, compared -# to naive implementation. -import fused_dense_lib as fused_dense_cuda -import torch -import torch.nn as nn -import torch.nn.functional as F - -from flash_attn.utils.torch import custom_fwd, custom_bwd -from flash_attn.ops.activations import sqrelu_bwd, sqrelu_fwd -from flash_attn.ops.triton.linear import triton_dgrad_act, triton_linear_act - - -class FusedDenseSqreluDenseFunc(torch.autograd.Function): - @staticmethod - @custom_fwd - def forward(ctx, x, weight1, bias1, weight2, bias2, checkpoint_lvl=0): - """checkpoint_lvl: - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute act_input and gelu_out in the bwd - """ - if torch.is_autocast_enabled(): - dtype = torch.get_autocast_gpu_dtype() - x, weight1, bias1, weight2, bias2 = [ - a.to(dtype=dtype) for a in [x, weight1, bias1, weight2, bias2] - ] - is_bf16 = x.dtype == torch.bfloat16 - assert checkpoint_lvl in [0, 1, 2] - x = x.contiguous() - weight1 = weight1.contiguous() - bias1 = bias1.contiguous() - weight2 = weight2.contiguous() - bias2 = bias2.contiguous() - batch_shape, n = x.shape[:-1], x.shape[-1] - batch_dim = batch_shape.numel() - if is_bf16: - act_input = fused_dense_cuda.linear_bias_forward( - x.reshape(batch_dim, n), weight1, bias1 - ) - output1 = sqrelu_fwd(act_input) - else: - save_act_input = checkpoint_lvl != 2 - result = triton_linear_act( - x.reshape(batch_dim, n), - weight1, - bias1, - activation="squared_relu", - save_act_input=save_act_input, - ) - if save_act_input: - output1, act_input = result - else: - output1 = result - output2 = fused_dense_cuda.linear_bias_forward(output1, weight2, bias2) - ctx.checkpoint_lvl = checkpoint_lvl - if checkpoint_lvl == 0: - ctx.save_for_backward(x, weight1, bias1, weight2, act_input, output1) - elif checkpoint_lvl == 1: - ctx.save_for_backward(x, weight1, bias1, weight2, act_input) - elif checkpoint_lvl == 2: - ctx.save_for_backward(x, weight1, bias1, weight2) - return output2.reshape(*batch_shape, output2.shape[-1]) - - @staticmethod - @custom_bwd - def backward(ctx, grad_output): - grad_output = grad_output.contiguous() - checkpoint_lvl = ctx.checkpoint_lvl - x, weight1, bias1, weight2, *rest = ctx.saved_tensors - batch_shape, n = x.shape[:-1], x.shape[-1] - batch_dim = batch_shape.numel() - is_bf16 = x.dtype == torch.bfloat16 - if checkpoint_lvl == 0: - act_input, output1 = rest - elif checkpoint_lvl == 1: - (act_input,) = rest - output1 = sqrelu_fwd(act_input) - elif checkpoint_lvl == 2: - if is_bf16: - act_input = fused_dense_cuda.linear_bias_forward( - x.reshape(batch_dim, n), weight1, bias1 - ) - output1 = sqrelu_fwd(act_input) - else: - output1, act_input = triton_linear_act( - x.reshape(batch_dim, n), - weight1, - bias1, - activation="squared_relu", - save_act_input=True, - ) - - if is_bf16: - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad(output1, grad_output) - grad_output1 = grad_output @ weight2 - grad_act_input = sqrelu_bwd(grad_output1, act_input) - grad_input, grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_backward( - x.reshape(batch_dim, n), weight1, grad_act_input - ) - else: - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad(output1, grad_output) - grad_act_input = triton_dgrad_act( - grad_output, weight2, activation="squared_relu", act_input=act_input - ) - grad_input, grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_backward( - x.reshape(batch_dim, n), weight1, grad_act_input - ) - return grad_input.reshape_as(x), grad_weight1, grad_bias1, grad_weight2, grad_bias2, None - - -fused_dense_sqrelu_dense_function = FusedDenseSqreluDenseFunc.apply - - -class FusedDenseSqreluDense(nn.Module): - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - bias1=True, - bias2=True, - checkpoint_lvl=0, - device=None, - dtype=None, - ): - """ - checkpoint_lvl (increasing lvl means slower but more memory saving): - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute gelu_in and gelu_out in the bwd - """ - assert checkpoint_lvl in [0, 1, 2] - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features * 4 - assert bias1 == True, "DenseSqreluDense module without bias is currently not supported" - assert bias2 == True, "DenseSqreluDense module without bias is currently not supported" - self.checkpoint_lvl = checkpoint_lvl - self.fc1 = nn.Linear(in_features, hidden_features, bias=bias1, **factory_kwargs) - self.fc2 = nn.Linear(hidden_features, out_features, bias=bias2, **factory_kwargs) - - def forward(self, x): - assert x.is_cuda - return fused_dense_sqrelu_dense_function( - x, self.fc1.weight, self.fc1.bias, self.fc2.weight, self.fc2.bias, self.checkpoint_lvl - ) diff --git a/build/torch212-cxx11-cu126-aarch64-linux/ops/triton/rotary.py b/build/torch212-cxx11-cu126-aarch64-linux/ops/triton/rotary.py deleted file mode 100644 index 6ad3ab412a084270120167ad32b67ad1720f2c37..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu126-aarch64-linux/ops/triton/rotary.py +++ /dev/null @@ -1,186 +0,0 @@ -# Copyright (c) 2025, Tri Dao. -# As of 2025-04-23, we require triton >= 3.0 - -from typing import Optional, Union - -import torch - -import triton -import triton.language as tl - - -@triton.jit -def rotary_kernel( - OUT, # Pointers to matrices - X, - COS, - SIN, - CU_SEQLENS, - SEQLEN_OFFSETS, # this could be int or a pointer - # Matrix dimensions - seqlen, - nheads, - seqlen_ro, - # strides - stride_out_batch, - stride_out_seqlen, - stride_out_nheads, - stride_out_headdim, - stride_x_batch, - stride_x_seqlen, - stride_x_nheads, - stride_x_headdim, - # Meta-parameters - # We want ROTARY_DIM to be constexpr, otherwise the triton compiler doesn't know that - # the mask is constant every 8 elements, and it will generate LDG.16 instead of LDG.128 - ROTARY_DIM: tl.constexpr, - IS_SEQLEN_OFFSETS_TENSOR: tl.constexpr, - IS_VARLEN: tl.constexpr, - INTERLEAVED: tl.constexpr, - CONJUGATE: tl.constexpr, - BLOCK_H: tl.constexpr, - BLOCK_M: tl.constexpr, -): - BLOCK_K: tl.constexpr = triton.next_power_of_2(ROTARY_DIM) - ROTARY_DIM_HALF = ROTARY_DIM // 2 - pid_head = tl.program_id(axis=0) - pid_m = tl.program_id(axis=1) - pid_batch = tl.program_id(axis=2) - - if not IS_VARLEN: - X = X + pid_batch * stride_x_batch - OUT = OUT + pid_batch * stride_out_batch - else: - start_idx = tl.load(CU_SEQLENS + pid_batch) - seqlen = tl.load(CU_SEQLENS + pid_batch + 1) - start_idx - X = X + start_idx * stride_x_seqlen - OUT = OUT + start_idx * stride_out_seqlen - - if pid_m * BLOCK_M >= seqlen: - return - - rh = pid_head * BLOCK_H + tl.arange(0, BLOCK_H) - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - if not IS_SEQLEN_OFFSETS_TENSOR: - rm_cs = rm + SEQLEN_OFFSETS - else: - rm_cs = rm + tl.load(SEQLEN_OFFSETS + pid_batch) - - rk_half = tl.arange(0, BLOCK_K // 2) - COS = COS + (rm_cs[:, None] * ROTARY_DIM_HALF + rk_half[None, :]) - SIN = SIN + (rm_cs[:, None] * ROTARY_DIM_HALF + rk_half[None, :]) - mask_cs = (rm_cs[:, None] < seqlen_ro) & (rk_half[None, :] < ROTARY_DIM_HALF) - cos = tl.load(COS, mask=mask_cs, other=1.0).to(tl.float32) - sin = tl.load(SIN, mask=mask_cs, other=0.0).to(tl.float32) - if CONJUGATE: - sin = -sin - - if not INTERLEAVED: - # Load the 1st and 2nd halves of X, do calculation, then store to 1st and 2nd halves of OUT - X = X + (rh[:, None, None] * stride_x_nheads + rm[None, :, None] * stride_x_seqlen + rk_half[None, None, :] * stride_x_headdim) - OUT = OUT + (rh[:, None, None] * stride_out_nheads + rm[None, :, None] * stride_out_seqlen + rk_half[None, None, :] * stride_out_headdim) - mask = (rh[:, None, None] < nheads) & (rm[None, :, None] < seqlen) & (rk_half[None, None, :] < ROTARY_DIM_HALF) - x0 = tl.load(X, mask=mask, other=0.0).to(tl.float32) - x1 = tl.load(X + ROTARY_DIM_HALF * stride_x_headdim, mask=mask, other=0.0,).to(tl.float32) - o0 = x0 * cos - x1 * sin - o1 = x0 * sin + x1 * cos - tl.store(OUT, o0, mask=mask) - tl.store(OUT + ROTARY_DIM_HALF * stride_out_headdim, o1, mask=mask) - else: - rk = tl.arange(0, BLOCK_K) - X = X + (rh[:, None, None] * stride_x_nheads + rm[None, :, None] * stride_x_seqlen + rk[None, None, :] * stride_x_headdim) - OUT = OUT + (rh[:, None, None] * stride_out_nheads + rm[None, :, None] * stride_out_seqlen + rk[None, None, :] * stride_out_headdim) - mask = (rh[:, None, None] < nheads) & (rm[None, :, None] < seqlen) & (rk[None, None, :] < ROTARY_DIM) - x = tl.load(X, mask=mask, other=0.0).to(tl.float32) - x0, x1 = tl.split(tl.reshape(x, [BLOCK_H, BLOCK_M, BLOCK_K // 2, 2])) - o0 = x0 * cos - x1 * sin - o1 = x0 * sin + x1 * cos - o = tl.reshape(tl.join(o0, o1), [BLOCK_H, BLOCK_M, BLOCK_K]) - tl.store(OUT, o, mask=mask) - - -def apply_rotary( - x: torch.Tensor, - cos: torch.Tensor, - sin: torch.Tensor, - seqlen_offsets: Union[int, torch.Tensor] = 0, - cu_seqlens: Optional[torch.Tensor] = None, - max_seqlen: Optional[int] = None, - interleaved=False, - inplace=False, - conjugate=False, -) -> torch.Tensor: - """ - Arguments: - x: (batch, seqlen, nheads, headdim) if cu_seqlens is None - else (total_seqlen, nheads, headdim). - cos: (seqlen_ro, rotary_dim / 2) - sin: (seqlen_ro, rotary_dim / 2) - seqlen_offsets: integer or integer tensor of size (batch,) - cu_seqlens: (batch + 1,) or None - max_seqlen: int - Returns: - y: (batch, seqlen, nheads, headdim) - """ - is_varlen = cu_seqlens is not None - if not is_varlen: - batch, seqlen, nheads, headdim = x.shape - else: - assert max_seqlen is not None, "If cu_seqlens is passed in, then max_seqlen must be passed" - total_seqlen, nheads, headdim = x.shape - batch_p_1 = cu_seqlens.shape[0] - batch = batch_p_1 - 1 - seqlen = max_seqlen - seqlen_ro, rotary_dim = cos.shape - assert sin.shape == cos.shape - rotary_dim *= 2 - assert rotary_dim <= headdim, "rotary_dim must be <= headdim" - assert headdim <= 256, "Only support headdim <= 256" - assert seqlen_ro >= seqlen, "seqlen_ro must be >= seqlen" - - cos, sin = cos.contiguous(), sin.contiguous() - if isinstance(seqlen_offsets, torch.Tensor): - assert seqlen_offsets.shape == (batch,) - assert seqlen_offsets.dtype in [torch.int32, torch.int64] - seqlen_offsets = seqlen_offsets.contiguous() - else: - assert seqlen_offsets + seqlen <= seqlen_ro - - output = torch.empty_like(x) if not inplace else x - if rotary_dim < headdim and not inplace: - output[..., rotary_dim:].copy_(x[..., rotary_dim:]) - - grid = lambda META: (triton.cdiv(nheads, META["BLOCK_H"]), triton.cdiv(seqlen, META["BLOCK_M"]), batch) # noqa - BLOCK_M = 8 if rotary_dim <= 128 else 4 - - # Need this, otherwise Triton tries to launch from cuda:0 and we get - # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?) - device_ctx = torch.cuda.device(x.device.index) if x.device.type == 'cuda' else torch.xpu.device(x.device.index) - with device_ctx: - torch.library.wrap_triton(rotary_kernel)[grid]( - output, # data ptrs - x, - cos, - sin, - cu_seqlens, - seqlen_offsets, - seqlen, # shapes - nheads, - seqlen_ro, - output.stride(0) if not is_varlen else 0, # batch_strides if not varlen else 0 - output.stride(-3), # seqlen_stride or total_seqlen_stride - output.stride(-2), # nheads_stride - output.stride(-1), # headdim_stride - x.stride(0) if not is_varlen else 0, # batch_strides if not varlen else 0 - x.stride(-3), # seqlen stride or total_seqlen_stride - x.stride(-2), # nheads stride - x.stride(-1), # headdim stride - rotary_dim, - isinstance(seqlen_offsets, torch.Tensor), - is_varlen, - interleaved, - conjugate, - BLOCK_M=BLOCK_M, - BLOCK_H=2, - ) - return output diff --git a/build/torch212-cxx11-cu126-x86_64-linux/__init__.py b/build/torch212-cxx11-cu126-x86_64-linux/__init__.py deleted file mode 100644 index ecc2f9d896b6c93f90b0a1499856dc0612177422..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu126-x86_64-linux/__init__.py +++ /dev/null @@ -1,393 +0,0 @@ -from typing import Optional, List -import torch -from ._ops import ops as flash_attn_ops -from .flash_attn_interface import ( - flash_attn_func, - flash_attn_kvpacked_func, - flash_attn_qkvpacked_func, - flash_attn_varlen_func, - flash_attn_varlen_kvpacked_func, - flash_attn_varlen_qkvpacked_func, - flash_attn_with_kvcache, -) - - -def fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - return_softmax: bool = False, - gen: Optional[torch.Generator] = None, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention. - - Args: - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Optional output tensor, same shape as q - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - return_softmax: Whether to return softmax weights - gen: Optional random number generator - - Returns: - List of tensors: [output, softmax_lse, (softmax if return_softmax)] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.fwd( - q, - k, - v, - out, - alibi_slopes, - p_dropout, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - gen, - ) - - -def varlen_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - out: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - block_table: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - max_seqlen_q: int = 0, - max_seqlen_k: int = 0, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - zero_tensors: bool = False, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - return_softmax: bool = False, - gen: Optional[torch.Generator] = None, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention with variable sequence lengths. - - Args: - q: Query tensor of shape [total_q, num_heads, head_size] - k: Key tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - v: Value tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1] - cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1] - out: Optional output tensor of shape [total_q, num_heads, head_size] - seqused_k: Optional tensor specifying how many keys to use per batch element [batch_size] - leftpad_k: Optional left padding for keys of shape [batch_size] - block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq] - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - max_seqlen_q: Maximum sequence length for queries - max_seqlen_k: Maximum sequence length for keys - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - zero_tensors: Whether to zero tensors before computation - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - return_softmax: Whether to return softmax weights - gen: Optional random number generator - - Returns: - List of tensors: [output, softmax_lse, (softmax if return_softmax)] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.varlen_fwd( - q, - k, - v, - out, - cu_seqlens_q, - cu_seqlens_k, - seqused_k, - leftpad_k, - block_table, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - p_dropout, - softmax_scale, - zero_tensors, - is_causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - gen, - ) - - -def bwd( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - p_dropout: float = 0.0, - 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, - gen: Optional[torch.Generator] = None, - rng_state: Optional[torch.Tensor] = None, -) -> List[torch.Tensor]: - """ - Backward pass for multi-head attention. - - Args: - dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size] - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size] - softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q] - dq: Optional gradient tensor for queries, same shape as q - dk: Optional gradient tensor for keys, same shape as k - dv: Optional gradient tensor for values, same shape as v - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - deterministic: Whether to use deterministic algorithms - gen: Optional random number generator - rng_state: Optional RNG state from forward pass - - Returns: - List of tensors: [dq, dk, dv] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - alibi_slopes, - p_dropout, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - deterministic, - gen, - rng_state, - ) - - -def varlen_bwd( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - max_seqlen_q: int = 0, - max_seqlen_k: int = 0, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - zero_tensors: bool = False, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - deterministic: bool = False, - gen: Optional[torch.Generator] = None, - rng_state: Optional[torch.Tensor] = None, -) -> List[torch.Tensor]: - """ - Backward pass for multi-head attention with variable sequence lengths. - - Args: - dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size] - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size] - softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q] - cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1] - cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1] - dq: Optional gradient tensor for queries, same shape as q - dk: Optional gradient tensor for keys, same shape as k - dv: Optional gradient tensor for values, same shape as v - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - max_seqlen_q: Maximum sequence length for queries - max_seqlen_k: Maximum sequence length for keys - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - zero_tensors: Whether to zero tensors before computation - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - deterministic: Whether to use deterministic algorithms - gen: Optional random number generator - rng_state: Optional RNG state from forward pass - - Returns: - List of tensors: [dq, dk, dv] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.varlen_bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - p_dropout, - softmax_scale, - zero_tensors, - is_causal, - window_size_left, - window_size_right, - softcap, - deterministic, - gen, - rng_state, - ) - - -def fwd_kvcache( - q: torch.Tensor, - kcache: torch.Tensor, - vcache: torch.Tensor, - k: Optional[torch.Tensor] = None, - v: Optional[torch.Tensor] = None, - seqlens_k: Optional[torch.Tensor] = None, - rotary_cos: Optional[torch.Tensor] = None, - rotary_sin: Optional[torch.Tensor] = None, - cache_batch_idx: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - block_table: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - out: 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, - is_rotary_interleaved: bool = False, - num_splits: int = 1, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention with KV cache. - - Args: - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - kcache: Key cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - vcache: Value cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - k: Optional new keys tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size] - v: Optional new values tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size] - seqlens_k: Optional sequence lengths for keys of shape [batch_size] - rotary_cos: Optional rotary cosine tensor of shape [seqlen_ro, rotary_dim/2] - rotary_sin: Optional rotary sine tensor of shape [seqlen_ro, rotary_dim/2] - cache_batch_idx: Optional indices to index into the KV cache - leftpad_k: Optional left padding for keys of shape [batch_size] - block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq] - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - out: Optional output tensor, same shape as q - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - is_rotary_interleaved: Whether rotary embeddings are interleaved - num_splits: Number of splits for computation - - Returns: - List of tensors: [output, softmax_lse] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.fwd_kvcache( - q, - kcache, - vcache, - k, - v, - seqlens_k, - rotary_cos, - rotary_sin, - cache_batch_idx, - leftpad_k, - block_table, - alibi_slopes, - out, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - is_rotary_interleaved, - num_splits, - ) diff --git a/build/torch212-cxx11-cu126-x86_64-linux/_flash_attn2_cuda_f12afc9.abi3.so b/build/torch212-cxx11-cu126-x86_64-linux/_flash_attn2_cuda_f12afc9.abi3.so deleted file mode 100644 index e7eb45da33a5ac3ba32a638dd7c8ca515d90d98f..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu126-x86_64-linux/_flash_attn2_cuda_f12afc9.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:e9803828f4321163636f495216d326bf0dac1d95ee90ff584d52aa90462d5487 -size 446835240 diff --git a/build/torch212-cxx11-cu126-x86_64-linux/_ops.py b/build/torch212-cxx11-cu126-x86_64-linux/_ops.py deleted file mode 100644 index af5b731cb04b5e6ce54e862db0501f3683c3163f..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu126-x86_64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _flash_attn2_cuda_f12afc9 -ops = torch.ops._flash_attn2_cuda_f12afc9 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_flash_attn2_cuda_f12afc9::{op_name}" diff --git a/build/torch212-cxx11-cu126-x86_64-linux/bert_padding.py b/build/torch212-cxx11-cu126-x86_64-linux/bert_padding.py deleted file mode 100644 index 3c2d35159a014a9d03aabead9e52e009168696ea..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu126-x86_64-linux/bert_padding.py +++ /dev/null @@ -1,218 +0,0 @@ -# Adapted from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/padding.py - -import torch -import torch.nn.functional as F -from einops import rearrange, repeat - - -class IndexFirstAxis(torch.autograd.Function): - @staticmethod - def forward(ctx, input, indices): - ctx.save_for_backward(indices) - assert input.ndim >= 2 - ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] - second_dim = other_shape.numel() - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - # return input[indices] - return torch.gather( - rearrange(input, "b ... -> b (...)"), 0, repeat(indices, "z -> z d", d=second_dim) - ).reshape(-1, *other_shape) - - @staticmethod - def backward(ctx, grad_output): - (indices,) = ctx.saved_tensors - assert grad_output.ndim >= 2 - other_shape = grad_output.shape[1:] - grad_output = rearrange(grad_output, "b ... -> b (...)") - grad_input = torch.zeros( - [ctx.first_axis_dim, grad_output.shape[1]], - device=grad_output.device, - dtype=grad_output.dtype, - ) - # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. - # grad_input[indices] = grad_output - grad_input.scatter_(0, repeat(indices, "z -> z d", d=grad_output.shape[1]), grad_output) - return grad_input.reshape(ctx.first_axis_dim, *other_shape), None - - -index_first_axis = IndexFirstAxis.apply - - -class IndexPutFirstAxis(torch.autograd.Function): - @staticmethod - def forward(ctx, values, indices, first_axis_dim): - ctx.save_for_backward(indices) - assert indices.ndim == 1 - assert values.ndim >= 2 - output = torch.zeros( - first_axis_dim, *values.shape[1:], device=values.device, dtype=values.dtype - ) - # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. - output[indices] = values - # output.scatter_(0, repeat(indices, 'z -> z d', d=values.shape[1]), values) - return output - - @staticmethod - def backward(ctx, grad_output): - (indices,) = ctx.saved_tensors - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - grad_values = grad_output[indices] - # grad_values = torch.gather(grad_output, 0, repeat(indices, 'z -> z d', d=grad_output.shape[1])) - return grad_values, None, None - - -index_put_first_axis = IndexPutFirstAxis.apply - - -class IndexFirstAxisResidual(torch.autograd.Function): - @staticmethod - def forward(ctx, input, indices): - ctx.save_for_backward(indices) - assert input.ndim >= 2 - ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] - second_dim = other_shape.numel() - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - output = input[indices] - # We don't want to reshape input (b ... -> b (...)) since it could change the channel_last - # memory format to channel_first. In other words, input might not be contiguous. - # If we don't detach, Pytorch complains about output being a view and is being modified inplace - return output, input.detach() - - @staticmethod - def backward(ctx, grad_output, grad_residual): - (indices,) = ctx.saved_tensors - assert grad_output.ndim >= 2 - other_shape = grad_output.shape[1:] - assert grad_residual.shape[1:] == other_shape - grad_input = grad_residual - # grad_input[indices] += grad_output - indices = indices.reshape(indices.shape[0], *((1,) * (grad_output.ndim - 1))) - indices = indices.expand_as(grad_output) - grad_input.scatter_add_(0, indices, grad_output) - return grad_input.reshape(ctx.first_axis_dim, *other_shape), None - - -index_first_axis_residual = IndexFirstAxisResidual.apply - - -def unpad_input(hidden_states, attention_mask, unused_mask=None): - """ - Arguments: - hidden_states: (batch, seqlen, ...) - attention_mask: (batch, seqlen), bool / int, 1 means valid and 0 means not valid. - unused_mask: (batch, seqlen), bool / int, 1 means the element is allocated but unused. - Return: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens selected in attention_mask + unused_mask. - indices: (total_nnz), the indices of masked tokens from the flattened input sequence. - cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. - max_seqlen_in_batch: int - seqused: (batch), returns the number of tokens selected in attention_mask + unused_mask. - """ - all_masks = (attention_mask + unused_mask) if unused_mask is not None else attention_mask - seqlens_in_batch = all_masks.sum(dim=-1, dtype=torch.int32) - used_seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) - indices = torch.nonzero(all_masks.flatten(), as_tuple=False).flatten() - max_seqlen_in_batch = seqlens_in_batch.max().item() - cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) - # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the - # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim - # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to - # index with integer indices. Moreover, torch's index is a bit slower than it needs to be, - # so we write custom forward and backward to make it a bit faster. - return ( - index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices), - indices, - cu_seqlens, - max_seqlen_in_batch, - used_seqlens_in_batch, - ) - - -def unpad_input_for_concatenated_sequences(hidden_states, attention_mask_in_length): - """ - Supports concatenating short samples in one sequence. The attention_mask_in_length is utilized to mask other short samples. It helps efficient training of variant lengths-based samples (e.g., the supervised fine-tuning task in large language model). - The motivation for this function is explained [here](https://github.com/Dao-AILab/flash-attention/issues/432#issuecomment-1668822286). - - For example, if batch = 3 and seqlen = 6, the attention_mask_in_length is: - ``` - [ - [2, 3, 0, 0, 0, 0], - [3, 2, 0, 0, 0, 0], - [6, 0, 0, 0, 0, 0] - ] - ``` - , which refers to the 3D-attention mask: - ``` - [ - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [0, 0, 1, 0, 0, 0], - [0, 0, 1, 1, 0, 0], - [0, 0, 1, 1, 1, 0], - [0, 0, 0, 0, 0, 1] - ], - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [0, 0, 0, 1, 0, 0], - [0, 0, 0, 1, 1, 0], - [0, 0, 0, 0, 0, 1] - ], - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [1, 1, 1, 1, 0, 0], - [1, 1, 1, 1, 1, 0], - [1, 1, 1, 1, 1, 1] - ] - ] - ```. - - Arguments: - hidden_states: (batch, seqlen, ...) - attention_mask_in_length: (batch, seqlen), int, a nonzero number (e.g., 1, 2, 3, etc.) means length of concatenated sequence in b-th batch, and 0 means none. - Return: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. - indices: (total_nnz), the indices of non-masked tokens from the flattened input sequence. - cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. - max_seqlen_in_batch: int - """ - length = attention_mask_in_length.sum(dim=-1) - seqlen = attention_mask_in_length.size(-1) - attention_mask_2d = torch.arange(seqlen, device=length.device, dtype=length.dtype).expand(len(length), seqlen) < length.unsqueeze(1) - real_indices_idx = torch.nonzero(attention_mask_in_length.flatten(), as_tuple=False).flatten() - seqlens_in_batch = attention_mask_in_length.flatten()[real_indices_idx] - indices = torch.nonzero(attention_mask_2d.flatten(), as_tuple=False).flatten() - max_seqlen_in_batch = seqlens_in_batch.max().item() - cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) - # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the - # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim - # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to - # index with integer indices. Moreover, torch's index is a bit slower than it needs to be, - # so we write custom forward and backward to make it a bit faster. - return ( - index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices), - indices, - cu_seqlens, - max_seqlen_in_batch, - ) - - -def pad_input(hidden_states, indices, batch, seqlen): - """ - Arguments: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. - indices: (total_nnz), the indices that represent the non-masked tokens of the original padded input sequence. - batch: int, batch size for the padded sequence. - seqlen: int, maximum sequence length for the padded sequence. - Return: - hidden_states: (batch, seqlen, ...) - """ - dim = hidden_states.shape[-1] - # output = torch.zeros((batch * seqlen), dim, device=hidden_states.device, dtype=hidden_states.dtype) - # output[indices] = hidden_states - output = index_put_first_axis(hidden_states, indices, batch * seqlen) - return rearrange(output, "(b s) ... -> b s ...", b=batch) diff --git a/build/torch212-cxx11-cu126-x86_64-linux/flash_attn2/__init__.py b/build/torch212-cxx11-cu126-x86_64-linux/flash_attn2/__init__.py deleted file mode 100644 index a9b2672c1cd85b74c1b3ded0fc0b2100e1aeac23..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu126-x86_64-linux/flash_attn2/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -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/torch212-cxx11-cu126-x86_64-linux/flash_attn_interface.py b/build/torch212-cxx11-cu126-x86_64-linux/flash_attn_interface.py deleted file mode 100644 index b014e569dc9812a9c840d1043f2cebe767b4d477..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu126-x86_64-linux/flash_attn_interface.py +++ /dev/null @@ -1,1631 +0,0 @@ -# Copyright (c) 2023, Tri Dao. - -from typing import Optional, Sequence, Tuple, Union - -import torch -import torch.nn as nn -import os - -# # isort: off -# # We need to import the CUDA kernels after importing torch -# USE_TRITON_ROCM = os.getenv("FLASH_ATTENTION_TRITON_AMD_ENABLE", "FALSE") == "TRUE" -# if USE_TRITON_ROCM: -# from .flash_attn_triton_amd import interface_fa as flash_attn -# else: -# import flash_attn_2_cuda as flash_attn - - -from . import _ops -from ._ops import ops as flash_attn -from ._ops import add_op_namespace_prefix - -# # isort: on - -def maybe_contiguous(x): - return x.contiguous() if x is not None and x.stride(-1) != 1 else x - - -def _get_device(): - if torch.xpu.is_available(): - return "xpu" - elif torch.cuda.is_available(): - return "cuda" - else: - return "cpu" - - -def _get_block_size_n(device, head_dim, is_dropout, is_causal): - assert head_dim <= 256 - - if device.type == "xpu": - if head_dim <= 96: - return 64 - elif head_dim <= 128: - return 32 - elif head_dim <= 256: - return 64 - else: - return 32 - - # This should match the block sizes in the CUDA kernel - major, minor = torch.cuda.get_device_capability(device) - is_sm8x = major == 8 and minor > 0 # Only include sm86 and sm89, exclude sm80 (A100) - is_sm80 = major == 8 and minor == 0 - is_sm90 = major == 9 and minor == 0 - if head_dim <= 32: - return 128 - if head_dim <= 64: - return 128 if not is_dropout else 64 - elif head_dim <= 96: - return 64 - elif head_dim <= 128: - if is_sm8x: - return 64 if (not is_dropout and is_causal) else 32 - else: - return 64 if not is_dropout else 32 - elif head_dim <= 192: - return 64 - elif head_dim <= 224: - return 64 - elif head_dim <= 256: - return 64 - - -def round_multiple(x, m): - return (x + m - 1) // m * m - - -# torch.compile() support is only enabled for pytorch >= 2.4 -# The reason for this is that we are using the new custom_op and register_fake -# APIs, which support inplace modification of inputs in the function itself -if torch.__version__ >= "2.4.0": - _torch_custom_op_wrapper = torch.library.custom_op - _torch_register_fake_wrapper = torch.library.register_fake -else: - def noop_custom_op_wrapper(name, fn=None, /, *, mutates_args, device_types=None, schema=None): - def wrap(func): - return func - if fn is None: - return wrap - return fn - def noop_register_fake_wrapper(op, fn=None, /, *, lib=None, _stacklevel=1): - def wrap(func): - return func - if fn is None: - return wrap - return fn - _torch_custom_op_wrapper = noop_custom_op_wrapper - _torch_register_fake_wrapper = noop_register_fake_wrapper - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_forward"), mutates_args=(), device_types=_get_device()) -def _flash_attn_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - return_softmax: bool -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - out, softmax_lse, S_dmask, rng_state = flash_attn.fwd( - q, - k, - v, - None, - alibi_slopes, - dropout_p, - softmax_scale, - causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - None, - ) - return out, softmax_lse, S_dmask, rng_state - - -@_torch_register_fake_wrapper(add_op_namespace_prefix("_flash_attn_forward")) -def _flash_attn_forward_fake( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - return_softmax: bool -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - batch_size, seqlen_q, num_heads, head_size = q.shape - seqlen_k = k.shape[1] - out = torch.empty_like(q) - softmax_lse = torch.empty((batch_size, num_heads, seqlen_q), dtype=torch.float32, device=q.device, layout=q.layout) - p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) - if return_softmax: - p = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128), round_multiple(seqlen_k, 128)), dtype=q.dtype, device=q.device, layout=q.layout) - rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) - - return out, softmax_lse, p, rng_state - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_forward = _ops.ops._flash_attn_forward -else: - _wrapped_flash_attn_forward = _flash_attn_forward - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_varlen_forward"), mutates_args=(), device_types=_get_device()) -def _flash_attn_varlen_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - alibi_slopes: Optional[torch.Tensor] = None, - return_softmax: bool = False, - block_table: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - out, softmax_lse, S_dmask, rng_state = flash_attn.varlen_fwd( - q, - k, - v, - None, - cu_seqlens_q, - cu_seqlens_k, - seqused_k, - leftpad_k, - block_table, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - zero_tensors, - causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - None, - ) - # if out.isnan().any() or softmax_lse.isnan().any(): - # breakpoint() - return out, softmax_lse, S_dmask, rng_state - - -@_torch_register_fake_wrapper(add_op_namespace_prefix("_flash_attn_varlen_forward")) -def _flash_attn_varlen_forward_fake( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - alibi_slopes: Optional[torch.Tensor] = None, - return_softmax: bool = False, - block_table: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - paged_kv = block_table is not None - batch_size = cu_seqlens_q.numel() - 1 - total_q, num_heads, _ = q.shape - - out = torch.empty_like(q) - softmax_lse = torch.empty((num_heads, total_q), dtype=torch.float32, device=q.device, layout=q.layout) - p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) - seqlen_q_rounded = round_multiple(max_seqlen_q, 128) - seqlen_k_rounded = round_multiple(max_seqlen_k, 128) - if return_softmax: - p = torch.empty((batch_size, num_heads, seqlen_q_rounded, seqlen_k_rounded), dtype=q.dtype, device=q.device, layout=q.layout) - rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) - return out, softmax_lse, p, rng_state - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_varlen_forward = _ops.ops._flash_attn_varlen_forward -else: - _wrapped_flash_attn_varlen_forward = _flash_attn_varlen_forward - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_backward"), mutates_args=("dq", "dk", "dv"), device_types=_get_device()) -def _flash_attn_backward( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, -) -> 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)] - ( - dq, - dk, - dv, - softmax_d, - ) = flash_attn.bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - alibi_slopes, - dropout_p, - softmax_scale, - causal, - window_size_left, - window_size_right, - softcap, - deterministic, - None, - rng_state, - ) - return softmax_d - - -@_torch_register_fake_wrapper(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, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, -) -> torch.Tensor: - dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] - if dq is None: - dq = torch.empty_like(q) - if dk is None: - dk = torch.empty_like(k) - if dv is None: - dv = torch.empty_like(v) - batch_size, seqlen_q, num_heads, _ = q.shape - softmax_d = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128)), device=q.device, dtype=torch.float32) - - return softmax_d - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_backward = _ops.ops._flash_attn_backward -else: - _wrapped_flash_attn_backward = _flash_attn_backward - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_varlen_backward"), mutates_args=("dq", "dk", "dv"), device_types=_get_device()) -def _flash_attn_varlen_backward( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> 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)] - ( - dq, - dk, - dv, - softmax_d, - ) = flash_attn.varlen_bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - zero_tensors, - causal, - window_size_left, - window_size_right, - softcap, - deterministic, - None, - rng_state, - ) - # if dk.isnan().any() or dk.isnan().any() or dv.isnan().any() or softmax_d.isnan().any(): - # breakpoint() - return softmax_d - - -@_torch_register_fake_wrapper(add_op_namespace_prefix("_flash_attn_varlen_backward")) -def _flash_attn_varlen_backward_fake( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> torch.Tensor: - dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] - batch_size = cu_seqlens_q.numel() - 1 - total_q, num_heads, _ = q.shape - - if dq is None: - dq = torch.empty_like(q) - if dk is None: - dk = torch.empty_like(k) - if dv is None: - dv = torch.empty_like(v) - softmax_d = torch.empty((num_heads, total_q + 128 * batch_size), device=q.device, dtype=torch.float32) - - return softmax_d - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_varlen_backward = _ops.ops._flash_attn_varlen_backward -else: - _wrapped_flash_attn_varlen_backward = _flash_attn_varlen_backward - - -class FlashAttnQKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and qkv.requires_grad - if softmax_scale is None: - softmax_scale = qkv.shape[-1] ** (-0.5) - q, k, v = qkv[:, :, 0].detach(), qkv[:, :, 1].detach(), qkv[:, :, 2].detach() - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) - dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dqkv[:, :, 0], - dqkv[:, :, 1], - dqkv[:, :, 2], - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dqkv = dqkv[..., : dout.shape[-1]] # We could have padded the head dimension - return dqkv, None, None, None, None, None, None, None, None, None - - -class FlashAttnVarlenQKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - cu_seqlens, - max_seqlen, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and qkv.requires_grad - if softmax_scale is None: - softmax_scale = qkv.shape[-1] ** (-0.5) - q, k, v = qkv[:, 0].detach(), qkv[:, 1].detach(), qkv[:, 2].detach() - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens, - cu_seqlens, - max_seqlen, - max_seqlen, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=None, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens, rng_state) - ctx.dropout_p = dropout_p - ctx.max_seqlen = max_seqlen - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens, rng_state = ctx.saved_tensors - qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) - dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dqkv[:, 0], - dqkv[:, 1], - dqkv[:, 2], - cu_seqlens, - cu_seqlens, - ctx.max_seqlen, - ctx.max_seqlen, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - 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 FlashAttnKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - kv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, kv] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - k, v = kv[:, :, 0].detach(), kv[:, :, 1].detach() - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - dq = torch.empty_like(q) - kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) - dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dkv[:, :, 0], - dkv[:, :, 1], - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dkv = dkv[..., : dout.shape[-1]] - return dq, dkv, None, None, None, None, None, None, None, None, None - - -class FlashAttnVarlenKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, kv] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - k, v = kv[:, 0].detach(), kv[:, 1].detach() - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=None, - ) - if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state - ) - ctx.dropout_p = dropout_p - 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.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors - dq = torch.empty_like(q) - kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) - dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dkv[:, 0], - dkv[:, 1], - cu_seqlens_q, - cu_seqlens_k, - ctx.max_seqlen_q, - ctx.max_seqlen_k, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dkv = dkv[..., : dout.shape[-1]] - return dq, dkv, None, None, None, None, None, None, None, None, None, None, None, None, None - - -class FlashAttnFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - k, - v, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, k, v] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dk = dk[..., : dout.shape[-1]] - dv = dv[..., : dout.shape[-1]] - return dq, dk, dv, 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, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - block_table, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, k, v] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=block_table, - ) - if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state - ) - ctx.dropout_p = dropout_p - 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.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors - dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - ctx.max_seqlen_q, - ctx.max_seqlen_k, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dk = dk[..., : dout.shape[-1]] - dv = dv[..., : dout.shape[-1]] - return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None - - -def flash_attn_qkvpacked_func( - qkv, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # <=0.0 means deactivate - alibi_slopes=None, - deterministic=False, - 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, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_kvpacked_func( - q, - kv, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, -): - """dropout_p should be set to 0.0 during evaluation - If 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 K, V. - 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) - kv: (batch_size, seqlen, 2, 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. - softcap: float. Anything > 0 activates softcapping 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). - 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 FlashAttnKVPackedFunc.apply( - q, - kv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_func( - q, - k, - v, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - 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). - 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 FlashAttnFunc.apply( - q, - k, - v, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_qkvpacked_func( - qkv, - cu_seqlens, - max_seqlen, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - 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_varlen_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_varlen_kvpacked_func and flash_attn_varlen_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: (total, 3, nheads, headdim), where total = total number of tokens in the batch. - cu_seqlens: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into qkv. - max_seqlen: int. Maximum sequence length in the batch. - 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenQKVPackedFunc.apply( - qkv, - cu_seqlens, - max_seqlen, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_kvpacked_func( - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, -): - """dropout_p should be set to 0.0 during evaluation - If 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 K, V. - 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: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch. - kv: (total_k, 2, nheads_k, headdim), where total_k = total number of key tokens in the batch. - cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into q. - cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into kv. - max_seqlen_q: int. Maximum query sequence length in the batch. - max_seqlen_k: int. Maximum key sequence length in the batch. - 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 + 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenKVPackedFunc.apply( - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_func( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, - block_table=None, -): - """dropout_p should be set to 0.0 during evaluation - Supports multi-query and grouped-query attention (MQA/GQA) by passing in K, V 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: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch. - k: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. - v: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. - cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into q. - cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into kv. - max_seqlen_q: int. Maximum query sequence length in the batch. - max_seqlen_k: int. Maximum key sequence length in the batch. - 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 + 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenFunc.apply( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - block_table, - False if q.device.type == "cpu" else torch.is_grad_enabled(), - ) - - -def flash_attn_with_kvcache( - q, - k_cache, - v_cache, - k=None, - v=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, - block_table: Optional[torch.Tensor] = None, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - rotary_interleaved=True, - alibi_slopes=None, - num_splits=0, - 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 block_table, - or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_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) if there's no block_table, - or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_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). Similar to k. - 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. - block_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). - 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. - 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" - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-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) - cache_batch_idx = maybe_contiguous(cache_batch_idx) - block_table = maybe_contiguous(block_table) - out, softmax_lse = flash_attn.fwd_kvcache( - q, - k_cache, - v_cache, - k, - v, - cache_seqlens, - rotary_cos, - rotary_sin, - cache_batch_idx, - cache_leftpad, - block_table, - alibi_slopes, - None, - softmax_scale, - causal, - window_size[0], - window_size[1], - softcap, - rotary_interleaved, - num_splits, - ) - return (out, softmax_lse) if return_softmax_lse else out diff --git a/build/torch212-cxx11-cu126-x86_64-linux/layers/__init__.py b/build/torch212-cxx11-cu126-x86_64-linux/layers/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/build/torch212-cxx11-cu126-x86_64-linux/layers/patch_embed.py b/build/torch212-cxx11-cu126-x86_64-linux/layers/patch_embed.py deleted file mode 100644 index 05562f8e8bcdb58e947c6f402a49eacd2d031871..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu126-x86_64-linux/layers/patch_embed.py +++ /dev/null @@ -1,67 +0,0 @@ -# We use the same API as https://github.com/rwightman/pytorch-image-models/blob/v0.6.11/timm/models/layers/patch_embed.py -# But we use nn.Linear instead of Conv2d and it's about 8x faster. - -from functools import partial - -import torch.nn as nn -from einops import rearrange -from torch import _assert -from torch.nn.modules.utils import _pair - -try: - from flash_attn.ops.fused_dense import FusedDense -except ImportError: - FusedDense = None - - -class PatchEmbed(nn.Module): - """2D Image to Patch Embedding""" - - def __init__( - self, - img_size=224, - patch_size=16, - in_chans=3, - embed_dim=768, - norm_layer=None, - flatten=True, - bias=True, - fused_bias_fc=False, - ): - super().__init__() - img_size = _pair(img_size) - patch_size = _pair(patch_size) - self.img_size = img_size - self.patch_size = patch_size - self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1]) - self.num_patches = self.grid_size[0] * self.grid_size[1] - self.flatten = flatten - if fused_bias_fc and FusedDense is None: - raise ImportError("fused_dense is not installed") - - linear_cls = nn.Linear if not fused_bias_fc or not bias else FusedDense - self.proj = linear_cls(in_chans * patch_size[0] * patch_size[1], embed_dim, bias=bias) - self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() - - def forward(self, x): - _, _, H, W = x.shape - _assert( - H == self.img_size[0], - f"Input image height ({H}) doesn't match model ({self.img_size[0]}).", - ) - _assert( - W == self.img_size[1], - f"Input image width ({W}) doesn't match model ({self.img_size[1]}).", - ) - x = self.proj( - rearrange( - x, - "b c (h p1) (w p2) -> b h w (c p1 p2)", - p1=self.patch_size[0], - p2=self.patch_size[1], - ) - ) - if self.flatten: - x = rearrange(x, "b h w c -> b (h w) c") - x = self.norm(x) - return x diff --git a/build/torch212-cxx11-cu126-x86_64-linux/layers/rotary.py b/build/torch212-cxx11-cu126-x86_64-linux/layers/rotary.py deleted file mode 100644 index d1bfc21fc7de1dd287e8f382847b194a48075981..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu126-x86_64-linux/layers/rotary.py +++ /dev/null @@ -1,483 +0,0 @@ -# Copyright (c) 2025, Tri Dao - -import math -from functools import partial -from typing import Optional, Tuple, Union - -import torch -from torch import Tensor - -from einops import rearrange, repeat -# from flash_attn.ops.triton.rotary import apply_rotary -from ..ops.triton.rotary import apply_rotary - - -def rotate_half(x, interleaved=False): - if not interleaved: - x1, x2 = x.chunk(2, dim=-1) - return torch.cat((-x2, x1), dim=-1) - else: - x1, x2 = x[..., ::2], x[..., 1::2] - return rearrange(torch.stack((-x2, x1), dim=-1), "... d two -> ... (d two)", two=2) - - -def apply_rotary_emb_torch(x, cos, sin, interleaved=False): - """ - x: (batch_size, seqlen, nheads, headdim) - cos, sin: (seqlen, rotary_dim / 2) or (batch_size, seqlen, rotary_dim / 2) - """ - ro_dim = cos.shape[-1] * 2 - assert ro_dim <= x.shape[-1] - cos = repeat(cos, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)") - sin = repeat(sin, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)") - return torch.cat( - [x[..., :ro_dim] * cos + rotate_half(x[..., :ro_dim], interleaved) * sin, x[..., ro_dim:]], - dim=-1, - ) - - -class ApplyRotaryEmb(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x, - cos, - sin, - interleaved=False, - inplace=False, - seqlen_offsets: Union[int, Tensor] = 0, - cu_seqlens: Optional[Tensor] = None, - max_seqlen: Optional[int] = None, - ): - out = apply_rotary( - x, - cos, - sin, - seqlen_offsets=seqlen_offsets, - cu_seqlens=cu_seqlens, - max_seqlen=max_seqlen, - interleaved=interleaved, - inplace=inplace, - ) - if isinstance(seqlen_offsets, int): - ctx.save_for_backward(cos, sin, cu_seqlens) # Can't save int with save_for_backward - ctx.seqlen_offsets = seqlen_offsets - else: - ctx.save_for_backward(cos, sin, cu_seqlens, seqlen_offsets) - ctx.seqlen_offsets = None - ctx.interleaved = interleaved - ctx.inplace = inplace - ctx.max_seqlen = max_seqlen - return out if not inplace else x - - @staticmethod - def backward(ctx, do): - seqlen_offsets = ctx.seqlen_offsets - if seqlen_offsets is None: - cos, sin, cu_seqlens, seqlen_offsets = ctx.saved_tensors - else: - cos, sin, cu_seqlens = ctx.saved_tensors - dx = apply_rotary( - do, - cos, - sin, - seqlen_offsets=seqlen_offsets, - cu_seqlens=cu_seqlens, - max_seqlen=ctx.max_seqlen, - interleaved=ctx.interleaved, - inplace=ctx.inplace, - conjugate=True, - ) - return dx, None, None, None, None, None, None, None - - -def apply_rotary_emb( - x, - cos, - sin, - interleaved=False, - inplace=False, - seqlen_offsets: Union[int, Tensor] = 0, - cu_seqlens: Optional[Tensor] = None, - max_seqlen: Optional[int] = None, -): - """ - Arguments: - x: (batch_size, seqlen, nheads, headdim) if cu_seqlens is None - else (total_seqlen, nheads, headdim) - cos, sin: (seqlen_rotary, rotary_dim / 2) - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead - of 1st half and 2nd half (GPT-NeoX style). - inplace: if True, apply rotary embedding in-place. - seqlen_offsets: (batch_size,) or int. Each sequence in x is shifted by this amount. - Most commonly used in inference when we have KV cache. - cu_seqlens: (batch + 1,) or None - max_seqlen: int - Return: - out: (batch_size, seqlen, nheads, headdim) if cu_seqlens is None - else (total_seqlen, nheads, headdim) - rotary_dim must be <= headdim - Apply rotary embedding to the first rotary_dim of x. - """ - return ApplyRotaryEmb.apply( - x, cos, sin, interleaved, inplace, seqlen_offsets, cu_seqlens, max_seqlen - ) - - -# For backward compatibility -apply_rotary_emb_func = apply_rotary_emb - - -def _apply_rotary_emb_qkv( - qkv, - cos, - sin, - cos_k=None, - sin_k=None, - interleaved=False, - inplace=False, - conjugate=False, - seqlen_offsets: Union[int, Tensor] = 0, - num_heads_q: Optional[int] = None, -): - apply_rotary_fn = partial( - apply_rotary, - interleaved=interleaved, - inplace=inplace, - conjugate=conjugate, - seqlen_offsets=seqlen_offsets - ) - if cos_k is None and sin_k is None and qkv.is_contiguous(): - # Call 1 kernel instead of 2 kernels - # We need qkv to be contiguous so that when we reshape to combine (3, nheads) - # dimensions, we get the same tensor - if qkv.dim() == 5: - batch, seqlen, three, nheads, headdim = qkv.shape - assert three == 3 - # qk = rearrange(qkv[:, :, :2], "b s t h d -> b s (t h) d") - qk = qkv[:, :, :2].reshape(batch, seqlen, -1, headdim) - qk = apply_rotary_fn(qk, cos, sin) - else: - assert qkv.dim() == 4 - assert num_heads_q is not None - num_heads_k = (qkv.shape[2] - num_heads_q) // 2 - assert qkv.shape[2] == num_heads_q + 2 * num_heads_k - qk = qkv[:, :, :num_heads_q + num_heads_k] - qk = apply_rotary_fn(qk, cos, sin) - if not inplace: - if qkv.dim() == 5: - qkv = torch.cat([rearrange(qk, "b s (t h) d -> b s t h d", t=2), qkv[:, :, 2:]], dim=2) - else: - qkv = torch.cat([qk, qkv[:, :, num_heads_q + num_heads_k :]], dim=2) - else: - cos_k = cos if cos_k is None else cos_k - sin_k = sin if sin_k is None else sin_k - if qkv.dim() == 5: - batch, seqlen, three, nheads, headdim = qkv.shape - assert three == 3 - q, k = qkv[:, :, 0], qkv[:, :, 1] - else: - assert qkv.dim() == 4 - assert num_heads_q is not None - num_heads_k = (qkv.shape[2] - num_heads_q) // 2 - assert qkv.shape[2] == num_heads_q + 2 * num_heads_k - q, k = qkv[:, :, :num_heads_q], qkv[:, :, num_heads_q : num_heads_q + num_heads_k] - q = apply_rotary_fn(q, cos, sin) - k = apply_rotary_fn(k, cos_k, sin_k) - if not inplace: - if qkv.dim() == 5: - qkv = torch.stack([q, k, qkv[:, :, 2]], dim=2) - else: - qkv = torch.cat([q, k, qkv[:, :, num_heads_q + num_heads_k:]], dim=2) - return qkv - - -class ApplyRotaryEmbQKV_(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - cos, - sin, - cos_k=None, - sin_k=None, - interleaved=False, - seqlen_offsets: Union[int, torch.Tensor] = 0, - num_heads_q: Optional[int] = None, - ): - # apply_rotary_emb_qkv_inplace( - qkv = _apply_rotary_emb_qkv( - qkv, cos, sin, cos_k, sin_k, interleaved=interleaved, inplace=True, - seqlen_offsets=seqlen_offsets, num_heads_q=num_heads_q, - ) - if isinstance(seqlen_offsets, int): - ctx.save_for_backward(cos, sin, cos_k, sin_k) - ctx.seqlen_offsets = seqlen_offsets - else: - ctx.save_for_backward(cos, sin, cos_k, sin_k, seqlen_offsets) - ctx.seqlen_offsets = None - ctx.interleaved = interleaved - ctx.num_heads_q = num_heads_q - return qkv - - @staticmethod - def backward(ctx, dqkv): - seqlen_offsets = ctx.seqlen_offsets - if seqlen_offsets is None: - cos, sin, cos_k, sin_k, seqlen_offsets = ctx.saved_tensors - else: - cos, sin, cos_k, sin_k = ctx.saved_tensors - dqkv = _apply_rotary_emb_qkv( - dqkv, cos, sin, cos_k, sin_k, interleaved=ctx.interleaved, inplace=True, - seqlen_offsets=seqlen_offsets, num_heads_q=ctx.num_heads_q, conjugate=True, - ) - return dqkv, None, None, None, None, None, None, None - - -def apply_rotary_emb_qkv_( - qkv, - cos, - sin, - cos_k=None, - sin_k=None, - interleaved=False, - seqlen_offsets: Union[int, torch.Tensor] = 0, - num_heads_q: Optional[int] = None, -): - """ - Arguments: - qkv: (batch_size, seqlen, 3, nheads, headdim) or (batch_size, seqlen, num_heads_q + 2 * num_heads_k, headdim). - If qkv has shape (batch_size, seqlen, num_heads_q + 2 * num_heads_k, headdim) (e.g. MQA / GQA), - then num_heads_q must be provided. - cos, sin: (seqlen, rotary_dim / 2) - cos_k, sin_k: (seqlen, rotary_dim / 2), optional - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead of - 1st half and 2nd half (GPT-NeoX style). - seqlen_offsets: (batch_size,) or int. Each sequence in Q and K is shifted by this amount. - Most commonly used in inference when we have KV cache. - Return: - qkv: (batch_size, seqlen, 3, nheads, headdim) or (batch_size, seqlen, num_heads_q + 2 * num_heads_k, headdim) - rotary_dim must be <= headdim - Apply rotary embedding *inplace* to the first rotary_dim of Q and K. - """ - return ApplyRotaryEmbQKV_.apply( - qkv, cos, sin, cos_k, sin_k, interleaved, seqlen_offsets, num_heads_q - ) - - -class ApplyRotaryEmbKV_(torch.autograd.Function): - - @staticmethod - def forward(ctx, kv, cos, sin, interleaved=False, seqlen_offsets: Union[int, torch.Tensor] = 0): - batch, seqlen, two, nheads, headdim = kv.shape - assert two == 2 - k = kv[:, :, 0] - apply_rotary( - k, cos, sin, seqlen_offsets=seqlen_offsets, interleaved=interleaved, inplace=True - ) - if isinstance(seqlen_offsets, int): - ctx.save_for_backward(cos, sin) # Can't save int with save_for_backward - ctx.seqlen_offsets = seqlen_offsets - else: - ctx.save_for_backward(cos, sin, seqlen_offsets) - ctx.seqlen_offsets = None - ctx.interleaved = interleaved - return kv - - @staticmethod - def backward(ctx, dkv): - seqlen_offsets = ctx.seqlen_offsets - if seqlen_offsets is None: - cos, sin, seqlen_offsets = ctx.saved_tensors - else: - cos, sin = ctx.saved_tensors - apply_rotary( - dkv[:, :, 0], - cos, - sin, - seqlen_offsets=seqlen_offsets, - interleaved=ctx.interleaved, - inplace=True, - conjugate=True, - ) - return dkv, None, None, None, None - - -apply_rotary_emb_kv_ = ApplyRotaryEmbKV_.apply - - -def apply_rotary_emb_kv_( - kv, - cos, - sin, - interleaved=False, - seqlen_offsets: Union[int, torch.Tensor] = 0, -): - """ - Arguments: - kv: (batch_size, seqlen, 2, nheads, headdim) - cos, sin: (seqlen, rotary_dim / 2) - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead of - 1st half and 2nd half (GPT-NeoX style). - seqlen_offsets: (batch_size,) or int. Each sequence in Q and K is shifted by this amount. - Most commonly used in inference when we have KV cache. - Return: - kv: (batch_size, seqlen, 2, nheads, headdim) - rotary_dim must be <= headdim - Apply rotary embedding *inplace* to the first rotary_dim of K. - """ - return ApplyRotaryEmbKV_.apply(kv, cos, sin, interleaved, seqlen_offsets) - - -class RotaryEmbedding(torch.nn.Module): - """ - The rotary position embeddings from RoFormer_ (Su et. al). - A crucial insight from the method is that the query and keys are - transformed by rotation matrices which depend on the relative positions. - - Other implementations are available in the Rotary Transformer repo_ and in - GPT-NeoX_, GPT-NeoX was an inspiration - - .. _RoFormer: https://arxiv.org/abs/2104.09864 - .. _repo: https://github.com/ZhuiyiTechnology/roformer - .. _GPT-NeoX: https://github.com/EleutherAI/gpt-neox - - If scale_base is not None, this implements XPos (Sun et al., https://arxiv.org/abs/2212.10554). - A recommended value for scale_base is 512: https://github.com/HazyResearch/flash-attention/issues/96 - Reference: https://github.com/sunyt32/torchscale/blob/main/torchscale/component/xpos_relative_position.py - """ - - def __init__( - self, - dim: int, - base=10000.0, - interleaved=False, - scale_base=None, - device=None, - ): - """ - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead - of 1st half and 2nd half (GPT-NeoX style). - """ - super().__init__() - self.dim = dim - self.base = float(base) - # Generate and save the inverse frequency buffer (non trainable) - inv_freq = self._compute_inv_freq(device) - self.register_buffer("inv_freq", inv_freq, persistent=False) - self.interleaved = interleaved - self.scale_base = scale_base - scale = ( - (torch.arange(0, dim, 2, device=device, dtype=torch.float32) + 0.4 * dim) / (1.4 * dim) - if scale_base is not None - else None - ) - self.register_buffer("scale", scale, persistent=False) - - self._seq_len_cached = 0 - self._cos_cached = None - self._sin_cached = None - self._cos_k_cached = None - self._sin_k_cached = None - - def _compute_inv_freq(self, device=None): - return 1.0 / ( - self.base - ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim) - ) - - def _update_cos_sin_cache(self, seqlen, device=None, dtype=None): - # Reset the tables if the sequence length has changed, - # if we're on a new device (possibly due to tracing for instance), - # or if we're switching from inference mode to training - if ( - seqlen > self._seq_len_cached - or self._cos_cached is None - or self._cos_cached.device != device - or self._cos_cached.dtype != dtype - or (self.training and self._cos_cached.is_inference()) - ): - self._seq_len_cached = seqlen - # We want fp32 here, not self.inv_freq.dtype, since the model could be loaded in bf16 - # And the output of arange can be quite large, so bf16 would lose a lot of precision. - t = torch.arange(seqlen, device=device, dtype=torch.float32) - # We want fp32 here as well since inv_freq will be multiplied with t, and the output - # will be large. Having it in bf16 will lose a lot of precision and cause the - # cos & sin output to change significantly. - # We want to recompute self.inv_freq if it was not loaded in fp32 - if self.inv_freq.dtype != torch.float32: - inv_freq = self._compute_inv_freq(device=device) - else: - inv_freq = self.inv_freq - # Don't do einsum, it converts fp32 to bf16 under AMP - # freqs = torch.einsum("i,j->ij", t, self.inv_freq) - freqs = torch.outer(t, inv_freq) - if self.scale is None: - self._cos_cached = torch.cos(freqs).to(dtype) - self._sin_cached = torch.sin(freqs).to(dtype) - else: - power = ( - torch.arange(seqlen, dtype=self.scale.dtype, device=self.scale.device) - - seqlen // 2 - ) / self.scale_base - scale = self.scale.to(device=power.device) ** rearrange(power, "s -> s 1") - # We want the multiplication by scale to happen in fp32 - self._cos_cached = (torch.cos(freqs) * scale).to(dtype) - self._sin_cached = (torch.sin(freqs) * scale).to(dtype) - self._cos_k_cached = (torch.cos(freqs) / scale).to(dtype) - self._sin_k_cached = (torch.sin(freqs) / scale).to(dtype) - - def forward( - self, - qkv: torch.Tensor, - kv: Optional[torch.Tensor] = None, - seqlen_offset: Union[int, torch.Tensor] = 0, - max_seqlen: Optional[int] = None, - num_heads_q: Optional[int] = None, - ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: - """ - qkv: (batch, seqlen, 3, nheads, headdim) or (batch, seqlen, num_heads_q + 2 * num_heads_k, headdim) - if kv is none, else it's just q of shape (batch, seqlen, nheads, headdim). - If qkv has shape (batch, seqlen, num_heads_q + 2 * num_heads_k, headdim) (e.g. MQA / GQA), - then num_heads_q must be provided. - kv: (batch, seqlen, 2, nheads, headdim) - seqlen_offset: (batch_size,) or int. Each sequence in x is shifted by this amount. - Most commonly used in inference when we have KV cache. - If it's a tensor of shape (batch_size,), then to update the cos / sin cache, one - should pass in max_seqlen, which will update the cos / sin cache up to that length. - Apply rotary embedding *inplace* to qkv and / or kv. - """ - seqlen = qkv.shape[1] - if max_seqlen is not None: - self._update_cos_sin_cache(max_seqlen, device=qkv.device, dtype=qkv.dtype) - elif isinstance(seqlen_offset, int): - self._update_cos_sin_cache(seqlen + seqlen_offset, device=qkv.device, dtype=qkv.dtype) - if kv is None: - return apply_rotary_emb_qkv_( - qkv, - self._cos_cached, - self._sin_cached, - self._cos_k_cached if self.scale is not None else None, - self._sin_k_cached if self.scale is not None else None, - interleaved=self.interleaved, - seqlen_offsets=seqlen_offset, - num_heads_q=num_heads_q, - ) - else: - q = qkv - q = apply_rotary_emb_func( - q, - self._cos_cached, - self._sin_cached, - interleaved=self.interleaved, - inplace=True, - seqlen_offsets=seqlen_offset, - ) - kv = apply_rotary_emb_kv_( - kv, - self._cos_cached if self.scale is None else self._cos_k_cached, - self._sin_cached if self.scale is None else self._sin_k_cached, - interleaved=self.interleaved, - seqlen_offsets=seqlen_offset, - ) - return q, kv diff --git a/build/torch212-cxx11-cu126-x86_64-linux/metadata.json b/build/torch212-cxx11-cu126-x86_64-linux/metadata.json deleted file mode 100644 index 3db84691c6190cc1f22cd18d028913713a8f3cba..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu126-x86_64-linux/metadata.json +++ /dev/null @@ -1,40 +0,0 @@ -{ - "name": "flash-attn2", - "id": "_flash_attn2_cuda_f12afc9", - "version": 1, - "license": "BSD-3-Clause", - "python-depends": [], - "backend": { - "type": "cuda", - "archs": [ - "8.0", - "9.0" - ] - }, - "digest": { - "algorithm": "sha256", - "files": { - "__init__.py": "Cd9MEdd/WUYnBmfiQpC4SyPqAxrTilTjvsvswNju4zc=", - "_flash_attn2_cuda_f12afc9.abi3.so": "6YA4KPQyEWNjb0lSFtMmvw2sHZXukP9YTVKqkEYtVIc=", - "_ops.py": "GQbmTTeSOU+7Q2nI68Hz+1RvEogyfEFccx1ubUYw+2s=", - "bert_padding.py": "gF1EmsdJ+HpQ86MRQ4VxDw+Sb/RVISdQALdNnoByHlw=", - "flash_attn2/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY=", - "flash_attn_interface.py": "W7a9bHR2ly6QL3aKqsBqxWA7M7ZWYcZPmtsPJu6rAfY=", - "layers/__init__.py": "47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU=", - "layers/patch_embed.py": "H58CgME/qSOPTZLOG08wFgrQS1j34pvNwMPrkTj3Ek4=", - "layers/rotary.py": "rFIdNrZQFNawwLD6d7t5LE2m+M93wwDZ8EuExw+gfD8=", - "ops/__init__.py": "47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU=", - "ops/activations.py": "t5lzNg1In8LP6bKeTnyeMizwqjv27JGbJ6ylPdGvZYg=", - "ops/fused_dense.py": "BqT8VTSCGdLJAjh6WBMkCRFYrdVodPeoElWBvH6AIzA=", - "ops/layer_norm.py": "zr7NXIm+2mtEynTp1CS0fbFGI2Mqdp41dY4AfDWF6EQ=", - "ops/rms_norm.py": "XEnihcj0a4aSz4LO55m5iKGVn4HKTeKN8TIyHjuDgxI=", - "ops/triton/__init__.py": "AbpHGcgLb+kRsJGnwFEktk7uzpZOCcBY74+YBdrKVGs=", - "ops/triton/cross_entropy.py": "y56nmuOtpSNxM7sDZcUSgXKIsIEayfuvrU8bAVog/nw=", - "ops/triton/k_activations.py": "+Z3vIyO4JkqBMipKsPvhzmxljtBdIhJCsl/M+/ESqBo=", - "ops/triton/layer_norm.py": "BBQdJiqopbIS12viUqzx3SvOQCNbDSAYLJXhpyMGyC8=", - "ops/triton/linear.py": "OtRvKz8xdpl+7v3q/ZTaS9fdBt9XrzMyapgRr50uBbM=", - "ops/triton/mlp.py": "uWA+D7HyuRUsOT+4nhi3vvCDNZnNEXwYcHMQ02UREfQ=", - "ops/triton/rotary.py": "0YSrlbep0lJmzLpTiANs0slTFeGcQQqR+el6uzbSLvI=" - } - } -} \ No newline at end of file diff --git a/build/torch212-cxx11-cu126-x86_64-linux/metadata.json.sigstore b/build/torch212-cxx11-cu126-x86_64-linux/metadata.json.sigstore deleted file mode 100644 index dbe484f95858a4dae6e21740acd984af19e534de..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu126-x86_64-linux/metadata.json.sigstore +++ /dev/null @@ -1 +0,0 @@ -{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json","verificationMaterial":{"certificate":{"rawBytes":"MIIHTTCCBtKgAwIBAgIUK4fnA/nvQAudm9xfLL7/VOq1LbIwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwNjI0MDEwMTIzWhcNMjYwNjI0MDExMTIzWjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEZlYfgO7GoSuZ3IbIzxC80gNtpSRdLfhIkKKQoT7vyhGlwBozYBeSkgCaUFZ0wEY/NKFMkdUGX8rQICg/Do5QrKOCBfEwggXtMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUj2FmdwC3DZvQiRZqqevM+94sU20wHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoZjEyYWZjOTJjNDBiYmQ3NmQ5OTBkOTExZWJhYmYyOGU0OTBhN2JkZTATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoZjEyYWZjOTJjNDBiYmQ3NmQ5OTBkOTExZWJhYmYyOGU0OTBhN2JkZTAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoZjEyYWZjOTJjNDBiYmQ3NmQ5OTBkOTExZWJhYmYyOGU0OTBhN2JkZTAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKGYxMmFmYzkyYzQwYmJkNzZkOTkwZDkxMWViYWJmMjhlNDkwYTdiZGUwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMjgwMzYyNjg5MzUvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBiwYKKwYBBAHWeQIEAgR9BHsAeQB3AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABnvclY5oAAAQDAEgwRgIhANYu47H6cNp3MuCxYz28cLLh4+PatzBEiCwSl5+S0EMmAiEAzIr2zTn6aUWwFemjpKAy5ICas0demazRrcvusEcaJZ8wCgYIKoZIzj0EAwMDaQAwZgIxAPOrSSG8SO/xPBGUkAbw/PDg1BXEeS20XCVqwIPd3GwAyVxM86GsKUvXq30uufx9jwIxAP41BEK3vmUVb81lOCwQc7JEEfP8mRvM+SoYzvyxfEKrH8Fr1YeDFqqGE3LfKmH46A=="},"tlogEntries":[{"logIndex":"1933772387","logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="},"kindVersion":{"kind":"hashedrekord","version":"0.0.1"},"integratedTime":"1782262883","inclusionPromise":{"signedEntryTimestamp":"MEUCIQD3sCY+elM+eX6kJ8JbDgg9xdlrDkbJSeUBTvBvnR6veQIgCPk2DNC6RcCUzBUJzkUXt4P5+OPxq54Q3Gts44f70JI="},"inclusionProof":{"logIndex":"1811868125","rootHash":"MmBm/7hFBeKeYCc/Y2SZwxh3p6rJTyaxoaD3uuaQMvY=","treeSize":"1811868129","hashes":["ud5QOn88yLqB3zp46TffgxU1FMIJpOcO31m4U1vOkDM=","MGPJfI+A7HZjdW0CSfJb97tt5VVQh3IWzZdnrdGNS80=","0vXgJAQPrOyXvmEDQvoybj/pzIM12/9O0y4ww+wF+E8=","8m7e3pVGJ345Snj6dRFdzy4fO6/wHDiheV/IJtCN4Wo=","h+8bwxOpbD8Nqhnuuytjq0OfXe9xLeAAwji87SYBI38=","3QGOY25zSQowdYrSLZ78L4Z41pexJOpYImKsUcRWE/k=","LE3BdXo2HKOOHG1CHHaxGsOSzOULgaGOaRKySauKQmE=","btQJuvlhowysbKp6G5CvwXXeoaFFUnfwtXLzOdGvBCE=","oTwvuS73/Mr8Eznxu3zK7mutJv/d67rW15lLY1XiQKs=","fqtytk2if0R+M3kWRpiPUrBEhf2z8y5Li6MeOuWmiks=","mzp9m+2At7GZdgfx0m7RwtO9YzltCZ6UhZBy4p+lUdM=","ipWWLDkqay5JAYe6P4Nrqsck/1C0C5uwqT9QbFsiY1E=","COKvK4TO+dkpqUDGbTbSKewajXzAonx1B1YaRb/ymLw=","dQlR+DXGDKXm5V1fupxaMSOwNZ9F1MybXyyxB96yj80=","JKJdxSZwKPO6QFzLTxHCrUyI5sDU/CwY57zDlA9ni6Q=","OgfnQCh7xr8xUTErIw//YAKaKMDwRJWiThvYExQvLJM=","S8Oe9ZicoJqbVuL1I0dIdTgVpLc82Cm2CMOcV83sKVs=","E1O2ilWlpFRLWGdbmJUMDZ2qbGu/eDZ12m6YjQ9/JPA=","nyv2TDlfkz0XfRwjY98zsbUbOkBZ+Tgz7PML0c2uN14=","OMfDh4k9YD0E+t4sjqJsh+vg8D1ueo1cC07exyxmr6s=","+vHKm4wNwGJ46ceyZ8mePD7Vn1SediStNHCStRoTSWQ=","5HwWjuGNDisetbiWM0cWIXM65TIM3EJlqk13ntCQlcU=","NDIRlqxO0IaGr94q86MIU7VCAN3yEEgOHS4loya5lu4=","daxmZaajRpZV+JxHiOYZhJBiSKN5ucqjh2WnGbHhirw=","DOCeoSMovIvLExkhIvisow9AuNXgeWs4ECkyR6EcqYU="],"checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n1811868129\nMmBm/7hFBeKeYCc/Y2SZwxh3p6rJTyaxoaD3uuaQMvY=\n\n— rekor.sigstore.dev wNI9ajBFAiEA7nYHEoaa7T1I5loHv2ir8SriLadsSOB6dJB0yxO2hIwCICQ3A/QrpJJdOssmriE4BtvS2WEDEvYc0+kC5MOngjF3\n"}},"canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiI1MTM0NzkwYTQ2YjdmMzIzMzY5MWQ1OWU3NDJiNjM0NDkyNWUzYTU2YzBiZjFjNDJmNjU5MjFkNDA1ZGY1Njg3In19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FWUNJUUNaVzhaMTc4YXp2aEYvY2hGeDRVejFoc1p3OUFTb0hiUUNKOUZBV01ZTWpRSWhBSVU0UUF3eEpWd3ZEYittSVRyeWpsbjROZ3hCNTRjM2sxZTRjRVM5SEcxUCIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFVWRU5EUW5STFowRjNTVUpCWjBsVlN6Um1ia0V2Ym5aUlFYVmtiVGw0Wmt4TU55OVdUM0V4VEdKSmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDVxU1RCTlJFVjNUVlJKZWxkb1kwNU5hbGwzVG1wSk1FMUVSWGhOVkVsNlYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZhYkZsbVowODNSMjlUZFZvelNXSkplbmhET0RCblRuUndVMUprVEdab1NXdExTMUVLYjFRM2RubG9SMngzUW05NldVSmxVMnRuUTJGVlJsb3dkMFZaTDA1TFJrMXJaRlZIV0RoeVVVbERaeTlFYnpWUmNrdFBRMEptUlhkbloxaDBUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlZxTWtadENtUjNRek5FV25aUmFWSmFjWEZsZGswck9UUnpWVEl3ZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOWFha1Y1V1ZkYWFrOVVTbXBPUkVKcFdXMVJNMDV0VVRWUFZFSnJDazlVUlhoYVYwcG9XVzFaZVU5SFZUQlBWRUpvVGpKS2ExcFVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMXBxUlhsWlYxcHFUMVJLYWs1RVFtbFpiVkV6VG0xUk5VOVVRbXRQVkVWNFdsZEthRmx0V1hrS1QwZFZNRTlVUW1oT01rcHJXbFJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMXBxUlhsWlYxcHFUMVJLYWs1RVFtbFpiVkV6VG0xUk5VOVVRbXNLVDFSRmVGcFhTbWhaYlZsNVQwZFZNRTlVUW1oT01rcHJXbFJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwZFplRTF0Um0wS1dYcHJlVmw2VVhkWmJVcHJUbnBhYTA5VWEzZGFSR3Q0VFZkV2FWbFhTbTFOYW1oc1RrUnJkMWxVWkdsYVIxVjNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOYW1kM1RYcFplVTVxWnpWTmVsVjJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwZDFsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTT1VKSWMwRUtaVkZDTTBGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtNTJZMnhaTlc5QlFVRlJSQXBCUldkM1VtZEphRUZPV1hVME4wZzJZMDV3TTAxMVEzaFplakk0WTB4TWFEUXJVR0YwZWtKRmFVTjNVMncxSzFNd1JVMXRRV2xGUVhwSmNqSjZWRzQyQ21GVlYzZEdaVzFxY0V0QmVUVkpRMkZ6TUdSbGJXRjZVbkpqZG5WelJXTmhTbG80ZDBObldVbExiMXBKZW1vd1JVRjNUVVJoVVVGM1dtZEplRUZRVDNJS1UxTkhPRk5QTDNoUVFrZFZhMEZpZHk5UVJHY3hRbGhGWlZNeU1GaERWbkYzU1ZCa00wZDNRWGxXZUUwNE5rZHpTMVYyV0hFek1IVjFabmc1YW5kSmVBcEJVRFF4UWtWTE0zWnRWVlppT0RGc1QwTjNVV00zU2tWRlpsQTRiVkoyVFN0VGIxbDZkbmw0WmtWTGNrZzRSbkl4V1dWRVJuRnhSMFV6VEdaTGJVZzBDalpCUFQwS0xTMHRMUzFGVGtRZ1EwVlNWRWxHU1VOQlZFVXRMUzB0TFFvPSJ9fX19"}],"timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyTADAgEAMIICwAYJKoZIhvcNAQcCoIICsTCCAq0CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQg/I/1ZI/P2HpaVVo4oT9Cmu4UEQ3JqG41miXZ78KmDoQCFQDioItfDuR58ZayA+Pai5mIMPY4kRgPMjAyNjA2MjQwMTAxMjNaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHaMIIB1gIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MDYyNDAxMDEyM1owLwYJKoZIhvcNAQkEMSIEII7s6EUFrndKcjFlRog3xjdOaAobzknDVc2mDeJ+Y825MIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRmMGQCMCYtD1m7xgTWHVHXiKVzGiu0QrnvF9faXxxQnsTDT0/fuLi3w4yROIMQad4dFbZVPwIwWJflxwHHYPd1veyOur/tzY/oBfyyriAG4PnkvHY+f4UZJAs8AiX2tOyGvkANcGzD"}]}},"messageSignature":{"messageDigest":{"algorithm":"SHA2_256","digest":"UTR5Cka38yM2kdWedCtjRJJeOlbAvxxC9lkh1AXfVoc="},"signature":"MEYCIQCZW8Z178azvhF/chFx4Uz1hsZw9ASoHbQCJ9FAWMYMjQIhAIU4QAwxJVwvDb+mITryjln4NgxB54c3k1e4cES9HG1P"}} \ No newline at end of file diff --git a/build/torch212-cxx11-cu126-x86_64-linux/ops/__init__.py b/build/torch212-cxx11-cu126-x86_64-linux/ops/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/build/torch212-cxx11-cu126-x86_64-linux/ops/activations.py b/build/torch212-cxx11-cu126-x86_64-linux/ops/activations.py deleted file mode 100644 index 7c09649fc41e12d5a360c5672825d8380bc7ec80..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu126-x86_64-linux/ops/activations.py +++ /dev/null @@ -1,135 +0,0 @@ -# Copied from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/model/layers/activations.py -import math - -import torch -import torch.nn as nn -import torch.nn.functional as F - -# 1/sqrt(2*pi)-> 0.3989423 -# 1/sqrt(2) -> 0.70710678 -# sqrt(2/pi) -> 0.79788456 - -# this function is tanh approximation of gelu -# actual gelu is: -# x * 0.5 * (1.0 + torch.erf(x * 0.70710678)) -@torch.jit.script -def bias_gelu(y, bias): - x = bias + y - return (x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))).to(dtype=y.dtype) - - -# gradient of tanh approximation of gelu -# gradient of actual gelu is: -# 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x) -@torch.jit.script -def bias_gelu_back(g, y, bias): - """Assume that y has shape (B, D) and bias has shape (D)""" - x = bias + y - tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)) - # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243 - ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( - 1 + tanh_out - ) - grad_y = ff * g - return grad_y.to(dtype=y.dtype), grad_y.sum(dim=(0), dtype=bias.dtype) - - -class GeLUFunction(torch.autograd.Function): - @staticmethod - # bias is an optional argument - def forward(ctx, input, bias): - ctx.save_for_backward(input, bias) - return bias_gelu(input, bias) - - @staticmethod - def backward(ctx, grad_output): - input, bias = ctx.saved_tensors - tmp = bias_gelu_back(grad_output, input, bias) - return tmp, tmp - - -bias_gelu_impl = GeLUFunction.apply - -# this function is tanh approximation of gelu -# actual gelu is: -# x * 0.5 * (1.0 + torch.erf(x * 0.70710678)) -@torch.jit.script -def gelu_fwd(x): - return (x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))).to(dtype=x.dtype) - - -# gradient of tanh approximation of gelu -# gradient of actual gelu is: -# 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x) -@torch.jit.script -def gelu_bwd(g, x): - tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)) - # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243 - ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( - 1 + tanh_out - ) - return (ff * g).to(dtype=x.dtype) - - -class FastGeLUFunction(torch.autograd.Function): - @staticmethod - # bias is an optional argument - def forward(ctx, input): - ctx.save_for_backward(input) - return gelu_fwd(input) - - @staticmethod - def backward(ctx, grad_output): - (input,) = ctx.saved_tensors - tmp = gelu_bwd(grad_output, input) - return tmp - - -fast_gelu_impl = FastGeLUFunction.apply - - -@torch.jit.script -def relu_bwd(g, x): - return torch.where(x >= 0, g, 0.0).to(dtype=x.dtype) - - -@torch.jit.script -def sqrelu_fwd(x): - r = F.relu(x) - return (r * r).to(dtype=x.dtype) - - -@torch.jit.script -def sqrelu_bwd(g, x): - return (2.0 * g * F.relu(x)).to(dtype=x.dtype) - - -swiglu_fwd_codestring = """ -template T swiglu_fwd(T x, T y) { - return float(x) * float(y) / (1.0f + ::exp(-float(x))); -} -""" -swiglu_bwd_codestring = """ -template void swiglu_bwd(T x, T y, T g, T& dx, T& dy) { - float x_sigmoid = 1.0f / (1.0f + ::exp(-float(x))); - dx = x_sigmoid * (1 + float(x) * (1.0f - x_sigmoid)) * float(g) * float(y); - dy = float(x) * x_sigmoid * float(g); -} -""" -swiglu_fwd = torch.cuda.jiterator._create_jit_fn(swiglu_fwd_codestring) -swiglu_bwd = torch.cuda.jiterator._create_multi_output_jit_fn(swiglu_bwd_codestring, num_outputs=2) - - -class SwiGLUFunction(torch.autograd.Function): - - @staticmethod - def forward(ctx, x, y): - ctx.save_for_backward(x, y) - return swiglu_fwd(x, y) - - @staticmethod - def backward(ctx, dout): - x, y = ctx.saved_tensors - return swiglu_bwd(x, y, dout) - -swiglu = SwiGLUFunction.apply diff --git a/build/torch212-cxx11-cu126-x86_64-linux/ops/fused_dense.py b/build/torch212-cxx11-cu126-x86_64-linux/ops/fused_dense.py deleted file mode 100644 index 6b4033d134e4093fe278f7b3f8c7d3128ce9f36d..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu126-x86_64-linux/ops/fused_dense.py +++ /dev/null @@ -1,688 +0,0 @@ -# Copyright (c) 2023, Tri Dao. -# Inspired by https://github.com/NVIDIA/apex/blob/master/apex/fused_dense/fused_dense.py -# We make it work with pytorch amp and with bfloat16. -# The TensorParallel linear modules are inspired by https://github.com/NVIDIA/apex/blob/master/apex/transformer/tensor_parallel/layers.py -from functools import partial -from typing import Optional - -# import fused_dense_cuda # from apex -import fused_dense_lib as fused_dense_cuda -import torch -import torch.nn as nn -import torch.nn.functional as F -from torch import Tensor -from torch.distributed import ProcessGroup - -from flash_attn.utils.torch import custom_fwd, custom_bwd -from flash_attn.ops.activations import gelu_bwd, relu_bwd, sqrelu_bwd, sqrelu_fwd -from flash_attn.utils.distributed import ( - all_gather_raw, - all_reduce, - all_reduce_raw, - reduce_scatter, - reduce_scatter_raw, -) - - -class FusedDenseFunc(torch.autograd.Function): - @staticmethod - @custom_fwd - def forward( - ctx, x, weight, bias, return_residual=False, process_group=None, sequence_parallel=True - ): - """ - If process_group is not None and sequence_parallel=True, we're doing Tensor Parallel - with sequence parallelism: we do an all_gather_raw of x before doing the matmul. - """ - ctx.compute_weight_gradient = weight.requires_grad - ctx.return_residual = return_residual - ctx.process_group = process_group - ctx.sequence_parallel = sequence_parallel - - if torch.is_autocast_enabled(): - x = x.to(dtype=torch.get_autocast_gpu_dtype()) - x = x.contiguous() - if process_group is not None and sequence_parallel: - # We want to kick off the all_gather early, before weight dtype conversion - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - else: - total_x = x - - if torch.is_autocast_enabled(): - weight = weight.to(dtype=torch.get_autocast_gpu_dtype()) - bias = bias.to(dtype=torch.get_autocast_gpu_dtype()) if bias is not None else None - weight = weight.contiguous() - if process_group is not None and sequence_parallel: - handle_x.wait() - batch_shape, n = total_x.shape[:-1], total_x.shape[-1] - batch_dim = batch_shape.numel() - # https://github.com/pytorch/pytorch/blob/5b51849b48a7dbccd297286cc0110def4706f9e7/aten/src/ATen/native/cuda/Blas.cpp#L174 - if min(batch_dim, n, *weight.shape) > 65535 * 32: - raise RuntimeError("fused_dense only supports matrix dims <= 2M") - output = F.linear(total_x, weight, bias) - if ctx.compute_weight_gradient: - ctx.save_for_backward(x, weight) - else: - ctx.save_for_backward(weight) - return output if not return_residual else (output, x) - - @staticmethod - @custom_bwd - def backward(ctx, grad_output, *args): - grad_output = grad_output.contiguous() - if ctx.return_residual: - (grad_input,) = args - grad_input = grad_input.contiguous() - process_group = ctx.process_group - sequence_parallel = ctx.sequence_parallel - if ctx.compute_weight_gradient: - x, weight = ctx.saved_tensors - if process_group is not None and sequence_parallel: - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - else: - total_x = x - else: - (weight,) = ctx.saved_tensors - total_x = None - batch_shape = grad_output.shape[:-1] - batch_dim = batch_shape.numel() - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - if ctx.needs_input_grad[0]: - if not ctx.return_residual: - grad_input = F.linear(grad_output, weight.t()) - else: - grad_input = torch.addmm( - grad_input.reshape(batch_dim, grad_input.shape[-1]), grad_output, weight - ) - grad_input = grad_input.reshape(*batch_shape, grad_input.shape[-1]) - if process_group is not None: - reduce_fn = reduce_scatter_raw if sequence_parallel else all_reduce_raw - grad_input, handle_grad_input = reduce_fn(grad_input, process_group, async_op=True) - else: - grad_input = None - if ctx.needs_input_grad[1]: - assert ctx.compute_weight_gradient - if process_group is not None and sequence_parallel: - handle_x.wait() - grad_weight, grad_bias = fused_dense_cuda.linear_bias_wgrad( - total_x.reshape(batch_dim, total_x.shape[-1]), grad_output, ctx.needs_input_grad[2] - ) - else: - grad_weight = None - grad_bias = grad_output if ctx.needs_input_grad[2] else None - if process_group is not None and ctx.needs_input_grad[0]: - handle_grad_input.wait() - return grad_input, grad_weight, grad_bias, None, None, None - - -def fused_dense_func( - x: Tensor, - weight: Tensor, - bias: Optional[Tensor] = None, - return_residual: bool = False, - process_group: Optional[ProcessGroup] = None, - sequence_parallel: bool = True, -): - dtype_eligible = x.dtype in [torch.float16, torch.bfloat16] or ( - x.dtype == torch.float32 and torch.is_autocast_enabled() - ) - if x.is_cuda and weight.is_cuda and (bias is None or bias.is_cuda) and dtype_eligible: - return FusedDenseFunc.apply( - x, weight, bias, return_residual, process_group, sequence_parallel - ) - else: - assert process_group is None - out = F.linear(x, weight, bias) - return out if not return_residual else (out, x) - - -class FusedDense(nn.Linear): - def __init__( - self, - in_features: int, - out_features: int, - bias: bool = True, - return_residual: bool = False, - device=None, - dtype=None, - ) -> None: - super().__init__(in_features, out_features, bias=bias, device=device, dtype=dtype) - self.return_residual = return_residual - - def forward(self, x, process_group=None): - """ - If process_group is not None, we're doing Tensor Parallel with sequence parallelism: - we do an all_gather of x before doing the matmul. - """ - return fused_dense_func( - x, - self.weight, - self.bias, - return_residual=self.return_residual, - process_group=process_group, - ) - - -class ColumnParallelLinear(nn.Linear): - def __init__( - self, - in_features: int, - out_features: int, - process_group: ProcessGroup, - bias: bool = True, - sequence_parallel=True, - multiple_of=1, - device=None, - dtype=None, - ) -> None: - world_size = torch.distributed.get_world_size(process_group) - if out_features % multiple_of: - raise ValueError(f"out_features ({out_features}) must be a multiple of {multiple_of}") - multiple = out_features // multiple_of - # We want to split @multiple across world_size, but it could be an uneven split - div = multiple // world_size - mod = multiple % world_size - # The first @mod ranks get @div + 1 copies, the rest get @div copies - local_multiple = div + int(torch.distributed.get_rank(process_group) < mod) - super().__init__( - in_features, local_multiple * multiple_of, bias=bias, device=device, dtype=dtype - ) - self.process_group = process_group - self.sequence_parallel = sequence_parallel - - def forward(self, x): - # If self.sequence_parallel is True, we're doing Tensor Parallel with sequence parallelism: - # we do an all_gather of x before doing the matmul. - # If not, then the input is already gathered. - return fused_dense_func( - x, - self.weight, - self.bias, - process_group=self.process_group, - sequence_parallel=self.sequence_parallel, - ) - - -class RowParallelLinear(nn.Linear): - def __init__( - self, - in_features: int, - out_features: int, - process_group: ProcessGroup, - bias: bool = True, - sequence_parallel=True, - multiple_of=1, - device=None, - dtype=None, - ) -> None: - world_size = torch.distributed.get_world_size(process_group) - rank = torch.distributed.get_rank(process_group) - if in_features % multiple_of: - raise ValueError(f"in_features ({in_features}) must be a multiple of {multiple_of}") - multiple = in_features // multiple_of - # We want to split @multiple across world_size, but it could be an uneven split - div = multiple // world_size - mod = multiple % world_size - # The first @mod ranks get @div + 1 copies, the rest get @div copies - local_multiple = div + int(torch.distributed.get_rank(process_group) < mod) - # Only rank 0 will have bias - super().__init__( - local_multiple * multiple_of, - out_features, - bias=bias and rank == 0, - device=device, - dtype=dtype, - ) - self.process_group = process_group - self.sequence_parallel = sequence_parallel - - def forward(self, x): - """ - We're doing Tensor Parallel with sequence parallelism: we do the matmul and then - a reduce_scatter of the result. - """ - out = fused_dense_func(x, self.weight, self.bias) - reduce_fn = reduce_scatter if self.sequence_parallel else all_reduce - return reduce_fn(out, self.process_group) - - -class FusedMLPFunc(torch.autograd.Function): - @staticmethod - @custom_fwd - def forward( - ctx, - x, - weight1, - bias1, - weight2, - bias2, - activation="gelu_approx", - save_pre_act=True, - return_residual=False, - checkpoint_lvl=0, - heuristic=0, - process_group=None, - sequence_parallel=True, - ): - """ - If process_group is not None and sequence_parallel=True, we're doing Tensor Parallel - with sequence parallelism: we do an all_gather of x before doing the matmul. - If sequence_parallel=False, then the input is already gathered. - - checkpoint_lvl: - 0: no recomputation in the bwd - 1: recompute gelu_out / relu_out in the bwd - 2: recompute pre_act and gelu_out / relu_out in the bwd - """ - assert -1 <= heuristic <= 4 - assert activation in ["gelu_approx", "relu", "sqrelu"] - if activation == "sqrelu": - assert heuristic == -1 - if not save_pre_act: - checkpoint_lvl = 2 - assert checkpoint_lvl in [0, 1, 2] - ctx.return_residual = return_residual - ctx.process_group = process_group - ctx.sequence_parallel = sequence_parallel - ctx.checkpoint_lvl = checkpoint_lvl - ctx.activation = activation - ctx.heuristic = heuristic - - if torch.is_autocast_enabled(): - x = x.to(dtype=torch.get_autocast_gpu_dtype()) - x = x.contiguous() - if process_group is not None and sequence_parallel: - # We want to kick off the all_gather early, before weight dtype conversion - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - else: - total_x = x - - if torch.is_autocast_enabled(): - dtype = torch.get_autocast_gpu_dtype() - weight1, weight2 = [a.to(dtype=dtype) for a in [weight1, weight2]] - bias1 = bias1.to(dtype=dtype) if bias1 is not None else None - bias2 = bias2.to(dtype=dtype) if bias2 is not None else None - weight1 = weight1.contiguous() - bias1 = bias1.contiguous() if bias1 is not None else None - weight2 = weight2.contiguous() - bias2 = bias2.contiguous() if bias2 is not None else None - if process_group is not None and sequence_parallel: - handle_x.wait() - batch_shape, n = total_x.shape[:-1], total_x.shape[-1] - batch_dim = batch_shape.numel() - # https://github.com/pytorch/pytorch/blob/5b51849b48a7dbccd297286cc0110def4706f9e7/aten/src/ATen/native/cuda/Blas.cpp#L174 - if min(batch_dim, n, *weight1.shape, *weight2.shape) > 65535 * 32: - raise RuntimeError("fused_dense only supports matrix dims <= 2M") - if heuristic == -1: - pre_act = F.linear(total_x, weight1, bias1) - activation_fn = ( - partial(F.gelu, approximate="tanh") - if activation == "gelu_approx" - else (sqrelu_fwd if activation == "sqrelu" else F.relu) - ) - with torch.jit.fuser("fuser2"): - output1 = activation_fn(pre_act) - # This is before adding bias1 - # pre_act = F.linear(total_x.reshape(batch_dim, n), weight1) - # with torch.jit.fuser('fuser2'): - # output1 = bias_gelu(pre_act, bias1) - else: - is_gelu = activation == "gelu_approx" - output1, *rest = fused_dense_cuda.linear_act_forward( - total_x.reshape(batch_dim, n), weight1, bias1, is_gelu, save_pre_act, heuristic - ) - if save_pre_act: - pre_act = rest[0] - output2 = F.linear(output1, weight2, bias2) - if checkpoint_lvl == 0 or (checkpoint_lvl == 1 and activation == "relu"): - # For RELU the pre_act is very small (just a bit-mask) so we just save it - ctx.save_for_backward(x, weight1, weight2, pre_act, output1) - elif checkpoint_lvl == 1: - ctx.save_for_backward(x, weight1, weight2, pre_act) - elif checkpoint_lvl == 2: - ctx.save_for_backward(x, weight1, weight2, bias1) - output2 = output2.reshape(*batch_shape, output2.shape[-1]) - return output2 if not return_residual else (output2, x) - - @staticmethod - @custom_bwd - def backward(ctx, grad_output, *args): - grad_output = grad_output.contiguous() - checkpoint_lvl = ctx.checkpoint_lvl - activation = ctx.activation - activation_fn = ( - partial(F.gelu, approximate="tanh") - if activation == "gelu_approx" - else (sqrelu_fwd if activation == "sqrelu" else F.relu) - ) - if ctx.return_residual: - (grad_input,) = args - grad_input = grad_input.contiguous() - process_group = ctx.process_group - sequence_parallel = ctx.sequence_parallel - x, weight1, weight2, *rest = ctx.saved_tensors - if process_group is None or not sequence_parallel: - total_x = x - batch_shape = grad_output.shape[:-1] - batch_dim = batch_shape.numel() - if checkpoint_lvl in [0, 1]: - if process_group is not None and sequence_parallel: - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - if checkpoint_lvl == 0 or (checkpoint_lvl == 1 and activation == "relu"): - pre_act, output1 = rest - elif checkpoint_lvl == 1: - (pre_act,) = rest - with torch.jit.fuser("fuser2"): - output1 = activation_fn(pre_act) - elif checkpoint_lvl == 2: - (bias1,) = rest - if process_group is not None and sequence_parallel: - total_x, _ = all_gather_raw(x, process_group) - if ctx.heuristic == -1: - pre_act = F.linear(total_x, weight1, bias1) - with torch.jit.fuser("fuser2"): - output1 = activation_fn(pre_act) - else: - output1, pre_act = fused_dense_cuda.linear_act_forward( - total_x.reshape(batch_dim, total_x.shape[-1]), - weight1, - bias1, - activation == "gelu_approx", - True, - ctx.heuristic, - ) - - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - output1 = output1.reshape(batch_dim, output1.shape[-1]) - pre_act = pre_act.reshape(batch_dim, pre_act.shape[-1]) - if ctx.needs_input_grad[3]: - grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad( - output1, grad_output, ctx.needs_input_grad[4] - ) - else: - grad_weight2 = None - grad_bias2 = grad_output if ctx.needs_input_grad[4] else None - if ctx.heuristic == -1: - # grad_pre_act = matmul_dgelu(grad_output, weight2, pre_act) - grad_output1 = F.linear(grad_output, weight2.t()) - activation_grad_fn = ( - gelu_bwd - if activation == "gelu_approx" - else (sqrelu_bwd if activation == "sqrelu" else relu_bwd) - ) - with torch.jit.fuser("fuser2"): - grad_pre_act = activation_grad_fn(grad_output1, pre_act) - else: - # The cublasLt epilogue has to compute both gelu/relu grad and bias grad, we can't - # just compute gelu/relu grad - grad_pre_act, grad_bias1 = fused_dense_cuda.bias_act_linear_dgrad_bgrad( - weight2, grad_output, pre_act, activation == "gelu_approx", ctx.heuristic - ) - if not ctx.needs_input_grad[2]: - grad_bias1 = None - if ctx.needs_input_grad[0]: - if not ctx.return_residual: - grad_input = F.linear(grad_pre_act, weight1.t()) - else: - grad_input = torch.addmm( - grad_input.reshape(batch_dim, grad_input.shape[-1]), grad_pre_act, weight1 - ) - grad_input = grad_input.reshape(*batch_shape, grad_input.shape[-1]) - if process_group is not None: - reduce_fn = reduce_scatter_raw if sequence_parallel else all_reduce_raw - grad_input, handle_grad_input = reduce_fn(grad_input, process_group, async_op=True) - else: - grad_input = None - if ctx.heuristic == -1: - if ctx.needs_input_grad[1]: - if process_group is not None and sequence_parallel and checkpoint_lvl != 2: - handle_x.wait() - grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_wgrad( - total_x.reshape(batch_dim, total_x.shape[-1]), - grad_pre_act, - ctx.needs_input_grad[2], - ) - else: - grad_weight1 = None - grad_bias1 = grad_pre_act if ctx.needs_input_grad[2] else None - else: - if ctx.needs_input_grad[1]: - if process_group is not None and sequence_parallel and checkpoint_lvl != 2: - handle_x.wait() - grad_weight1 = F.linear( - grad_pre_act.t(), total_x.reshape(batch_dim, total_x.shape[-1]).t() - ) - else: - grad_weight1 = None - if process_group is not None and ctx.needs_input_grad[0]: - handle_grad_input.wait() - return ( - grad_input, - grad_weight1, - grad_bias1, - grad_weight2, - grad_bias2, - None, - None, - None, - None, - None, - None, - None, - ) - - -def fused_mlp_func( - x: Tensor, - weight1: Tensor, - weight2: Tensor, - bias1: Optional[Tensor] = None, - bias2: Optional[Tensor] = None, - activation: str = "gelu_approx", - save_pre_act: bool = True, - return_residual: bool = False, - checkpoint_lvl: int = 0, - heuristic: int = 0, - process_group: Optional[ProcessGroup] = None, - sequence_parallel: bool = True, -): - assert activation in ["gelu_approx", "relu", "sqrelu"] - dtype_eligible = x.dtype in [torch.float16, torch.bfloat16] or ( - x.dtype == torch.float32 and torch.is_autocast_enabled() - ) - # If we save pre-activation, dimension must be divisible by 128 (relu) or 8 (gelu) - dim_eligible = not save_pre_act or (x.shape[-1] % (128 if activation == "relu" else 8) == 0) - if ( - x.is_cuda - and weight1.is_cuda - and weight2.is_cuda - and (bias1 is None or bias1.is_cuda) - and (bias2 is None or bias2.is_cuda) - and dtype_eligible - and dim_eligible - ): - return FusedMLPFunc.apply( - x, - weight1, - bias1, - weight2, - bias2, - activation, - save_pre_act, - return_residual, - checkpoint_lvl, - heuristic, - process_group, - sequence_parallel, - ) - else: - assert process_group is None - pre_act = F.linear(x, weight1, bias1) - activation_fn = ( - partial(F.gelu, approximate="tanh") - if activation == "gelu_approx" - else partial(F.relu, inplace=True) - ) - output1 = activation_fn(pre_act) - output2 = F.linear(output1, weight2, bias2) - return output2 if not return_residual else (output2, x) - - -class FusedMLP(nn.Module): - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - bias1=True, - bias2=True, - activation="gelu_approx", - return_residual=False, - checkpoint_lvl=0, - heuristic="auto", - device=None, - dtype=None, - ): - """ - If process_group is not None, we're doing Tensor Parallel with sequence parallelism: - we do an all_gather of x before doing the matmul, gelu, then matmul. - Finally we do a reduce_scatter of the output. - - checkpoint_lvl (increasing lvl means slower but more memory saving): - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute pre_act and gelu_out in the bwd - heuristic: - -1: don't fuse gemm + gelu (separate kernel) - 0..4: use this heuristic for the algo section in the fused gemm + gelu - 'auto': heuristic will be picked automatically: - For CUDA >= 11.8, we set heuristic=0 for both fp16 and bf16 for best perf. - For CUDA <= 11.7, we set heuristic=1 for fp16 and heuristic=-1 for bf16. - For H100, we set heuristic=-1 for both fp16 and bf16 as the fused cuBlasLt implementation - is slower than the unfused version. - return_residual: whether to return the input x along with the output. This is for - performance reason: for post-norm architecture, returning the input allows us - to fuse the backward of nn.Linear with the residual connection. - """ - assert checkpoint_lvl in [0, 1, 2] - assert activation in ["gelu_approx", "relu", "sqrelu"] - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features * 4 - self.activation = activation - self.return_residual = return_residual - self.checkpoint_lvl = checkpoint_lvl - self.heuristic = heuristic if activation != "sqrelu" else -1 - self.fc1 = nn.Linear(in_features, hidden_features, bias=bias1, **factory_kwargs) - self.fc2 = nn.Linear(hidden_features, out_features, bias=bias2, **factory_kwargs) - - def forward(self, x, process_group=None): - dtype = x.dtype if not torch.is_autocast_enabled() else torch.get_autocast_gpu_dtype() - if self.heuristic == "auto": - if self.activation == "gelu_approx": - if torch.cuda.get_device_capability("cuda") == (9, 0): - heuristic = -1 - else: - cuda_ver = tuple(map(int, torch.version.cuda.split("."))) - heuristic = 0 if cuda_ver >= (11, 8) else (1 if dtype == torch.float16 else -1) - else: - heuristic = 0 - else: - heuristic = self.heuristic - out = fused_mlp_func( - x, - self.fc1.weight, - self.fc2.weight, - self.fc1.bias, - self.fc2.bias, - activation=self.activation, - save_pre_act=self.training, - return_residual=self.return_residual, - checkpoint_lvl=self.checkpoint_lvl, - heuristic=heuristic, - process_group=process_group, - ) - if self.return_residual: - out, x = out - if process_group is not None: - out = reduce_scatter(out, process_group) - return out if not self.return_residual else (out, x) - - -class ParallelFusedMLP(nn.Module): - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - activation="gelu_approx", - process_group: ProcessGroup = None, - bias1=True, - bias2=True, - sequence_parallel=True, - checkpoint_lvl=0, - heuristic="auto", - device=None, - dtype=None, - ): - """ - process_group is required. We're doing Tensor Parallel with sequence parallelism: - we do an all_gather of x before doing the matmul, gelu, then matmul. - Finally we do a reduce_scatter of the output. - - checkpoint_lvl (increasing lvl means slower but more memory saving): - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute pre_act and gelu_out in the bwd - heuristic: - -1: don't fuse gemm + gelu (separate kernel) - 0..4: use this heuristic for the algo section in the fused gemm + gelu - 'auto': heuristic will be picked automatically: - For CUDA >= 11.8, we set heuristic=0 for both fp16 and bf16 for best perf. - For CUDA <= 11.7, we set heuristic=1 for fp16 and heuristic=-1 for bf16. - """ - assert checkpoint_lvl in [0, 1, 2] - assert activation in ["gelu_approx", "relu", "sqrelu"] - assert process_group is not None - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features * 4 - self.activation = activation - self.process_group = process_group - self.sequence_parallel = sequence_parallel - self.checkpoint_lvl = checkpoint_lvl - self.heuristic = heuristic if activation != "sqrelu" else -1 - self.fc1 = ColumnParallelLinear( - in_features, hidden_features, process_group, bias=bias1, **factory_kwargs - ) - self.fc2 = RowParallelLinear( - hidden_features, out_features, process_group, bias=bias2, **factory_kwargs - ) - - def forward(self, x): - dtype = x.dtype if not torch.is_autocast_enabled() else torch.get_autocast_gpu_dtype() - if self.heuristic == "auto": - if self.activation == "gelu_approx": - cuda_ver = tuple(map(int, torch.version.cuda.split("."))) - heuristic = 0 if cuda_ver >= (11, 8) else (1 if dtype == torch.float16 else -1) - else: - heuristic = 0 - else: - heuristic = self.heuristic - out = fused_mlp_func( - x, - self.fc1.weight, - self.fc2.weight, - self.fc1.bias, - self.fc2.bias, - activation=self.activation, - save_pre_act=self.training, - checkpoint_lvl=self.checkpoint_lvl, - heuristic=heuristic, - process_group=self.process_group, - sequence_parallel=self.sequence_parallel, - ) - reduce_fn = reduce_scatter if self.sequence_parallel else all_reduce - return reduce_fn(out, self.process_group) diff --git a/build/torch212-cxx11-cu126-x86_64-linux/ops/layer_norm.py b/build/torch212-cxx11-cu126-x86_64-linux/ops/layer_norm.py deleted file mode 100644 index 4b6cd798fd02844ef9cd3897f8ab95e490e638bf..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu126-x86_64-linux/ops/layer_norm.py +++ /dev/null @@ -1,800 +0,0 @@ -# Copyright (c) 2022, Tri Dao. -# Adapted from https://github.com/NVIDIA/apex/blob/master/apex/contrib/layer_norm/layer_norm.py - -import dropout_layer_norm -import torch -from torch.nn import init - - -def maybe_align(x, alignment_in_bytes=16): - """Assume that x already has last dim divisible by alignment_in_bytes""" - # TD [2023-07-04] I'm not 100% sure that clone will align the memory - # https://discuss.pytorch.org/t/how-to-ensure-that-tensor-data-ptr-is-aligned-to-16-bytes/183440 - return x if x.data_ptr() % alignment_in_bytes == 0 else x.clone() - - -def _dropout_add_layer_norm_forward( - x0, - residual, - gamma, - beta, - rowscale, - colscale, - dropout_p, - epsilon, - residual_in_fp32=False, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes""" - hidden_size = gamma.numel() - x0mat = x0.view((-1, hidden_size)) - residualmat = residual.view((-1, hidden_size)) if residual is not None else None - rowscale = rowscale.view(-1) if rowscale is not None else None - zmat, xmat, dmask, mu, rsigma = dropout_layer_norm.dropout_add_ln_fwd( - x0mat, - residualmat, - gamma, - beta, - rowscale, - colscale, - None, - None, - dropout_p, - epsilon, - 1.0, - 0, - None, - residual_in_fp32, - is_rms_norm, - ) - # dmask is None if dropout_p == 0.0 - # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype - return zmat, xmat if xmat is not None else x0mat, dmask, mu, rsigma - - -def _dropout_add_layer_norm_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - rowscale, - colscale, - dropout_p, - has_residual, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes - dx == None means that it was a post-norm architecture - (x = drop(x0) + residual was not returned in the fwd). - x0 must not be None if we have colscale. - """ - hidden_size = gamma.numel() - xmat = x.view((-1, hidden_size)) - dzmat = dz.view(xmat.shape) - dxmat = dx.view(xmat.shape) if dx is not None else None - x0mat = x0.view((-1, hidden_size)) if x0 is not None else None - rowscale = rowscale.view(-1) if rowscale is not None else None - if colscale is not None: - assert x0 is not None, "x0 is required to compute the gradient of colscale" - dx0mat, dresidualmat, dgamma, dbeta, _, _, *rest = dropout_layer_norm.dropout_add_ln_bwd( - dzmat, - dxmat, - xmat, - x0mat, - dmask, - mu, - rsigma, - gamma, - rowscale, - colscale, - None, - None, - dropout_p, - 1.0, - 0, - has_residual, - is_rms_norm, - ) - # dresidualmat is None if not has_residual - if colscale is None: - return dx0mat, dresidualmat, dgamma, dbeta - else: - dcolscale = rest[0] - return dx0mat, dresidualmat, dgamma, dbeta, dcolscale - - -def _dropout_add_layer_norm_subset_forward( - x0, - residual, - gamma, - beta, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32=False, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes""" - hidden_size = gamma.numel() - x0mat = x0.view((-1, hidden_size)) - residualmat = residual.view((-1, hidden_size)) if residual is not None else None - x0_subset = x0_subset.view(-1) if x0_subset is not None else None - out_subset = out_subset.view(-1) if out_subset is not None else None - zmat, xmat, dmask, mu, rsigma = dropout_layer_norm.dropout_add_ln_fwd( - x0mat, - residualmat, - gamma, - beta, - None, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - None, - residual_in_fp32, - is_rms_norm, - ) - # dmask is None if dropout_p == 0.0 - # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype - return zmat, xmat if xmat is not None else x0mat, dmask, mu, rsigma - - -def _dropout_add_layer_norm_subset_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - colscale, - x0_subset, - out_subset, - dropout_p, - rowscale_const, - x0_numrows, - has_residual, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes - dx == None means that it was a post-norm architecture - (x = drop(x0) + residual was not returned in the fwd). - x0 must not be None if we have colscale. - """ - hidden_size = gamma.numel() - xmat = x.view((-1, hidden_size)) - dzmat = dz.view(-1, hidden_size) - dxmat = dx.view(xmat.shape) if dx is not None else None - x0mat = x0.view((-1, hidden_size)) if x0 is not None else None - x0_subset = x0_subset.view(-1) if x0_subset is not None else None - out_subset = out_subset.view(-1) if out_subset is not None else None - if colscale is not None: - assert x0 is not None, "x0 is required to compute the gradient of colscale" - dx0mat, dresidualmat, dgamma, dbeta, _, _, *rest = dropout_layer_norm.dropout_add_ln_bwd( - dzmat, - dxmat, - xmat, - x0mat, - dmask, - mu, - rsigma, - gamma, - None, - colscale, - x0_subset, - out_subset, - dropout_p, - rowscale_const, - x0_numrows, - has_residual, - is_rms_norm, - ) - # dresidualmat is None if not has_residual - if colscale is None: - return dx0mat, dresidualmat, dgamma, dbeta - else: - dcolscale = rest[0] - return dx0mat, dresidualmat, dgamma, dbeta, dcolscale - - -def _dropout_add_layer_norm_parallel_residual_forward( - x0, - x1, - residual, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - residual_in_fp32=False, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes""" - hidden_size = gamma0.numel() - x0mat = x0.view((-1, hidden_size)) - x1mat = x1.view((-1, hidden_size)) if x1 is not None else None - residualmat = residual.view((-1, hidden_size)) if residual is not None else None - ( - z0mat, - z1mat, - xmat, - dmask0, - dmask1, - mu, - rsigma, - ) = dropout_layer_norm.dropout_add_ln_parallel_residual_fwd( - x0mat, - x1mat, - residualmat, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - None, - residual_in_fp32, - is_rms_norm, - ) - # dmask0 and dmask1 are None if dropout_p == 0.0 - # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype - return z0mat, z1mat, xmat if xmat is not None else x0mat, dmask0, dmask1, mu, rsigma - - -def _dropout_add_layer_norm_parallel_residual_backward( - dz0, - dz1, - dx, - x, - dmask0, - dmask1, - mu, - rsigma, - gamma0, - gamma1, - dropout_p, - has_x1, - has_residual, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes - dx == None means that it was a post-norm architecture - (x = drop(x0) + residual was not returned in the fwd). - """ - hidden_size = gamma0.numel() - xmat = x.view((-1, hidden_size)) - dz0mat = dz0.view(xmat.shape) - dz1mat = dz1.view(xmat.shape) if dz1 is not None else None - dxmat = dx.view(xmat.shape) if dx is not None else None - ( - dx0mat, - dx1mat, - dresidualmat, - dgamma0, - dbeta0, - dgamma1, - dbeta1, - *rest, - ) = dropout_layer_norm.dropout_add_ln_parallel_residual_bwd( - dz0mat, - dz1mat, - dxmat, - xmat, - dmask0, - dmask1, - mu, - rsigma, - gamma0, - gamma1, - dropout_p, - has_x1, - has_residual, - is_rms_norm, - ) - # dresidualmat is None if not has_residual - return dx0mat, dx1mat, dresidualmat, dgamma0, dbeta0, dgamma1, dbeta1 - - -class DropoutAddLayerNormFn(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x0, - residual, - gamma, - beta, - rowscale, - colscale, - dropout_p, - epsilon, - residual_in_fp32=False, - prenorm=False, - is_rms_norm=False, - return_dmask=False, - ): - x0 = maybe_align(x0.contiguous(), 16) - residual = maybe_align(residual.contiguous(), 16) if residual is not None else None - gamma = maybe_align(gamma.contiguous(), 16) - beta = maybe_align(beta.contiguous(), 16) if beta is not None else None - rowscale = maybe_align(rowscale.contiguous(), 16) if rowscale is not None else None - colscale = maybe_align(colscale.contiguous(), 16) if colscale is not None else None - zmat, xmat, dmask, mu, rsigma = _dropout_add_layer_norm_forward( - x0, - residual, - gamma, - beta, - rowscale, - colscale, - dropout_p, - epsilon, - residual_in_fp32, - is_rms_norm, - ) - # Only need to save x0 if we need to compute gradient wrt colscale - x0_saved = x0 if colscale is not None else None - ctx.save_for_backward( - xmat.view(x0.shape), x0_saved, dmask, gamma, mu, rsigma, rowscale, colscale - ) - ctx.prenorm = prenorm - ctx.dropout_p = dropout_p - ctx.has_residual = residual is not None - ctx.is_rms_norm = is_rms_norm - ctx.has_beta = beta is not None - if not return_dmask: - return ( - zmat.view(x0.shape) if not prenorm else (zmat.view(x0.shape), xmat.view(x0.shape)) - ) - else: - dmask = ( - dmask.view(x0.shape) - if dropout_p > 0.0 - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - ctx.mark_non_differentiable(dmask) - return ( - (zmat.view(x0.shape), dmask) - if not prenorm - else (zmat.view(x0.shape), xmat.view(x0.shape), dmask) - ) - - @staticmethod - def backward(ctx, dz, *args): - # assert dz.is_contiguous() - dz = maybe_align(dz.contiguous(), 16) # this happens! - dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None - x, x0, dmask, gamma, mu, rsigma, rowscale, colscale = ctx.saved_tensors - # x0 is None if colscale is None - dropout_p = ctx.dropout_p - has_residual = ctx.has_residual - dx0mat, dresidualmat, dgamma, dbeta, *rest = _dropout_add_layer_norm_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - rowscale, - colscale, - dropout_p, - has_residual, - ctx.is_rms_norm, - ) - dx0 = dx0mat.view(x.shape) - dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None - dcolscale = rest[0] if colscale is not None else None - return ( - dx0, - dresidual, - dgamma, - dbeta if ctx.has_beta else None, - None, - dcolscale, - None, - None, - None, - None, - None, - None, - ) - - -class DropoutAddLayerNormSubsetFn(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x0, - residual, - gamma, - beta, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32=False, - prenorm=False, - is_rms_norm=False, - return_dmask=False, - ): - x0 = maybe_align(x0.contiguous(), 16) - residual = maybe_align(residual.contiguous(), 16) if residual is not None else None - gamma = maybe_align(gamma.contiguous(), 16) - beta = maybe_align(beta.contiguous(), 16) if beta is not None else None - colscale = maybe_align(colscale.contiguous(), 16) if colscale is not None else None - zmat, xmat, dmask, mu, rsigma = _dropout_add_layer_norm_subset_forward( - x0, - residual, - gamma, - beta, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32, - is_rms_norm, - ) - # Only need to save x0 if we need to compute gradient wrt colscale - x0_saved = x0 if colscale is not None else None - x_shape = (-1, *x0.shape[1:]) - ctx.save_for_backward( - xmat.view(x_shape), x0_saved, dmask, gamma, mu, rsigma, colscale, x0_subset, out_subset - ) - ctx.prenorm = prenorm - ctx.dropout_p = dropout_p - ctx.rowscale_const = rowscale_const - ctx.x0_numrows = x0.shape[:-1].numel() - ctx.has_residual = residual is not None - ctx.is_rms_norm = is_rms_norm - ctx.has_beta = beta is not None - z_shape = (-1, *x0.shape[1:]) - if not return_dmask: - return zmat.view(z_shape) if not prenorm else (zmat.view(z_shape), xmat.view(x0.shape)) - else: - z = zmat.view(z_shape) - dmask = ( - dmask.view(x0.shape) - if dropout_p > 0.0 - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - ctx.mark_non_differentiable(dmask) - return (z, dmask) if not prenorm else (z, xmat.view(x_shape), dmask) - - @staticmethod - def backward(ctx, dz, *args): - # assert dz.is_contiguous() - dz = maybe_align(dz.contiguous(), 16) # this happens! - dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None - x, x0, dmask, gamma, mu, rsigma, colscale, x0_subset, out_subset = ctx.saved_tensors - # x0 is None if colscale is None - dropout_p = ctx.dropout_p - has_residual = ctx.has_residual - dx0mat, dresidualmat, dgamma, dbeta, *rest = _dropout_add_layer_norm_subset_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - colscale, - x0_subset, - out_subset, - dropout_p, - ctx.rowscale_const, - ctx.x0_numrows, - has_residual, - ctx.is_rms_norm, - ) - dx0 = dx0mat.view(-1, *x.shape[1:]) - dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None - dcolscale = rest[0] if colscale is not None else None - return ( - dx0, - dresidual, - dgamma, - dbeta if ctx.has_beta else None, - dcolscale, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) - - -class DropoutAddLayerNormParallelResidualFn(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x0, - x1, - residual, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - residual_in_fp32=False, - prenorm=False, - is_rms_norm=False, - return_dmask=False, - ): - x0 = maybe_align(x0.contiguous(), 16) - x1 = maybe_align(x1.contiguous(), 16) if x1 is not None else None - residual = maybe_align(residual.contiguous(), 16) if residual is not None else None - gamma0 = maybe_align(gamma0.contiguous(), 16) - beta0 = maybe_align(beta0.contiguous(), 16) if beta0 is not None else None - gamma1 = maybe_align(gamma1.contiguous(), 16) if gamma1 is not None else None - beta1 = maybe_align(beta1.contiguous(), 16) if beta1 is not None else None - ( - z0mat, - z1mat, - xmat, - dmask0, - dmask1, - mu, - rsigma, - ) = _dropout_add_layer_norm_parallel_residual_forward( - x0, - x1, - residual, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - residual_in_fp32, - is_rms_norm, - ) - ctx.save_for_backward(xmat.view(x0.shape), dmask0, dmask1, gamma0, gamma1, mu, rsigma) - ctx.prenorm = prenorm - ctx.dropout_p = dropout_p - ctx.has_x1 = x1 is not None - ctx.has_residual = residual is not None - ctx.is_rms_norm = is_rms_norm - ctx.has_beta = beta0 is not None - z = (z0mat.view(x0.shape), z1mat.view(x0.shape) if z1mat is not None else None) - if not return_dmask: - return z if not prenorm else (*z, xmat.view(x0.shape)) - else: - dmask0 = ( - dmask0.view(x0.shape) - if dropout_p > 0.0 - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - dmask1 = ( - dmask1.view(x0.shape) - if dropout_p > 0.0 and x1 is not None - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - ctx.mark_non_differentiable(dmask0) - ctx.mark_non_differentiable(dmask1) - return ( - (*z, dmask0, dmask1) if not prenorm else (*z, xmat.view(x0.shape), dmask0, dmask1) - ) - - @staticmethod - def backward(ctx, dz0, dz1, *args): - dz0 = maybe_align(dz0.contiguous(), 16) # this happens! - dz1 = maybe_align(dz1.contiguous(), 16) if dz1 is not None else None - dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None - x, dmask0, dmask1, gamma0, gamma1, mu, rsigma = ctx.saved_tensors - dropout_p = ctx.dropout_p - has_x1 = ctx.has_x1 - has_residual = ctx.has_residual - ( - dx0mat, - dx1mat, - dresidualmat, - dgamma0, - dbeta0, - dgamma1, - dbeta1, - ) = _dropout_add_layer_norm_parallel_residual_backward( - dz0, - dz1, - dx, - x, - dmask0, - dmask1, - mu, - rsigma, - gamma0, - gamma1, - dropout_p, - has_x1, - has_residual, - ctx.is_rms_norm, - ) - dx0 = dx0mat.view(x.shape) - dx1 = dx1mat.view(x.shape) if dx1mat is not None else None - dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None - return ( - dx0, - dx1, - dresidual, - dgamma0, - dbeta0 if ctx.has_beta else None, - dgamma1, - dbeta1 if ctx.has_beta else None, - None, - None, - None, - None, - None, - None, - ) - - -def layer_norm(x, weight, bias, epsilon): - return DropoutAddLayerNormFn.apply(x, None, weight, bias, None, None, 0.0, epsilon, False) - - -def dropout_add_layer_norm( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - rowscale=None, - layerscale=None, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormFn.apply( - x0, - residual, - weight, - bias, - rowscale, - layerscale, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - False, - return_dropout_mask, - ) - - -def dropout_add_layer_norm_subset( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - layerscale=None, - x0_subset=None, - out_subset=None, - rowscale_const=1.0, - out_numrows=0, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormSubsetFn.apply( - x0, - residual, - weight, - bias, - layerscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32, - prenorm, - False, - return_dropout_mask, - ) - - -def dropout_add_layer_norm_parallel_residual( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormParallelResidualFn.apply( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - False, - return_dropout_mask, - ) - - -class DropoutAddLayerNorm(torch.nn.Module): - def __init__( - self, - hidden_size, - prenorm=False, - p=0.0, - eps=1e-5, - residual_in_fp32=False, - device=None, - dtype=None, - ): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.prenorm = prenorm - self.p = p - self.eps = eps - self.residual_in_fp32 = residual_in_fp32 - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.bias = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.reset_parameters() - - def reset_parameters(self): - init.ones_(self.weight) - init.zeros_(self.bias) - - def forward(self, x0, residual=None): - return dropout_add_layer_norm( - x0, - residual, - self.weight, - self.bias, - self.p if self.training else 0.0, - self.eps, - prenorm=self.prenorm, - residual_in_fp32=self.residual_in_fp32, - ) diff --git a/build/torch212-cxx11-cu126-x86_64-linux/ops/rms_norm.py b/build/torch212-cxx11-cu126-x86_64-linux/ops/rms_norm.py deleted file mode 100644 index 068348d61290e3839dd082b540d898578ba1e8e2..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu126-x86_64-linux/ops/rms_norm.py +++ /dev/null @@ -1,174 +0,0 @@ -# Copyright (c) 2022, Tri Dao. -# Adapted from https://github.com/NVIDIA/apex/blob/master/apex/contrib/layer_norm/layer_norm.py - -import torch -from torch.nn import init - -from flash_attn.ops.layer_norm import ( - DropoutAddLayerNormFn, - DropoutAddLayerNormParallelResidualFn, - DropoutAddLayerNormSubsetFn, -) - - -def rms_norm(x, weight, epsilon): - return DropoutAddLayerNormFn.apply( - x, None, weight, None, None, None, 0.0, epsilon, False, False, True - ) - - -def dropout_add_rms_norm( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - rowscale=None, - layerscale=None, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormFn.apply( - x0, - residual, - weight, - bias, - rowscale, - layerscale, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - True, - return_dropout_mask, - ) - - -def dropout_add_rms_norm_subset( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - layerscale=None, - x0_subset=None, - out_subset=None, - rowscale_const=1.0, - out_numrows=0, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormSubsetFn.apply( - x0, - residual, - weight, - bias, - layerscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32, - prenorm, - True, - return_dropout_mask, - ) - - -def dropout_add_rms_norm_parallel_residual( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormParallelResidualFn.apply( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - True, - return_dropout_mask, - ) - - -class RMSNorm(torch.nn.Module): - def __init__(self, hidden_size, eps=1e-5, device=None, dtype=None): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.eps = eps - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.register_parameter("bias", None) - self.reset_parameters() - - def reset_parameters(self): - init.ones_(self.weight) - - def forward(self, x): - return rms_norm(x, self.weight, self.eps) - - -class DropoutAddRMSNorm(torch.nn.Module): - def __init__( - self, - hidden_size, - prenorm=False, - p=0.0, - eps=1e-5, - residual_in_fp32=False, - device=None, - dtype=None, - ): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.prenorm = prenorm - self.p = p - self.eps = eps - self.residual_in_fp32 = residual_in_fp32 - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.register_parameter("bias", None) - self.reset_parameters() - - def reset_parameters(self): - init.ones_(self.weight) - - def forward(self, x0, residual=None): - return dropout_add_rms_norm( - x0, - residual, - self.weight, - None, - self.p if self.training else 0.0, - self.eps, - prenorm=self.prenorm, - residual_in_fp32=self.residual_in_fp32, - ) diff --git a/build/torch212-cxx11-cu126-x86_64-linux/ops/triton/__init__.py b/build/torch212-cxx11-cu126-x86_64-linux/ops/triton/__init__.py deleted file mode 100644 index 8b137891791fe96927ad78e64b0aad7bded08bdc..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu126-x86_64-linux/ops/triton/__init__.py +++ /dev/null @@ -1 +0,0 @@ - diff --git a/build/torch212-cxx11-cu126-x86_64-linux/ops/triton/cross_entropy.py b/build/torch212-cxx11-cu126-x86_64-linux/ops/triton/cross_entropy.py deleted file mode 100644 index 1b5a415b73f236f3e05fb14b9141959559e18526..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu126-x86_64-linux/ops/triton/cross_entropy.py +++ /dev/null @@ -1,330 +0,0 @@ -# Copyright (c) 2023, Tri Dao. - -from typing import Tuple, Optional, Union - -import torch -import torch.nn.functional as F - -import triton -import triton.language as tl - -# `all_gather_into_tensor` and `reduce_scatter_tensor` are new placeholders for -# `_all_gather_base` and `_reduce_scatter_base`. They require the most recent -# version of PyTorch. The following 2 lines are for backward compatibility with -# older PyTorch. -if "all_gather_into_tensor" not in dir(torch.distributed): - torch.distributed.all_gather_into_tensor = torch.distributed._all_gather_base - - -@triton.heuristics( - { - "HAS_SMOOTHING": lambda args: args["smoothing"] > 0.0, - } -) -@triton.jit -def cross_entropy_fwd_kernel( - loss_ptr, # data ptrs - lse_ptr, - z_loss_ptr, - logits_ptr, - labels_ptr, - smoothing, - logit_scale, - lse_square_scale, - ignore_index, - total_classes, - class_start_idx, # Useful for tensor parallel when each rank only has a subset of classes - n_cols, # shapes - logits_row_stride, # strides - BLOCK_SIZE: tl.constexpr, - HAS_SMOOTHING: tl.constexpr, - # if SPLIT (e.g. tensor parallel), don't include the LSE in the loss since it's not the final LSE - SPLIT: tl.constexpr, - PRECOMPUTED_LSE: tl.constexpr, # If LSE is already computed (also no smoothing and logit_scale == 1.0) -): - row_idx = tl.program_id(0) - logits_ptr = logits_ptr + row_idx * logits_row_stride.to(tl.int64) - sum_logits = 0.0 # For smoothing - if not PRECOMPUTED_LSE: - # Statistics for online softmax - m_i = -float("inf") - l_i = 0.0 - for col_offset in range(0, n_cols, BLOCK_SIZE): - cols = col_offset + tl.arange(0, BLOCK_SIZE) - logits = tl.load(logits_ptr + cols, mask=cols < n_cols, other=-float("inf")).to( - tl.float32 - ) * logit_scale - if HAS_SMOOTHING: - sum_logits += tl.sum(tl.where(cols < n_cols, logits, 0.0)) - m_i_new = tl.maximum(m_i, tl.max(logits)) - l_i = tl.exp(m_i - m_i_new) * l_i + tl.sum(tl.exp(logits - m_i_new)) - m_i = m_i_new - lse = tl.log(l_i) + m_i - tl.store(lse_ptr + row_idx, lse) - else: - lse = tl.load(lse_ptr + row_idx) - label_idx = tl.load(labels_ptr + row_idx) - if label_idx == ignore_index: - loss = 0.0 - z_loss = 0.0 - else: - label_idx -= class_start_idx - if label_idx >= 0 and label_idx < n_cols: - logits_label = tl.load(logits_ptr + label_idx) * logit_scale - if HAS_SMOOTHING: - loss = ( - (lse if not SPLIT else 0.0) - - smoothing * sum_logits / total_classes - - (1 - smoothing) * logits_label - ) - else: - loss = (lse if not SPLIT else 0.0) - logits_label - else: - # If label is out of bounds, we set the CE loss to 0.0. But we still want the smoothing loss - if HAS_SMOOTHING: - loss = smoothing * ((lse if not SPLIT else 0.0) - sum_logits / total_classes) - else: - loss = 0.0 - if not SPLIT: - z_loss = lse_square_scale * lse * lse - loss += z_loss - else: - z_loss = 0.0 - tl.store(loss_ptr + row_idx, loss) - if not SPLIT: - tl.store(z_loss_ptr + row_idx, z_loss) - - -@triton.heuristics( - { - "HAS_SMOOTHING": lambda args: args["smoothing"] > 0.0, - } -) -@triton.jit -def cross_entropy_bwd_kernel( - dlogits_ptr, # data ptrs - dloss_ptr, - logits_ptr, - lse_ptr, - labels_ptr, - smoothing, - logit_scale, - lse_square_scale, - ignore_index, - total_classes, - class_start_idx, # Useful for tensor parallel when each rank only has a subset of classes - n_cols, # shapes - logits_row_stride, # strides - dlogits_row_stride, - dloss_row_stride, - BLOCK_SIZE: tl.constexpr, - HAS_SMOOTHING: tl.constexpr, -): - row_idx = tl.program_id(0) - col_block_idx = tl.program_id(1) - logits_ptr = logits_ptr + row_idx * logits_row_stride.to(tl.int64) - dlogits_ptr = dlogits_ptr + row_idx * dlogits_row_stride.to(tl.int64) - col_offsets = col_block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - label_idx = tl.load(labels_ptr + row_idx) - if label_idx != ignore_index: - dloss = tl.load(dloss_ptr + row_idx * dloss_row_stride) - else: - dloss = 0.0 - logits = tl.load(logits_ptr + col_offsets, mask=col_offsets < n_cols, other=-float("inf")).to( - tl.float32 - ) * logit_scale - lse = tl.load(lse_ptr + row_idx) - probs = tl.exp(logits - lse) - probs += 2.0 * lse_square_scale * lse * probs - label_idx -= class_start_idx - if HAS_SMOOTHING: - smooth_positive = 1.0 - smoothing - smooth_negative = smoothing / total_classes - probs = tl.where(col_offsets == label_idx, probs - smooth_positive, probs) - smooth_negative - else: - probs = tl.where(col_offsets == label_idx, probs - 1.0, probs) - tl.store(dlogits_ptr + col_offsets, (dloss * logit_scale) * probs, mask=col_offsets < n_cols) - - -class CrossEntropyLoss(torch.autograd.Function): - - @staticmethod - def forward( - ctx, - logits, - labels, - precomputed_lse=None, - smoothing=0.0, - logit_scale=1.0, - lse_square_scale=0.0, - ignore_index=-100, - inplace_backward=False, - process_group=None, - ): - # For some reason Triton generates wrong code when labels has dtype long and its address - # is not aligned to 16 bytes. The ld.global.b64 seems to load the wrong label index. - if labels.dtype == torch.long and labels.data_ptr() % 16 != 0: - labels = F.pad(labels, (0, 1))[..., :-1] - assert labels.data_ptr() % 16 == 0 - assert logit_scale > 0.0 - n_rows, n_cols = logits.shape - assert labels.shape == (n_rows,) - world_size = 1 if process_group is None else torch.distributed.get_world_size(process_group) - total_classes = world_size * n_cols - rank = 0 if process_group is None else torch.distributed.get_rank(process_group) - class_start_idx = rank * n_cols - use_precomputed_lse = precomputed_lse is not None and logit_scale == 1.0 and smoothing == 0.0 - - if logits.stride(-1) != 1: - logits = logits.contiguous() - MAX_BLOCK_SIZE = 16 * 1024 - BLOCK_SIZE = min(triton.next_power_of_2(n_cols), MAX_BLOCK_SIZE) - num_warps = ( - 4 - if BLOCK_SIZE < 2048 - else (8 if BLOCK_SIZE < 8192 else (16 if BLOCK_SIZE < 128 * 1024 else 32)) - ) - losses = torch.empty(n_rows, dtype=torch.float, device=logits.device) - if use_precomputed_lse: - assert precomputed_lse.shape == (n_rows,) - lse = precomputed_lse.contiguous() - else: - lse = torch.empty(n_rows, dtype=torch.float, device=logits.device) - z_losses = torch.empty(n_rows, dtype=torch.float, device=logits.device) - # Need this, otherwise Triton tries to launch from cuda:0 and we get - # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?) - with torch.cuda.device(logits.device.index): - cross_entropy_fwd_kernel[(n_rows,)]( - losses, # data ptrs - lse, - z_losses, - logits, - labels, - smoothing, - logit_scale, - lse_square_scale, - ignore_index, - total_classes, - class_start_idx, - n_cols, # shapes - logits.stride(0), # strides - BLOCK_SIZE=BLOCK_SIZE, # constants - SPLIT=world_size > 1, - PRECOMPUTED_LSE=use_precomputed_lse, - num_warps=num_warps, - ) - - if world_size > 1: - # If there's no smoothing, if labels are in the vocab of this partition, losses contains - # - predicted logit, and 0 otherwise. - # If there's smoothing=0.1, for labels in the vocab of this partition, losses contains - # -0.9 * predicted logit - 0.1 * sum logit / total_classes. - # For labels not in the vocab of this partition, losses contains - # -0.1 * sum logit / total_classes. - if world_size > 1: - lse_allgather = torch.empty(world_size, n_rows, dtype=lse.dtype, device=lse.device) - torch.distributed.all_gather_into_tensor(lse_allgather, lse, group=process_group) - handle_losses = torch.distributed.all_reduce( - losses, op=torch.distributed.ReduceOp.SUM, group=process_group, async_op=True - ) - lse = torch.logsumexp(lse_allgather, dim=0) - handle_losses.wait() - # After the allreduce, if there's no smoothing, the total losses are - predicted_logit, - # we just have to add the (global) lse. - # If there's smoothing=0.1, the total losses are - # -0.9 * predicted_logit - 0.1 * sum logit / total_classes. - # Again, we just have to add the (global) lse. - losses += lse - if lse_square_scale != 0.0: - z_losses = lse_square_scale * lse.square() - z_losses.masked_fill_(labels == ignore_index, 0.0) - losses += z_losses - else: - z_losses = torch.zeros_like(losses) - losses.masked_fill_(labels == ignore_index, 0.0) - - ctx.save_for_backward(logits, lse, labels) - ctx.mark_non_differentiable(z_losses) - ctx.smoothing = smoothing - ctx.logit_scale = logit_scale - ctx.lse_square_scale = lse_square_scale - ctx.ignore_index = ignore_index - ctx.total_classes = total_classes - ctx.class_start_idx = class_start_idx - ctx.inplace_backward = inplace_backward - return losses, z_losses - - @staticmethod - def backward(ctx, grad_losses, grad_z_losses): - del grad_z_losses # z_losses are only for logging. - - logits, lse, labels = ctx.saved_tensors - dlogits = logits if ctx.inplace_backward else torch.empty_like(logits) - n_rows, n_cols = logits.shape - BLOCK_SIZE = min(triton.next_power_of_2(n_cols), 4 * 1024) - num_warps = 4 if BLOCK_SIZE < 2048 else (8 if BLOCK_SIZE < 8192 else 16) - grid = lambda META: (n_rows, triton.cdiv(n_cols, META["BLOCK_SIZE"])) # noqa - # Need this, otherwise Triton tries to launch from cuda:0 and we get - # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?) - with torch.cuda.device(logits.device.index): - cross_entropy_bwd_kernel[grid]( - dlogits, # data ptrs - grad_losses, - logits, - lse, - labels, - ctx.smoothing, - ctx.logit_scale, - ctx.lse_square_scale, - ctx.ignore_index, - ctx.total_classes, - ctx.class_start_idx, - n_cols, # shapes - logits.stride(0), # strides - dlogits.stride(0), - grad_losses.stride(0), - BLOCK_SIZE=BLOCK_SIZE, # constants - num_warps=num_warps, - ) - return dlogits, None, None, None, None, None, None, None, None, None - - -def cross_entropy_loss( - logits: torch.Tensor, - labels: torch.Tensor, - precomputed_lse: Optional[torch.Tensor] = None, - label_smoothing: float = 0.0, - logit_scale: float = 1.0, - lse_square_scale: float = 0.0, - ignore_index=-100, - inplace_backward: bool = False, - process_group=None, -) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Arguments: - logits: (batch, vocab_size) - labels: (batch,) - label_smoothing: float - logit_scale: float. Multiply logits by this scale before calculating the loss. - lse_square_scale: float. If > 0, we add lse_square_scale * lse(logits) ^ 2 to the loss. - This is also referred to as "z-loss". - ignore_index: int. If labels == ignore_index, the loss is set to 0.0. - inplace_backward: bool. If True, we do the backward pass in-place by modifying the logits. - This saves memory. - process_group: if not None, we're doing Tensor Parallel: each process is responsible for - one part of the vocab. The loss will be aggregated across processes. - Returns: - losses: (batch,), float - z_losses: (batch,), float - """ - return CrossEntropyLoss.apply( - logits, - labels, - precomputed_lse, - label_smoothing, - logit_scale, - lse_square_scale, - ignore_index, - inplace_backward, - process_group, - ) diff --git a/build/torch212-cxx11-cu126-x86_64-linux/ops/triton/k_activations.py b/build/torch212-cxx11-cu126-x86_64-linux/ops/triton/k_activations.py deleted file mode 100644 index efb83c358eb4a85d069ee340a3c83f418f9a805b..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu126-x86_64-linux/ops/triton/k_activations.py +++ /dev/null @@ -1,162 +0,0 @@ -# Adapted from https://github.com/facebookresearch/xformers/blob/main/xformers/triton/k_activations.py -# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved. -# -# This source code is licensed under the BSD license found in the -# LICENSE file in the root directory of this source tree. - -import math -from enum import Enum -from typing import Optional - -import triton -import triton.language as tl - -_sqrt2pi = math.sqrt(2.0 / math.pi) -_sqrt1_2 = math.sqrt(1.0 / 2) -_gaussian_pdf_normalization = 1.0 / math.sqrt(2 * math.pi) - - -class Activation(str, Enum): - SquaredReLU = "squared_relu" - GeLU = "gelu" - GeLUApprox = "gelu_approx" - LeakyReLU = "leaky_relu" - ReLU = "relu" - - -def get_triton_activation_kernel(activation: Optional[Activation]): - return ( - { - Activation.ReLU: relu, - Activation.LeakyReLU: leaky_relu, - Activation.GeLU: gelu, - Activation.GeLUApprox: gelu_approx, - Activation.SquaredReLU: squared_relu, - }[activation] - if activation - else None - ) - - -def get_triton_activation_bwd_kernel(activation: Optional[Activation]): - return ( - { - Activation.ReLU: relu_grad, - Activation.LeakyReLU: leaky_relu_grad, - Activation.GeLU: gelu_grad, - Activation.GeLUApprox: gelu_approx_grad, - Activation.SquaredReLU: squared_relu_grad, - }[activation] - if activation - else None - ) - - -@triton.jit -def tanh(x): - # Tanh is just a scaled sigmoid - return 2 * tl.sigmoid(2 * x) - 1 - - -@triton.jit -def cosh(x): - exp_x = tl.exp(x) - return (exp_x + 1.0 / exp_x) * 0.5 - - -# a Triton implementation of the most used activations -# See for instance http://arxiv.org/abs/1606.08415 for an overview - -# ReLU -@triton.jit -def relu(x): - """ - ReLU_ activation function - - .. _ReLU: https://pytorch.org/docs/stable/generated/torch.nn.ReLU.html - """ - zero = 0.0 - return tl.where(x >= 0, x, zero.to(x.dtype)) - - -@triton.jit -def relu_grad(x): - # ReLU is different from other activations - # in that it does not require the input to retrospectively compute its gradient - # here the input is the downstream gradient, and we return the upstream gradient directly - zero = 0.0 - one = 1.0 - return tl.where(x >= 0, one.to(x.dtype), zero.to(x.dtype)) - - -@triton.jit -def squared_relu(x): - """ - Squared ReLU activation, as proposed in the Primer_ paper. - - .. _Primer: https://arxiv.org/abs/2109.08668 - """ - x_ = relu(x) - return (x_ * x_).to(x.dtype) - - -@triton.jit -def squared_relu_grad(x): - return tl.where(x >= 0, 2.0 * x, 0.0) - - -# Leaky ReLU -@triton.jit -def leaky_relu(x): - """ - LeakyReLU_ activation - - .. _LeakyReLU: https://pytorch.org/docs/stable/generated/torch.nn.LeakyReLU.html - """ - scale = 0.01 + 0.0 - scale = scale.to(x.dtype) - return tl.where(x >= 0, x, scale * x) - - -@triton.jit -def leaky_relu_grad(x): - min_grad = 0.01 - max_grad = 1 - - min_grad = min_grad.to(x.dtype) - max_grad = max_grad.to(x.dtype) - - return tl.where(x >= 0, max_grad, min_grad) - - -@triton.jit -def gelu(x): - """Gaussian Error Linear Unit (GELU)""" - return x * 0.5 * (1.0 + tl.libdevice.erf(x * _sqrt1_2)) - - -@triton.jit -def gelu_grad(x): - cdf = 0.5 * (1.0 + tl.libdevice.erf(x * _sqrt1_2)) - pdf = tl.exp(-0.5 * x * x) * _gaussian_pdf_normalization - return cdf + x * pdf - - -@triton.jit -def gelu_approx(x): - """ - GeLU_ activation - Gaussian error linear unit, with tanh approximation - - .. _GeLU: https://arxiv.org/pdf/1606.08415.pdf - """ - return 0.5 * x * (1.0 + tanh(_sqrt2pi * x * (1.0 + 0.044715 * x * x))) - - -@triton.jit -def gelu_approx_grad(x): - # CREDITS: Fast implementation proposed in - # https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/fused_bias_gelu.py#L30 - tanh_out = tanh(0.79788456 * x * (1 + 0.044715 * x * x)) - return 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( - 1 + tanh_out - ) diff --git a/build/torch212-cxx11-cu126-x86_64-linux/ops/triton/layer_norm.py b/build/torch212-cxx11-cu126-x86_64-linux/ops/triton/layer_norm.py deleted file mode 100644 index 5e6f0b3697b612e6a19031bbc55ac76ad233dc85..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu126-x86_64-linux/ops/triton/layer_norm.py +++ /dev/null @@ -1,1254 +0,0 @@ -# Copyright (c) 2024, Tri Dao. -# Implement dropout + residual + layer_norm / rms_norm. - -# Based on the Triton LayerNorm tutorial: https://triton-lang.org/main/getting-started/tutorials/05-layer-norm.html -# For the backward pass, we keep weight_grad and bias_grad in registers and accumulate. -# This is faster for dimensions up to 8k, but after that it's much slower due to register spilling. -# The models we train have hidden dim up to 8k anyway (e.g. Llama 70B), so this is fine. - -import math -from typing import Optional, List - -import torch -import torch.nn.functional as F -from torch import Tensor - -import triton -import triton.language as tl - -from flash_attn.utils.torch import custom_fwd, custom_bwd -from flash_attn.utils.library import triton_op - -from ..._ops import add_op_namespace_prefix - - -def maybe_contiguous_lastdim(x): - return x.contiguous() if x is not None and x.stride(-1) != 1 else x - - -def maybe_contiguous(x): - return x.contiguous() if x is not None else None - - -def triton_autotune_configs(): - # Return configs with a valid warp count for the current device - configs = [] - # Maximum threads per block is architecture-dependent in theory, but in reality all are 1024 - max_threads_per_block = 1024 - # Default to warp size 32 if not defined by device - warp_size = getattr(torch.cuda.get_device_properties(torch.cuda.current_device()), "warp_size", 32) - # Autotune for warp counts which are powers of 2 and do not exceed thread per block limit - return [triton.Config({}, num_warps=warp_count) for warp_count in [1, 2, 4, 8, 16, 32] - if warp_count * warp_size <= max_threads_per_block] - # return [triton.Config({}, num_warps=8)] - - -def layer_norm_ref( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - zero_centered_weight=False, - dropout_mask=None, - dropout_mask1=None, - upcast=False, -): - dtype = x.dtype - if upcast: - x = x.float() - weight = weight.float() - bias = bias.float() if bias is not None else None - residual = residual.float() if residual is not None else residual - x1 = x1.float() if x1 is not None else None - weight1 = weight1.float() if weight1 is not None else None - bias1 = bias1.float() if bias1 is not None else None - if zero_centered_weight: - weight = weight + 1.0 - if weight1 is not None: - weight1 = weight1 + 1.0 - if x1 is not None: - assert rowscale is None, "rowscale is not supported with parallel LayerNorm" - if rowscale is not None: - x = x * rowscale[..., None] - if dropout_p > 0.0: - if dropout_mask is not None: - x = x.masked_fill(~dropout_mask, 0.0) / (1.0 - dropout_p) - else: - x = F.dropout(x, p=dropout_p) - if x1 is not None: - if dropout_mask1 is not None: - x1 = x1.masked_fill(~dropout_mask1, 0.0) / (1.0 - dropout_p) - else: - x1 = F.dropout(x1, p=dropout_p) - if x1 is not None: - x = x + x1 - if residual is not None: - x = (x + residual).to(x.dtype) - out = F.layer_norm(x.to(weight.dtype), x.shape[-1:], weight=weight, bias=bias, eps=eps).to( - dtype - ) - if weight1 is None: - return out if not prenorm else (out, x) - else: - out1 = F.layer_norm( - x.to(weight1.dtype), x.shape[-1:], weight=weight1, bias=bias1, eps=eps - ).to(dtype) - return (out, out1) if not prenorm else (out, out1, x) - - -def rms_norm_ref( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - zero_centered_weight=False, - dropout_mask=None, - dropout_mask1=None, - upcast=False, -): - dtype = x.dtype - if upcast: - x = x.float() - weight = weight.float() - bias = bias.float() if bias is not None else None - residual = residual.float() if residual is not None else residual - x1 = x1.float() if x1 is not None else None - weight1 = weight1.float() if weight1 is not None else None - bias1 = bias1.float() if bias1 is not None else None - if zero_centered_weight: - weight = weight + 1.0 - if weight1 is not None: - weight1 = weight1 + 1.0 - if x1 is not None: - assert rowscale is None, "rowscale is not supported with parallel LayerNorm" - if rowscale is not None: - x = x * rowscale[..., None] - if dropout_p > 0.0: - if dropout_mask is not None: - x = x.masked_fill(~dropout_mask, 0.0) / (1.0 - dropout_p) - else: - x = F.dropout(x, p=dropout_p) - if x1 is not None: - if dropout_mask1 is not None: - x1 = x1.masked_fill(~dropout_mask1, 0.0) / (1.0 - dropout_p) - else: - x1 = F.dropout(x1, p=dropout_p) - if x1 is not None: - x = x + x1 - if residual is not None: - x = (x + residual).to(x.dtype) - rstd = 1 / torch.sqrt((x.square()).mean(dim=-1, keepdim=True) + eps) - out = ((x * rstd * weight) + bias if bias is not None else (x * rstd * weight)).to(dtype) - if weight1 is None: - return out if not prenorm else (out, x) - else: - out1 = ((x * rstd * weight1) + bias1 if bias1 is not None else (x * rstd * weight1)).to( - dtype - ) - return (out, out1) if not prenorm else (out, out1, x) - - -@triton.autotune( - configs=triton_autotune_configs(), - key=["N", "HAS_RESIDUAL", "STORE_RESIDUAL_OUT", "IS_RMS_NORM", "HAS_BIAS", "HAS_X1", "HAS_W1", "HAS_B1"], -) -# torch compile doesn't like triton.heuristics, so we set these manually when calling the kernel -# @triton.heuristics({"HAS_BIAS": lambda args: args["B"] is not None}) -# @triton.heuristics({"HAS_RESIDUAL": lambda args: args["RESIDUAL"] is not None}) -# @triton.heuristics({"HAS_X1": lambda args: args["X1"] is not None}) -# @triton.heuristics({"HAS_W1": lambda args: args["W1"] is not None}) -# @triton.heuristics({"HAS_B1": lambda args: args["B1"] is not None}) -@triton.jit -def _layer_norm_fwd_1pass_kernel( - X, # pointer to the input - Y, # pointer to the output - W, # pointer to the weights - B, # pointer to the biases - RESIDUAL, # pointer to the residual - X1, - W1, - B1, - Y1, - RESIDUAL_OUT, # pointer to the residual - ROWSCALE, - SEEDS, # Dropout seeds for each row - DROPOUT_MASK, - DROPOUT_MASK1, - Mean, # pointer to the mean - Rstd, # pointer to the 1/std - stride_x_row, # how much to increase the pointer when moving by 1 row - stride_y_row, - stride_res_row, - stride_res_out_row, - stride_x1_row, - stride_y1_row, - M, # number of rows in X - N, # number of columns in X - eps, # epsilon to avoid division by zero - dropout_p, # Dropout probability - zero_centered_weight, # If true, add 1.0 to the weight - IS_RMS_NORM: tl.constexpr, - BLOCK_N: tl.constexpr, - HAS_RESIDUAL: tl.constexpr, - STORE_RESIDUAL_OUT: tl.constexpr, - HAS_BIAS: tl.constexpr, - HAS_DROPOUT: tl.constexpr, - STORE_DROPOUT_MASK: tl.constexpr, - HAS_ROWSCALE: tl.constexpr, - HAS_X1: tl.constexpr, - HAS_W1: tl.constexpr, - HAS_B1: tl.constexpr, -): - # Map the program id to the row of X and Y it should compute. - row = tl.program_id(0) - X += row * stride_x_row - Y += row * stride_y_row - if HAS_RESIDUAL: - RESIDUAL += row * stride_res_row - if STORE_RESIDUAL_OUT: - RESIDUAL_OUT += row * stride_res_out_row - if HAS_X1: - X1 += row * stride_x1_row - if HAS_W1: - Y1 += row * stride_y1_row - # Compute mean and variance - cols = tl.arange(0, BLOCK_N) - x = tl.load(X + cols, mask=cols < N, other=0.0).to(tl.float32) - if HAS_ROWSCALE: - rowscale = tl.load(ROWSCALE + row).to(tl.float32) - x *= rowscale - if HAS_DROPOUT: - # Compute dropout mask - # 7 rounds is good enough, and reduces register pressure - keep_mask = tl.rand(tl.load(SEEDS + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - x = tl.where(keep_mask, x / (1.0 - dropout_p), 0.0) - if STORE_DROPOUT_MASK: - tl.store(DROPOUT_MASK + row * N + cols, keep_mask, mask=cols < N) - if HAS_X1: - x1 = tl.load(X1 + cols, mask=cols < N, other=0.0).to(tl.float32) - if HAS_ROWSCALE: - rowscale = tl.load(ROWSCALE + M + row).to(tl.float32) - x1 *= rowscale - if HAS_DROPOUT: - # Compute dropout mask - # 7 rounds is good enough, and reduces register pressure - keep_mask = ( - tl.rand(tl.load(SEEDS + M + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - ) - x1 = tl.where(keep_mask, x1 / (1.0 - dropout_p), 0.0) - if STORE_DROPOUT_MASK: - tl.store(DROPOUT_MASK1 + row * N + cols, keep_mask, mask=cols < N) - x += x1 - if HAS_RESIDUAL: - residual = tl.load(RESIDUAL + cols, mask=cols < N, other=0.0).to(tl.float32) - x += residual - if STORE_RESIDUAL_OUT: - tl.store(RESIDUAL_OUT + cols, x, mask=cols < N) - if not IS_RMS_NORM: - mean = tl.sum(x, axis=0) / N - tl.store(Mean + row, mean) - xbar = tl.where(cols < N, x - mean, 0.0) - var = tl.sum(xbar * xbar, axis=0) / N - else: - xbar = tl.where(cols < N, x, 0.0) - var = tl.sum(xbar * xbar, axis=0) / N - rstd = 1 / tl.sqrt(var + eps) - tl.store(Rstd + row, rstd) - # Normalize and apply linear transformation - mask = cols < N - w = tl.load(W + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w += 1.0 - if HAS_BIAS: - b = tl.load(B + cols, mask=mask).to(tl.float32) - x_hat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd - y = x_hat * w + b if HAS_BIAS else x_hat * w - # Write output - tl.store(Y + cols, y, mask=mask) - if HAS_W1: - w1 = tl.load(W1 + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w1 += 1.0 - if HAS_B1: - b1 = tl.load(B1 + cols, mask=mask).to(tl.float32) - y1 = x_hat * w1 + b1 if HAS_B1 else x_hat * w1 - tl.store(Y1 + cols, y1, mask=mask) - - -def _layer_norm_fwd( - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - residual: Optional[Tensor] = None, - x1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - out_dtype: Optional[torch.dtype] = None, - residual_dtype: Optional[torch.dtype] = None, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - return_dropout_mask: bool = False, - out: Optional[Tensor] = None, - residual_out: Optional[Tensor] = None -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - # Need to wrap to handle the case where residual_out is a alias of x, which makes torch.library - # and torch.compile unhappy. Also allocate memory for out and residual_out if they are None - # so that _layer_norm_fwd_impl doesn't have to return them. - if out is None: - out = torch.empty_like(x, dtype=x.dtype if out_dtype is None else out_dtype) - if residual is not None: - residual_dtype = residual.dtype - if residual_out is None and ( - residual is not None - or (residual_dtype is not None and residual_dtype != x.dtype) - or dropout_p > 0.0 - or rowscale is not None - or x1 is not None - ): - residual_out = torch.empty_like( - x, dtype=residual_dtype if residual_dtype is not None else x.dtype - ) - else: - residual_out = None - y1, mean, rstd, seeds, dropout_mask, dropout_mask1 = _layer_norm_fwd_impl( - x, - weight, - bias, - eps, - out, - residual=residual, - x1=x1, - weight1=weight1, - bias1=bias1, - dropout_p=dropout_p, - rowscale=rowscale, - zero_centered_weight=zero_centered_weight, - is_rms_norm=is_rms_norm, - return_dropout_mask=return_dropout_mask, - residual_out=residual_out, - ) - # residual_out is None if residual is None and residual_dtype == input_dtype and dropout_p == 0.0 - if residual_out is None: - residual_out = x - return out, y1, mean, rstd, residual_out, seeds, dropout_mask, dropout_mask1 - - -# [2025-04-28] torch.library.triton_op ignores the schema argument, but here we need the schema -# since we're returning a tuple of tensors -@triton_op(add_op_namespace_prefix("layer_norm_fwd_impl"), mutates_args={"out", "residual_out"}, - schema="(Tensor x, Tensor weight, Tensor bias, float eps, Tensor(a!) out, Tensor? residual, Tensor? x1, Tensor? weight1, Tensor? bias1, float dropout_p, Tensor? rowscale, bool zero_centered_weight, bool is_rms_norm, bool return_dropout_mask, Tensor(a!)? residual_out) -> (Tensor y1, Tensor mean, Tensor rstd, Tensor seeds, Tensor dropout_mask, Tensor dropout_mask1)") -def _layer_norm_fwd_impl( - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - out: Tensor, - residual: Optional[Tensor] = None, - x1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - return_dropout_mask: bool = False, - residual_out: Optional[Tensor] = None -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - M, N = x.shape - assert x.stride(-1) == 1 - if residual is not None: - assert residual.stride(-1) == 1 - assert residual.shape == (M, N) - assert weight.shape == (N,) - assert weight.stride(-1) == 1 - if bias is not None: - assert bias.stride(-1) == 1 - assert bias.shape == (N,) - if x1 is not None: - assert x1.shape == x.shape - assert rowscale is None - assert x1.stride(-1) == 1 - if weight1 is not None: - assert weight1.shape == (N,) - assert weight1.stride(-1) == 1 - if bias1 is not None: - assert bias1.shape == (N,) - assert bias1.stride(-1) == 1 - if rowscale is not None: - assert rowscale.is_contiguous() - assert rowscale.shape == (M,) - assert out.shape == x.shape - assert out.stride(-1) == 1 - if residual_out is not None: - assert residual_out.shape == x.shape - assert residual_out.stride(-1) == 1 - if weight1 is not None: - y1 = torch.empty_like(out) - assert y1.stride(-1) == 1 - else: - y1 = None - mean = torch.empty((M,), dtype=torch.float32, device=x.device) if not is_rms_norm else None - rstd = torch.empty((M,), dtype=torch.float32, device=x.device) - if dropout_p > 0.0: - seeds = torch.randint( - 2**32, (M if x1 is None else 2 * M,), device=x.device, dtype=torch.int64 - ) - else: - seeds = None - if return_dropout_mask and dropout_p > 0.0: - dropout_mask = torch.empty(M, N, device=x.device, dtype=torch.bool) - if x1 is not None: - dropout_mask1 = torch.empty(M, N, device=x.device, dtype=torch.bool) - else: - dropout_mask1 = None - else: - dropout_mask, dropout_mask1 = None, None - # Less than 64KB per feature: enqueue fused kernel - MAX_FUSED_SIZE = 65536 // x.element_size() - BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) - if N > BLOCK_N: - raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") - with torch.cuda.device(x.device.index): - torch.library.wrap_triton(_layer_norm_fwd_1pass_kernel)[(M,)]( - x, - out, - weight, - bias, - residual, - x1, - weight1, - bias1, - y1, - residual_out, - rowscale, - seeds, - dropout_mask, - dropout_mask1, - mean, - rstd, - x.stride(0), - out.stride(0), - residual.stride(0) if residual is not None else 0, - residual_out.stride(0) if residual_out is not None else 0, - x1.stride(0) if x1 is not None else 0, - y1.stride(0) if y1 is not None else 0, - M, - N, - eps, - dropout_p, - # Passing bool make torch inductor very unhappy since it then tries to compare to int_max - int(zero_centered_weight), - is_rms_norm, - BLOCK_N, - residual is not None, - residual_out is not None, - bias is not None, - dropout_p > 0.0, - dropout_mask is not None, - rowscale is not None, - HAS_X1=x1 is not None, - HAS_W1=weight1 is not None, - HAS_B1=bias1 is not None, - ) - return y1, mean, rstd, seeds, dropout_mask, dropout_mask1 - - -@triton.autotune( - configs=triton_autotune_configs(), - key=["N", "HAS_DRESIDUAL", "STORE_DRESIDUAL", "IS_RMS_NORM", "HAS_BIAS", "HAS_DROPOUT"], -) -# torch compile doesn't like triton.heuristics, so we set these manually when calling the kernel -# @triton.heuristics({"HAS_BIAS": lambda args: args["B"] is not None}) -# @triton.heuristics({"HAS_DRESIDUAL": lambda args: args["DRESIDUAL"] is not None}) -# @triton.heuristics({"STORE_DRESIDUAL": lambda args: args["DRESIDUAL_IN"] is not None}) -# @triton.heuristics({"HAS_ROWSCALE": lambda args: args["ROWSCALE"] is not None}) -# @triton.heuristics({"HAS_DY1": lambda args: args["DY1"] is not None}) -# @triton.heuristics({"HAS_DX1": lambda args: args["DX1"] is not None}) -# @triton.heuristics({"HAS_B1": lambda args: args["DB1"] is not None}) -# @triton.heuristics({"RECOMPUTE_OUTPUT": lambda args: args["Y"] is not None}) -@triton.jit -def _layer_norm_bwd_kernel( - X, # pointer to the input - W, # pointer to the weights - B, # pointer to the biases - Y, # pointer to the output to be recomputed - DY, # pointer to the output gradient - DX, # pointer to the input gradient - DW, # pointer to the partial sum of weights gradient - DB, # pointer to the partial sum of biases gradient - DRESIDUAL, - W1, - DY1, - DX1, - DW1, - DB1, - DRESIDUAL_IN, - ROWSCALE, - SEEDS, - Mean, # pointer to the mean - Rstd, # pointer to the 1/std - stride_x_row, # how much to increase the pointer when moving by 1 row - stride_y_row, - stride_dy_row, - stride_dx_row, - stride_dres_row, - stride_dy1_row, - stride_dx1_row, - stride_dres_in_row, - M, # number of rows in X - N, # number of columns in X - eps, # epsilon to avoid division by zero - dropout_p, - zero_centered_weight, - rows_per_program, - IS_RMS_NORM: tl.constexpr, - BLOCK_N: tl.constexpr, - HAS_DRESIDUAL: tl.constexpr, - STORE_DRESIDUAL: tl.constexpr, - HAS_BIAS: tl.constexpr, - HAS_DROPOUT: tl.constexpr, - HAS_ROWSCALE: tl.constexpr, - HAS_DY1: tl.constexpr, - HAS_DX1: tl.constexpr, - HAS_B1: tl.constexpr, - RECOMPUTE_OUTPUT: tl.constexpr, -): - # Map the program id to the elements of X, DX, and DY it should compute. - row_block_id = tl.program_id(0) - row_start = row_block_id * rows_per_program - # Do not early exit if row_start >= M, because we need to write DW and DB - cols = tl.arange(0, BLOCK_N) - mask = cols < N - X += row_start * stride_x_row - if HAS_DRESIDUAL: - DRESIDUAL += row_start * stride_dres_row - if STORE_DRESIDUAL: - DRESIDUAL_IN += row_start * stride_dres_in_row - DY += row_start * stride_dy_row - DX += row_start * stride_dx_row - if HAS_DY1: - DY1 += row_start * stride_dy1_row - if HAS_DX1: - DX1 += row_start * stride_dx1_row - if RECOMPUTE_OUTPUT: - Y += row_start * stride_y_row - w = tl.load(W + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w += 1.0 - if RECOMPUTE_OUTPUT and HAS_BIAS: - b = tl.load(B + cols, mask=mask, other=0.0).to(tl.float32) - if HAS_DY1: - w1 = tl.load(W1 + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w1 += 1.0 - dw = tl.zeros((BLOCK_N,), dtype=tl.float32) - if HAS_BIAS: - db = tl.zeros((BLOCK_N,), dtype=tl.float32) - if HAS_DY1: - dw1 = tl.zeros((BLOCK_N,), dtype=tl.float32) - if HAS_B1: - db1 = tl.zeros((BLOCK_N,), dtype=tl.float32) - row_end = min((row_block_id + 1) * rows_per_program, M) - for row in range(row_start, row_end): - # Load data to SRAM - x = tl.load(X + cols, mask=mask, other=0).to(tl.float32) - dy = tl.load(DY + cols, mask=mask, other=0).to(tl.float32) - if HAS_DY1: - dy1 = tl.load(DY1 + cols, mask=mask, other=0).to(tl.float32) - if not IS_RMS_NORM: - mean = tl.load(Mean + row) - rstd = tl.load(Rstd + row) - # Compute dx - xhat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd - xhat = tl.where(mask, xhat, 0.0) - if RECOMPUTE_OUTPUT: - y = xhat * w + b if HAS_BIAS else xhat * w - tl.store(Y + cols, y, mask=mask) - wdy = w * dy - dw += dy * xhat - if HAS_BIAS: - db += dy - if HAS_DY1: - wdy += w1 * dy1 - dw1 += dy1 * xhat - if HAS_B1: - db1 += dy1 - if not IS_RMS_NORM: - c1 = tl.sum(xhat * wdy, axis=0) / N - c2 = tl.sum(wdy, axis=0) / N - dx = (wdy - (xhat * c1 + c2)) * rstd - else: - c1 = tl.sum(xhat * wdy, axis=0) / N - dx = (wdy - xhat * c1) * rstd - if HAS_DRESIDUAL: - dres = tl.load(DRESIDUAL + cols, mask=mask, other=0).to(tl.float32) - dx += dres - # Write dx - if STORE_DRESIDUAL: - tl.store(DRESIDUAL_IN + cols, dx, mask=mask) - if HAS_DX1: - if HAS_DROPOUT: - keep_mask = ( - tl.rand(tl.load(SEEDS + M + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - ) - dx1 = tl.where(keep_mask, dx / (1.0 - dropout_p), 0.0) - else: - dx1 = dx - tl.store(DX1 + cols, dx1, mask=mask) - if HAS_DROPOUT: - keep_mask = tl.rand(tl.load(SEEDS + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - dx = tl.where(keep_mask, dx / (1.0 - dropout_p), 0.0) - if HAS_ROWSCALE: - rowscale = tl.load(ROWSCALE + row).to(tl.float32) - dx *= rowscale - tl.store(DX + cols, dx, mask=mask) - - X += stride_x_row - if HAS_DRESIDUAL: - DRESIDUAL += stride_dres_row - if STORE_DRESIDUAL: - DRESIDUAL_IN += stride_dres_in_row - if RECOMPUTE_OUTPUT: - Y += stride_y_row - DY += stride_dy_row - DX += stride_dx_row - if HAS_DY1: - DY1 += stride_dy1_row - if HAS_DX1: - DX1 += stride_dx1_row - tl.store(DW + row_block_id * N + cols, dw, mask=mask) - if HAS_BIAS: - tl.store(DB + row_block_id * N + cols, db, mask=mask) - if HAS_DY1: - tl.store(DW1 + row_block_id * N + cols, dw1, mask=mask) - if HAS_B1: - tl.store(DB1 + row_block_id * N + cols, db1, mask=mask) - - -def _layer_norm_bwd( - dy: Tensor, - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - mean: Tensor, - rstd: Tensor, - dresidual: Optional[Tensor] = None, - dy1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - seeds: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - has_residual: bool = False, - has_x1: bool = False, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - x_dtype: Optional[torch.dtype] = None, - recompute_output: bool = False, -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - # Need to wrap to handle the case where dresidual_in or dx1 are aliases of x, - # which makes torch.library unhappy - dx, dw, db, dresidual_in, dx1, dw1, db1, y = _layer_norm_bwd_impl( - dy, - x, - weight, - bias, - eps, - mean, - rstd, - dresidual, - dy1, - weight1, - bias1, - seeds, - dropout_p, - rowscale, - has_residual, - has_x1, - zero_centered_weight, - is_rms_norm, - x_dtype=x_dtype, - recompute_output=recompute_output, - ) - # Don't need to compute dresidual_in separately in this case - if has_residual and dx.dtype == x.dtype and dropout_p == 0.0 and rowscale is None: - dresidual_in = dx - if has_x1 and dropout_p == 0.0: - dx1 = dx - return dx, dw, db, dresidual_in, dx1, dw1, db1, y - - - -@triton_op(add_op_namespace_prefix("layer_norm_bwd_impl"), mutates_args={}, - schema="(Tensor dy, Tensor x, Tensor weight, Tensor bias, float eps, Tensor mean, Tensor rstd, Tensor? dresidual, Tensor? dy1, Tensor? weight1, Tensor? bias1, Tensor? seeds, float dropout_p, Tensor? rowscale, bool has_residual, bool has_x1, bool zero_centered_weight, bool is_rms_norm, ScalarType? x_dtype, bool recompute_output) -> (Tensor dx, Tensor dw, Tensor db, Tensor dresidual_in, Tensor dx1, Tensor dw1, Tensor db1, Tensor y)", - allow_decomposition=False, # Don't let torch.compile trace inside - ) -def _layer_norm_bwd_impl( - dy: Tensor, - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - mean: Tensor, - rstd: Tensor, - dresidual: Optional[Tensor] = None, - dy1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - seeds: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - has_residual: bool = False, - has_x1: bool = False, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - x_dtype: Optional[torch.dtype] = None, - recompute_output: bool = False, -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - M, N = x.shape - assert x.stride(-1) == 1 - dy = maybe_contiguous_lastdim(dy) - assert dy.stride(-1) == 1 - assert dy.shape == (M, N) - if dresidual is not None: - dresidual = maybe_contiguous_lastdim(dresidual) - assert dresidual.stride(-1) == 1 - assert dresidual.shape == (M, N) - assert weight.shape == (N,) - assert weight.stride(-1) == 1 - if bias is not None: - assert bias.stride(-1) == 1 - assert bias.shape == (N,) - if dy1 is not None: - dy1 = maybe_contiguous_lastdim(dy1) - assert weight1 is not None - assert dy1.shape == dy.shape - assert dy1.stride(-1) == 1 - if weight1 is not None: - assert weight1.shape == (N,) - assert weight1.stride(-1) == 1 - if bias1 is not None: - assert bias1.shape == (N,) - assert bias1.stride(-1) == 1 - if seeds is not None: - assert seeds.is_contiguous() - assert seeds.shape == (M if not has_x1 else M * 2,) - if rowscale is not None: - assert rowscale.is_contiguous() - assert rowscale.shape == (M,) - # allocate output - dx = ( - torch.empty_like(x) - if x_dtype is None - else torch.empty(M, N, dtype=x_dtype, device=x.device) - ) - dresidual_in = ( - torch.empty_like(x) - if has_residual - and (dx.dtype != x.dtype or dropout_p > 0.0 or rowscale is not None or has_x1) - else None - ) - dx1 = torch.empty_like(dx) if (has_x1 and dropout_p > 0.0) else None - y = torch.empty(M, N, dtype=dy.dtype, device=dy.device) if recompute_output else None - if recompute_output: - assert weight1 is None, "recompute_output is not supported with parallel LayerNorm" - - # Less than 64KB per feature: enqueue fused kernel - MAX_FUSED_SIZE = 65536 // x.element_size() - BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) - if N > BLOCK_N: - raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") - # Increasing the multiple (e.g. 8) will allow more thread blocks to be launched and hide the - # latency of the gmem reads/writes, but will increase the time of summing up dw / db. - sm_count = torch.cuda.get_device_properties(x.device).multi_processor_count * 8 - _dw = torch.empty((sm_count, N), dtype=torch.float32, device=weight.device) - _db = ( - torch.empty((sm_count, N), dtype=torch.float32, device=bias.device) - if bias is not None - else None - ) - _dw1 = torch.empty_like(_dw) if weight1 is not None else None - _db1 = torch.empty_like(_db) if bias1 is not None else None - rows_per_program = math.ceil(M / sm_count) - grid = (sm_count,) - with torch.cuda.device(x.device.index): - torch.library.wrap_triton(_layer_norm_bwd_kernel)[grid]( - x, - weight, - bias, - y, - dy, - dx, - _dw, - _db, - dresidual, - weight1, - dy1, - dx1, - _dw1, - _db1, - dresidual_in, - rowscale, - seeds, - mean, - rstd, - x.stride(0), - 0 if not recompute_output else y.stride(0), - dy.stride(0), - dx.stride(0), - dresidual.stride(0) if dresidual is not None else 0, - dy1.stride(0) if dy1 is not None else 0, - dx1.stride(0) if dx1 is not None else 0, - dresidual_in.stride(0) if dresidual_in is not None else 0, - M, - N, - eps, - dropout_p, - # Passing bool make torch inductor very unhappy since it then tries to compare to int_max - int(zero_centered_weight), - rows_per_program, - is_rms_norm, - BLOCK_N, - dresidual is not None, - dresidual_in is not None, - bias is not None, - dropout_p > 0.0, - HAS_ROWSCALE=rowscale is not None, - HAS_DY1=dy1 is not None, - HAS_DX1=dx1 is not None, - HAS_B1=bias1 is not None, - RECOMPUTE_OUTPUT=y is not None, - ) - dw = _dw.sum(0).to(weight.dtype) - db = _db.sum(0).to(bias.dtype) if bias is not None else None - dw1 = _dw1.sum(0).to(weight1.dtype) if weight1 is not None else None - db1 = _db1.sum(0).to(bias1.dtype) if bias1 is not None else None - # dresidual_in and dx1 could be None, the wrapper will handle assigning them from dx - return dx, dw, db, dresidual_in, dx1, dw1, db1, y - - -class LayerNormFn(torch.autograd.Function): - - @staticmethod - def forward( - ctx, - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - residual_in_fp32=False, - zero_centered_weight=False, - is_rms_norm=False, - return_dropout_mask=False, - out_dtype=None, - out=None, - residual_out=None - ): - x_shape_og = x.shape - # reshape input data into 2D tensor - x = maybe_contiguous_lastdim(x.reshape(-1, x.shape[-1])) - if residual is not None: - assert residual.shape == x_shape_og - residual = maybe_contiguous_lastdim(residual.reshape(-1, residual.shape[-1])) - if x1 is not None: - assert x1.shape == x_shape_og - assert rowscale is None, "rowscale is not supported with parallel LayerNorm" - x1 = maybe_contiguous_lastdim(x1.reshape(-1, x1.shape[-1])) - weight = weight.contiguous() - bias = maybe_contiguous(bias) - weight1 = maybe_contiguous(weight1) - bias1 = maybe_contiguous(bias1) - if rowscale is not None: - rowscale = rowscale.reshape(-1).contiguous() - residual_dtype = ( - residual.dtype - if residual is not None - else (torch.float32 if residual_in_fp32 else None) - ) - if out is not None: - out = out.reshape(-1, out.shape[-1]) - if residual_out is not None: - residual_out = residual_out.reshape(-1, residual_out.shape[-1]) - y, y1, mean, rstd, residual_out, seeds, dropout_mask, dropout_mask1 = _layer_norm_fwd( - x, - weight, - bias, - eps, - residual, - x1, - weight1, - bias1, - dropout_p=dropout_p, - rowscale=rowscale, - out_dtype=out_dtype, - residual_dtype=residual_dtype, - zero_centered_weight=zero_centered_weight, - is_rms_norm=is_rms_norm, - return_dropout_mask=return_dropout_mask, - out=out, - residual_out=residual_out, - ) - ctx.save_for_backward( - residual_out, weight, bias, weight1, bias1, rowscale, seeds, mean, rstd - ) - ctx.x_shape_og = x_shape_og - ctx.eps = eps - ctx.dropout_p = dropout_p - ctx.is_rms_norm = is_rms_norm - ctx.has_residual = residual is not None - ctx.has_x1 = x1 is not None - ctx.prenorm = prenorm - ctx.x_dtype = x.dtype - ctx.zero_centered_weight = zero_centered_weight - y = y.reshape(x_shape_og) - y1 = y1.reshape(x_shape_og) if y1 is not None else None - residual_out = residual_out.reshape(x_shape_og) if residual_out is not None else None - dropout_mask = dropout_mask.reshape(x_shape_og) if dropout_mask is not None else None - dropout_mask1 = dropout_mask1.reshape(x_shape_og) if dropout_mask1 is not None else None - if not return_dropout_mask: - if weight1 is None: - return y if not prenorm else (y, residual_out) - else: - return (y, y1) if not prenorm else (y, y1, residual_out) - else: - if weight1 is None: - return ( - (y, dropout_mask, dropout_mask1) - if not prenorm - else (y, residual_out, dropout_mask, dropout_mask1) - ) - else: - return ( - (y, y1, dropout_mask, dropout_mask1) - if not prenorm - else (y, y1, residual_out, dropout_mask, dropout_mask1) - ) - - @staticmethod - def backward(ctx, dy, *args): - x, weight, bias, weight1, bias1, rowscale, seeds, mean, rstd = ctx.saved_tensors - dy = dy.reshape(-1, dy.shape[-1]) - if weight1 is not None: - dy1, args = args[0], args[1:] - dy1 = dy1.reshape(-1, dy1.shape[-1]) - assert dy1.shape == x.shape - else: - dy1 = None - if ctx.prenorm: - dresidual = args[0] - dresidual = dresidual.reshape(-1, dresidual.shape[-1]) - assert dresidual.shape == x.shape - else: - dresidual = None - dx, dw, db, dresidual_in, dx1, dw1, db1, _ = _layer_norm_bwd( - dy, - x, - weight, - bias, - ctx.eps, - mean, - rstd, - dresidual, - dy1, - weight1, - bias1, - seeds, - ctx.dropout_p, - rowscale, - ctx.has_residual, - ctx.has_x1, - ctx.zero_centered_weight, - ctx.is_rms_norm, - x_dtype=ctx.x_dtype, - recompute_output=False, - ) - return ( - dx.reshape(ctx.x_shape_og), - dw, - db, - dresidual_in.reshape(ctx.x_shape_og) if ctx.has_residual else None, - dx1.reshape(ctx.x_shape_og) if dx1 is not None else None, - dw1, - db1, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) - - -def layer_norm_fn( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - residual_in_fp32=False, - zero_centered_weight=False, - is_rms_norm=False, - return_dropout_mask=False, - out_dtype=None, - out=None, - residual_out=None -): - return LayerNormFn.apply( - x, - weight, - bias, - residual, - x1, - weight1, - bias1, - eps, - dropout_p, - rowscale, - prenorm, - residual_in_fp32, - zero_centered_weight, - is_rms_norm, - return_dropout_mask, - out_dtype, - out, - residual_out - ) - - -def rms_norm_fn( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - residual_in_fp32=False, - zero_centered_weight=False, - return_dropout_mask=False, - out_dtype=None, - out=None, - residual_out=None -): - return LayerNormFn.apply( - x, - weight, - bias, - residual, - x1, - weight1, - bias1, - eps, - dropout_p, - rowscale, - prenorm, - residual_in_fp32, - zero_centered_weight, - True, - return_dropout_mask, - out_dtype, - out, - residual_out - ) - - -class RMSNorm(torch.nn.Module): - - def __init__(self, hidden_size, eps=1e-5, dropout_p=0.0, zero_centered_weight=False, - device=None, dtype=None): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.eps = eps - if dropout_p > 0.0: - self.drop = torch.nn.Dropout(dropout_p) - else: - self.drop = None - self.zero_centered_weight = zero_centered_weight - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.register_parameter("bias", None) - self.reset_parameters() - - def reset_parameters(self): - if not self.zero_centered_weight: - torch.nn.init.ones_(self.weight) - else: - torch.nn.init.zeros_(self.weight) - - def forward(self, x, residual=None, prenorm=False, residual_in_fp32=False): - return rms_norm_fn( - x, - self.weight, - self.bias, - residual=residual, - eps=self.eps, - dropout_p=self.drop.p if self.drop is not None and self.training else 0.0, - prenorm=prenorm, - residual_in_fp32=residual_in_fp32, - zero_centered_weight=self.zero_centered_weight, - ) - - -class LayerNormLinearFn(torch.autograd.Function): - - @staticmethod - @custom_fwd - def forward( - ctx, - x, - norm_weight, - norm_bias, - linear_weight, - linear_bias, - residual=None, - eps=1e-6, - prenorm=False, - residual_in_fp32=False, - is_rms_norm=False, - ): - x_shape_og = x.shape - # reshape input data into 2D tensor - x = maybe_contiguous_lastdim(x.reshape(-1, x.shape[-1])) - if residual is not None: - assert residual.shape == x_shape_og - residual = maybe_contiguous_lastdim(residual.reshape(-1, residual.shape[-1])) - norm_weight = norm_weight.contiguous() - norm_bias = maybe_contiguous(norm_bias) - residual_dtype = ( - residual.dtype - if residual is not None - else (torch.float32 if residual_in_fp32 else None) - ) - y, _, mean, rstd, residual_out, *rest = _layer_norm_fwd( - x, - norm_weight, - norm_bias, - eps, - residual, - out_dtype=None if not torch.is_autocast_enabled() else torch.get_autocast_dtype("cuda"), - residual_dtype=residual_dtype, - is_rms_norm=is_rms_norm, - ) - y = y.reshape(x_shape_og) - dtype = torch.get_autocast_dtype("cuda") if torch.is_autocast_enabled() else y.dtype - linear_weight = linear_weight.to(dtype) - linear_bias = linear_bias.to(dtype) if linear_bias is not None else None - out = F.linear(y.to(linear_weight.dtype), linear_weight, linear_bias) - # We don't store y, will be recomputed in the backward pass to save memory - ctx.save_for_backward(residual_out, norm_weight, norm_bias, linear_weight, mean, rstd) - ctx.x_shape_og = x_shape_og - ctx.eps = eps - ctx.is_rms_norm = is_rms_norm - ctx.has_residual = residual is not None - ctx.prenorm = prenorm - ctx.x_dtype = x.dtype - ctx.linear_bias_is_none = linear_bias is None - return out if not prenorm else (out, residual_out.reshape(x_shape_og)) - - @staticmethod - @custom_bwd - def backward(ctx, dout, *args): - x, norm_weight, norm_bias, linear_weight, mean, rstd = ctx.saved_tensors - dout = dout.reshape(-1, dout.shape[-1]) - dy = F.linear(dout, linear_weight.t()) - dlinear_bias = None if ctx.linear_bias_is_none else dout.sum(0) - dy = maybe_contiguous_lastdim(dy) - assert dy.shape == x.shape - if ctx.prenorm: - dresidual = args[0] - dresidual = maybe_contiguous_lastdim(dresidual.reshape(-1, dresidual.shape[-1])) - assert dresidual.shape == x.shape - else: - dresidual = None - dx, dnorm_weight, dnorm_bias, dresidual_in, _, _, _, y = _layer_norm_bwd( - dy, - x, - norm_weight, - norm_bias, - ctx.eps, - mean, - rstd, - dresidual=dresidual, - has_residual=ctx.has_residual, - is_rms_norm=ctx.is_rms_norm, - x_dtype=ctx.x_dtype, - recompute_output=True, - ) - dlinear_weight = torch.einsum("bo,bi->oi", dout, y) - return ( - dx.reshape(ctx.x_shape_og), - dnorm_weight, - dnorm_bias, - dlinear_weight, - dlinear_bias, - dresidual_in.reshape(ctx.x_shape_og) if ctx.has_residual else None, - None, - None, - None, - None, - ) - - -def layer_norm_linear_fn( - x, - norm_weight, - norm_bias, - linear_weight, - linear_bias, - residual=None, - eps=1e-6, - prenorm=False, - residual_in_fp32=False, - is_rms_norm=False, -): - return LayerNormLinearFn.apply( - x, - norm_weight, - norm_bias, - linear_weight, - linear_bias, - residual, - eps, - prenorm, - residual_in_fp32, - is_rms_norm, - ) diff --git a/build/torch212-cxx11-cu126-x86_64-linux/ops/triton/linear.py b/build/torch212-cxx11-cu126-x86_64-linux/ops/triton/linear.py deleted file mode 100644 index a8966dbc345ab0e593df0124451ee7be3dae131a..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu126-x86_64-linux/ops/triton/linear.py +++ /dev/null @@ -1,594 +0,0 @@ -# Adapted from https://github.com/ELS-RD/kernl/blob/main/src/kernl/implementations/linear_layer.py -# and https://github.com/openai/triton/blob/master/python/triton/ops/matmul.py -from typing import Optional - -import torch -import triton -import triton.language as tl -from triton.ops.matmul_perf_model import early_config_prune, estimate_matmul_time - -from flash_attn.ops.triton.k_activations import ( - gelu, - gelu_approx, - gelu_approx_grad, - gelu_grad, - squared_relu, - squared_relu_grad, -) - -# CREDITS: Initially inspired by the Triton tutorial on matrix multiplications - - -def init_to_zero(name): - return lambda nargs: nargs[name].zero_() - - -def get_configs_io_bound(): - configs = [] - for num_stages in [2, 3, 4, 5, 6]: - for block_m in [16, 32]: - for block_k in [32, 64]: - for block_n in [32, 64, 128, 256]: - num_warps = 2 if block_n <= 64 else 4 - configs.append( - triton.Config( - { - "BLOCK_M": block_m, - "BLOCK_N": block_n, - "BLOCK_K": block_k, - "SPLIT_K": 1, - }, - num_stages=num_stages, - num_warps=num_warps, - ) - ) - # split_k not used - # for split_k in [2, 4, 8, 16]: - # configs.append(triton.Config( - # {'BLOCK_M': block_m, 'BLOCK_N': block_n, 'BLOCK_K': block_k, 'SPLIT_K': split_k}, - # num_stages=num_stages, num_warps=num_warps, pre_hook=init_to_zero('C'))) - return configs - - -@triton.autotune( - configs=[ - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - # good for int8 - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=4, - num_warps=4, - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - ] - + get_configs_io_bound(), - key=["CACHE_KEY_M", "CACHE_KEY_N", "CACHE_KEY_K"], - prune_configs_by={ - "early_config_prune": early_config_prune, - "perf_model": estimate_matmul_time, - "top_k": 10, - }, -) -@triton.heuristics( - { - "EVEN_K": lambda args: args["K"] % (args["BLOCK_K"] * args["SPLIT_K"]) == 0, - } -) -@triton.jit -def kernel_fwd( - C, # Pointers to matrices - ACT_INPUT, - A, - B, - bias, - # Matrix dimensions - M, - N, - K, - CACHE_KEY_M, - CACHE_KEY_N, - CACHE_KEY_K, - # The stride variables represent how much to increase the ptr by when moving by 1 - # element in a particular dimension. E.g. stride_am is how much to increase a_ptr - # by to get the element one row down (A has M rows) - stride_cm, - # stride_cn, # Assume that stride_cn == 1 - stride_am, - stride_ak, - stride_bn, - stride_bk, - # Meta-parameters - BLOCK_M: tl.constexpr, - GROUP_M: tl.constexpr, - BLOCK_N: tl.constexpr, - BLOCK_K: tl.constexpr, - # split k not used, not performant with activation, kept because early_config_prune is expecting it - SPLIT_K: tl.constexpr, - EVEN_K: tl.constexpr, - A_ROWMAJOR: tl.constexpr, - B_COLMAJOR: tl.constexpr, - BIAS: tl.constexpr, - SAVE_ACT_INPUT: tl.constexpr, - ACTIVATION: tl.constexpr, -): - - """ - Kernel for computing Out = activation(A x W + C) - - Input has shape (M, K) - - Weight has shape (K, N) - - Bias has shape (N,) - - Output has shape (M, N) - - ActInputs (optional) has shape (M, N) - 'ActInputs' optionally saves the A x W + C intermediate for backward computations - This kernel will consolidate over K - """ - - pid = tl.program_id(axis=0) - - grid_m = (M + BLOCK_M - 1) // BLOCK_M - grid_n = (N + BLOCK_N - 1) // BLOCK_N - # re-order program ID for better L2 performance - width = GROUP_M * grid_n - group_id = pid // width - group_size = min(grid_m - group_id * GROUP_M, GROUP_M) - pid_m = group_id * GROUP_M + (pid % group_size) - pid_n = (pid % width) // (group_size) - - # now compute the block that each program will go through - # rm (resp. rn) denotes a range of indices - # for rows (resp. col) of C - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - # trick to avoid masking on M and N axis - ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M) - rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N) - rk = tl.arange(0, BLOCK_K) - - if A_ROWMAJOR: - A = A + (ram[:, None] * stride_am + rk[None, :]) - else: - A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak) - if B_COLMAJOR: - B = B + (rk[:, None] + rbn[None, :] * stride_bn) - else: - B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn) - - acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) - - for k in range(K, 0, -BLOCK_K): - if EVEN_K: - a = tl.load(A) - b = tl.load(B) - else: - a = tl.load(A, mask=rk[None, :] < k, other=0.0) - b = tl.load(B, mask=rk[:, None] < k, other=0.0) - acc += tl.dot(a, b) - - if A_ROWMAJOR: - A += BLOCK_K - else: - A += BLOCK_K * stride_ak - if B_COLMAJOR: - B += BLOCK_K - else: - B += BLOCK_K * stride_bk - - # Putting bias after the matmul (instead of before) is faster, idk why - if BIAS: - bias = tl.load(bias + rn, mask=rn < N, other=0.0).to(tl.float32) - acc += bias[None, :] - - # optional: save the activation inputs - if SAVE_ACT_INPUT: - # act_in_ptrs = ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :] * stride_cn - act_in_ptrs = ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :] - tl.store(act_in_ptrs, acc) - - # optional: fused activation (while the data is in shared memory) - if ACTIVATION == "gelu": - acc = gelu(acc) - elif ACTIVATION == "gelu_approx": - acc = gelu_approx(acc) - elif ACTIVATION == "squared_relu": - acc = squared_relu(acc) - # rematerialize rm and rn to save registers - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - - # write back result - # C = C + rm[:, None] * stride_cm + rn[None, :] * stride_cn - C = C + rm[:, None] * stride_cm + rn[None, :] - mask = (rm < M)[:, None] & (rn < N)[None, :] - tl.store(C, acc) - - -def triton_linear_act( - x: torch.Tensor, - weight: torch.Tensor, - bias: Optional[torch.Tensor] = None, - activation: str = "id", - save_act_input: bool = False, -) -> torch.Tensor: - """ - Compute e = activation(x @ weight.T + bias). - This wrapper kicks the `kernel_fwd` Triton kernel - :param x: input tensor - :param weight: weight matrix - :param bias: an optional bias tensor - :param activation: Activation name. Needs to be a Triton kernel. - :param act_input: an optional tensor to save the activation inputs (for backward) - :return: result tensor - """ - # if torch.is_autocast_enabled(): - # dtype = torch.get_autocast_gpu_dtype() - # x, weight, bias = [a.to(dtype=dtype) for a in [x, weight, bias]] - - assert activation in ["id", "gelu", "gelu_approx", "squared_relu"] - - batch_shape, n = x.shape[:-1], x.shape[-1] - batch_dim = batch_shape.numel() - x_reshaped = x.reshape(batch_dim, n) - - if x_reshaped.stride(0) > 1 and x_reshaped.stride(1) > 1: - x_reshaped = x_reshaped.contiguous() - if weight.stride(0) > 1 and weight.stride(1) > 1: - weight = weight.contiguous() - bias = bias.contiguous() if bias is not None else None - - assert ( - x.dtype == weight.dtype - ), f"Input and weight must have the same dtype, got {x.dtype} and {weight.dtype}" - if bias is not None: - assert ( - x.dtype == bias.dtype - ), f"Input and bias must have the same dtype, got {x.dtype} and {bias.dtype}" - assert ( - x_reshaped.shape[1] == weight.shape[1] - ), f"Incompatible dimensions: {x_reshaped.shape} - {weight.shape}" - - assert ( - bias is None or bias.shape[0] == weight.shape[0] - ), "Incompatible dimensions in between weight and bias" - - M, K = x_reshaped.shape - N, K = weight.shape - - output = torch.empty((M, N), device=x.device, dtype=x.dtype) - act_input = torch.empty_like(output) if save_act_input else None - - # 1D launch kernel where each block gets its own program. - grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),) # noqa - - kernel_fwd[grid]( - output, - act_input, - x_reshaped, - weight, # data ptrs - bias if bias is not None else x, # auto skip bias if not present - M, # shapes - N, - K, - M // 32, # key for triton cache (limit number of compilations) - N // 32, - K // 32, - stride_cm=output.stride(0), # strides - # stride_cn=output.stride(1), - stride_am=x_reshaped.stride(0), - stride_ak=x_reshaped.stride(1), - stride_bk=weight.stride(1), - stride_bn=weight.stride(0), - BIAS=bias is not None, # optional fused bias - SAVE_ACT_INPUT=save_act_input, # optional save activation inputs - ACTIVATION=activation, # optional fused activation - A_ROWMAJOR=x_reshaped.stride(1) == 1, - B_COLMAJOR=weight.stride(1) == 1, - GROUP_M=8, # speed optimization: group the programs - ) - - if not save_act_input: - return output.reshape(*batch_shape, output.shape[-1]) - else: - return ( - output.reshape(*batch_shape, output.shape[-1]), - act_input.reshape(*batch_shape, act_input.shape[-1]), - ) - - -@triton.autotune( - configs=[ - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - # good for int8 - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=4, - num_warps=4, - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - ] - + get_configs_io_bound(), - key=["CACHE_KEY_M", "CACHE_KEY_N", "CACHE_KEY_K"], - prune_configs_by={ - "early_config_prune": early_config_prune, - "perf_model": estimate_matmul_time, - "top_k": 10, - }, -) -@triton.heuristics( - { - "EVEN_K": lambda args: args["K"] % (args["BLOCK_K"] * args["SPLIT_K"]) == 0, - } -) -@triton.jit -def kernel_bwd( - C, # Pointers to matrices - ACT_INPUT, - A, - B, - # Matrix dimensions - M, - N, - K, - CACHE_KEY_M, - CACHE_KEY_N, - CACHE_KEY_K, - # The stride variables represent how much to increase the ptr by when moving by 1 - # element in a particular dimension. E.g. stride_am is how much to increase a_ptr - # by to get the element one row down (A has M rows) - stride_cm, - # stride_cn, # Assume that stride_cn == 1 - stride_am, - stride_ak, - stride_bk, - stride_bn, - # Meta-parameters - BLOCK_M: tl.constexpr, - GROUP_M: tl.constexpr, - BLOCK_N: tl.constexpr, - BLOCK_K: tl.constexpr, - # split k not used, not performant with activation, kept because early_config_prune is expecting it - SPLIT_K: tl.constexpr, - EVEN_K: tl.constexpr, - ACTIVATION: tl.constexpr, -): - - """ - Kernel for computing Out = activation(A x W + C) - - Input has shape (M, K) - - Weight has shape (K, N) - - Output has shape (M, N) - - ActInputs (optional) has shape (M, N) - 'ActInputs' optionally saves the A x W + C intermediate for backward computations - This kernel will consolidate over K - """ - - pid = tl.program_id(axis=0) - - grid_m = (M + BLOCK_M - 1) // BLOCK_M - grid_n = (N + BLOCK_N - 1) // BLOCK_N - # re-order program ID for better L2 performance - width = GROUP_M * grid_n - group_id = pid // width - group_size = min(grid_m - group_id * GROUP_M, GROUP_M) - pid_m = group_id * GROUP_M + (pid % group_size) - pid_n = (pid % width) // (group_size) - - # now compute the block that each program will go through - # rm (resp. rn) denotes a range of indices - # for rows (resp. col) of C - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - # trick to avoid masking on M and N axis - ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M) - rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N) - rk = tl.arange(0, BLOCK_K) - - A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak) - B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn) - - acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) - - for k in range(K, 0, -BLOCK_K): - if EVEN_K: - a = tl.load(A) - b = tl.load(B) - else: - a = tl.load(A, mask=rk[None, :] < k, other=0.0) - b = tl.load(B, mask=rk[:, None] < k, other=0.0) - acc += tl.dot(a, b) - - A += BLOCK_K * stride_ak - B += BLOCK_K * stride_bk - - # optional: fused activation (while the data is in shared memory) - if ACTIVATION != "id": - act_in_ptrs = ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :] - act_input = tl.load(act_in_ptrs).to(acc.dtype) - if ACTIVATION == "gelu": - acc *= gelu_grad(act_input) - elif ACTIVATION == "gelu_approx": - acc *= gelu_approx_grad(act_input) - elif ACTIVATION == "squared_relu": - acc *= squared_relu_grad(act_input) - - # rematerialize rm and rn to save registers - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - - # write back result - C = C + rm[:, None] * stride_cm + rn[None, :] - mask = (rm < M)[:, None] & (rn < N)[None, :] - tl.store(C, acc, mask=mask) - - -def triton_dgrad_act( - grad_output: torch.Tensor, - weight: torch.Tensor, - activation: str = "id", - act_input: Optional[torch.Tensor] = None, -) -> torch.Tensor: - """ - Compute e = activation(grad_output @ weight + bias). - This wrapper kicks the `kernel_fwd` Triton kernel - :param grad_output: input tensor - :param weight: weight matrix - :param activation: Activation name. Needs to be a Triton kernel. - :param act_input: an optional tensor to save the activation inputs (for backward) - :return: result tensor - """ - assert activation in ["id", "gelu", "gelu_approx", "squared_relu"] - - batch_shape, n = grad_output.shape[:-1], grad_output.shape[-1] - batch_dim = batch_shape.numel() - grad_output_reshaped = grad_output.reshape(batch_dim, n) - - if grad_output_reshaped.stride(0) > 1 and grad_output_reshaped.stride(1) > 1: - grad_output_reshaped = grad_output_reshaped.contiguous() - if weight.stride(0) > 1 and weight.stride(1) > 1: - weight = weight.contiguous() - - assert ( - grad_output.dtype == weight.dtype - ), f"grad_output and weight must have the same dtype, got {grad_output.dtype} and {weight.dtype}" - assert ( - grad_output_reshaped.shape[1] == weight.shape[0] - ), f"Incompatible dimensions: {grad_output_reshaped.shape} - {weight.shape}" - if activation != "id": - assert act_input is not None, f"act_input is required for activation {activation}" - - # M, N, K in bwd are different from M, N, K in fwd - M, K = grad_output_reshaped.shape - K, N = weight.shape - - grad_input = torch.empty((M, N), device=grad_output.device, dtype=grad_output.dtype) - - # 1D launch kernel where each block gets its own program. - grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),) # noqa - - kernel_bwd[grid]( - grad_input, - act_input, - grad_output_reshaped, - weight, # data ptrs - M, # shapes - N, - K, - M // 32, # key for triton cache (limit number of compilations) - N // 32, - K // 32, - stride_cm=grad_input.stride(0), # strides - # stride_cn=grad_input.stride(1), - stride_am=grad_output_reshaped.stride(0), - stride_ak=grad_output_reshaped.stride(1), - stride_bk=weight.stride(0), - stride_bn=weight.stride(1), - ACTIVATION=activation, # optional fused activation - GROUP_M=8, # speed optimization: group the programs - ) - - return grad_input.reshape(*batch_shape, grad_input.shape[-1]) diff --git a/build/torch212-cxx11-cu126-x86_64-linux/ops/triton/mlp.py b/build/torch212-cxx11-cu126-x86_64-linux/ops/triton/mlp.py deleted file mode 100644 index 059f4f8a5e174c1f4824e43d313fca18eaa799b8..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu126-x86_64-linux/ops/triton/mlp.py +++ /dev/null @@ -1,149 +0,0 @@ -# The triton fused matmul + sqrelu is faster for fp16 but slower for bf16, compared -# to naive implementation. -import fused_dense_lib as fused_dense_cuda -import torch -import torch.nn as nn -import torch.nn.functional as F - -from flash_attn.utils.torch import custom_fwd, custom_bwd -from flash_attn.ops.activations import sqrelu_bwd, sqrelu_fwd -from flash_attn.ops.triton.linear import triton_dgrad_act, triton_linear_act - - -class FusedDenseSqreluDenseFunc(torch.autograd.Function): - @staticmethod - @custom_fwd - def forward(ctx, x, weight1, bias1, weight2, bias2, checkpoint_lvl=0): - """checkpoint_lvl: - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute act_input and gelu_out in the bwd - """ - if torch.is_autocast_enabled(): - dtype = torch.get_autocast_gpu_dtype() - x, weight1, bias1, weight2, bias2 = [ - a.to(dtype=dtype) for a in [x, weight1, bias1, weight2, bias2] - ] - is_bf16 = x.dtype == torch.bfloat16 - assert checkpoint_lvl in [0, 1, 2] - x = x.contiguous() - weight1 = weight1.contiguous() - bias1 = bias1.contiguous() - weight2 = weight2.contiguous() - bias2 = bias2.contiguous() - batch_shape, n = x.shape[:-1], x.shape[-1] - batch_dim = batch_shape.numel() - if is_bf16: - act_input = fused_dense_cuda.linear_bias_forward( - x.reshape(batch_dim, n), weight1, bias1 - ) - output1 = sqrelu_fwd(act_input) - else: - save_act_input = checkpoint_lvl != 2 - result = triton_linear_act( - x.reshape(batch_dim, n), - weight1, - bias1, - activation="squared_relu", - save_act_input=save_act_input, - ) - if save_act_input: - output1, act_input = result - else: - output1 = result - output2 = fused_dense_cuda.linear_bias_forward(output1, weight2, bias2) - ctx.checkpoint_lvl = checkpoint_lvl - if checkpoint_lvl == 0: - ctx.save_for_backward(x, weight1, bias1, weight2, act_input, output1) - elif checkpoint_lvl == 1: - ctx.save_for_backward(x, weight1, bias1, weight2, act_input) - elif checkpoint_lvl == 2: - ctx.save_for_backward(x, weight1, bias1, weight2) - return output2.reshape(*batch_shape, output2.shape[-1]) - - @staticmethod - @custom_bwd - def backward(ctx, grad_output): - grad_output = grad_output.contiguous() - checkpoint_lvl = ctx.checkpoint_lvl - x, weight1, bias1, weight2, *rest = ctx.saved_tensors - batch_shape, n = x.shape[:-1], x.shape[-1] - batch_dim = batch_shape.numel() - is_bf16 = x.dtype == torch.bfloat16 - if checkpoint_lvl == 0: - act_input, output1 = rest - elif checkpoint_lvl == 1: - (act_input,) = rest - output1 = sqrelu_fwd(act_input) - elif checkpoint_lvl == 2: - if is_bf16: - act_input = fused_dense_cuda.linear_bias_forward( - x.reshape(batch_dim, n), weight1, bias1 - ) - output1 = sqrelu_fwd(act_input) - else: - output1, act_input = triton_linear_act( - x.reshape(batch_dim, n), - weight1, - bias1, - activation="squared_relu", - save_act_input=True, - ) - - if is_bf16: - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad(output1, grad_output) - grad_output1 = grad_output @ weight2 - grad_act_input = sqrelu_bwd(grad_output1, act_input) - grad_input, grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_backward( - x.reshape(batch_dim, n), weight1, grad_act_input - ) - else: - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad(output1, grad_output) - grad_act_input = triton_dgrad_act( - grad_output, weight2, activation="squared_relu", act_input=act_input - ) - grad_input, grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_backward( - x.reshape(batch_dim, n), weight1, grad_act_input - ) - return grad_input.reshape_as(x), grad_weight1, grad_bias1, grad_weight2, grad_bias2, None - - -fused_dense_sqrelu_dense_function = FusedDenseSqreluDenseFunc.apply - - -class FusedDenseSqreluDense(nn.Module): - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - bias1=True, - bias2=True, - checkpoint_lvl=0, - device=None, - dtype=None, - ): - """ - checkpoint_lvl (increasing lvl means slower but more memory saving): - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute gelu_in and gelu_out in the bwd - """ - assert checkpoint_lvl in [0, 1, 2] - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features * 4 - assert bias1 == True, "DenseSqreluDense module without bias is currently not supported" - assert bias2 == True, "DenseSqreluDense module without bias is currently not supported" - self.checkpoint_lvl = checkpoint_lvl - self.fc1 = nn.Linear(in_features, hidden_features, bias=bias1, **factory_kwargs) - self.fc2 = nn.Linear(hidden_features, out_features, bias=bias2, **factory_kwargs) - - def forward(self, x): - assert x.is_cuda - return fused_dense_sqrelu_dense_function( - x, self.fc1.weight, self.fc1.bias, self.fc2.weight, self.fc2.bias, self.checkpoint_lvl - ) diff --git a/build/torch212-cxx11-cu126-x86_64-linux/ops/triton/rotary.py b/build/torch212-cxx11-cu126-x86_64-linux/ops/triton/rotary.py deleted file mode 100644 index 6ad3ab412a084270120167ad32b67ad1720f2c37..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu126-x86_64-linux/ops/triton/rotary.py +++ /dev/null @@ -1,186 +0,0 @@ -# Copyright (c) 2025, Tri Dao. -# As of 2025-04-23, we require triton >= 3.0 - -from typing import Optional, Union - -import torch - -import triton -import triton.language as tl - - -@triton.jit -def rotary_kernel( - OUT, # Pointers to matrices - X, - COS, - SIN, - CU_SEQLENS, - SEQLEN_OFFSETS, # this could be int or a pointer - # Matrix dimensions - seqlen, - nheads, - seqlen_ro, - # strides - stride_out_batch, - stride_out_seqlen, - stride_out_nheads, - stride_out_headdim, - stride_x_batch, - stride_x_seqlen, - stride_x_nheads, - stride_x_headdim, - # Meta-parameters - # We want ROTARY_DIM to be constexpr, otherwise the triton compiler doesn't know that - # the mask is constant every 8 elements, and it will generate LDG.16 instead of LDG.128 - ROTARY_DIM: tl.constexpr, - IS_SEQLEN_OFFSETS_TENSOR: tl.constexpr, - IS_VARLEN: tl.constexpr, - INTERLEAVED: tl.constexpr, - CONJUGATE: tl.constexpr, - BLOCK_H: tl.constexpr, - BLOCK_M: tl.constexpr, -): - BLOCK_K: tl.constexpr = triton.next_power_of_2(ROTARY_DIM) - ROTARY_DIM_HALF = ROTARY_DIM // 2 - pid_head = tl.program_id(axis=0) - pid_m = tl.program_id(axis=1) - pid_batch = tl.program_id(axis=2) - - if not IS_VARLEN: - X = X + pid_batch * stride_x_batch - OUT = OUT + pid_batch * stride_out_batch - else: - start_idx = tl.load(CU_SEQLENS + pid_batch) - seqlen = tl.load(CU_SEQLENS + pid_batch + 1) - start_idx - X = X + start_idx * stride_x_seqlen - OUT = OUT + start_idx * stride_out_seqlen - - if pid_m * BLOCK_M >= seqlen: - return - - rh = pid_head * BLOCK_H + tl.arange(0, BLOCK_H) - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - if not IS_SEQLEN_OFFSETS_TENSOR: - rm_cs = rm + SEQLEN_OFFSETS - else: - rm_cs = rm + tl.load(SEQLEN_OFFSETS + pid_batch) - - rk_half = tl.arange(0, BLOCK_K // 2) - COS = COS + (rm_cs[:, None] * ROTARY_DIM_HALF + rk_half[None, :]) - SIN = SIN + (rm_cs[:, None] * ROTARY_DIM_HALF + rk_half[None, :]) - mask_cs = (rm_cs[:, None] < seqlen_ro) & (rk_half[None, :] < ROTARY_DIM_HALF) - cos = tl.load(COS, mask=mask_cs, other=1.0).to(tl.float32) - sin = tl.load(SIN, mask=mask_cs, other=0.0).to(tl.float32) - if CONJUGATE: - sin = -sin - - if not INTERLEAVED: - # Load the 1st and 2nd halves of X, do calculation, then store to 1st and 2nd halves of OUT - X = X + (rh[:, None, None] * stride_x_nheads + rm[None, :, None] * stride_x_seqlen + rk_half[None, None, :] * stride_x_headdim) - OUT = OUT + (rh[:, None, None] * stride_out_nheads + rm[None, :, None] * stride_out_seqlen + rk_half[None, None, :] * stride_out_headdim) - mask = (rh[:, None, None] < nheads) & (rm[None, :, None] < seqlen) & (rk_half[None, None, :] < ROTARY_DIM_HALF) - x0 = tl.load(X, mask=mask, other=0.0).to(tl.float32) - x1 = tl.load(X + ROTARY_DIM_HALF * stride_x_headdim, mask=mask, other=0.0,).to(tl.float32) - o0 = x0 * cos - x1 * sin - o1 = x0 * sin + x1 * cos - tl.store(OUT, o0, mask=mask) - tl.store(OUT + ROTARY_DIM_HALF * stride_out_headdim, o1, mask=mask) - else: - rk = tl.arange(0, BLOCK_K) - X = X + (rh[:, None, None] * stride_x_nheads + rm[None, :, None] * stride_x_seqlen + rk[None, None, :] * stride_x_headdim) - OUT = OUT + (rh[:, None, None] * stride_out_nheads + rm[None, :, None] * stride_out_seqlen + rk[None, None, :] * stride_out_headdim) - mask = (rh[:, None, None] < nheads) & (rm[None, :, None] < seqlen) & (rk[None, None, :] < ROTARY_DIM) - x = tl.load(X, mask=mask, other=0.0).to(tl.float32) - x0, x1 = tl.split(tl.reshape(x, [BLOCK_H, BLOCK_M, BLOCK_K // 2, 2])) - o0 = x0 * cos - x1 * sin - o1 = x0 * sin + x1 * cos - o = tl.reshape(tl.join(o0, o1), [BLOCK_H, BLOCK_M, BLOCK_K]) - tl.store(OUT, o, mask=mask) - - -def apply_rotary( - x: torch.Tensor, - cos: torch.Tensor, - sin: torch.Tensor, - seqlen_offsets: Union[int, torch.Tensor] = 0, - cu_seqlens: Optional[torch.Tensor] = None, - max_seqlen: Optional[int] = None, - interleaved=False, - inplace=False, - conjugate=False, -) -> torch.Tensor: - """ - Arguments: - x: (batch, seqlen, nheads, headdim) if cu_seqlens is None - else (total_seqlen, nheads, headdim). - cos: (seqlen_ro, rotary_dim / 2) - sin: (seqlen_ro, rotary_dim / 2) - seqlen_offsets: integer or integer tensor of size (batch,) - cu_seqlens: (batch + 1,) or None - max_seqlen: int - Returns: - y: (batch, seqlen, nheads, headdim) - """ - is_varlen = cu_seqlens is not None - if not is_varlen: - batch, seqlen, nheads, headdim = x.shape - else: - assert max_seqlen is not None, "If cu_seqlens is passed in, then max_seqlen must be passed" - total_seqlen, nheads, headdim = x.shape - batch_p_1 = cu_seqlens.shape[0] - batch = batch_p_1 - 1 - seqlen = max_seqlen - seqlen_ro, rotary_dim = cos.shape - assert sin.shape == cos.shape - rotary_dim *= 2 - assert rotary_dim <= headdim, "rotary_dim must be <= headdim" - assert headdim <= 256, "Only support headdim <= 256" - assert seqlen_ro >= seqlen, "seqlen_ro must be >= seqlen" - - cos, sin = cos.contiguous(), sin.contiguous() - if isinstance(seqlen_offsets, torch.Tensor): - assert seqlen_offsets.shape == (batch,) - assert seqlen_offsets.dtype in [torch.int32, torch.int64] - seqlen_offsets = seqlen_offsets.contiguous() - else: - assert seqlen_offsets + seqlen <= seqlen_ro - - output = torch.empty_like(x) if not inplace else x - if rotary_dim < headdim and not inplace: - output[..., rotary_dim:].copy_(x[..., rotary_dim:]) - - grid = lambda META: (triton.cdiv(nheads, META["BLOCK_H"]), triton.cdiv(seqlen, META["BLOCK_M"]), batch) # noqa - BLOCK_M = 8 if rotary_dim <= 128 else 4 - - # Need this, otherwise Triton tries to launch from cuda:0 and we get - # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?) - device_ctx = torch.cuda.device(x.device.index) if x.device.type == 'cuda' else torch.xpu.device(x.device.index) - with device_ctx: - torch.library.wrap_triton(rotary_kernel)[grid]( - output, # data ptrs - x, - cos, - sin, - cu_seqlens, - seqlen_offsets, - seqlen, # shapes - nheads, - seqlen_ro, - output.stride(0) if not is_varlen else 0, # batch_strides if not varlen else 0 - output.stride(-3), # seqlen_stride or total_seqlen_stride - output.stride(-2), # nheads_stride - output.stride(-1), # headdim_stride - x.stride(0) if not is_varlen else 0, # batch_strides if not varlen else 0 - x.stride(-3), # seqlen stride or total_seqlen_stride - x.stride(-2), # nheads stride - x.stride(-1), # headdim stride - rotary_dim, - isinstance(seqlen_offsets, torch.Tensor), - is_varlen, - interleaved, - conjugate, - BLOCK_M=BLOCK_M, - BLOCK_H=2, - ) - return output diff --git a/build/torch212-cxx11-cu130-aarch64-linux/__init__.py b/build/torch212-cxx11-cu130-aarch64-linux/__init__.py deleted file mode 100644 index ecc2f9d896b6c93f90b0a1499856dc0612177422..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu130-aarch64-linux/__init__.py +++ /dev/null @@ -1,393 +0,0 @@ -from typing import Optional, List -import torch -from ._ops import ops as flash_attn_ops -from .flash_attn_interface import ( - flash_attn_func, - flash_attn_kvpacked_func, - flash_attn_qkvpacked_func, - flash_attn_varlen_func, - flash_attn_varlen_kvpacked_func, - flash_attn_varlen_qkvpacked_func, - flash_attn_with_kvcache, -) - - -def fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - return_softmax: bool = False, - gen: Optional[torch.Generator] = None, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention. - - Args: - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Optional output tensor, same shape as q - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - return_softmax: Whether to return softmax weights - gen: Optional random number generator - - Returns: - List of tensors: [output, softmax_lse, (softmax if return_softmax)] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.fwd( - q, - k, - v, - out, - alibi_slopes, - p_dropout, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - gen, - ) - - -def varlen_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - out: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - block_table: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - max_seqlen_q: int = 0, - max_seqlen_k: int = 0, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - zero_tensors: bool = False, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - return_softmax: bool = False, - gen: Optional[torch.Generator] = None, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention with variable sequence lengths. - - Args: - q: Query tensor of shape [total_q, num_heads, head_size] - k: Key tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - v: Value tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1] - cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1] - out: Optional output tensor of shape [total_q, num_heads, head_size] - seqused_k: Optional tensor specifying how many keys to use per batch element [batch_size] - leftpad_k: Optional left padding for keys of shape [batch_size] - block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq] - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - max_seqlen_q: Maximum sequence length for queries - max_seqlen_k: Maximum sequence length for keys - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - zero_tensors: Whether to zero tensors before computation - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - return_softmax: Whether to return softmax weights - gen: Optional random number generator - - Returns: - List of tensors: [output, softmax_lse, (softmax if return_softmax)] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.varlen_fwd( - q, - k, - v, - out, - cu_seqlens_q, - cu_seqlens_k, - seqused_k, - leftpad_k, - block_table, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - p_dropout, - softmax_scale, - zero_tensors, - is_causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - gen, - ) - - -def bwd( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - p_dropout: float = 0.0, - 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, - gen: Optional[torch.Generator] = None, - rng_state: Optional[torch.Tensor] = None, -) -> List[torch.Tensor]: - """ - Backward pass for multi-head attention. - - Args: - dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size] - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size] - softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q] - dq: Optional gradient tensor for queries, same shape as q - dk: Optional gradient tensor for keys, same shape as k - dv: Optional gradient tensor for values, same shape as v - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - deterministic: Whether to use deterministic algorithms - gen: Optional random number generator - rng_state: Optional RNG state from forward pass - - Returns: - List of tensors: [dq, dk, dv] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - alibi_slopes, - p_dropout, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - deterministic, - gen, - rng_state, - ) - - -def varlen_bwd( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - max_seqlen_q: int = 0, - max_seqlen_k: int = 0, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - zero_tensors: bool = False, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - deterministic: bool = False, - gen: Optional[torch.Generator] = None, - rng_state: Optional[torch.Tensor] = None, -) -> List[torch.Tensor]: - """ - Backward pass for multi-head attention with variable sequence lengths. - - Args: - dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size] - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size] - softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q] - cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1] - cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1] - dq: Optional gradient tensor for queries, same shape as q - dk: Optional gradient tensor for keys, same shape as k - dv: Optional gradient tensor for values, same shape as v - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - max_seqlen_q: Maximum sequence length for queries - max_seqlen_k: Maximum sequence length for keys - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - zero_tensors: Whether to zero tensors before computation - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - deterministic: Whether to use deterministic algorithms - gen: Optional random number generator - rng_state: Optional RNG state from forward pass - - Returns: - List of tensors: [dq, dk, dv] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.varlen_bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - p_dropout, - softmax_scale, - zero_tensors, - is_causal, - window_size_left, - window_size_right, - softcap, - deterministic, - gen, - rng_state, - ) - - -def fwd_kvcache( - q: torch.Tensor, - kcache: torch.Tensor, - vcache: torch.Tensor, - k: Optional[torch.Tensor] = None, - v: Optional[torch.Tensor] = None, - seqlens_k: Optional[torch.Tensor] = None, - rotary_cos: Optional[torch.Tensor] = None, - rotary_sin: Optional[torch.Tensor] = None, - cache_batch_idx: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - block_table: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - out: 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, - is_rotary_interleaved: bool = False, - num_splits: int = 1, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention with KV cache. - - Args: - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - kcache: Key cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - vcache: Value cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - k: Optional new keys tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size] - v: Optional new values tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size] - seqlens_k: Optional sequence lengths for keys of shape [batch_size] - rotary_cos: Optional rotary cosine tensor of shape [seqlen_ro, rotary_dim/2] - rotary_sin: Optional rotary sine tensor of shape [seqlen_ro, rotary_dim/2] - cache_batch_idx: Optional indices to index into the KV cache - leftpad_k: Optional left padding for keys of shape [batch_size] - block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq] - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - out: Optional output tensor, same shape as q - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - is_rotary_interleaved: Whether rotary embeddings are interleaved - num_splits: Number of splits for computation - - Returns: - List of tensors: [output, softmax_lse] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.fwd_kvcache( - q, - kcache, - vcache, - k, - v, - seqlens_k, - rotary_cos, - rotary_sin, - cache_batch_idx, - leftpad_k, - block_table, - alibi_slopes, - out, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - is_rotary_interleaved, - num_splits, - ) diff --git a/build/torch212-cxx11-cu130-aarch64-linux/_flash_attn2_cuda_f12afc9.abi3.so b/build/torch212-cxx11-cu130-aarch64-linux/_flash_attn2_cuda_f12afc9.abi3.so deleted file mode 100644 index ee2d1868747684c3f5a986c9cd5782b609721fa0..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu130-aarch64-linux/_flash_attn2_cuda_f12afc9.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:08fa0c232c7064b8d87d431bb8aadf2d4e92f79120648f6d392725ce522e3a3c -size 1006697936 diff --git a/build/torch212-cxx11-cu130-aarch64-linux/_ops.py b/build/torch212-cxx11-cu130-aarch64-linux/_ops.py deleted file mode 100644 index af5b731cb04b5e6ce54e862db0501f3683c3163f..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu130-aarch64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _flash_attn2_cuda_f12afc9 -ops = torch.ops._flash_attn2_cuda_f12afc9 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_flash_attn2_cuda_f12afc9::{op_name}" diff --git a/build/torch212-cxx11-cu130-aarch64-linux/bert_padding.py b/build/torch212-cxx11-cu130-aarch64-linux/bert_padding.py deleted file mode 100644 index 3c2d35159a014a9d03aabead9e52e009168696ea..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu130-aarch64-linux/bert_padding.py +++ /dev/null @@ -1,218 +0,0 @@ -# Adapted from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/padding.py - -import torch -import torch.nn.functional as F -from einops import rearrange, repeat - - -class IndexFirstAxis(torch.autograd.Function): - @staticmethod - def forward(ctx, input, indices): - ctx.save_for_backward(indices) - assert input.ndim >= 2 - ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] - second_dim = other_shape.numel() - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - # return input[indices] - return torch.gather( - rearrange(input, "b ... -> b (...)"), 0, repeat(indices, "z -> z d", d=second_dim) - ).reshape(-1, *other_shape) - - @staticmethod - def backward(ctx, grad_output): - (indices,) = ctx.saved_tensors - assert grad_output.ndim >= 2 - other_shape = grad_output.shape[1:] - grad_output = rearrange(grad_output, "b ... -> b (...)") - grad_input = torch.zeros( - [ctx.first_axis_dim, grad_output.shape[1]], - device=grad_output.device, - dtype=grad_output.dtype, - ) - # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. - # grad_input[indices] = grad_output - grad_input.scatter_(0, repeat(indices, "z -> z d", d=grad_output.shape[1]), grad_output) - return grad_input.reshape(ctx.first_axis_dim, *other_shape), None - - -index_first_axis = IndexFirstAxis.apply - - -class IndexPutFirstAxis(torch.autograd.Function): - @staticmethod - def forward(ctx, values, indices, first_axis_dim): - ctx.save_for_backward(indices) - assert indices.ndim == 1 - assert values.ndim >= 2 - output = torch.zeros( - first_axis_dim, *values.shape[1:], device=values.device, dtype=values.dtype - ) - # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. - output[indices] = values - # output.scatter_(0, repeat(indices, 'z -> z d', d=values.shape[1]), values) - return output - - @staticmethod - def backward(ctx, grad_output): - (indices,) = ctx.saved_tensors - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - grad_values = grad_output[indices] - # grad_values = torch.gather(grad_output, 0, repeat(indices, 'z -> z d', d=grad_output.shape[1])) - return grad_values, None, None - - -index_put_first_axis = IndexPutFirstAxis.apply - - -class IndexFirstAxisResidual(torch.autograd.Function): - @staticmethod - def forward(ctx, input, indices): - ctx.save_for_backward(indices) - assert input.ndim >= 2 - ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] - second_dim = other_shape.numel() - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - output = input[indices] - # We don't want to reshape input (b ... -> b (...)) since it could change the channel_last - # memory format to channel_first. In other words, input might not be contiguous. - # If we don't detach, Pytorch complains about output being a view and is being modified inplace - return output, input.detach() - - @staticmethod - def backward(ctx, grad_output, grad_residual): - (indices,) = ctx.saved_tensors - assert grad_output.ndim >= 2 - other_shape = grad_output.shape[1:] - assert grad_residual.shape[1:] == other_shape - grad_input = grad_residual - # grad_input[indices] += grad_output - indices = indices.reshape(indices.shape[0], *((1,) * (grad_output.ndim - 1))) - indices = indices.expand_as(grad_output) - grad_input.scatter_add_(0, indices, grad_output) - return grad_input.reshape(ctx.first_axis_dim, *other_shape), None - - -index_first_axis_residual = IndexFirstAxisResidual.apply - - -def unpad_input(hidden_states, attention_mask, unused_mask=None): - """ - Arguments: - hidden_states: (batch, seqlen, ...) - attention_mask: (batch, seqlen), bool / int, 1 means valid and 0 means not valid. - unused_mask: (batch, seqlen), bool / int, 1 means the element is allocated but unused. - Return: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens selected in attention_mask + unused_mask. - indices: (total_nnz), the indices of masked tokens from the flattened input sequence. - cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. - max_seqlen_in_batch: int - seqused: (batch), returns the number of tokens selected in attention_mask + unused_mask. - """ - all_masks = (attention_mask + unused_mask) if unused_mask is not None else attention_mask - seqlens_in_batch = all_masks.sum(dim=-1, dtype=torch.int32) - used_seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) - indices = torch.nonzero(all_masks.flatten(), as_tuple=False).flatten() - max_seqlen_in_batch = seqlens_in_batch.max().item() - cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) - # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the - # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim - # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to - # index with integer indices. Moreover, torch's index is a bit slower than it needs to be, - # so we write custom forward and backward to make it a bit faster. - return ( - index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices), - indices, - cu_seqlens, - max_seqlen_in_batch, - used_seqlens_in_batch, - ) - - -def unpad_input_for_concatenated_sequences(hidden_states, attention_mask_in_length): - """ - Supports concatenating short samples in one sequence. The attention_mask_in_length is utilized to mask other short samples. It helps efficient training of variant lengths-based samples (e.g., the supervised fine-tuning task in large language model). - The motivation for this function is explained [here](https://github.com/Dao-AILab/flash-attention/issues/432#issuecomment-1668822286). - - For example, if batch = 3 and seqlen = 6, the attention_mask_in_length is: - ``` - [ - [2, 3, 0, 0, 0, 0], - [3, 2, 0, 0, 0, 0], - [6, 0, 0, 0, 0, 0] - ] - ``` - , which refers to the 3D-attention mask: - ``` - [ - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [0, 0, 1, 0, 0, 0], - [0, 0, 1, 1, 0, 0], - [0, 0, 1, 1, 1, 0], - [0, 0, 0, 0, 0, 1] - ], - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [0, 0, 0, 1, 0, 0], - [0, 0, 0, 1, 1, 0], - [0, 0, 0, 0, 0, 1] - ], - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [1, 1, 1, 1, 0, 0], - [1, 1, 1, 1, 1, 0], - [1, 1, 1, 1, 1, 1] - ] - ] - ```. - - Arguments: - hidden_states: (batch, seqlen, ...) - attention_mask_in_length: (batch, seqlen), int, a nonzero number (e.g., 1, 2, 3, etc.) means length of concatenated sequence in b-th batch, and 0 means none. - Return: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. - indices: (total_nnz), the indices of non-masked tokens from the flattened input sequence. - cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. - max_seqlen_in_batch: int - """ - length = attention_mask_in_length.sum(dim=-1) - seqlen = attention_mask_in_length.size(-1) - attention_mask_2d = torch.arange(seqlen, device=length.device, dtype=length.dtype).expand(len(length), seqlen) < length.unsqueeze(1) - real_indices_idx = torch.nonzero(attention_mask_in_length.flatten(), as_tuple=False).flatten() - seqlens_in_batch = attention_mask_in_length.flatten()[real_indices_idx] - indices = torch.nonzero(attention_mask_2d.flatten(), as_tuple=False).flatten() - max_seqlen_in_batch = seqlens_in_batch.max().item() - cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) - # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the - # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim - # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to - # index with integer indices. Moreover, torch's index is a bit slower than it needs to be, - # so we write custom forward and backward to make it a bit faster. - return ( - index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices), - indices, - cu_seqlens, - max_seqlen_in_batch, - ) - - -def pad_input(hidden_states, indices, batch, seqlen): - """ - Arguments: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. - indices: (total_nnz), the indices that represent the non-masked tokens of the original padded input sequence. - batch: int, batch size for the padded sequence. - seqlen: int, maximum sequence length for the padded sequence. - Return: - hidden_states: (batch, seqlen, ...) - """ - dim = hidden_states.shape[-1] - # output = torch.zeros((batch * seqlen), dim, device=hidden_states.device, dtype=hidden_states.dtype) - # output[indices] = hidden_states - output = index_put_first_axis(hidden_states, indices, batch * seqlen) - return rearrange(output, "(b s) ... -> b s ...", b=batch) diff --git a/build/torch212-cxx11-cu130-aarch64-linux/flash_attn2/__init__.py b/build/torch212-cxx11-cu130-aarch64-linux/flash_attn2/__init__.py deleted file mode 100644 index a9b2672c1cd85b74c1b3ded0fc0b2100e1aeac23..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu130-aarch64-linux/flash_attn2/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -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/torch212-cxx11-cu130-aarch64-linux/flash_attn_interface.py b/build/torch212-cxx11-cu130-aarch64-linux/flash_attn_interface.py deleted file mode 100644 index b014e569dc9812a9c840d1043f2cebe767b4d477..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu130-aarch64-linux/flash_attn_interface.py +++ /dev/null @@ -1,1631 +0,0 @@ -# Copyright (c) 2023, Tri Dao. - -from typing import Optional, Sequence, Tuple, Union - -import torch -import torch.nn as nn -import os - -# # isort: off -# # We need to import the CUDA kernels after importing torch -# USE_TRITON_ROCM = os.getenv("FLASH_ATTENTION_TRITON_AMD_ENABLE", "FALSE") == "TRUE" -# if USE_TRITON_ROCM: -# from .flash_attn_triton_amd import interface_fa as flash_attn -# else: -# import flash_attn_2_cuda as flash_attn - - -from . import _ops -from ._ops import ops as flash_attn -from ._ops import add_op_namespace_prefix - -# # isort: on - -def maybe_contiguous(x): - return x.contiguous() if x is not None and x.stride(-1) != 1 else x - - -def _get_device(): - if torch.xpu.is_available(): - return "xpu" - elif torch.cuda.is_available(): - return "cuda" - else: - return "cpu" - - -def _get_block_size_n(device, head_dim, is_dropout, is_causal): - assert head_dim <= 256 - - if device.type == "xpu": - if head_dim <= 96: - return 64 - elif head_dim <= 128: - return 32 - elif head_dim <= 256: - return 64 - else: - return 32 - - # This should match the block sizes in the CUDA kernel - major, minor = torch.cuda.get_device_capability(device) - is_sm8x = major == 8 and minor > 0 # Only include sm86 and sm89, exclude sm80 (A100) - is_sm80 = major == 8 and minor == 0 - is_sm90 = major == 9 and minor == 0 - if head_dim <= 32: - return 128 - if head_dim <= 64: - return 128 if not is_dropout else 64 - elif head_dim <= 96: - return 64 - elif head_dim <= 128: - if is_sm8x: - return 64 if (not is_dropout and is_causal) else 32 - else: - return 64 if not is_dropout else 32 - elif head_dim <= 192: - return 64 - elif head_dim <= 224: - return 64 - elif head_dim <= 256: - return 64 - - -def round_multiple(x, m): - return (x + m - 1) // m * m - - -# torch.compile() support is only enabled for pytorch >= 2.4 -# The reason for this is that we are using the new custom_op and register_fake -# APIs, which support inplace modification of inputs in the function itself -if torch.__version__ >= "2.4.0": - _torch_custom_op_wrapper = torch.library.custom_op - _torch_register_fake_wrapper = torch.library.register_fake -else: - def noop_custom_op_wrapper(name, fn=None, /, *, mutates_args, device_types=None, schema=None): - def wrap(func): - return func - if fn is None: - return wrap - return fn - def noop_register_fake_wrapper(op, fn=None, /, *, lib=None, _stacklevel=1): - def wrap(func): - return func - if fn is None: - return wrap - return fn - _torch_custom_op_wrapper = noop_custom_op_wrapper - _torch_register_fake_wrapper = noop_register_fake_wrapper - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_forward"), mutates_args=(), device_types=_get_device()) -def _flash_attn_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - return_softmax: bool -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - out, softmax_lse, S_dmask, rng_state = flash_attn.fwd( - q, - k, - v, - None, - alibi_slopes, - dropout_p, - softmax_scale, - causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - None, - ) - return out, softmax_lse, S_dmask, rng_state - - -@_torch_register_fake_wrapper(add_op_namespace_prefix("_flash_attn_forward")) -def _flash_attn_forward_fake( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - return_softmax: bool -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - batch_size, seqlen_q, num_heads, head_size = q.shape - seqlen_k = k.shape[1] - out = torch.empty_like(q) - softmax_lse = torch.empty((batch_size, num_heads, seqlen_q), dtype=torch.float32, device=q.device, layout=q.layout) - p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) - if return_softmax: - p = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128), round_multiple(seqlen_k, 128)), dtype=q.dtype, device=q.device, layout=q.layout) - rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) - - return out, softmax_lse, p, rng_state - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_forward = _ops.ops._flash_attn_forward -else: - _wrapped_flash_attn_forward = _flash_attn_forward - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_varlen_forward"), mutates_args=(), device_types=_get_device()) -def _flash_attn_varlen_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - alibi_slopes: Optional[torch.Tensor] = None, - return_softmax: bool = False, - block_table: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - out, softmax_lse, S_dmask, rng_state = flash_attn.varlen_fwd( - q, - k, - v, - None, - cu_seqlens_q, - cu_seqlens_k, - seqused_k, - leftpad_k, - block_table, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - zero_tensors, - causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - None, - ) - # if out.isnan().any() or softmax_lse.isnan().any(): - # breakpoint() - return out, softmax_lse, S_dmask, rng_state - - -@_torch_register_fake_wrapper(add_op_namespace_prefix("_flash_attn_varlen_forward")) -def _flash_attn_varlen_forward_fake( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - alibi_slopes: Optional[torch.Tensor] = None, - return_softmax: bool = False, - block_table: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - paged_kv = block_table is not None - batch_size = cu_seqlens_q.numel() - 1 - total_q, num_heads, _ = q.shape - - out = torch.empty_like(q) - softmax_lse = torch.empty((num_heads, total_q), dtype=torch.float32, device=q.device, layout=q.layout) - p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) - seqlen_q_rounded = round_multiple(max_seqlen_q, 128) - seqlen_k_rounded = round_multiple(max_seqlen_k, 128) - if return_softmax: - p = torch.empty((batch_size, num_heads, seqlen_q_rounded, seqlen_k_rounded), dtype=q.dtype, device=q.device, layout=q.layout) - rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) - return out, softmax_lse, p, rng_state - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_varlen_forward = _ops.ops._flash_attn_varlen_forward -else: - _wrapped_flash_attn_varlen_forward = _flash_attn_varlen_forward - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_backward"), mutates_args=("dq", "dk", "dv"), device_types=_get_device()) -def _flash_attn_backward( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, -) -> 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)] - ( - dq, - dk, - dv, - softmax_d, - ) = flash_attn.bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - alibi_slopes, - dropout_p, - softmax_scale, - causal, - window_size_left, - window_size_right, - softcap, - deterministic, - None, - rng_state, - ) - return softmax_d - - -@_torch_register_fake_wrapper(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, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, -) -> torch.Tensor: - dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] - if dq is None: - dq = torch.empty_like(q) - if dk is None: - dk = torch.empty_like(k) - if dv is None: - dv = torch.empty_like(v) - batch_size, seqlen_q, num_heads, _ = q.shape - softmax_d = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128)), device=q.device, dtype=torch.float32) - - return softmax_d - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_backward = _ops.ops._flash_attn_backward -else: - _wrapped_flash_attn_backward = _flash_attn_backward - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_varlen_backward"), mutates_args=("dq", "dk", "dv"), device_types=_get_device()) -def _flash_attn_varlen_backward( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> 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)] - ( - dq, - dk, - dv, - softmax_d, - ) = flash_attn.varlen_bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - zero_tensors, - causal, - window_size_left, - window_size_right, - softcap, - deterministic, - None, - rng_state, - ) - # if dk.isnan().any() or dk.isnan().any() or dv.isnan().any() or softmax_d.isnan().any(): - # breakpoint() - return softmax_d - - -@_torch_register_fake_wrapper(add_op_namespace_prefix("_flash_attn_varlen_backward")) -def _flash_attn_varlen_backward_fake( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> torch.Tensor: - dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] - batch_size = cu_seqlens_q.numel() - 1 - total_q, num_heads, _ = q.shape - - if dq is None: - dq = torch.empty_like(q) - if dk is None: - dk = torch.empty_like(k) - if dv is None: - dv = torch.empty_like(v) - softmax_d = torch.empty((num_heads, total_q + 128 * batch_size), device=q.device, dtype=torch.float32) - - return softmax_d - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_varlen_backward = _ops.ops._flash_attn_varlen_backward -else: - _wrapped_flash_attn_varlen_backward = _flash_attn_varlen_backward - - -class FlashAttnQKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and qkv.requires_grad - if softmax_scale is None: - softmax_scale = qkv.shape[-1] ** (-0.5) - q, k, v = qkv[:, :, 0].detach(), qkv[:, :, 1].detach(), qkv[:, :, 2].detach() - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) - dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dqkv[:, :, 0], - dqkv[:, :, 1], - dqkv[:, :, 2], - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dqkv = dqkv[..., : dout.shape[-1]] # We could have padded the head dimension - return dqkv, None, None, None, None, None, None, None, None, None - - -class FlashAttnVarlenQKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - cu_seqlens, - max_seqlen, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and qkv.requires_grad - if softmax_scale is None: - softmax_scale = qkv.shape[-1] ** (-0.5) - q, k, v = qkv[:, 0].detach(), qkv[:, 1].detach(), qkv[:, 2].detach() - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens, - cu_seqlens, - max_seqlen, - max_seqlen, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=None, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens, rng_state) - ctx.dropout_p = dropout_p - ctx.max_seqlen = max_seqlen - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens, rng_state = ctx.saved_tensors - qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) - dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dqkv[:, 0], - dqkv[:, 1], - dqkv[:, 2], - cu_seqlens, - cu_seqlens, - ctx.max_seqlen, - ctx.max_seqlen, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - 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 FlashAttnKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - kv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, kv] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - k, v = kv[:, :, 0].detach(), kv[:, :, 1].detach() - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - dq = torch.empty_like(q) - kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) - dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dkv[:, :, 0], - dkv[:, :, 1], - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dkv = dkv[..., : dout.shape[-1]] - return dq, dkv, None, None, None, None, None, None, None, None, None - - -class FlashAttnVarlenKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, kv] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - k, v = kv[:, 0].detach(), kv[:, 1].detach() - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=None, - ) - if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state - ) - ctx.dropout_p = dropout_p - 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.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors - dq = torch.empty_like(q) - kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) - dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dkv[:, 0], - dkv[:, 1], - cu_seqlens_q, - cu_seqlens_k, - ctx.max_seqlen_q, - ctx.max_seqlen_k, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dkv = dkv[..., : dout.shape[-1]] - return dq, dkv, None, None, None, None, None, None, None, None, None, None, None, None, None - - -class FlashAttnFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - k, - v, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, k, v] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dk = dk[..., : dout.shape[-1]] - dv = dv[..., : dout.shape[-1]] - return dq, dk, dv, 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, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - block_table, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, k, v] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=block_table, - ) - if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state - ) - ctx.dropout_p = dropout_p - 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.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors - dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - ctx.max_seqlen_q, - ctx.max_seqlen_k, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dk = dk[..., : dout.shape[-1]] - dv = dv[..., : dout.shape[-1]] - return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None - - -def flash_attn_qkvpacked_func( - qkv, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # <=0.0 means deactivate - alibi_slopes=None, - deterministic=False, - 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, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_kvpacked_func( - q, - kv, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, -): - """dropout_p should be set to 0.0 during evaluation - If 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 K, V. - 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) - kv: (batch_size, seqlen, 2, 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. - softcap: float. Anything > 0 activates softcapping 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). - 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 FlashAttnKVPackedFunc.apply( - q, - kv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_func( - q, - k, - v, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - 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). - 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 FlashAttnFunc.apply( - q, - k, - v, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_qkvpacked_func( - qkv, - cu_seqlens, - max_seqlen, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - 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_varlen_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_varlen_kvpacked_func and flash_attn_varlen_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: (total, 3, nheads, headdim), where total = total number of tokens in the batch. - cu_seqlens: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into qkv. - max_seqlen: int. Maximum sequence length in the batch. - 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenQKVPackedFunc.apply( - qkv, - cu_seqlens, - max_seqlen, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_kvpacked_func( - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, -): - """dropout_p should be set to 0.0 during evaluation - If 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 K, V. - 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: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch. - kv: (total_k, 2, nheads_k, headdim), where total_k = total number of key tokens in the batch. - cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into q. - cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into kv. - max_seqlen_q: int. Maximum query sequence length in the batch. - max_seqlen_k: int. Maximum key sequence length in the batch. - 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 + 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenKVPackedFunc.apply( - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_func( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, - block_table=None, -): - """dropout_p should be set to 0.0 during evaluation - Supports multi-query and grouped-query attention (MQA/GQA) by passing in K, V 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: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch. - k: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. - v: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. - cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into q. - cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into kv. - max_seqlen_q: int. Maximum query sequence length in the batch. - max_seqlen_k: int. Maximum key sequence length in the batch. - 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 + 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenFunc.apply( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - block_table, - False if q.device.type == "cpu" else torch.is_grad_enabled(), - ) - - -def flash_attn_with_kvcache( - q, - k_cache, - v_cache, - k=None, - v=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, - block_table: Optional[torch.Tensor] = None, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - rotary_interleaved=True, - alibi_slopes=None, - num_splits=0, - 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 block_table, - or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_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) if there's no block_table, - or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_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). Similar to k. - 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. - block_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). - 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. - 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" - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-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) - cache_batch_idx = maybe_contiguous(cache_batch_idx) - block_table = maybe_contiguous(block_table) - out, softmax_lse = flash_attn.fwd_kvcache( - q, - k_cache, - v_cache, - k, - v, - cache_seqlens, - rotary_cos, - rotary_sin, - cache_batch_idx, - cache_leftpad, - block_table, - alibi_slopes, - None, - softmax_scale, - causal, - window_size[0], - window_size[1], - softcap, - rotary_interleaved, - num_splits, - ) - return (out, softmax_lse) if return_softmax_lse else out diff --git a/build/torch212-cxx11-cu130-aarch64-linux/layers/__init__.py b/build/torch212-cxx11-cu130-aarch64-linux/layers/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/build/torch212-cxx11-cu130-aarch64-linux/layers/patch_embed.py b/build/torch212-cxx11-cu130-aarch64-linux/layers/patch_embed.py deleted file mode 100644 index 05562f8e8bcdb58e947c6f402a49eacd2d031871..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu130-aarch64-linux/layers/patch_embed.py +++ /dev/null @@ -1,67 +0,0 @@ -# We use the same API as https://github.com/rwightman/pytorch-image-models/blob/v0.6.11/timm/models/layers/patch_embed.py -# But we use nn.Linear instead of Conv2d and it's about 8x faster. - -from functools import partial - -import torch.nn as nn -from einops import rearrange -from torch import _assert -from torch.nn.modules.utils import _pair - -try: - from flash_attn.ops.fused_dense import FusedDense -except ImportError: - FusedDense = None - - -class PatchEmbed(nn.Module): - """2D Image to Patch Embedding""" - - def __init__( - self, - img_size=224, - patch_size=16, - in_chans=3, - embed_dim=768, - norm_layer=None, - flatten=True, - bias=True, - fused_bias_fc=False, - ): - super().__init__() - img_size = _pair(img_size) - patch_size = _pair(patch_size) - self.img_size = img_size - self.patch_size = patch_size - self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1]) - self.num_patches = self.grid_size[0] * self.grid_size[1] - self.flatten = flatten - if fused_bias_fc and FusedDense is None: - raise ImportError("fused_dense is not installed") - - linear_cls = nn.Linear if not fused_bias_fc or not bias else FusedDense - self.proj = linear_cls(in_chans * patch_size[0] * patch_size[1], embed_dim, bias=bias) - self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() - - def forward(self, x): - _, _, H, W = x.shape - _assert( - H == self.img_size[0], - f"Input image height ({H}) doesn't match model ({self.img_size[0]}).", - ) - _assert( - W == self.img_size[1], - f"Input image width ({W}) doesn't match model ({self.img_size[1]}).", - ) - x = self.proj( - rearrange( - x, - "b c (h p1) (w p2) -> b h w (c p1 p2)", - p1=self.patch_size[0], - p2=self.patch_size[1], - ) - ) - if self.flatten: - x = rearrange(x, "b h w c -> b (h w) c") - x = self.norm(x) - return x diff --git a/build/torch212-cxx11-cu130-aarch64-linux/layers/rotary.py b/build/torch212-cxx11-cu130-aarch64-linux/layers/rotary.py deleted file mode 100644 index d1bfc21fc7de1dd287e8f382847b194a48075981..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu130-aarch64-linux/layers/rotary.py +++ /dev/null @@ -1,483 +0,0 @@ -# Copyright (c) 2025, Tri Dao - -import math -from functools import partial -from typing import Optional, Tuple, Union - -import torch -from torch import Tensor - -from einops import rearrange, repeat -# from flash_attn.ops.triton.rotary import apply_rotary -from ..ops.triton.rotary import apply_rotary - - -def rotate_half(x, interleaved=False): - if not interleaved: - x1, x2 = x.chunk(2, dim=-1) - return torch.cat((-x2, x1), dim=-1) - else: - x1, x2 = x[..., ::2], x[..., 1::2] - return rearrange(torch.stack((-x2, x1), dim=-1), "... d two -> ... (d two)", two=2) - - -def apply_rotary_emb_torch(x, cos, sin, interleaved=False): - """ - x: (batch_size, seqlen, nheads, headdim) - cos, sin: (seqlen, rotary_dim / 2) or (batch_size, seqlen, rotary_dim / 2) - """ - ro_dim = cos.shape[-1] * 2 - assert ro_dim <= x.shape[-1] - cos = repeat(cos, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)") - sin = repeat(sin, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)") - return torch.cat( - [x[..., :ro_dim] * cos + rotate_half(x[..., :ro_dim], interleaved) * sin, x[..., ro_dim:]], - dim=-1, - ) - - -class ApplyRotaryEmb(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x, - cos, - sin, - interleaved=False, - inplace=False, - seqlen_offsets: Union[int, Tensor] = 0, - cu_seqlens: Optional[Tensor] = None, - max_seqlen: Optional[int] = None, - ): - out = apply_rotary( - x, - cos, - sin, - seqlen_offsets=seqlen_offsets, - cu_seqlens=cu_seqlens, - max_seqlen=max_seqlen, - interleaved=interleaved, - inplace=inplace, - ) - if isinstance(seqlen_offsets, int): - ctx.save_for_backward(cos, sin, cu_seqlens) # Can't save int with save_for_backward - ctx.seqlen_offsets = seqlen_offsets - else: - ctx.save_for_backward(cos, sin, cu_seqlens, seqlen_offsets) - ctx.seqlen_offsets = None - ctx.interleaved = interleaved - ctx.inplace = inplace - ctx.max_seqlen = max_seqlen - return out if not inplace else x - - @staticmethod - def backward(ctx, do): - seqlen_offsets = ctx.seqlen_offsets - if seqlen_offsets is None: - cos, sin, cu_seqlens, seqlen_offsets = ctx.saved_tensors - else: - cos, sin, cu_seqlens = ctx.saved_tensors - dx = apply_rotary( - do, - cos, - sin, - seqlen_offsets=seqlen_offsets, - cu_seqlens=cu_seqlens, - max_seqlen=ctx.max_seqlen, - interleaved=ctx.interleaved, - inplace=ctx.inplace, - conjugate=True, - ) - return dx, None, None, None, None, None, None, None - - -def apply_rotary_emb( - x, - cos, - sin, - interleaved=False, - inplace=False, - seqlen_offsets: Union[int, Tensor] = 0, - cu_seqlens: Optional[Tensor] = None, - max_seqlen: Optional[int] = None, -): - """ - Arguments: - x: (batch_size, seqlen, nheads, headdim) if cu_seqlens is None - else (total_seqlen, nheads, headdim) - cos, sin: (seqlen_rotary, rotary_dim / 2) - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead - of 1st half and 2nd half (GPT-NeoX style). - inplace: if True, apply rotary embedding in-place. - seqlen_offsets: (batch_size,) or int. Each sequence in x is shifted by this amount. - Most commonly used in inference when we have KV cache. - cu_seqlens: (batch + 1,) or None - max_seqlen: int - Return: - out: (batch_size, seqlen, nheads, headdim) if cu_seqlens is None - else (total_seqlen, nheads, headdim) - rotary_dim must be <= headdim - Apply rotary embedding to the first rotary_dim of x. - """ - return ApplyRotaryEmb.apply( - x, cos, sin, interleaved, inplace, seqlen_offsets, cu_seqlens, max_seqlen - ) - - -# For backward compatibility -apply_rotary_emb_func = apply_rotary_emb - - -def _apply_rotary_emb_qkv( - qkv, - cos, - sin, - cos_k=None, - sin_k=None, - interleaved=False, - inplace=False, - conjugate=False, - seqlen_offsets: Union[int, Tensor] = 0, - num_heads_q: Optional[int] = None, -): - apply_rotary_fn = partial( - apply_rotary, - interleaved=interleaved, - inplace=inplace, - conjugate=conjugate, - seqlen_offsets=seqlen_offsets - ) - if cos_k is None and sin_k is None and qkv.is_contiguous(): - # Call 1 kernel instead of 2 kernels - # We need qkv to be contiguous so that when we reshape to combine (3, nheads) - # dimensions, we get the same tensor - if qkv.dim() == 5: - batch, seqlen, three, nheads, headdim = qkv.shape - assert three == 3 - # qk = rearrange(qkv[:, :, :2], "b s t h d -> b s (t h) d") - qk = qkv[:, :, :2].reshape(batch, seqlen, -1, headdim) - qk = apply_rotary_fn(qk, cos, sin) - else: - assert qkv.dim() == 4 - assert num_heads_q is not None - num_heads_k = (qkv.shape[2] - num_heads_q) // 2 - assert qkv.shape[2] == num_heads_q + 2 * num_heads_k - qk = qkv[:, :, :num_heads_q + num_heads_k] - qk = apply_rotary_fn(qk, cos, sin) - if not inplace: - if qkv.dim() == 5: - qkv = torch.cat([rearrange(qk, "b s (t h) d -> b s t h d", t=2), qkv[:, :, 2:]], dim=2) - else: - qkv = torch.cat([qk, qkv[:, :, num_heads_q + num_heads_k :]], dim=2) - else: - cos_k = cos if cos_k is None else cos_k - sin_k = sin if sin_k is None else sin_k - if qkv.dim() == 5: - batch, seqlen, three, nheads, headdim = qkv.shape - assert three == 3 - q, k = qkv[:, :, 0], qkv[:, :, 1] - else: - assert qkv.dim() == 4 - assert num_heads_q is not None - num_heads_k = (qkv.shape[2] - num_heads_q) // 2 - assert qkv.shape[2] == num_heads_q + 2 * num_heads_k - q, k = qkv[:, :, :num_heads_q], qkv[:, :, num_heads_q : num_heads_q + num_heads_k] - q = apply_rotary_fn(q, cos, sin) - k = apply_rotary_fn(k, cos_k, sin_k) - if not inplace: - if qkv.dim() == 5: - qkv = torch.stack([q, k, qkv[:, :, 2]], dim=2) - else: - qkv = torch.cat([q, k, qkv[:, :, num_heads_q + num_heads_k:]], dim=2) - return qkv - - -class ApplyRotaryEmbQKV_(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - cos, - sin, - cos_k=None, - sin_k=None, - interleaved=False, - seqlen_offsets: Union[int, torch.Tensor] = 0, - num_heads_q: Optional[int] = None, - ): - # apply_rotary_emb_qkv_inplace( - qkv = _apply_rotary_emb_qkv( - qkv, cos, sin, cos_k, sin_k, interleaved=interleaved, inplace=True, - seqlen_offsets=seqlen_offsets, num_heads_q=num_heads_q, - ) - if isinstance(seqlen_offsets, int): - ctx.save_for_backward(cos, sin, cos_k, sin_k) - ctx.seqlen_offsets = seqlen_offsets - else: - ctx.save_for_backward(cos, sin, cos_k, sin_k, seqlen_offsets) - ctx.seqlen_offsets = None - ctx.interleaved = interleaved - ctx.num_heads_q = num_heads_q - return qkv - - @staticmethod - def backward(ctx, dqkv): - seqlen_offsets = ctx.seqlen_offsets - if seqlen_offsets is None: - cos, sin, cos_k, sin_k, seqlen_offsets = ctx.saved_tensors - else: - cos, sin, cos_k, sin_k = ctx.saved_tensors - dqkv = _apply_rotary_emb_qkv( - dqkv, cos, sin, cos_k, sin_k, interleaved=ctx.interleaved, inplace=True, - seqlen_offsets=seqlen_offsets, num_heads_q=ctx.num_heads_q, conjugate=True, - ) - return dqkv, None, None, None, None, None, None, None - - -def apply_rotary_emb_qkv_( - qkv, - cos, - sin, - cos_k=None, - sin_k=None, - interleaved=False, - seqlen_offsets: Union[int, torch.Tensor] = 0, - num_heads_q: Optional[int] = None, -): - """ - Arguments: - qkv: (batch_size, seqlen, 3, nheads, headdim) or (batch_size, seqlen, num_heads_q + 2 * num_heads_k, headdim). - If qkv has shape (batch_size, seqlen, num_heads_q + 2 * num_heads_k, headdim) (e.g. MQA / GQA), - then num_heads_q must be provided. - cos, sin: (seqlen, rotary_dim / 2) - cos_k, sin_k: (seqlen, rotary_dim / 2), optional - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead of - 1st half and 2nd half (GPT-NeoX style). - seqlen_offsets: (batch_size,) or int. Each sequence in Q and K is shifted by this amount. - Most commonly used in inference when we have KV cache. - Return: - qkv: (batch_size, seqlen, 3, nheads, headdim) or (batch_size, seqlen, num_heads_q + 2 * num_heads_k, headdim) - rotary_dim must be <= headdim - Apply rotary embedding *inplace* to the first rotary_dim of Q and K. - """ - return ApplyRotaryEmbQKV_.apply( - qkv, cos, sin, cos_k, sin_k, interleaved, seqlen_offsets, num_heads_q - ) - - -class ApplyRotaryEmbKV_(torch.autograd.Function): - - @staticmethod - def forward(ctx, kv, cos, sin, interleaved=False, seqlen_offsets: Union[int, torch.Tensor] = 0): - batch, seqlen, two, nheads, headdim = kv.shape - assert two == 2 - k = kv[:, :, 0] - apply_rotary( - k, cos, sin, seqlen_offsets=seqlen_offsets, interleaved=interleaved, inplace=True - ) - if isinstance(seqlen_offsets, int): - ctx.save_for_backward(cos, sin) # Can't save int with save_for_backward - ctx.seqlen_offsets = seqlen_offsets - else: - ctx.save_for_backward(cos, sin, seqlen_offsets) - ctx.seqlen_offsets = None - ctx.interleaved = interleaved - return kv - - @staticmethod - def backward(ctx, dkv): - seqlen_offsets = ctx.seqlen_offsets - if seqlen_offsets is None: - cos, sin, seqlen_offsets = ctx.saved_tensors - else: - cos, sin = ctx.saved_tensors - apply_rotary( - dkv[:, :, 0], - cos, - sin, - seqlen_offsets=seqlen_offsets, - interleaved=ctx.interleaved, - inplace=True, - conjugate=True, - ) - return dkv, None, None, None, None - - -apply_rotary_emb_kv_ = ApplyRotaryEmbKV_.apply - - -def apply_rotary_emb_kv_( - kv, - cos, - sin, - interleaved=False, - seqlen_offsets: Union[int, torch.Tensor] = 0, -): - """ - Arguments: - kv: (batch_size, seqlen, 2, nheads, headdim) - cos, sin: (seqlen, rotary_dim / 2) - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead of - 1st half and 2nd half (GPT-NeoX style). - seqlen_offsets: (batch_size,) or int. Each sequence in Q and K is shifted by this amount. - Most commonly used in inference when we have KV cache. - Return: - kv: (batch_size, seqlen, 2, nheads, headdim) - rotary_dim must be <= headdim - Apply rotary embedding *inplace* to the first rotary_dim of K. - """ - return ApplyRotaryEmbKV_.apply(kv, cos, sin, interleaved, seqlen_offsets) - - -class RotaryEmbedding(torch.nn.Module): - """ - The rotary position embeddings from RoFormer_ (Su et. al). - A crucial insight from the method is that the query and keys are - transformed by rotation matrices which depend on the relative positions. - - Other implementations are available in the Rotary Transformer repo_ and in - GPT-NeoX_, GPT-NeoX was an inspiration - - .. _RoFormer: https://arxiv.org/abs/2104.09864 - .. _repo: https://github.com/ZhuiyiTechnology/roformer - .. _GPT-NeoX: https://github.com/EleutherAI/gpt-neox - - If scale_base is not None, this implements XPos (Sun et al., https://arxiv.org/abs/2212.10554). - A recommended value for scale_base is 512: https://github.com/HazyResearch/flash-attention/issues/96 - Reference: https://github.com/sunyt32/torchscale/blob/main/torchscale/component/xpos_relative_position.py - """ - - def __init__( - self, - dim: int, - base=10000.0, - interleaved=False, - scale_base=None, - device=None, - ): - """ - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead - of 1st half and 2nd half (GPT-NeoX style). - """ - super().__init__() - self.dim = dim - self.base = float(base) - # Generate and save the inverse frequency buffer (non trainable) - inv_freq = self._compute_inv_freq(device) - self.register_buffer("inv_freq", inv_freq, persistent=False) - self.interleaved = interleaved - self.scale_base = scale_base - scale = ( - (torch.arange(0, dim, 2, device=device, dtype=torch.float32) + 0.4 * dim) / (1.4 * dim) - if scale_base is not None - else None - ) - self.register_buffer("scale", scale, persistent=False) - - self._seq_len_cached = 0 - self._cos_cached = None - self._sin_cached = None - self._cos_k_cached = None - self._sin_k_cached = None - - def _compute_inv_freq(self, device=None): - return 1.0 / ( - self.base - ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim) - ) - - def _update_cos_sin_cache(self, seqlen, device=None, dtype=None): - # Reset the tables if the sequence length has changed, - # if we're on a new device (possibly due to tracing for instance), - # or if we're switching from inference mode to training - if ( - seqlen > self._seq_len_cached - or self._cos_cached is None - or self._cos_cached.device != device - or self._cos_cached.dtype != dtype - or (self.training and self._cos_cached.is_inference()) - ): - self._seq_len_cached = seqlen - # We want fp32 here, not self.inv_freq.dtype, since the model could be loaded in bf16 - # And the output of arange can be quite large, so bf16 would lose a lot of precision. - t = torch.arange(seqlen, device=device, dtype=torch.float32) - # We want fp32 here as well since inv_freq will be multiplied with t, and the output - # will be large. Having it in bf16 will lose a lot of precision and cause the - # cos & sin output to change significantly. - # We want to recompute self.inv_freq if it was not loaded in fp32 - if self.inv_freq.dtype != torch.float32: - inv_freq = self._compute_inv_freq(device=device) - else: - inv_freq = self.inv_freq - # Don't do einsum, it converts fp32 to bf16 under AMP - # freqs = torch.einsum("i,j->ij", t, self.inv_freq) - freqs = torch.outer(t, inv_freq) - if self.scale is None: - self._cos_cached = torch.cos(freqs).to(dtype) - self._sin_cached = torch.sin(freqs).to(dtype) - else: - power = ( - torch.arange(seqlen, dtype=self.scale.dtype, device=self.scale.device) - - seqlen // 2 - ) / self.scale_base - scale = self.scale.to(device=power.device) ** rearrange(power, "s -> s 1") - # We want the multiplication by scale to happen in fp32 - self._cos_cached = (torch.cos(freqs) * scale).to(dtype) - self._sin_cached = (torch.sin(freqs) * scale).to(dtype) - self._cos_k_cached = (torch.cos(freqs) / scale).to(dtype) - self._sin_k_cached = (torch.sin(freqs) / scale).to(dtype) - - def forward( - self, - qkv: torch.Tensor, - kv: Optional[torch.Tensor] = None, - seqlen_offset: Union[int, torch.Tensor] = 0, - max_seqlen: Optional[int] = None, - num_heads_q: Optional[int] = None, - ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: - """ - qkv: (batch, seqlen, 3, nheads, headdim) or (batch, seqlen, num_heads_q + 2 * num_heads_k, headdim) - if kv is none, else it's just q of shape (batch, seqlen, nheads, headdim). - If qkv has shape (batch, seqlen, num_heads_q + 2 * num_heads_k, headdim) (e.g. MQA / GQA), - then num_heads_q must be provided. - kv: (batch, seqlen, 2, nheads, headdim) - seqlen_offset: (batch_size,) or int. Each sequence in x is shifted by this amount. - Most commonly used in inference when we have KV cache. - If it's a tensor of shape (batch_size,), then to update the cos / sin cache, one - should pass in max_seqlen, which will update the cos / sin cache up to that length. - Apply rotary embedding *inplace* to qkv and / or kv. - """ - seqlen = qkv.shape[1] - if max_seqlen is not None: - self._update_cos_sin_cache(max_seqlen, device=qkv.device, dtype=qkv.dtype) - elif isinstance(seqlen_offset, int): - self._update_cos_sin_cache(seqlen + seqlen_offset, device=qkv.device, dtype=qkv.dtype) - if kv is None: - return apply_rotary_emb_qkv_( - qkv, - self._cos_cached, - self._sin_cached, - self._cos_k_cached if self.scale is not None else None, - self._sin_k_cached if self.scale is not None else None, - interleaved=self.interleaved, - seqlen_offsets=seqlen_offset, - num_heads_q=num_heads_q, - ) - else: - q = qkv - q = apply_rotary_emb_func( - q, - self._cos_cached, - self._sin_cached, - interleaved=self.interleaved, - inplace=True, - seqlen_offsets=seqlen_offset, - ) - kv = apply_rotary_emb_kv_( - kv, - self._cos_cached if self.scale is None else self._cos_k_cached, - self._sin_cached if self.scale is None else self._sin_k_cached, - interleaved=self.interleaved, - seqlen_offsets=seqlen_offset, - ) - return q, kv diff --git a/build/torch212-cxx11-cu130-aarch64-linux/metadata.json b/build/torch212-cxx11-cu130-aarch64-linux/metadata.json deleted file mode 100644 index 70274076109309fb354fe5f1a26665d21b9a8dbb..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu130-aarch64-linux/metadata.json +++ /dev/null @@ -1,42 +0,0 @@ -{ - "name": "flash-attn2", - "id": "_flash_attn2_cuda_f12afc9", - "version": 1, - "license": "BSD-3-Clause", - "python-depends": [], - "backend": { - "type": "cuda", - "archs": [ - "10.0", - "12.0", - "8.0", - "9.0" - ] - }, - "digest": { - "algorithm": "sha256", - "files": { - "__init__.py": "Cd9MEdd/WUYnBmfiQpC4SyPqAxrTilTjvsvswNju4zc=", - "_flash_attn2_cuda_f12afc9.abi3.so": "CPoMIyxwZLjYfUMbuKrfLU6S95EgZI9tOSclzlIuOjw=", - "_ops.py": "GQbmTTeSOU+7Q2nI68Hz+1RvEogyfEFccx1ubUYw+2s=", - "bert_padding.py": "gF1EmsdJ+HpQ86MRQ4VxDw+Sb/RVISdQALdNnoByHlw=", - "flash_attn2/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY=", - "flash_attn_interface.py": "W7a9bHR2ly6QL3aKqsBqxWA7M7ZWYcZPmtsPJu6rAfY=", - "layers/__init__.py": "47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU=", - "layers/patch_embed.py": "H58CgME/qSOPTZLOG08wFgrQS1j34pvNwMPrkTj3Ek4=", - "layers/rotary.py": "rFIdNrZQFNawwLD6d7t5LE2m+M93wwDZ8EuExw+gfD8=", - "ops/__init__.py": "47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU=", - "ops/activations.py": "t5lzNg1In8LP6bKeTnyeMizwqjv27JGbJ6ylPdGvZYg=", - "ops/fused_dense.py": "BqT8VTSCGdLJAjh6WBMkCRFYrdVodPeoElWBvH6AIzA=", - "ops/layer_norm.py": "zr7NXIm+2mtEynTp1CS0fbFGI2Mqdp41dY4AfDWF6EQ=", - "ops/rms_norm.py": "XEnihcj0a4aSz4LO55m5iKGVn4HKTeKN8TIyHjuDgxI=", - "ops/triton/__init__.py": "AbpHGcgLb+kRsJGnwFEktk7uzpZOCcBY74+YBdrKVGs=", - "ops/triton/cross_entropy.py": "y56nmuOtpSNxM7sDZcUSgXKIsIEayfuvrU8bAVog/nw=", - "ops/triton/k_activations.py": "+Z3vIyO4JkqBMipKsPvhzmxljtBdIhJCsl/M+/ESqBo=", - "ops/triton/layer_norm.py": "BBQdJiqopbIS12viUqzx3SvOQCNbDSAYLJXhpyMGyC8=", - "ops/triton/linear.py": "OtRvKz8xdpl+7v3q/ZTaS9fdBt9XrzMyapgRr50uBbM=", - "ops/triton/mlp.py": "uWA+D7HyuRUsOT+4nhi3vvCDNZnNEXwYcHMQ02UREfQ=", - "ops/triton/rotary.py": "0YSrlbep0lJmzLpTiANs0slTFeGcQQqR+el6uzbSLvI=" - } - } -} \ No newline at end of file diff --git a/build/torch212-cxx11-cu130-aarch64-linux/metadata.json.sigstore b/build/torch212-cxx11-cu130-aarch64-linux/metadata.json.sigstore deleted file mode 100644 index bb1f01244fbaeee334eee74fcd80171bf9fb2c8d..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu130-aarch64-linux/metadata.json.sigstore +++ /dev/null @@ -1 +0,0 @@ -{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHSzCCBtGgAwIBAgIUNKcOn66lXwzSJNxN8MDEs5DwkHswCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwNjIzMjM0MDQ2WhcNMjYwNjIzMjM1MDQ2WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEufXYt2XpReUz+6Q9NqWAXBYaReKMQfr8oQIpRsXCvhf1zsaNAuFdQGCcPvzcLqDOIxjFHyC/vE+w43IfXzfhJKOCBfAwggXsMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUYB63C+KvLi/bdXBMrfngzMonMqIwHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoZjEyYWZjOTJjNDBiYmQ3NmQ5OTBkOTExZWJhYmYyOGU0OTBhN2JkZTATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoZjEyYWZjOTJjNDBiYmQ3NmQ5OTBkOTExZWJhYmYyOGU0OTBhN2JkZTAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoZjEyYWZjOTJjNDBiYmQ3NmQ5OTBkOTExZWJhYmYyOGU0OTBhN2JkZTAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKGYxMmFmYzkyYzQwYmJkNzZkOTkwZDkxMWViYWJmMjhlNDkwYTdiZGUwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMjgwMzYyNjg5MzUvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBigYKKwYBBAHWeQIEAgR8BHoAeAB2AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABnvbbluoAAAQDAEcwRQIhAMP+6CP0OKl9/CdHZ4Q1k+rOe9C/IZ4gi8XwM4eoEG/BAiAy8VI3b8KcKJ1kewNoG7ceDBj7WLnp7S9581sjJLmFYzAKBggqhkjOPQQDAwNoADBlAjBx6SA85b+rY9eNtMX9Dhgzk6RsVeGGZAcU3DRFfY6QNo4GC7P/me17G1oBhmczkngCMQCi2iqn7WmymdUue38zIGUzxtlHSsHhSQshVsCGrqjnaeH7Cw7++bnByUbcpZU4TeY="}, "tlogEntries":[{"logIndex":"1933040231", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1782258047", "inclusionPromise":{"signedEntryTimestamp":"MEQCICh82RDzwljbIgqC3alMfkJRsQjAbWqSxtqaypuc3fRWAiBRV+mHDRC++NAzpyMc3DZfM8NRtHbBGoeDjVKN2D24Vg=="}, "inclusionProof":{"logIndex":"1811135969", "rootHash":"T4EnKTc/w5hw6SUeUfQ5Qs0oKPqBzBN8SYi1xfWFXmY=", "treeSize":"1811136052", "hashes":["t0vewF3dxH9DU0jULGrpK1540LbQ3hG9Kom5rGSBMaA=", "1BLIcg7DuacGIXy+KkFERC8N7eHulhpmmV6RCrFCK78=", "7TzwlcMQAYBZuRS//WqFisXpQ+BQnYGNr0RylCyiUiM=", "ci4ymBdX/0soolJXLbT79r2A/+L2qa76PFNRny5J9ZA=", "lDP8swxQjx/a5pYb2Z96Pyrf6OX7r/K9iSj6BYKD6dY=", "yOGrUSZI6pBF0Xbn/9j60RBR1q/uMt/7f6cbhLzE2L0=", "sOZhXdtH005zZeEhB3WSm31g3XhIBtHEc0wkelplZ5o=", "U4mkhZHvT25SXQTBLzez9yGo/3ze3i1Qb6sz7rAWTFI=", "RiSyDi6XzM92vKCyZ+HZ2MpfN2m2GsxRHoN79kocf9Y=", "eXVohfhx91i0ODjUsYF/ke04gJYA6i7oszgI5L4VC/E=", "LuDs4vs9YQGnXlArOp7boDhzWABrNJ1pOn55UTTP5B8=", "WD8KwmQq5R8K69koYj5k4K7EKo2VF3LqJcx2xoXcd7o=", "99JTcUz0aJeZZocenAppVOLl8v91qMn0v7CmEjq9lws=", "OSDrt6jTKV6LYAkaG63p2BFz4SmbNmGb3uerhF5pOpQ=", "yehxDZkBWOGIzNDb93qxmQa6VBVo7kacVSoHyXMm3ZI=", "oUh+N71k3jWcTmlDDB6Esa5uqlCnbmzxTmDdut2A16M=", "2BsqkdK2omLvu+CpsFjpwqGi6m22H/7zcwCVkPwslM8=", "S8Oe9ZicoJqbVuL1I0dIdTgVpLc82Cm2CMOcV83sKVs=", "E1O2ilWlpFRLWGdbmJUMDZ2qbGu/eDZ12m6YjQ9/JPA=", "nyv2TDlfkz0XfRwjY98zsbUbOkBZ+Tgz7PML0c2uN14=", "OMfDh4k9YD0E+t4sjqJsh+vg8D1ueo1cC07exyxmr6s=", "+vHKm4wNwGJ46ceyZ8mePD7Vn1SediStNHCStRoTSWQ=", "5HwWjuGNDisetbiWM0cWIXM65TIM3EJlqk13ntCQlcU=", "NDIRlqxO0IaGr94q86MIU7VCAN3yEEgOHS4loya5lu4=", "daxmZaajRpZV+JxHiOYZhJBiSKN5ucqjh2WnGbHhirw=", "DOCeoSMovIvLExkhIvisow9AuNXgeWs4ECkyR6EcqYU="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n1811136052\nT4EnKTc/w5hw6SUeUfQ5Qs0oKPqBzBN8SYi1xfWFXmY=\n\n— rekor.sigstore.dev wNI9ajBEAiBpyNdCsfFlbq6PEKCLQM7tGi6e8tpEsRSR9x6FbGZa0gIgHurF5Tgre9I4k7EG4Nf6Ru29qJoT8eHlcSfTqUvQlX8=\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiJjNjc3N2Q3ZDc4Nzg2YjM3NTljZjlhZjdhMjg5YTU2ZWQ4YmVhOTM4ODNjZTQxNjAzZDdkN2U1MWY3NDQ3NGFjIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FWUNJUUN5VGM2NjZPZ0NOaHltTmNXYnpnVXdHK3pJWmtRenBSRUJySFR6SldqWm5BSWhBT1I1eDRDZjRpVVBmR1phK2dCbUI5MzcwZm95bkhPK2ZhUzRhRXdJVkRjRyIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRla05EUW5SSFowRjNTVUpCWjBsVlRrdGpUMjQyTm14WWQzcFRTazU0VGpoTlJFVnpOVVIzYTBoemQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDVxU1hwTmFrMHdUVVJSTWxkb1kwNU5hbGwzVG1wSmVrMXFUVEZOUkZFeVYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVYxWmxoWmRESlljRkpsVlhvck5sRTVUbkZYUVZoQ1dXRlNaVXROVVdaeU9HOVJTWEFLVW5OWVEzWm9aakY2YzJGT1FYVkdaRkZIUTJOUWRucGpUSEZFVDBsNGFrWkllVU12ZGtVcmR6UXpTV1pZZW1ab1NrdFBRMEptUVhkbloxaHpUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlZaUWpZekNrTXJTM1pNYVM5aVpGaENUWEptYm1kNlRXOXVUWEZKZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOWFha1Y1V1ZkYWFrOVVTbXBPUkVKcFdXMVJNMDV0VVRWUFZFSnJDazlVUlhoYVYwcG9XVzFaZVU5SFZUQlBWRUpvVGpKS2ExcFVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMXBxUlhsWlYxcHFUMVJLYWs1RVFtbFpiVkV6VG0xUk5VOVVRbXRQVkVWNFdsZEthRmx0V1hrS1QwZFZNRTlVUW1oT01rcHJXbFJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMXBxUlhsWlYxcHFUMVJLYWs1RVFtbFpiVkV6VG0xUk5VOVVRbXNLVDFSRmVGcFhTbWhaYlZsNVQwZFZNRTlVUW1oT01rcHJXbFJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwZFplRTF0Um0wS1dYcHJlVmw2VVhkWmJVcHJUbnBhYTA5VWEzZGFSR3Q0VFZkV2FWbFhTbTFOYW1oc1RrUnJkMWxVWkdsYVIxVjNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOYW1kM1RYcFplVTVxWnpWTmVsVjJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwWjFsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTT0VKSWIwRUtaVUZDTWtGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtNTJZbUpzZFc5QlFVRlJSQXBCUldOM1VsRkphRUZOVUNzMlExQXdUMHRzT1M5RFpFaGFORkV4YXl0eVQyVTVReTlKV2pSbmFUaFlkMDAwWlc5RlJ5OUNRV2xCZVRoV1NUTmlPRXRqQ2t0S01XdGxkMDV2UnpkalpVUkNhamRYVEc1d04xTTVOVGd4YzJwS1RHMUdXWHBCUzBKblozRm9hMnBQVUZGUlJFRjNUbTlCUkVKc1FXcENlRFpUUVRnS05XSXJjbGs1WlU1MFRWZzVSR2huZW1zMlVuTldaVWRIV2tGalZUTkVVa1ptV1RaUlRtODBSME0zVUM5dFpURTNSekZ2UW1odFkzcHJibWREVFZGRGFRb3lhWEZ1TjFkdGVXMWtWWFZsTXpoNlNVZFZlbmgwYkVoVGMwaG9VMUZ6YUZaelEwZHljV3B1WVdWSU4wTjNOeXNyWW01Q2VWVmlZM0JhVlRSVVpWazlDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyjADAgEAMIICwQYJKoZIhvcNAQcCoIICsjCCAq4CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgY2fFEwVhrChD6w+tK5HiIcr/fpdA9Ys+g3ixM+fLCEsCFQCgv9br+0iwIiF8Q3IFWmSEalyF3BgPMjAyNjA2MjMyMzQwNDZaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHbMIIB1wIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MDYyMzIzNDA0NlowLwYJKoZIhvcNAQkEMSIEIPuRGEU1G/urW9GB21zs0PMw7ujvyvzLmmEVK9pU1RLcMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRnMGUCMQCcci0DzZoIj9xQ53aw6lEVL39Mt7Gu8RrFDiVTBQvsvSfTYt8HKrTE59GDXoGWihcCMAsvbsMOLnrzwx45YAicwmvYVvm4iCPjMZhZAxawQmdtrAe4OYyDiW/oeocz3JjRBA=="}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"xnd9fXh4azdZz5r3oomlbti+qTiDzkFgPX1+UfdEdKw="}, "signature":"MEYCIQCyTc666OgCNhymNcWbzgUwG+zIZkQzpREBrHTzJWjZnAIhAOR5x4Cf4iUPfGZa+gBmB9370foynHO+faS4aEwIVDcG"}} \ No newline at end of file diff --git a/build/torch212-cxx11-cu130-aarch64-linux/ops/__init__.py b/build/torch212-cxx11-cu130-aarch64-linux/ops/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/build/torch212-cxx11-cu130-aarch64-linux/ops/activations.py b/build/torch212-cxx11-cu130-aarch64-linux/ops/activations.py deleted file mode 100644 index 7c09649fc41e12d5a360c5672825d8380bc7ec80..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu130-aarch64-linux/ops/activations.py +++ /dev/null @@ -1,135 +0,0 @@ -# Copied from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/model/layers/activations.py -import math - -import torch -import torch.nn as nn -import torch.nn.functional as F - -# 1/sqrt(2*pi)-> 0.3989423 -# 1/sqrt(2) -> 0.70710678 -# sqrt(2/pi) -> 0.79788456 - -# this function is tanh approximation of gelu -# actual gelu is: -# x * 0.5 * (1.0 + torch.erf(x * 0.70710678)) -@torch.jit.script -def bias_gelu(y, bias): - x = bias + y - return (x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))).to(dtype=y.dtype) - - -# gradient of tanh approximation of gelu -# gradient of actual gelu is: -# 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x) -@torch.jit.script -def bias_gelu_back(g, y, bias): - """Assume that y has shape (B, D) and bias has shape (D)""" - x = bias + y - tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)) - # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243 - ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( - 1 + tanh_out - ) - grad_y = ff * g - return grad_y.to(dtype=y.dtype), grad_y.sum(dim=(0), dtype=bias.dtype) - - -class GeLUFunction(torch.autograd.Function): - @staticmethod - # bias is an optional argument - def forward(ctx, input, bias): - ctx.save_for_backward(input, bias) - return bias_gelu(input, bias) - - @staticmethod - def backward(ctx, grad_output): - input, bias = ctx.saved_tensors - tmp = bias_gelu_back(grad_output, input, bias) - return tmp, tmp - - -bias_gelu_impl = GeLUFunction.apply - -# this function is tanh approximation of gelu -# actual gelu is: -# x * 0.5 * (1.0 + torch.erf(x * 0.70710678)) -@torch.jit.script -def gelu_fwd(x): - return (x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))).to(dtype=x.dtype) - - -# gradient of tanh approximation of gelu -# gradient of actual gelu is: -# 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x) -@torch.jit.script -def gelu_bwd(g, x): - tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)) - # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243 - ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( - 1 + tanh_out - ) - return (ff * g).to(dtype=x.dtype) - - -class FastGeLUFunction(torch.autograd.Function): - @staticmethod - # bias is an optional argument - def forward(ctx, input): - ctx.save_for_backward(input) - return gelu_fwd(input) - - @staticmethod - def backward(ctx, grad_output): - (input,) = ctx.saved_tensors - tmp = gelu_bwd(grad_output, input) - return tmp - - -fast_gelu_impl = FastGeLUFunction.apply - - -@torch.jit.script -def relu_bwd(g, x): - return torch.where(x >= 0, g, 0.0).to(dtype=x.dtype) - - -@torch.jit.script -def sqrelu_fwd(x): - r = F.relu(x) - return (r * r).to(dtype=x.dtype) - - -@torch.jit.script -def sqrelu_bwd(g, x): - return (2.0 * g * F.relu(x)).to(dtype=x.dtype) - - -swiglu_fwd_codestring = """ -template T swiglu_fwd(T x, T y) { - return float(x) * float(y) / (1.0f + ::exp(-float(x))); -} -""" -swiglu_bwd_codestring = """ -template void swiglu_bwd(T x, T y, T g, T& dx, T& dy) { - float x_sigmoid = 1.0f / (1.0f + ::exp(-float(x))); - dx = x_sigmoid * (1 + float(x) * (1.0f - x_sigmoid)) * float(g) * float(y); - dy = float(x) * x_sigmoid * float(g); -} -""" -swiglu_fwd = torch.cuda.jiterator._create_jit_fn(swiglu_fwd_codestring) -swiglu_bwd = torch.cuda.jiterator._create_multi_output_jit_fn(swiglu_bwd_codestring, num_outputs=2) - - -class SwiGLUFunction(torch.autograd.Function): - - @staticmethod - def forward(ctx, x, y): - ctx.save_for_backward(x, y) - return swiglu_fwd(x, y) - - @staticmethod - def backward(ctx, dout): - x, y = ctx.saved_tensors - return swiglu_bwd(x, y, dout) - -swiglu = SwiGLUFunction.apply diff --git a/build/torch212-cxx11-cu130-aarch64-linux/ops/fused_dense.py b/build/torch212-cxx11-cu130-aarch64-linux/ops/fused_dense.py deleted file mode 100644 index 6b4033d134e4093fe278f7b3f8c7d3128ce9f36d..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu130-aarch64-linux/ops/fused_dense.py +++ /dev/null @@ -1,688 +0,0 @@ -# Copyright (c) 2023, Tri Dao. -# Inspired by https://github.com/NVIDIA/apex/blob/master/apex/fused_dense/fused_dense.py -# We make it work with pytorch amp and with bfloat16. -# The TensorParallel linear modules are inspired by https://github.com/NVIDIA/apex/blob/master/apex/transformer/tensor_parallel/layers.py -from functools import partial -from typing import Optional - -# import fused_dense_cuda # from apex -import fused_dense_lib as fused_dense_cuda -import torch -import torch.nn as nn -import torch.nn.functional as F -from torch import Tensor -from torch.distributed import ProcessGroup - -from flash_attn.utils.torch import custom_fwd, custom_bwd -from flash_attn.ops.activations import gelu_bwd, relu_bwd, sqrelu_bwd, sqrelu_fwd -from flash_attn.utils.distributed import ( - all_gather_raw, - all_reduce, - all_reduce_raw, - reduce_scatter, - reduce_scatter_raw, -) - - -class FusedDenseFunc(torch.autograd.Function): - @staticmethod - @custom_fwd - def forward( - ctx, x, weight, bias, return_residual=False, process_group=None, sequence_parallel=True - ): - """ - If process_group is not None and sequence_parallel=True, we're doing Tensor Parallel - with sequence parallelism: we do an all_gather_raw of x before doing the matmul. - """ - ctx.compute_weight_gradient = weight.requires_grad - ctx.return_residual = return_residual - ctx.process_group = process_group - ctx.sequence_parallel = sequence_parallel - - if torch.is_autocast_enabled(): - x = x.to(dtype=torch.get_autocast_gpu_dtype()) - x = x.contiguous() - if process_group is not None and sequence_parallel: - # We want to kick off the all_gather early, before weight dtype conversion - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - else: - total_x = x - - if torch.is_autocast_enabled(): - weight = weight.to(dtype=torch.get_autocast_gpu_dtype()) - bias = bias.to(dtype=torch.get_autocast_gpu_dtype()) if bias is not None else None - weight = weight.contiguous() - if process_group is not None and sequence_parallel: - handle_x.wait() - batch_shape, n = total_x.shape[:-1], total_x.shape[-1] - batch_dim = batch_shape.numel() - # https://github.com/pytorch/pytorch/blob/5b51849b48a7dbccd297286cc0110def4706f9e7/aten/src/ATen/native/cuda/Blas.cpp#L174 - if min(batch_dim, n, *weight.shape) > 65535 * 32: - raise RuntimeError("fused_dense only supports matrix dims <= 2M") - output = F.linear(total_x, weight, bias) - if ctx.compute_weight_gradient: - ctx.save_for_backward(x, weight) - else: - ctx.save_for_backward(weight) - return output if not return_residual else (output, x) - - @staticmethod - @custom_bwd - def backward(ctx, grad_output, *args): - grad_output = grad_output.contiguous() - if ctx.return_residual: - (grad_input,) = args - grad_input = grad_input.contiguous() - process_group = ctx.process_group - sequence_parallel = ctx.sequence_parallel - if ctx.compute_weight_gradient: - x, weight = ctx.saved_tensors - if process_group is not None and sequence_parallel: - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - else: - total_x = x - else: - (weight,) = ctx.saved_tensors - total_x = None - batch_shape = grad_output.shape[:-1] - batch_dim = batch_shape.numel() - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - if ctx.needs_input_grad[0]: - if not ctx.return_residual: - grad_input = F.linear(grad_output, weight.t()) - else: - grad_input = torch.addmm( - grad_input.reshape(batch_dim, grad_input.shape[-1]), grad_output, weight - ) - grad_input = grad_input.reshape(*batch_shape, grad_input.shape[-1]) - if process_group is not None: - reduce_fn = reduce_scatter_raw if sequence_parallel else all_reduce_raw - grad_input, handle_grad_input = reduce_fn(grad_input, process_group, async_op=True) - else: - grad_input = None - if ctx.needs_input_grad[1]: - assert ctx.compute_weight_gradient - if process_group is not None and sequence_parallel: - handle_x.wait() - grad_weight, grad_bias = fused_dense_cuda.linear_bias_wgrad( - total_x.reshape(batch_dim, total_x.shape[-1]), grad_output, ctx.needs_input_grad[2] - ) - else: - grad_weight = None - grad_bias = grad_output if ctx.needs_input_grad[2] else None - if process_group is not None and ctx.needs_input_grad[0]: - handle_grad_input.wait() - return grad_input, grad_weight, grad_bias, None, None, None - - -def fused_dense_func( - x: Tensor, - weight: Tensor, - bias: Optional[Tensor] = None, - return_residual: bool = False, - process_group: Optional[ProcessGroup] = None, - sequence_parallel: bool = True, -): - dtype_eligible = x.dtype in [torch.float16, torch.bfloat16] or ( - x.dtype == torch.float32 and torch.is_autocast_enabled() - ) - if x.is_cuda and weight.is_cuda and (bias is None or bias.is_cuda) and dtype_eligible: - return FusedDenseFunc.apply( - x, weight, bias, return_residual, process_group, sequence_parallel - ) - else: - assert process_group is None - out = F.linear(x, weight, bias) - return out if not return_residual else (out, x) - - -class FusedDense(nn.Linear): - def __init__( - self, - in_features: int, - out_features: int, - bias: bool = True, - return_residual: bool = False, - device=None, - dtype=None, - ) -> None: - super().__init__(in_features, out_features, bias=bias, device=device, dtype=dtype) - self.return_residual = return_residual - - def forward(self, x, process_group=None): - """ - If process_group is not None, we're doing Tensor Parallel with sequence parallelism: - we do an all_gather of x before doing the matmul. - """ - return fused_dense_func( - x, - self.weight, - self.bias, - return_residual=self.return_residual, - process_group=process_group, - ) - - -class ColumnParallelLinear(nn.Linear): - def __init__( - self, - in_features: int, - out_features: int, - process_group: ProcessGroup, - bias: bool = True, - sequence_parallel=True, - multiple_of=1, - device=None, - dtype=None, - ) -> None: - world_size = torch.distributed.get_world_size(process_group) - if out_features % multiple_of: - raise ValueError(f"out_features ({out_features}) must be a multiple of {multiple_of}") - multiple = out_features // multiple_of - # We want to split @multiple across world_size, but it could be an uneven split - div = multiple // world_size - mod = multiple % world_size - # The first @mod ranks get @div + 1 copies, the rest get @div copies - local_multiple = div + int(torch.distributed.get_rank(process_group) < mod) - super().__init__( - in_features, local_multiple * multiple_of, bias=bias, device=device, dtype=dtype - ) - self.process_group = process_group - self.sequence_parallel = sequence_parallel - - def forward(self, x): - # If self.sequence_parallel is True, we're doing Tensor Parallel with sequence parallelism: - # we do an all_gather of x before doing the matmul. - # If not, then the input is already gathered. - return fused_dense_func( - x, - self.weight, - self.bias, - process_group=self.process_group, - sequence_parallel=self.sequence_parallel, - ) - - -class RowParallelLinear(nn.Linear): - def __init__( - self, - in_features: int, - out_features: int, - process_group: ProcessGroup, - bias: bool = True, - sequence_parallel=True, - multiple_of=1, - device=None, - dtype=None, - ) -> None: - world_size = torch.distributed.get_world_size(process_group) - rank = torch.distributed.get_rank(process_group) - if in_features % multiple_of: - raise ValueError(f"in_features ({in_features}) must be a multiple of {multiple_of}") - multiple = in_features // multiple_of - # We want to split @multiple across world_size, but it could be an uneven split - div = multiple // world_size - mod = multiple % world_size - # The first @mod ranks get @div + 1 copies, the rest get @div copies - local_multiple = div + int(torch.distributed.get_rank(process_group) < mod) - # Only rank 0 will have bias - super().__init__( - local_multiple * multiple_of, - out_features, - bias=bias and rank == 0, - device=device, - dtype=dtype, - ) - self.process_group = process_group - self.sequence_parallel = sequence_parallel - - def forward(self, x): - """ - We're doing Tensor Parallel with sequence parallelism: we do the matmul and then - a reduce_scatter of the result. - """ - out = fused_dense_func(x, self.weight, self.bias) - reduce_fn = reduce_scatter if self.sequence_parallel else all_reduce - return reduce_fn(out, self.process_group) - - -class FusedMLPFunc(torch.autograd.Function): - @staticmethod - @custom_fwd - def forward( - ctx, - x, - weight1, - bias1, - weight2, - bias2, - activation="gelu_approx", - save_pre_act=True, - return_residual=False, - checkpoint_lvl=0, - heuristic=0, - process_group=None, - sequence_parallel=True, - ): - """ - If process_group is not None and sequence_parallel=True, we're doing Tensor Parallel - with sequence parallelism: we do an all_gather of x before doing the matmul. - If sequence_parallel=False, then the input is already gathered. - - checkpoint_lvl: - 0: no recomputation in the bwd - 1: recompute gelu_out / relu_out in the bwd - 2: recompute pre_act and gelu_out / relu_out in the bwd - """ - assert -1 <= heuristic <= 4 - assert activation in ["gelu_approx", "relu", "sqrelu"] - if activation == "sqrelu": - assert heuristic == -1 - if not save_pre_act: - checkpoint_lvl = 2 - assert checkpoint_lvl in [0, 1, 2] - ctx.return_residual = return_residual - ctx.process_group = process_group - ctx.sequence_parallel = sequence_parallel - ctx.checkpoint_lvl = checkpoint_lvl - ctx.activation = activation - ctx.heuristic = heuristic - - if torch.is_autocast_enabled(): - x = x.to(dtype=torch.get_autocast_gpu_dtype()) - x = x.contiguous() - if process_group is not None and sequence_parallel: - # We want to kick off the all_gather early, before weight dtype conversion - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - else: - total_x = x - - if torch.is_autocast_enabled(): - dtype = torch.get_autocast_gpu_dtype() - weight1, weight2 = [a.to(dtype=dtype) for a in [weight1, weight2]] - bias1 = bias1.to(dtype=dtype) if bias1 is not None else None - bias2 = bias2.to(dtype=dtype) if bias2 is not None else None - weight1 = weight1.contiguous() - bias1 = bias1.contiguous() if bias1 is not None else None - weight2 = weight2.contiguous() - bias2 = bias2.contiguous() if bias2 is not None else None - if process_group is not None and sequence_parallel: - handle_x.wait() - batch_shape, n = total_x.shape[:-1], total_x.shape[-1] - batch_dim = batch_shape.numel() - # https://github.com/pytorch/pytorch/blob/5b51849b48a7dbccd297286cc0110def4706f9e7/aten/src/ATen/native/cuda/Blas.cpp#L174 - if min(batch_dim, n, *weight1.shape, *weight2.shape) > 65535 * 32: - raise RuntimeError("fused_dense only supports matrix dims <= 2M") - if heuristic == -1: - pre_act = F.linear(total_x, weight1, bias1) - activation_fn = ( - partial(F.gelu, approximate="tanh") - if activation == "gelu_approx" - else (sqrelu_fwd if activation == "sqrelu" else F.relu) - ) - with torch.jit.fuser("fuser2"): - output1 = activation_fn(pre_act) - # This is before adding bias1 - # pre_act = F.linear(total_x.reshape(batch_dim, n), weight1) - # with torch.jit.fuser('fuser2'): - # output1 = bias_gelu(pre_act, bias1) - else: - is_gelu = activation == "gelu_approx" - output1, *rest = fused_dense_cuda.linear_act_forward( - total_x.reshape(batch_dim, n), weight1, bias1, is_gelu, save_pre_act, heuristic - ) - if save_pre_act: - pre_act = rest[0] - output2 = F.linear(output1, weight2, bias2) - if checkpoint_lvl == 0 or (checkpoint_lvl == 1 and activation == "relu"): - # For RELU the pre_act is very small (just a bit-mask) so we just save it - ctx.save_for_backward(x, weight1, weight2, pre_act, output1) - elif checkpoint_lvl == 1: - ctx.save_for_backward(x, weight1, weight2, pre_act) - elif checkpoint_lvl == 2: - ctx.save_for_backward(x, weight1, weight2, bias1) - output2 = output2.reshape(*batch_shape, output2.shape[-1]) - return output2 if not return_residual else (output2, x) - - @staticmethod - @custom_bwd - def backward(ctx, grad_output, *args): - grad_output = grad_output.contiguous() - checkpoint_lvl = ctx.checkpoint_lvl - activation = ctx.activation - activation_fn = ( - partial(F.gelu, approximate="tanh") - if activation == "gelu_approx" - else (sqrelu_fwd if activation == "sqrelu" else F.relu) - ) - if ctx.return_residual: - (grad_input,) = args - grad_input = grad_input.contiguous() - process_group = ctx.process_group - sequence_parallel = ctx.sequence_parallel - x, weight1, weight2, *rest = ctx.saved_tensors - if process_group is None or not sequence_parallel: - total_x = x - batch_shape = grad_output.shape[:-1] - batch_dim = batch_shape.numel() - if checkpoint_lvl in [0, 1]: - if process_group is not None and sequence_parallel: - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - if checkpoint_lvl == 0 or (checkpoint_lvl == 1 and activation == "relu"): - pre_act, output1 = rest - elif checkpoint_lvl == 1: - (pre_act,) = rest - with torch.jit.fuser("fuser2"): - output1 = activation_fn(pre_act) - elif checkpoint_lvl == 2: - (bias1,) = rest - if process_group is not None and sequence_parallel: - total_x, _ = all_gather_raw(x, process_group) - if ctx.heuristic == -1: - pre_act = F.linear(total_x, weight1, bias1) - with torch.jit.fuser("fuser2"): - output1 = activation_fn(pre_act) - else: - output1, pre_act = fused_dense_cuda.linear_act_forward( - total_x.reshape(batch_dim, total_x.shape[-1]), - weight1, - bias1, - activation == "gelu_approx", - True, - ctx.heuristic, - ) - - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - output1 = output1.reshape(batch_dim, output1.shape[-1]) - pre_act = pre_act.reshape(batch_dim, pre_act.shape[-1]) - if ctx.needs_input_grad[3]: - grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad( - output1, grad_output, ctx.needs_input_grad[4] - ) - else: - grad_weight2 = None - grad_bias2 = grad_output if ctx.needs_input_grad[4] else None - if ctx.heuristic == -1: - # grad_pre_act = matmul_dgelu(grad_output, weight2, pre_act) - grad_output1 = F.linear(grad_output, weight2.t()) - activation_grad_fn = ( - gelu_bwd - if activation == "gelu_approx" - else (sqrelu_bwd if activation == "sqrelu" else relu_bwd) - ) - with torch.jit.fuser("fuser2"): - grad_pre_act = activation_grad_fn(grad_output1, pre_act) - else: - # The cublasLt epilogue has to compute both gelu/relu grad and bias grad, we can't - # just compute gelu/relu grad - grad_pre_act, grad_bias1 = fused_dense_cuda.bias_act_linear_dgrad_bgrad( - weight2, grad_output, pre_act, activation == "gelu_approx", ctx.heuristic - ) - if not ctx.needs_input_grad[2]: - grad_bias1 = None - if ctx.needs_input_grad[0]: - if not ctx.return_residual: - grad_input = F.linear(grad_pre_act, weight1.t()) - else: - grad_input = torch.addmm( - grad_input.reshape(batch_dim, grad_input.shape[-1]), grad_pre_act, weight1 - ) - grad_input = grad_input.reshape(*batch_shape, grad_input.shape[-1]) - if process_group is not None: - reduce_fn = reduce_scatter_raw if sequence_parallel else all_reduce_raw - grad_input, handle_grad_input = reduce_fn(grad_input, process_group, async_op=True) - else: - grad_input = None - if ctx.heuristic == -1: - if ctx.needs_input_grad[1]: - if process_group is not None and sequence_parallel and checkpoint_lvl != 2: - handle_x.wait() - grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_wgrad( - total_x.reshape(batch_dim, total_x.shape[-1]), - grad_pre_act, - ctx.needs_input_grad[2], - ) - else: - grad_weight1 = None - grad_bias1 = grad_pre_act if ctx.needs_input_grad[2] else None - else: - if ctx.needs_input_grad[1]: - if process_group is not None and sequence_parallel and checkpoint_lvl != 2: - handle_x.wait() - grad_weight1 = F.linear( - grad_pre_act.t(), total_x.reshape(batch_dim, total_x.shape[-1]).t() - ) - else: - grad_weight1 = None - if process_group is not None and ctx.needs_input_grad[0]: - handle_grad_input.wait() - return ( - grad_input, - grad_weight1, - grad_bias1, - grad_weight2, - grad_bias2, - None, - None, - None, - None, - None, - None, - None, - ) - - -def fused_mlp_func( - x: Tensor, - weight1: Tensor, - weight2: Tensor, - bias1: Optional[Tensor] = None, - bias2: Optional[Tensor] = None, - activation: str = "gelu_approx", - save_pre_act: bool = True, - return_residual: bool = False, - checkpoint_lvl: int = 0, - heuristic: int = 0, - process_group: Optional[ProcessGroup] = None, - sequence_parallel: bool = True, -): - assert activation in ["gelu_approx", "relu", "sqrelu"] - dtype_eligible = x.dtype in [torch.float16, torch.bfloat16] or ( - x.dtype == torch.float32 and torch.is_autocast_enabled() - ) - # If we save pre-activation, dimension must be divisible by 128 (relu) or 8 (gelu) - dim_eligible = not save_pre_act or (x.shape[-1] % (128 if activation == "relu" else 8) == 0) - if ( - x.is_cuda - and weight1.is_cuda - and weight2.is_cuda - and (bias1 is None or bias1.is_cuda) - and (bias2 is None or bias2.is_cuda) - and dtype_eligible - and dim_eligible - ): - return FusedMLPFunc.apply( - x, - weight1, - bias1, - weight2, - bias2, - activation, - save_pre_act, - return_residual, - checkpoint_lvl, - heuristic, - process_group, - sequence_parallel, - ) - else: - assert process_group is None - pre_act = F.linear(x, weight1, bias1) - activation_fn = ( - partial(F.gelu, approximate="tanh") - if activation == "gelu_approx" - else partial(F.relu, inplace=True) - ) - output1 = activation_fn(pre_act) - output2 = F.linear(output1, weight2, bias2) - return output2 if not return_residual else (output2, x) - - -class FusedMLP(nn.Module): - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - bias1=True, - bias2=True, - activation="gelu_approx", - return_residual=False, - checkpoint_lvl=0, - heuristic="auto", - device=None, - dtype=None, - ): - """ - If process_group is not None, we're doing Tensor Parallel with sequence parallelism: - we do an all_gather of x before doing the matmul, gelu, then matmul. - Finally we do a reduce_scatter of the output. - - checkpoint_lvl (increasing lvl means slower but more memory saving): - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute pre_act and gelu_out in the bwd - heuristic: - -1: don't fuse gemm + gelu (separate kernel) - 0..4: use this heuristic for the algo section in the fused gemm + gelu - 'auto': heuristic will be picked automatically: - For CUDA >= 11.8, we set heuristic=0 for both fp16 and bf16 for best perf. - For CUDA <= 11.7, we set heuristic=1 for fp16 and heuristic=-1 for bf16. - For H100, we set heuristic=-1 for both fp16 and bf16 as the fused cuBlasLt implementation - is slower than the unfused version. - return_residual: whether to return the input x along with the output. This is for - performance reason: for post-norm architecture, returning the input allows us - to fuse the backward of nn.Linear with the residual connection. - """ - assert checkpoint_lvl in [0, 1, 2] - assert activation in ["gelu_approx", "relu", "sqrelu"] - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features * 4 - self.activation = activation - self.return_residual = return_residual - self.checkpoint_lvl = checkpoint_lvl - self.heuristic = heuristic if activation != "sqrelu" else -1 - self.fc1 = nn.Linear(in_features, hidden_features, bias=bias1, **factory_kwargs) - self.fc2 = nn.Linear(hidden_features, out_features, bias=bias2, **factory_kwargs) - - def forward(self, x, process_group=None): - dtype = x.dtype if not torch.is_autocast_enabled() else torch.get_autocast_gpu_dtype() - if self.heuristic == "auto": - if self.activation == "gelu_approx": - if torch.cuda.get_device_capability("cuda") == (9, 0): - heuristic = -1 - else: - cuda_ver = tuple(map(int, torch.version.cuda.split("."))) - heuristic = 0 if cuda_ver >= (11, 8) else (1 if dtype == torch.float16 else -1) - else: - heuristic = 0 - else: - heuristic = self.heuristic - out = fused_mlp_func( - x, - self.fc1.weight, - self.fc2.weight, - self.fc1.bias, - self.fc2.bias, - activation=self.activation, - save_pre_act=self.training, - return_residual=self.return_residual, - checkpoint_lvl=self.checkpoint_lvl, - heuristic=heuristic, - process_group=process_group, - ) - if self.return_residual: - out, x = out - if process_group is not None: - out = reduce_scatter(out, process_group) - return out if not self.return_residual else (out, x) - - -class ParallelFusedMLP(nn.Module): - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - activation="gelu_approx", - process_group: ProcessGroup = None, - bias1=True, - bias2=True, - sequence_parallel=True, - checkpoint_lvl=0, - heuristic="auto", - device=None, - dtype=None, - ): - """ - process_group is required. We're doing Tensor Parallel with sequence parallelism: - we do an all_gather of x before doing the matmul, gelu, then matmul. - Finally we do a reduce_scatter of the output. - - checkpoint_lvl (increasing lvl means slower but more memory saving): - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute pre_act and gelu_out in the bwd - heuristic: - -1: don't fuse gemm + gelu (separate kernel) - 0..4: use this heuristic for the algo section in the fused gemm + gelu - 'auto': heuristic will be picked automatically: - For CUDA >= 11.8, we set heuristic=0 for both fp16 and bf16 for best perf. - For CUDA <= 11.7, we set heuristic=1 for fp16 and heuristic=-1 for bf16. - """ - assert checkpoint_lvl in [0, 1, 2] - assert activation in ["gelu_approx", "relu", "sqrelu"] - assert process_group is not None - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features * 4 - self.activation = activation - self.process_group = process_group - self.sequence_parallel = sequence_parallel - self.checkpoint_lvl = checkpoint_lvl - self.heuristic = heuristic if activation != "sqrelu" else -1 - self.fc1 = ColumnParallelLinear( - in_features, hidden_features, process_group, bias=bias1, **factory_kwargs - ) - self.fc2 = RowParallelLinear( - hidden_features, out_features, process_group, bias=bias2, **factory_kwargs - ) - - def forward(self, x): - dtype = x.dtype if not torch.is_autocast_enabled() else torch.get_autocast_gpu_dtype() - if self.heuristic == "auto": - if self.activation == "gelu_approx": - cuda_ver = tuple(map(int, torch.version.cuda.split("."))) - heuristic = 0 if cuda_ver >= (11, 8) else (1 if dtype == torch.float16 else -1) - else: - heuristic = 0 - else: - heuristic = self.heuristic - out = fused_mlp_func( - x, - self.fc1.weight, - self.fc2.weight, - self.fc1.bias, - self.fc2.bias, - activation=self.activation, - save_pre_act=self.training, - checkpoint_lvl=self.checkpoint_lvl, - heuristic=heuristic, - process_group=self.process_group, - sequence_parallel=self.sequence_parallel, - ) - reduce_fn = reduce_scatter if self.sequence_parallel else all_reduce - return reduce_fn(out, self.process_group) diff --git a/build/torch212-cxx11-cu130-aarch64-linux/ops/layer_norm.py b/build/torch212-cxx11-cu130-aarch64-linux/ops/layer_norm.py deleted file mode 100644 index 4b6cd798fd02844ef9cd3897f8ab95e490e638bf..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu130-aarch64-linux/ops/layer_norm.py +++ /dev/null @@ -1,800 +0,0 @@ -# Copyright (c) 2022, Tri Dao. -# Adapted from https://github.com/NVIDIA/apex/blob/master/apex/contrib/layer_norm/layer_norm.py - -import dropout_layer_norm -import torch -from torch.nn import init - - -def maybe_align(x, alignment_in_bytes=16): - """Assume that x already has last dim divisible by alignment_in_bytes""" - # TD [2023-07-04] I'm not 100% sure that clone will align the memory - # https://discuss.pytorch.org/t/how-to-ensure-that-tensor-data-ptr-is-aligned-to-16-bytes/183440 - return x if x.data_ptr() % alignment_in_bytes == 0 else x.clone() - - -def _dropout_add_layer_norm_forward( - x0, - residual, - gamma, - beta, - rowscale, - colscale, - dropout_p, - epsilon, - residual_in_fp32=False, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes""" - hidden_size = gamma.numel() - x0mat = x0.view((-1, hidden_size)) - residualmat = residual.view((-1, hidden_size)) if residual is not None else None - rowscale = rowscale.view(-1) if rowscale is not None else None - zmat, xmat, dmask, mu, rsigma = dropout_layer_norm.dropout_add_ln_fwd( - x0mat, - residualmat, - gamma, - beta, - rowscale, - colscale, - None, - None, - dropout_p, - epsilon, - 1.0, - 0, - None, - residual_in_fp32, - is_rms_norm, - ) - # dmask is None if dropout_p == 0.0 - # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype - return zmat, xmat if xmat is not None else x0mat, dmask, mu, rsigma - - -def _dropout_add_layer_norm_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - rowscale, - colscale, - dropout_p, - has_residual, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes - dx == None means that it was a post-norm architecture - (x = drop(x0) + residual was not returned in the fwd). - x0 must not be None if we have colscale. - """ - hidden_size = gamma.numel() - xmat = x.view((-1, hidden_size)) - dzmat = dz.view(xmat.shape) - dxmat = dx.view(xmat.shape) if dx is not None else None - x0mat = x0.view((-1, hidden_size)) if x0 is not None else None - rowscale = rowscale.view(-1) if rowscale is not None else None - if colscale is not None: - assert x0 is not None, "x0 is required to compute the gradient of colscale" - dx0mat, dresidualmat, dgamma, dbeta, _, _, *rest = dropout_layer_norm.dropout_add_ln_bwd( - dzmat, - dxmat, - xmat, - x0mat, - dmask, - mu, - rsigma, - gamma, - rowscale, - colscale, - None, - None, - dropout_p, - 1.0, - 0, - has_residual, - is_rms_norm, - ) - # dresidualmat is None if not has_residual - if colscale is None: - return dx0mat, dresidualmat, dgamma, dbeta - else: - dcolscale = rest[0] - return dx0mat, dresidualmat, dgamma, dbeta, dcolscale - - -def _dropout_add_layer_norm_subset_forward( - x0, - residual, - gamma, - beta, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32=False, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes""" - hidden_size = gamma.numel() - x0mat = x0.view((-1, hidden_size)) - residualmat = residual.view((-1, hidden_size)) if residual is not None else None - x0_subset = x0_subset.view(-1) if x0_subset is not None else None - out_subset = out_subset.view(-1) if out_subset is not None else None - zmat, xmat, dmask, mu, rsigma = dropout_layer_norm.dropout_add_ln_fwd( - x0mat, - residualmat, - gamma, - beta, - None, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - None, - residual_in_fp32, - is_rms_norm, - ) - # dmask is None if dropout_p == 0.0 - # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype - return zmat, xmat if xmat is not None else x0mat, dmask, mu, rsigma - - -def _dropout_add_layer_norm_subset_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - colscale, - x0_subset, - out_subset, - dropout_p, - rowscale_const, - x0_numrows, - has_residual, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes - dx == None means that it was a post-norm architecture - (x = drop(x0) + residual was not returned in the fwd). - x0 must not be None if we have colscale. - """ - hidden_size = gamma.numel() - xmat = x.view((-1, hidden_size)) - dzmat = dz.view(-1, hidden_size) - dxmat = dx.view(xmat.shape) if dx is not None else None - x0mat = x0.view((-1, hidden_size)) if x0 is not None else None - x0_subset = x0_subset.view(-1) if x0_subset is not None else None - out_subset = out_subset.view(-1) if out_subset is not None else None - if colscale is not None: - assert x0 is not None, "x0 is required to compute the gradient of colscale" - dx0mat, dresidualmat, dgamma, dbeta, _, _, *rest = dropout_layer_norm.dropout_add_ln_bwd( - dzmat, - dxmat, - xmat, - x0mat, - dmask, - mu, - rsigma, - gamma, - None, - colscale, - x0_subset, - out_subset, - dropout_p, - rowscale_const, - x0_numrows, - has_residual, - is_rms_norm, - ) - # dresidualmat is None if not has_residual - if colscale is None: - return dx0mat, dresidualmat, dgamma, dbeta - else: - dcolscale = rest[0] - return dx0mat, dresidualmat, dgamma, dbeta, dcolscale - - -def _dropout_add_layer_norm_parallel_residual_forward( - x0, - x1, - residual, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - residual_in_fp32=False, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes""" - hidden_size = gamma0.numel() - x0mat = x0.view((-1, hidden_size)) - x1mat = x1.view((-1, hidden_size)) if x1 is not None else None - residualmat = residual.view((-1, hidden_size)) if residual is not None else None - ( - z0mat, - z1mat, - xmat, - dmask0, - dmask1, - mu, - rsigma, - ) = dropout_layer_norm.dropout_add_ln_parallel_residual_fwd( - x0mat, - x1mat, - residualmat, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - None, - residual_in_fp32, - is_rms_norm, - ) - # dmask0 and dmask1 are None if dropout_p == 0.0 - # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype - return z0mat, z1mat, xmat if xmat is not None else x0mat, dmask0, dmask1, mu, rsigma - - -def _dropout_add_layer_norm_parallel_residual_backward( - dz0, - dz1, - dx, - x, - dmask0, - dmask1, - mu, - rsigma, - gamma0, - gamma1, - dropout_p, - has_x1, - has_residual, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes - dx == None means that it was a post-norm architecture - (x = drop(x0) + residual was not returned in the fwd). - """ - hidden_size = gamma0.numel() - xmat = x.view((-1, hidden_size)) - dz0mat = dz0.view(xmat.shape) - dz1mat = dz1.view(xmat.shape) if dz1 is not None else None - dxmat = dx.view(xmat.shape) if dx is not None else None - ( - dx0mat, - dx1mat, - dresidualmat, - dgamma0, - dbeta0, - dgamma1, - dbeta1, - *rest, - ) = dropout_layer_norm.dropout_add_ln_parallel_residual_bwd( - dz0mat, - dz1mat, - dxmat, - xmat, - dmask0, - dmask1, - mu, - rsigma, - gamma0, - gamma1, - dropout_p, - has_x1, - has_residual, - is_rms_norm, - ) - # dresidualmat is None if not has_residual - return dx0mat, dx1mat, dresidualmat, dgamma0, dbeta0, dgamma1, dbeta1 - - -class DropoutAddLayerNormFn(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x0, - residual, - gamma, - beta, - rowscale, - colscale, - dropout_p, - epsilon, - residual_in_fp32=False, - prenorm=False, - is_rms_norm=False, - return_dmask=False, - ): - x0 = maybe_align(x0.contiguous(), 16) - residual = maybe_align(residual.contiguous(), 16) if residual is not None else None - gamma = maybe_align(gamma.contiguous(), 16) - beta = maybe_align(beta.contiguous(), 16) if beta is not None else None - rowscale = maybe_align(rowscale.contiguous(), 16) if rowscale is not None else None - colscale = maybe_align(colscale.contiguous(), 16) if colscale is not None else None - zmat, xmat, dmask, mu, rsigma = _dropout_add_layer_norm_forward( - x0, - residual, - gamma, - beta, - rowscale, - colscale, - dropout_p, - epsilon, - residual_in_fp32, - is_rms_norm, - ) - # Only need to save x0 if we need to compute gradient wrt colscale - x0_saved = x0 if colscale is not None else None - ctx.save_for_backward( - xmat.view(x0.shape), x0_saved, dmask, gamma, mu, rsigma, rowscale, colscale - ) - ctx.prenorm = prenorm - ctx.dropout_p = dropout_p - ctx.has_residual = residual is not None - ctx.is_rms_norm = is_rms_norm - ctx.has_beta = beta is not None - if not return_dmask: - return ( - zmat.view(x0.shape) if not prenorm else (zmat.view(x0.shape), xmat.view(x0.shape)) - ) - else: - dmask = ( - dmask.view(x0.shape) - if dropout_p > 0.0 - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - ctx.mark_non_differentiable(dmask) - return ( - (zmat.view(x0.shape), dmask) - if not prenorm - else (zmat.view(x0.shape), xmat.view(x0.shape), dmask) - ) - - @staticmethod - def backward(ctx, dz, *args): - # assert dz.is_contiguous() - dz = maybe_align(dz.contiguous(), 16) # this happens! - dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None - x, x0, dmask, gamma, mu, rsigma, rowscale, colscale = ctx.saved_tensors - # x0 is None if colscale is None - dropout_p = ctx.dropout_p - has_residual = ctx.has_residual - dx0mat, dresidualmat, dgamma, dbeta, *rest = _dropout_add_layer_norm_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - rowscale, - colscale, - dropout_p, - has_residual, - ctx.is_rms_norm, - ) - dx0 = dx0mat.view(x.shape) - dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None - dcolscale = rest[0] if colscale is not None else None - return ( - dx0, - dresidual, - dgamma, - dbeta if ctx.has_beta else None, - None, - dcolscale, - None, - None, - None, - None, - None, - None, - ) - - -class DropoutAddLayerNormSubsetFn(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x0, - residual, - gamma, - beta, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32=False, - prenorm=False, - is_rms_norm=False, - return_dmask=False, - ): - x0 = maybe_align(x0.contiguous(), 16) - residual = maybe_align(residual.contiguous(), 16) if residual is not None else None - gamma = maybe_align(gamma.contiguous(), 16) - beta = maybe_align(beta.contiguous(), 16) if beta is not None else None - colscale = maybe_align(colscale.contiguous(), 16) if colscale is not None else None - zmat, xmat, dmask, mu, rsigma = _dropout_add_layer_norm_subset_forward( - x0, - residual, - gamma, - beta, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32, - is_rms_norm, - ) - # Only need to save x0 if we need to compute gradient wrt colscale - x0_saved = x0 if colscale is not None else None - x_shape = (-1, *x0.shape[1:]) - ctx.save_for_backward( - xmat.view(x_shape), x0_saved, dmask, gamma, mu, rsigma, colscale, x0_subset, out_subset - ) - ctx.prenorm = prenorm - ctx.dropout_p = dropout_p - ctx.rowscale_const = rowscale_const - ctx.x0_numrows = x0.shape[:-1].numel() - ctx.has_residual = residual is not None - ctx.is_rms_norm = is_rms_norm - ctx.has_beta = beta is not None - z_shape = (-1, *x0.shape[1:]) - if not return_dmask: - return zmat.view(z_shape) if not prenorm else (zmat.view(z_shape), xmat.view(x0.shape)) - else: - z = zmat.view(z_shape) - dmask = ( - dmask.view(x0.shape) - if dropout_p > 0.0 - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - ctx.mark_non_differentiable(dmask) - return (z, dmask) if not prenorm else (z, xmat.view(x_shape), dmask) - - @staticmethod - def backward(ctx, dz, *args): - # assert dz.is_contiguous() - dz = maybe_align(dz.contiguous(), 16) # this happens! - dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None - x, x0, dmask, gamma, mu, rsigma, colscale, x0_subset, out_subset = ctx.saved_tensors - # x0 is None if colscale is None - dropout_p = ctx.dropout_p - has_residual = ctx.has_residual - dx0mat, dresidualmat, dgamma, dbeta, *rest = _dropout_add_layer_norm_subset_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - colscale, - x0_subset, - out_subset, - dropout_p, - ctx.rowscale_const, - ctx.x0_numrows, - has_residual, - ctx.is_rms_norm, - ) - dx0 = dx0mat.view(-1, *x.shape[1:]) - dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None - dcolscale = rest[0] if colscale is not None else None - return ( - dx0, - dresidual, - dgamma, - dbeta if ctx.has_beta else None, - dcolscale, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) - - -class DropoutAddLayerNormParallelResidualFn(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x0, - x1, - residual, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - residual_in_fp32=False, - prenorm=False, - is_rms_norm=False, - return_dmask=False, - ): - x0 = maybe_align(x0.contiguous(), 16) - x1 = maybe_align(x1.contiguous(), 16) if x1 is not None else None - residual = maybe_align(residual.contiguous(), 16) if residual is not None else None - gamma0 = maybe_align(gamma0.contiguous(), 16) - beta0 = maybe_align(beta0.contiguous(), 16) if beta0 is not None else None - gamma1 = maybe_align(gamma1.contiguous(), 16) if gamma1 is not None else None - beta1 = maybe_align(beta1.contiguous(), 16) if beta1 is not None else None - ( - z0mat, - z1mat, - xmat, - dmask0, - dmask1, - mu, - rsigma, - ) = _dropout_add_layer_norm_parallel_residual_forward( - x0, - x1, - residual, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - residual_in_fp32, - is_rms_norm, - ) - ctx.save_for_backward(xmat.view(x0.shape), dmask0, dmask1, gamma0, gamma1, mu, rsigma) - ctx.prenorm = prenorm - ctx.dropout_p = dropout_p - ctx.has_x1 = x1 is not None - ctx.has_residual = residual is not None - ctx.is_rms_norm = is_rms_norm - ctx.has_beta = beta0 is not None - z = (z0mat.view(x0.shape), z1mat.view(x0.shape) if z1mat is not None else None) - if not return_dmask: - return z if not prenorm else (*z, xmat.view(x0.shape)) - else: - dmask0 = ( - dmask0.view(x0.shape) - if dropout_p > 0.0 - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - dmask1 = ( - dmask1.view(x0.shape) - if dropout_p > 0.0 and x1 is not None - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - ctx.mark_non_differentiable(dmask0) - ctx.mark_non_differentiable(dmask1) - return ( - (*z, dmask0, dmask1) if not prenorm else (*z, xmat.view(x0.shape), dmask0, dmask1) - ) - - @staticmethod - def backward(ctx, dz0, dz1, *args): - dz0 = maybe_align(dz0.contiguous(), 16) # this happens! - dz1 = maybe_align(dz1.contiguous(), 16) if dz1 is not None else None - dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None - x, dmask0, dmask1, gamma0, gamma1, mu, rsigma = ctx.saved_tensors - dropout_p = ctx.dropout_p - has_x1 = ctx.has_x1 - has_residual = ctx.has_residual - ( - dx0mat, - dx1mat, - dresidualmat, - dgamma0, - dbeta0, - dgamma1, - dbeta1, - ) = _dropout_add_layer_norm_parallel_residual_backward( - dz0, - dz1, - dx, - x, - dmask0, - dmask1, - mu, - rsigma, - gamma0, - gamma1, - dropout_p, - has_x1, - has_residual, - ctx.is_rms_norm, - ) - dx0 = dx0mat.view(x.shape) - dx1 = dx1mat.view(x.shape) if dx1mat is not None else None - dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None - return ( - dx0, - dx1, - dresidual, - dgamma0, - dbeta0 if ctx.has_beta else None, - dgamma1, - dbeta1 if ctx.has_beta else None, - None, - None, - None, - None, - None, - None, - ) - - -def layer_norm(x, weight, bias, epsilon): - return DropoutAddLayerNormFn.apply(x, None, weight, bias, None, None, 0.0, epsilon, False) - - -def dropout_add_layer_norm( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - rowscale=None, - layerscale=None, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormFn.apply( - x0, - residual, - weight, - bias, - rowscale, - layerscale, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - False, - return_dropout_mask, - ) - - -def dropout_add_layer_norm_subset( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - layerscale=None, - x0_subset=None, - out_subset=None, - rowscale_const=1.0, - out_numrows=0, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormSubsetFn.apply( - x0, - residual, - weight, - bias, - layerscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32, - prenorm, - False, - return_dropout_mask, - ) - - -def dropout_add_layer_norm_parallel_residual( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormParallelResidualFn.apply( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - False, - return_dropout_mask, - ) - - -class DropoutAddLayerNorm(torch.nn.Module): - def __init__( - self, - hidden_size, - prenorm=False, - p=0.0, - eps=1e-5, - residual_in_fp32=False, - device=None, - dtype=None, - ): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.prenorm = prenorm - self.p = p - self.eps = eps - self.residual_in_fp32 = residual_in_fp32 - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.bias = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.reset_parameters() - - def reset_parameters(self): - init.ones_(self.weight) - init.zeros_(self.bias) - - def forward(self, x0, residual=None): - return dropout_add_layer_norm( - x0, - residual, - self.weight, - self.bias, - self.p if self.training else 0.0, - self.eps, - prenorm=self.prenorm, - residual_in_fp32=self.residual_in_fp32, - ) diff --git a/build/torch212-cxx11-cu130-aarch64-linux/ops/rms_norm.py b/build/torch212-cxx11-cu130-aarch64-linux/ops/rms_norm.py deleted file mode 100644 index 068348d61290e3839dd082b540d898578ba1e8e2..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu130-aarch64-linux/ops/rms_norm.py +++ /dev/null @@ -1,174 +0,0 @@ -# Copyright (c) 2022, Tri Dao. -# Adapted from https://github.com/NVIDIA/apex/blob/master/apex/contrib/layer_norm/layer_norm.py - -import torch -from torch.nn import init - -from flash_attn.ops.layer_norm import ( - DropoutAddLayerNormFn, - DropoutAddLayerNormParallelResidualFn, - DropoutAddLayerNormSubsetFn, -) - - -def rms_norm(x, weight, epsilon): - return DropoutAddLayerNormFn.apply( - x, None, weight, None, None, None, 0.0, epsilon, False, False, True - ) - - -def dropout_add_rms_norm( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - rowscale=None, - layerscale=None, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormFn.apply( - x0, - residual, - weight, - bias, - rowscale, - layerscale, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - True, - return_dropout_mask, - ) - - -def dropout_add_rms_norm_subset( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - layerscale=None, - x0_subset=None, - out_subset=None, - rowscale_const=1.0, - out_numrows=0, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormSubsetFn.apply( - x0, - residual, - weight, - bias, - layerscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32, - prenorm, - True, - return_dropout_mask, - ) - - -def dropout_add_rms_norm_parallel_residual( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormParallelResidualFn.apply( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - True, - return_dropout_mask, - ) - - -class RMSNorm(torch.nn.Module): - def __init__(self, hidden_size, eps=1e-5, device=None, dtype=None): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.eps = eps - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.register_parameter("bias", None) - self.reset_parameters() - - def reset_parameters(self): - init.ones_(self.weight) - - def forward(self, x): - return rms_norm(x, self.weight, self.eps) - - -class DropoutAddRMSNorm(torch.nn.Module): - def __init__( - self, - hidden_size, - prenorm=False, - p=0.0, - eps=1e-5, - residual_in_fp32=False, - device=None, - dtype=None, - ): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.prenorm = prenorm - self.p = p - self.eps = eps - self.residual_in_fp32 = residual_in_fp32 - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.register_parameter("bias", None) - self.reset_parameters() - - def reset_parameters(self): - init.ones_(self.weight) - - def forward(self, x0, residual=None): - return dropout_add_rms_norm( - x0, - residual, - self.weight, - None, - self.p if self.training else 0.0, - self.eps, - prenorm=self.prenorm, - residual_in_fp32=self.residual_in_fp32, - ) diff --git a/build/torch212-cxx11-cu130-aarch64-linux/ops/triton/__init__.py b/build/torch212-cxx11-cu130-aarch64-linux/ops/triton/__init__.py deleted file mode 100644 index 8b137891791fe96927ad78e64b0aad7bded08bdc..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu130-aarch64-linux/ops/triton/__init__.py +++ /dev/null @@ -1 +0,0 @@ - diff --git a/build/torch212-cxx11-cu130-aarch64-linux/ops/triton/cross_entropy.py b/build/torch212-cxx11-cu130-aarch64-linux/ops/triton/cross_entropy.py deleted file mode 100644 index 1b5a415b73f236f3e05fb14b9141959559e18526..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu130-aarch64-linux/ops/triton/cross_entropy.py +++ /dev/null @@ -1,330 +0,0 @@ -# Copyright (c) 2023, Tri Dao. - -from typing import Tuple, Optional, Union - -import torch -import torch.nn.functional as F - -import triton -import triton.language as tl - -# `all_gather_into_tensor` and `reduce_scatter_tensor` are new placeholders for -# `_all_gather_base` and `_reduce_scatter_base`. They require the most recent -# version of PyTorch. The following 2 lines are for backward compatibility with -# older PyTorch. -if "all_gather_into_tensor" not in dir(torch.distributed): - torch.distributed.all_gather_into_tensor = torch.distributed._all_gather_base - - -@triton.heuristics( - { - "HAS_SMOOTHING": lambda args: args["smoothing"] > 0.0, - } -) -@triton.jit -def cross_entropy_fwd_kernel( - loss_ptr, # data ptrs - lse_ptr, - z_loss_ptr, - logits_ptr, - labels_ptr, - smoothing, - logit_scale, - lse_square_scale, - ignore_index, - total_classes, - class_start_idx, # Useful for tensor parallel when each rank only has a subset of classes - n_cols, # shapes - logits_row_stride, # strides - BLOCK_SIZE: tl.constexpr, - HAS_SMOOTHING: tl.constexpr, - # if SPLIT (e.g. tensor parallel), don't include the LSE in the loss since it's not the final LSE - SPLIT: tl.constexpr, - PRECOMPUTED_LSE: tl.constexpr, # If LSE is already computed (also no smoothing and logit_scale == 1.0) -): - row_idx = tl.program_id(0) - logits_ptr = logits_ptr + row_idx * logits_row_stride.to(tl.int64) - sum_logits = 0.0 # For smoothing - if not PRECOMPUTED_LSE: - # Statistics for online softmax - m_i = -float("inf") - l_i = 0.0 - for col_offset in range(0, n_cols, BLOCK_SIZE): - cols = col_offset + tl.arange(0, BLOCK_SIZE) - logits = tl.load(logits_ptr + cols, mask=cols < n_cols, other=-float("inf")).to( - tl.float32 - ) * logit_scale - if HAS_SMOOTHING: - sum_logits += tl.sum(tl.where(cols < n_cols, logits, 0.0)) - m_i_new = tl.maximum(m_i, tl.max(logits)) - l_i = tl.exp(m_i - m_i_new) * l_i + tl.sum(tl.exp(logits - m_i_new)) - m_i = m_i_new - lse = tl.log(l_i) + m_i - tl.store(lse_ptr + row_idx, lse) - else: - lse = tl.load(lse_ptr + row_idx) - label_idx = tl.load(labels_ptr + row_idx) - if label_idx == ignore_index: - loss = 0.0 - z_loss = 0.0 - else: - label_idx -= class_start_idx - if label_idx >= 0 and label_idx < n_cols: - logits_label = tl.load(logits_ptr + label_idx) * logit_scale - if HAS_SMOOTHING: - loss = ( - (lse if not SPLIT else 0.0) - - smoothing * sum_logits / total_classes - - (1 - smoothing) * logits_label - ) - else: - loss = (lse if not SPLIT else 0.0) - logits_label - else: - # If label is out of bounds, we set the CE loss to 0.0. But we still want the smoothing loss - if HAS_SMOOTHING: - loss = smoothing * ((lse if not SPLIT else 0.0) - sum_logits / total_classes) - else: - loss = 0.0 - if not SPLIT: - z_loss = lse_square_scale * lse * lse - loss += z_loss - else: - z_loss = 0.0 - tl.store(loss_ptr + row_idx, loss) - if not SPLIT: - tl.store(z_loss_ptr + row_idx, z_loss) - - -@triton.heuristics( - { - "HAS_SMOOTHING": lambda args: args["smoothing"] > 0.0, - } -) -@triton.jit -def cross_entropy_bwd_kernel( - dlogits_ptr, # data ptrs - dloss_ptr, - logits_ptr, - lse_ptr, - labels_ptr, - smoothing, - logit_scale, - lse_square_scale, - ignore_index, - total_classes, - class_start_idx, # Useful for tensor parallel when each rank only has a subset of classes - n_cols, # shapes - logits_row_stride, # strides - dlogits_row_stride, - dloss_row_stride, - BLOCK_SIZE: tl.constexpr, - HAS_SMOOTHING: tl.constexpr, -): - row_idx = tl.program_id(0) - col_block_idx = tl.program_id(1) - logits_ptr = logits_ptr + row_idx * logits_row_stride.to(tl.int64) - dlogits_ptr = dlogits_ptr + row_idx * dlogits_row_stride.to(tl.int64) - col_offsets = col_block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - label_idx = tl.load(labels_ptr + row_idx) - if label_idx != ignore_index: - dloss = tl.load(dloss_ptr + row_idx * dloss_row_stride) - else: - dloss = 0.0 - logits = tl.load(logits_ptr + col_offsets, mask=col_offsets < n_cols, other=-float("inf")).to( - tl.float32 - ) * logit_scale - lse = tl.load(lse_ptr + row_idx) - probs = tl.exp(logits - lse) - probs += 2.0 * lse_square_scale * lse * probs - label_idx -= class_start_idx - if HAS_SMOOTHING: - smooth_positive = 1.0 - smoothing - smooth_negative = smoothing / total_classes - probs = tl.where(col_offsets == label_idx, probs - smooth_positive, probs) - smooth_negative - else: - probs = tl.where(col_offsets == label_idx, probs - 1.0, probs) - tl.store(dlogits_ptr + col_offsets, (dloss * logit_scale) * probs, mask=col_offsets < n_cols) - - -class CrossEntropyLoss(torch.autograd.Function): - - @staticmethod - def forward( - ctx, - logits, - labels, - precomputed_lse=None, - smoothing=0.0, - logit_scale=1.0, - lse_square_scale=0.0, - ignore_index=-100, - inplace_backward=False, - process_group=None, - ): - # For some reason Triton generates wrong code when labels has dtype long and its address - # is not aligned to 16 bytes. The ld.global.b64 seems to load the wrong label index. - if labels.dtype == torch.long and labels.data_ptr() % 16 != 0: - labels = F.pad(labels, (0, 1))[..., :-1] - assert labels.data_ptr() % 16 == 0 - assert logit_scale > 0.0 - n_rows, n_cols = logits.shape - assert labels.shape == (n_rows,) - world_size = 1 if process_group is None else torch.distributed.get_world_size(process_group) - total_classes = world_size * n_cols - rank = 0 if process_group is None else torch.distributed.get_rank(process_group) - class_start_idx = rank * n_cols - use_precomputed_lse = precomputed_lse is not None and logit_scale == 1.0 and smoothing == 0.0 - - if logits.stride(-1) != 1: - logits = logits.contiguous() - MAX_BLOCK_SIZE = 16 * 1024 - BLOCK_SIZE = min(triton.next_power_of_2(n_cols), MAX_BLOCK_SIZE) - num_warps = ( - 4 - if BLOCK_SIZE < 2048 - else (8 if BLOCK_SIZE < 8192 else (16 if BLOCK_SIZE < 128 * 1024 else 32)) - ) - losses = torch.empty(n_rows, dtype=torch.float, device=logits.device) - if use_precomputed_lse: - assert precomputed_lse.shape == (n_rows,) - lse = precomputed_lse.contiguous() - else: - lse = torch.empty(n_rows, dtype=torch.float, device=logits.device) - z_losses = torch.empty(n_rows, dtype=torch.float, device=logits.device) - # Need this, otherwise Triton tries to launch from cuda:0 and we get - # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?) - with torch.cuda.device(logits.device.index): - cross_entropy_fwd_kernel[(n_rows,)]( - losses, # data ptrs - lse, - z_losses, - logits, - labels, - smoothing, - logit_scale, - lse_square_scale, - ignore_index, - total_classes, - class_start_idx, - n_cols, # shapes - logits.stride(0), # strides - BLOCK_SIZE=BLOCK_SIZE, # constants - SPLIT=world_size > 1, - PRECOMPUTED_LSE=use_precomputed_lse, - num_warps=num_warps, - ) - - if world_size > 1: - # If there's no smoothing, if labels are in the vocab of this partition, losses contains - # - predicted logit, and 0 otherwise. - # If there's smoothing=0.1, for labels in the vocab of this partition, losses contains - # -0.9 * predicted logit - 0.1 * sum logit / total_classes. - # For labels not in the vocab of this partition, losses contains - # -0.1 * sum logit / total_classes. - if world_size > 1: - lse_allgather = torch.empty(world_size, n_rows, dtype=lse.dtype, device=lse.device) - torch.distributed.all_gather_into_tensor(lse_allgather, lse, group=process_group) - handle_losses = torch.distributed.all_reduce( - losses, op=torch.distributed.ReduceOp.SUM, group=process_group, async_op=True - ) - lse = torch.logsumexp(lse_allgather, dim=0) - handle_losses.wait() - # After the allreduce, if there's no smoothing, the total losses are - predicted_logit, - # we just have to add the (global) lse. - # If there's smoothing=0.1, the total losses are - # -0.9 * predicted_logit - 0.1 * sum logit / total_classes. - # Again, we just have to add the (global) lse. - losses += lse - if lse_square_scale != 0.0: - z_losses = lse_square_scale * lse.square() - z_losses.masked_fill_(labels == ignore_index, 0.0) - losses += z_losses - else: - z_losses = torch.zeros_like(losses) - losses.masked_fill_(labels == ignore_index, 0.0) - - ctx.save_for_backward(logits, lse, labels) - ctx.mark_non_differentiable(z_losses) - ctx.smoothing = smoothing - ctx.logit_scale = logit_scale - ctx.lse_square_scale = lse_square_scale - ctx.ignore_index = ignore_index - ctx.total_classes = total_classes - ctx.class_start_idx = class_start_idx - ctx.inplace_backward = inplace_backward - return losses, z_losses - - @staticmethod - def backward(ctx, grad_losses, grad_z_losses): - del grad_z_losses # z_losses are only for logging. - - logits, lse, labels = ctx.saved_tensors - dlogits = logits if ctx.inplace_backward else torch.empty_like(logits) - n_rows, n_cols = logits.shape - BLOCK_SIZE = min(triton.next_power_of_2(n_cols), 4 * 1024) - num_warps = 4 if BLOCK_SIZE < 2048 else (8 if BLOCK_SIZE < 8192 else 16) - grid = lambda META: (n_rows, triton.cdiv(n_cols, META["BLOCK_SIZE"])) # noqa - # Need this, otherwise Triton tries to launch from cuda:0 and we get - # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?) - with torch.cuda.device(logits.device.index): - cross_entropy_bwd_kernel[grid]( - dlogits, # data ptrs - grad_losses, - logits, - lse, - labels, - ctx.smoothing, - ctx.logit_scale, - ctx.lse_square_scale, - ctx.ignore_index, - ctx.total_classes, - ctx.class_start_idx, - n_cols, # shapes - logits.stride(0), # strides - dlogits.stride(0), - grad_losses.stride(0), - BLOCK_SIZE=BLOCK_SIZE, # constants - num_warps=num_warps, - ) - return dlogits, None, None, None, None, None, None, None, None, None - - -def cross_entropy_loss( - logits: torch.Tensor, - labels: torch.Tensor, - precomputed_lse: Optional[torch.Tensor] = None, - label_smoothing: float = 0.0, - logit_scale: float = 1.0, - lse_square_scale: float = 0.0, - ignore_index=-100, - inplace_backward: bool = False, - process_group=None, -) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Arguments: - logits: (batch, vocab_size) - labels: (batch,) - label_smoothing: float - logit_scale: float. Multiply logits by this scale before calculating the loss. - lse_square_scale: float. If > 0, we add lse_square_scale * lse(logits) ^ 2 to the loss. - This is also referred to as "z-loss". - ignore_index: int. If labels == ignore_index, the loss is set to 0.0. - inplace_backward: bool. If True, we do the backward pass in-place by modifying the logits. - This saves memory. - process_group: if not None, we're doing Tensor Parallel: each process is responsible for - one part of the vocab. The loss will be aggregated across processes. - Returns: - losses: (batch,), float - z_losses: (batch,), float - """ - return CrossEntropyLoss.apply( - logits, - labels, - precomputed_lse, - label_smoothing, - logit_scale, - lse_square_scale, - ignore_index, - inplace_backward, - process_group, - ) diff --git a/build/torch212-cxx11-cu130-aarch64-linux/ops/triton/k_activations.py b/build/torch212-cxx11-cu130-aarch64-linux/ops/triton/k_activations.py deleted file mode 100644 index efb83c358eb4a85d069ee340a3c83f418f9a805b..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu130-aarch64-linux/ops/triton/k_activations.py +++ /dev/null @@ -1,162 +0,0 @@ -# Adapted from https://github.com/facebookresearch/xformers/blob/main/xformers/triton/k_activations.py -# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved. -# -# This source code is licensed under the BSD license found in the -# LICENSE file in the root directory of this source tree. - -import math -from enum import Enum -from typing import Optional - -import triton -import triton.language as tl - -_sqrt2pi = math.sqrt(2.0 / math.pi) -_sqrt1_2 = math.sqrt(1.0 / 2) -_gaussian_pdf_normalization = 1.0 / math.sqrt(2 * math.pi) - - -class Activation(str, Enum): - SquaredReLU = "squared_relu" - GeLU = "gelu" - GeLUApprox = "gelu_approx" - LeakyReLU = "leaky_relu" - ReLU = "relu" - - -def get_triton_activation_kernel(activation: Optional[Activation]): - return ( - { - Activation.ReLU: relu, - Activation.LeakyReLU: leaky_relu, - Activation.GeLU: gelu, - Activation.GeLUApprox: gelu_approx, - Activation.SquaredReLU: squared_relu, - }[activation] - if activation - else None - ) - - -def get_triton_activation_bwd_kernel(activation: Optional[Activation]): - return ( - { - Activation.ReLU: relu_grad, - Activation.LeakyReLU: leaky_relu_grad, - Activation.GeLU: gelu_grad, - Activation.GeLUApprox: gelu_approx_grad, - Activation.SquaredReLU: squared_relu_grad, - }[activation] - if activation - else None - ) - - -@triton.jit -def tanh(x): - # Tanh is just a scaled sigmoid - return 2 * tl.sigmoid(2 * x) - 1 - - -@triton.jit -def cosh(x): - exp_x = tl.exp(x) - return (exp_x + 1.0 / exp_x) * 0.5 - - -# a Triton implementation of the most used activations -# See for instance http://arxiv.org/abs/1606.08415 for an overview - -# ReLU -@triton.jit -def relu(x): - """ - ReLU_ activation function - - .. _ReLU: https://pytorch.org/docs/stable/generated/torch.nn.ReLU.html - """ - zero = 0.0 - return tl.where(x >= 0, x, zero.to(x.dtype)) - - -@triton.jit -def relu_grad(x): - # ReLU is different from other activations - # in that it does not require the input to retrospectively compute its gradient - # here the input is the downstream gradient, and we return the upstream gradient directly - zero = 0.0 - one = 1.0 - return tl.where(x >= 0, one.to(x.dtype), zero.to(x.dtype)) - - -@triton.jit -def squared_relu(x): - """ - Squared ReLU activation, as proposed in the Primer_ paper. - - .. _Primer: https://arxiv.org/abs/2109.08668 - """ - x_ = relu(x) - return (x_ * x_).to(x.dtype) - - -@triton.jit -def squared_relu_grad(x): - return tl.where(x >= 0, 2.0 * x, 0.0) - - -# Leaky ReLU -@triton.jit -def leaky_relu(x): - """ - LeakyReLU_ activation - - .. _LeakyReLU: https://pytorch.org/docs/stable/generated/torch.nn.LeakyReLU.html - """ - scale = 0.01 + 0.0 - scale = scale.to(x.dtype) - return tl.where(x >= 0, x, scale * x) - - -@triton.jit -def leaky_relu_grad(x): - min_grad = 0.01 - max_grad = 1 - - min_grad = min_grad.to(x.dtype) - max_grad = max_grad.to(x.dtype) - - return tl.where(x >= 0, max_grad, min_grad) - - -@triton.jit -def gelu(x): - """Gaussian Error Linear Unit (GELU)""" - return x * 0.5 * (1.0 + tl.libdevice.erf(x * _sqrt1_2)) - - -@triton.jit -def gelu_grad(x): - cdf = 0.5 * (1.0 + tl.libdevice.erf(x * _sqrt1_2)) - pdf = tl.exp(-0.5 * x * x) * _gaussian_pdf_normalization - return cdf + x * pdf - - -@triton.jit -def gelu_approx(x): - """ - GeLU_ activation - Gaussian error linear unit, with tanh approximation - - .. _GeLU: https://arxiv.org/pdf/1606.08415.pdf - """ - return 0.5 * x * (1.0 + tanh(_sqrt2pi * x * (1.0 + 0.044715 * x * x))) - - -@triton.jit -def gelu_approx_grad(x): - # CREDITS: Fast implementation proposed in - # https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/fused_bias_gelu.py#L30 - tanh_out = tanh(0.79788456 * x * (1 + 0.044715 * x * x)) - return 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( - 1 + tanh_out - ) diff --git a/build/torch212-cxx11-cu130-aarch64-linux/ops/triton/layer_norm.py b/build/torch212-cxx11-cu130-aarch64-linux/ops/triton/layer_norm.py deleted file mode 100644 index 5e6f0b3697b612e6a19031bbc55ac76ad233dc85..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu130-aarch64-linux/ops/triton/layer_norm.py +++ /dev/null @@ -1,1254 +0,0 @@ -# Copyright (c) 2024, Tri Dao. -# Implement dropout + residual + layer_norm / rms_norm. - -# Based on the Triton LayerNorm tutorial: https://triton-lang.org/main/getting-started/tutorials/05-layer-norm.html -# For the backward pass, we keep weight_grad and bias_grad in registers and accumulate. -# This is faster for dimensions up to 8k, but after that it's much slower due to register spilling. -# The models we train have hidden dim up to 8k anyway (e.g. Llama 70B), so this is fine. - -import math -from typing import Optional, List - -import torch -import torch.nn.functional as F -from torch import Tensor - -import triton -import triton.language as tl - -from flash_attn.utils.torch import custom_fwd, custom_bwd -from flash_attn.utils.library import triton_op - -from ..._ops import add_op_namespace_prefix - - -def maybe_contiguous_lastdim(x): - return x.contiguous() if x is not None and x.stride(-1) != 1 else x - - -def maybe_contiguous(x): - return x.contiguous() if x is not None else None - - -def triton_autotune_configs(): - # Return configs with a valid warp count for the current device - configs = [] - # Maximum threads per block is architecture-dependent in theory, but in reality all are 1024 - max_threads_per_block = 1024 - # Default to warp size 32 if not defined by device - warp_size = getattr(torch.cuda.get_device_properties(torch.cuda.current_device()), "warp_size", 32) - # Autotune for warp counts which are powers of 2 and do not exceed thread per block limit - return [triton.Config({}, num_warps=warp_count) for warp_count in [1, 2, 4, 8, 16, 32] - if warp_count * warp_size <= max_threads_per_block] - # return [triton.Config({}, num_warps=8)] - - -def layer_norm_ref( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - zero_centered_weight=False, - dropout_mask=None, - dropout_mask1=None, - upcast=False, -): - dtype = x.dtype - if upcast: - x = x.float() - weight = weight.float() - bias = bias.float() if bias is not None else None - residual = residual.float() if residual is not None else residual - x1 = x1.float() if x1 is not None else None - weight1 = weight1.float() if weight1 is not None else None - bias1 = bias1.float() if bias1 is not None else None - if zero_centered_weight: - weight = weight + 1.0 - if weight1 is not None: - weight1 = weight1 + 1.0 - if x1 is not None: - assert rowscale is None, "rowscale is not supported with parallel LayerNorm" - if rowscale is not None: - x = x * rowscale[..., None] - if dropout_p > 0.0: - if dropout_mask is not None: - x = x.masked_fill(~dropout_mask, 0.0) / (1.0 - dropout_p) - else: - x = F.dropout(x, p=dropout_p) - if x1 is not None: - if dropout_mask1 is not None: - x1 = x1.masked_fill(~dropout_mask1, 0.0) / (1.0 - dropout_p) - else: - x1 = F.dropout(x1, p=dropout_p) - if x1 is not None: - x = x + x1 - if residual is not None: - x = (x + residual).to(x.dtype) - out = F.layer_norm(x.to(weight.dtype), x.shape[-1:], weight=weight, bias=bias, eps=eps).to( - dtype - ) - if weight1 is None: - return out if not prenorm else (out, x) - else: - out1 = F.layer_norm( - x.to(weight1.dtype), x.shape[-1:], weight=weight1, bias=bias1, eps=eps - ).to(dtype) - return (out, out1) if not prenorm else (out, out1, x) - - -def rms_norm_ref( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - zero_centered_weight=False, - dropout_mask=None, - dropout_mask1=None, - upcast=False, -): - dtype = x.dtype - if upcast: - x = x.float() - weight = weight.float() - bias = bias.float() if bias is not None else None - residual = residual.float() if residual is not None else residual - x1 = x1.float() if x1 is not None else None - weight1 = weight1.float() if weight1 is not None else None - bias1 = bias1.float() if bias1 is not None else None - if zero_centered_weight: - weight = weight + 1.0 - if weight1 is not None: - weight1 = weight1 + 1.0 - if x1 is not None: - assert rowscale is None, "rowscale is not supported with parallel LayerNorm" - if rowscale is not None: - x = x * rowscale[..., None] - if dropout_p > 0.0: - if dropout_mask is not None: - x = x.masked_fill(~dropout_mask, 0.0) / (1.0 - dropout_p) - else: - x = F.dropout(x, p=dropout_p) - if x1 is not None: - if dropout_mask1 is not None: - x1 = x1.masked_fill(~dropout_mask1, 0.0) / (1.0 - dropout_p) - else: - x1 = F.dropout(x1, p=dropout_p) - if x1 is not None: - x = x + x1 - if residual is not None: - x = (x + residual).to(x.dtype) - rstd = 1 / torch.sqrt((x.square()).mean(dim=-1, keepdim=True) + eps) - out = ((x * rstd * weight) + bias if bias is not None else (x * rstd * weight)).to(dtype) - if weight1 is None: - return out if not prenorm else (out, x) - else: - out1 = ((x * rstd * weight1) + bias1 if bias1 is not None else (x * rstd * weight1)).to( - dtype - ) - return (out, out1) if not prenorm else (out, out1, x) - - -@triton.autotune( - configs=triton_autotune_configs(), - key=["N", "HAS_RESIDUAL", "STORE_RESIDUAL_OUT", "IS_RMS_NORM", "HAS_BIAS", "HAS_X1", "HAS_W1", "HAS_B1"], -) -# torch compile doesn't like triton.heuristics, so we set these manually when calling the kernel -# @triton.heuristics({"HAS_BIAS": lambda args: args["B"] is not None}) -# @triton.heuristics({"HAS_RESIDUAL": lambda args: args["RESIDUAL"] is not None}) -# @triton.heuristics({"HAS_X1": lambda args: args["X1"] is not None}) -# @triton.heuristics({"HAS_W1": lambda args: args["W1"] is not None}) -# @triton.heuristics({"HAS_B1": lambda args: args["B1"] is not None}) -@triton.jit -def _layer_norm_fwd_1pass_kernel( - X, # pointer to the input - Y, # pointer to the output - W, # pointer to the weights - B, # pointer to the biases - RESIDUAL, # pointer to the residual - X1, - W1, - B1, - Y1, - RESIDUAL_OUT, # pointer to the residual - ROWSCALE, - SEEDS, # Dropout seeds for each row - DROPOUT_MASK, - DROPOUT_MASK1, - Mean, # pointer to the mean - Rstd, # pointer to the 1/std - stride_x_row, # how much to increase the pointer when moving by 1 row - stride_y_row, - stride_res_row, - stride_res_out_row, - stride_x1_row, - stride_y1_row, - M, # number of rows in X - N, # number of columns in X - eps, # epsilon to avoid division by zero - dropout_p, # Dropout probability - zero_centered_weight, # If true, add 1.0 to the weight - IS_RMS_NORM: tl.constexpr, - BLOCK_N: tl.constexpr, - HAS_RESIDUAL: tl.constexpr, - STORE_RESIDUAL_OUT: tl.constexpr, - HAS_BIAS: tl.constexpr, - HAS_DROPOUT: tl.constexpr, - STORE_DROPOUT_MASK: tl.constexpr, - HAS_ROWSCALE: tl.constexpr, - HAS_X1: tl.constexpr, - HAS_W1: tl.constexpr, - HAS_B1: tl.constexpr, -): - # Map the program id to the row of X and Y it should compute. - row = tl.program_id(0) - X += row * stride_x_row - Y += row * stride_y_row - if HAS_RESIDUAL: - RESIDUAL += row * stride_res_row - if STORE_RESIDUAL_OUT: - RESIDUAL_OUT += row * stride_res_out_row - if HAS_X1: - X1 += row * stride_x1_row - if HAS_W1: - Y1 += row * stride_y1_row - # Compute mean and variance - cols = tl.arange(0, BLOCK_N) - x = tl.load(X + cols, mask=cols < N, other=0.0).to(tl.float32) - if HAS_ROWSCALE: - rowscale = tl.load(ROWSCALE + row).to(tl.float32) - x *= rowscale - if HAS_DROPOUT: - # Compute dropout mask - # 7 rounds is good enough, and reduces register pressure - keep_mask = tl.rand(tl.load(SEEDS + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - x = tl.where(keep_mask, x / (1.0 - dropout_p), 0.0) - if STORE_DROPOUT_MASK: - tl.store(DROPOUT_MASK + row * N + cols, keep_mask, mask=cols < N) - if HAS_X1: - x1 = tl.load(X1 + cols, mask=cols < N, other=0.0).to(tl.float32) - if HAS_ROWSCALE: - rowscale = tl.load(ROWSCALE + M + row).to(tl.float32) - x1 *= rowscale - if HAS_DROPOUT: - # Compute dropout mask - # 7 rounds is good enough, and reduces register pressure - keep_mask = ( - tl.rand(tl.load(SEEDS + M + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - ) - x1 = tl.where(keep_mask, x1 / (1.0 - dropout_p), 0.0) - if STORE_DROPOUT_MASK: - tl.store(DROPOUT_MASK1 + row * N + cols, keep_mask, mask=cols < N) - x += x1 - if HAS_RESIDUAL: - residual = tl.load(RESIDUAL + cols, mask=cols < N, other=0.0).to(tl.float32) - x += residual - if STORE_RESIDUAL_OUT: - tl.store(RESIDUAL_OUT + cols, x, mask=cols < N) - if not IS_RMS_NORM: - mean = tl.sum(x, axis=0) / N - tl.store(Mean + row, mean) - xbar = tl.where(cols < N, x - mean, 0.0) - var = tl.sum(xbar * xbar, axis=0) / N - else: - xbar = tl.where(cols < N, x, 0.0) - var = tl.sum(xbar * xbar, axis=0) / N - rstd = 1 / tl.sqrt(var + eps) - tl.store(Rstd + row, rstd) - # Normalize and apply linear transformation - mask = cols < N - w = tl.load(W + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w += 1.0 - if HAS_BIAS: - b = tl.load(B + cols, mask=mask).to(tl.float32) - x_hat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd - y = x_hat * w + b if HAS_BIAS else x_hat * w - # Write output - tl.store(Y + cols, y, mask=mask) - if HAS_W1: - w1 = tl.load(W1 + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w1 += 1.0 - if HAS_B1: - b1 = tl.load(B1 + cols, mask=mask).to(tl.float32) - y1 = x_hat * w1 + b1 if HAS_B1 else x_hat * w1 - tl.store(Y1 + cols, y1, mask=mask) - - -def _layer_norm_fwd( - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - residual: Optional[Tensor] = None, - x1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - out_dtype: Optional[torch.dtype] = None, - residual_dtype: Optional[torch.dtype] = None, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - return_dropout_mask: bool = False, - out: Optional[Tensor] = None, - residual_out: Optional[Tensor] = None -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - # Need to wrap to handle the case where residual_out is a alias of x, which makes torch.library - # and torch.compile unhappy. Also allocate memory for out and residual_out if they are None - # so that _layer_norm_fwd_impl doesn't have to return them. - if out is None: - out = torch.empty_like(x, dtype=x.dtype if out_dtype is None else out_dtype) - if residual is not None: - residual_dtype = residual.dtype - if residual_out is None and ( - residual is not None - or (residual_dtype is not None and residual_dtype != x.dtype) - or dropout_p > 0.0 - or rowscale is not None - or x1 is not None - ): - residual_out = torch.empty_like( - x, dtype=residual_dtype if residual_dtype is not None else x.dtype - ) - else: - residual_out = None - y1, mean, rstd, seeds, dropout_mask, dropout_mask1 = _layer_norm_fwd_impl( - x, - weight, - bias, - eps, - out, - residual=residual, - x1=x1, - weight1=weight1, - bias1=bias1, - dropout_p=dropout_p, - rowscale=rowscale, - zero_centered_weight=zero_centered_weight, - is_rms_norm=is_rms_norm, - return_dropout_mask=return_dropout_mask, - residual_out=residual_out, - ) - # residual_out is None if residual is None and residual_dtype == input_dtype and dropout_p == 0.0 - if residual_out is None: - residual_out = x - return out, y1, mean, rstd, residual_out, seeds, dropout_mask, dropout_mask1 - - -# [2025-04-28] torch.library.triton_op ignores the schema argument, but here we need the schema -# since we're returning a tuple of tensors -@triton_op(add_op_namespace_prefix("layer_norm_fwd_impl"), mutates_args={"out", "residual_out"}, - schema="(Tensor x, Tensor weight, Tensor bias, float eps, Tensor(a!) out, Tensor? residual, Tensor? x1, Tensor? weight1, Tensor? bias1, float dropout_p, Tensor? rowscale, bool zero_centered_weight, bool is_rms_norm, bool return_dropout_mask, Tensor(a!)? residual_out) -> (Tensor y1, Tensor mean, Tensor rstd, Tensor seeds, Tensor dropout_mask, Tensor dropout_mask1)") -def _layer_norm_fwd_impl( - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - out: Tensor, - residual: Optional[Tensor] = None, - x1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - return_dropout_mask: bool = False, - residual_out: Optional[Tensor] = None -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - M, N = x.shape - assert x.stride(-1) == 1 - if residual is not None: - assert residual.stride(-1) == 1 - assert residual.shape == (M, N) - assert weight.shape == (N,) - assert weight.stride(-1) == 1 - if bias is not None: - assert bias.stride(-1) == 1 - assert bias.shape == (N,) - if x1 is not None: - assert x1.shape == x.shape - assert rowscale is None - assert x1.stride(-1) == 1 - if weight1 is not None: - assert weight1.shape == (N,) - assert weight1.stride(-1) == 1 - if bias1 is not None: - assert bias1.shape == (N,) - assert bias1.stride(-1) == 1 - if rowscale is not None: - assert rowscale.is_contiguous() - assert rowscale.shape == (M,) - assert out.shape == x.shape - assert out.stride(-1) == 1 - if residual_out is not None: - assert residual_out.shape == x.shape - assert residual_out.stride(-1) == 1 - if weight1 is not None: - y1 = torch.empty_like(out) - assert y1.stride(-1) == 1 - else: - y1 = None - mean = torch.empty((M,), dtype=torch.float32, device=x.device) if not is_rms_norm else None - rstd = torch.empty((M,), dtype=torch.float32, device=x.device) - if dropout_p > 0.0: - seeds = torch.randint( - 2**32, (M if x1 is None else 2 * M,), device=x.device, dtype=torch.int64 - ) - else: - seeds = None - if return_dropout_mask and dropout_p > 0.0: - dropout_mask = torch.empty(M, N, device=x.device, dtype=torch.bool) - if x1 is not None: - dropout_mask1 = torch.empty(M, N, device=x.device, dtype=torch.bool) - else: - dropout_mask1 = None - else: - dropout_mask, dropout_mask1 = None, None - # Less than 64KB per feature: enqueue fused kernel - MAX_FUSED_SIZE = 65536 // x.element_size() - BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) - if N > BLOCK_N: - raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") - with torch.cuda.device(x.device.index): - torch.library.wrap_triton(_layer_norm_fwd_1pass_kernel)[(M,)]( - x, - out, - weight, - bias, - residual, - x1, - weight1, - bias1, - y1, - residual_out, - rowscale, - seeds, - dropout_mask, - dropout_mask1, - mean, - rstd, - x.stride(0), - out.stride(0), - residual.stride(0) if residual is not None else 0, - residual_out.stride(0) if residual_out is not None else 0, - x1.stride(0) if x1 is not None else 0, - y1.stride(0) if y1 is not None else 0, - M, - N, - eps, - dropout_p, - # Passing bool make torch inductor very unhappy since it then tries to compare to int_max - int(zero_centered_weight), - is_rms_norm, - BLOCK_N, - residual is not None, - residual_out is not None, - bias is not None, - dropout_p > 0.0, - dropout_mask is not None, - rowscale is not None, - HAS_X1=x1 is not None, - HAS_W1=weight1 is not None, - HAS_B1=bias1 is not None, - ) - return y1, mean, rstd, seeds, dropout_mask, dropout_mask1 - - -@triton.autotune( - configs=triton_autotune_configs(), - key=["N", "HAS_DRESIDUAL", "STORE_DRESIDUAL", "IS_RMS_NORM", "HAS_BIAS", "HAS_DROPOUT"], -) -# torch compile doesn't like triton.heuristics, so we set these manually when calling the kernel -# @triton.heuristics({"HAS_BIAS": lambda args: args["B"] is not None}) -# @triton.heuristics({"HAS_DRESIDUAL": lambda args: args["DRESIDUAL"] is not None}) -# @triton.heuristics({"STORE_DRESIDUAL": lambda args: args["DRESIDUAL_IN"] is not None}) -# @triton.heuristics({"HAS_ROWSCALE": lambda args: args["ROWSCALE"] is not None}) -# @triton.heuristics({"HAS_DY1": lambda args: args["DY1"] is not None}) -# @triton.heuristics({"HAS_DX1": lambda args: args["DX1"] is not None}) -# @triton.heuristics({"HAS_B1": lambda args: args["DB1"] is not None}) -# @triton.heuristics({"RECOMPUTE_OUTPUT": lambda args: args["Y"] is not None}) -@triton.jit -def _layer_norm_bwd_kernel( - X, # pointer to the input - W, # pointer to the weights - B, # pointer to the biases - Y, # pointer to the output to be recomputed - DY, # pointer to the output gradient - DX, # pointer to the input gradient - DW, # pointer to the partial sum of weights gradient - DB, # pointer to the partial sum of biases gradient - DRESIDUAL, - W1, - DY1, - DX1, - DW1, - DB1, - DRESIDUAL_IN, - ROWSCALE, - SEEDS, - Mean, # pointer to the mean - Rstd, # pointer to the 1/std - stride_x_row, # how much to increase the pointer when moving by 1 row - stride_y_row, - stride_dy_row, - stride_dx_row, - stride_dres_row, - stride_dy1_row, - stride_dx1_row, - stride_dres_in_row, - M, # number of rows in X - N, # number of columns in X - eps, # epsilon to avoid division by zero - dropout_p, - zero_centered_weight, - rows_per_program, - IS_RMS_NORM: tl.constexpr, - BLOCK_N: tl.constexpr, - HAS_DRESIDUAL: tl.constexpr, - STORE_DRESIDUAL: tl.constexpr, - HAS_BIAS: tl.constexpr, - HAS_DROPOUT: tl.constexpr, - HAS_ROWSCALE: tl.constexpr, - HAS_DY1: tl.constexpr, - HAS_DX1: tl.constexpr, - HAS_B1: tl.constexpr, - RECOMPUTE_OUTPUT: tl.constexpr, -): - # Map the program id to the elements of X, DX, and DY it should compute. - row_block_id = tl.program_id(0) - row_start = row_block_id * rows_per_program - # Do not early exit if row_start >= M, because we need to write DW and DB - cols = tl.arange(0, BLOCK_N) - mask = cols < N - X += row_start * stride_x_row - if HAS_DRESIDUAL: - DRESIDUAL += row_start * stride_dres_row - if STORE_DRESIDUAL: - DRESIDUAL_IN += row_start * stride_dres_in_row - DY += row_start * stride_dy_row - DX += row_start * stride_dx_row - if HAS_DY1: - DY1 += row_start * stride_dy1_row - if HAS_DX1: - DX1 += row_start * stride_dx1_row - if RECOMPUTE_OUTPUT: - Y += row_start * stride_y_row - w = tl.load(W + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w += 1.0 - if RECOMPUTE_OUTPUT and HAS_BIAS: - b = tl.load(B + cols, mask=mask, other=0.0).to(tl.float32) - if HAS_DY1: - w1 = tl.load(W1 + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w1 += 1.0 - dw = tl.zeros((BLOCK_N,), dtype=tl.float32) - if HAS_BIAS: - db = tl.zeros((BLOCK_N,), dtype=tl.float32) - if HAS_DY1: - dw1 = tl.zeros((BLOCK_N,), dtype=tl.float32) - if HAS_B1: - db1 = tl.zeros((BLOCK_N,), dtype=tl.float32) - row_end = min((row_block_id + 1) * rows_per_program, M) - for row in range(row_start, row_end): - # Load data to SRAM - x = tl.load(X + cols, mask=mask, other=0).to(tl.float32) - dy = tl.load(DY + cols, mask=mask, other=0).to(tl.float32) - if HAS_DY1: - dy1 = tl.load(DY1 + cols, mask=mask, other=0).to(tl.float32) - if not IS_RMS_NORM: - mean = tl.load(Mean + row) - rstd = tl.load(Rstd + row) - # Compute dx - xhat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd - xhat = tl.where(mask, xhat, 0.0) - if RECOMPUTE_OUTPUT: - y = xhat * w + b if HAS_BIAS else xhat * w - tl.store(Y + cols, y, mask=mask) - wdy = w * dy - dw += dy * xhat - if HAS_BIAS: - db += dy - if HAS_DY1: - wdy += w1 * dy1 - dw1 += dy1 * xhat - if HAS_B1: - db1 += dy1 - if not IS_RMS_NORM: - c1 = tl.sum(xhat * wdy, axis=0) / N - c2 = tl.sum(wdy, axis=0) / N - dx = (wdy - (xhat * c1 + c2)) * rstd - else: - c1 = tl.sum(xhat * wdy, axis=0) / N - dx = (wdy - xhat * c1) * rstd - if HAS_DRESIDUAL: - dres = tl.load(DRESIDUAL + cols, mask=mask, other=0).to(tl.float32) - dx += dres - # Write dx - if STORE_DRESIDUAL: - tl.store(DRESIDUAL_IN + cols, dx, mask=mask) - if HAS_DX1: - if HAS_DROPOUT: - keep_mask = ( - tl.rand(tl.load(SEEDS + M + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - ) - dx1 = tl.where(keep_mask, dx / (1.0 - dropout_p), 0.0) - else: - dx1 = dx - tl.store(DX1 + cols, dx1, mask=mask) - if HAS_DROPOUT: - keep_mask = tl.rand(tl.load(SEEDS + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - dx = tl.where(keep_mask, dx / (1.0 - dropout_p), 0.0) - if HAS_ROWSCALE: - rowscale = tl.load(ROWSCALE + row).to(tl.float32) - dx *= rowscale - tl.store(DX + cols, dx, mask=mask) - - X += stride_x_row - if HAS_DRESIDUAL: - DRESIDUAL += stride_dres_row - if STORE_DRESIDUAL: - DRESIDUAL_IN += stride_dres_in_row - if RECOMPUTE_OUTPUT: - Y += stride_y_row - DY += stride_dy_row - DX += stride_dx_row - if HAS_DY1: - DY1 += stride_dy1_row - if HAS_DX1: - DX1 += stride_dx1_row - tl.store(DW + row_block_id * N + cols, dw, mask=mask) - if HAS_BIAS: - tl.store(DB + row_block_id * N + cols, db, mask=mask) - if HAS_DY1: - tl.store(DW1 + row_block_id * N + cols, dw1, mask=mask) - if HAS_B1: - tl.store(DB1 + row_block_id * N + cols, db1, mask=mask) - - -def _layer_norm_bwd( - dy: Tensor, - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - mean: Tensor, - rstd: Tensor, - dresidual: Optional[Tensor] = None, - dy1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - seeds: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - has_residual: bool = False, - has_x1: bool = False, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - x_dtype: Optional[torch.dtype] = None, - recompute_output: bool = False, -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - # Need to wrap to handle the case where dresidual_in or dx1 are aliases of x, - # which makes torch.library unhappy - dx, dw, db, dresidual_in, dx1, dw1, db1, y = _layer_norm_bwd_impl( - dy, - x, - weight, - bias, - eps, - mean, - rstd, - dresidual, - dy1, - weight1, - bias1, - seeds, - dropout_p, - rowscale, - has_residual, - has_x1, - zero_centered_weight, - is_rms_norm, - x_dtype=x_dtype, - recompute_output=recompute_output, - ) - # Don't need to compute dresidual_in separately in this case - if has_residual and dx.dtype == x.dtype and dropout_p == 0.0 and rowscale is None: - dresidual_in = dx - if has_x1 and dropout_p == 0.0: - dx1 = dx - return dx, dw, db, dresidual_in, dx1, dw1, db1, y - - - -@triton_op(add_op_namespace_prefix("layer_norm_bwd_impl"), mutates_args={}, - schema="(Tensor dy, Tensor x, Tensor weight, Tensor bias, float eps, Tensor mean, Tensor rstd, Tensor? dresidual, Tensor? dy1, Tensor? weight1, Tensor? bias1, Tensor? seeds, float dropout_p, Tensor? rowscale, bool has_residual, bool has_x1, bool zero_centered_weight, bool is_rms_norm, ScalarType? x_dtype, bool recompute_output) -> (Tensor dx, Tensor dw, Tensor db, Tensor dresidual_in, Tensor dx1, Tensor dw1, Tensor db1, Tensor y)", - allow_decomposition=False, # Don't let torch.compile trace inside - ) -def _layer_norm_bwd_impl( - dy: Tensor, - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - mean: Tensor, - rstd: Tensor, - dresidual: Optional[Tensor] = None, - dy1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - seeds: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - has_residual: bool = False, - has_x1: bool = False, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - x_dtype: Optional[torch.dtype] = None, - recompute_output: bool = False, -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - M, N = x.shape - assert x.stride(-1) == 1 - dy = maybe_contiguous_lastdim(dy) - assert dy.stride(-1) == 1 - assert dy.shape == (M, N) - if dresidual is not None: - dresidual = maybe_contiguous_lastdim(dresidual) - assert dresidual.stride(-1) == 1 - assert dresidual.shape == (M, N) - assert weight.shape == (N,) - assert weight.stride(-1) == 1 - if bias is not None: - assert bias.stride(-1) == 1 - assert bias.shape == (N,) - if dy1 is not None: - dy1 = maybe_contiguous_lastdim(dy1) - assert weight1 is not None - assert dy1.shape == dy.shape - assert dy1.stride(-1) == 1 - if weight1 is not None: - assert weight1.shape == (N,) - assert weight1.stride(-1) == 1 - if bias1 is not None: - assert bias1.shape == (N,) - assert bias1.stride(-1) == 1 - if seeds is not None: - assert seeds.is_contiguous() - assert seeds.shape == (M if not has_x1 else M * 2,) - if rowscale is not None: - assert rowscale.is_contiguous() - assert rowscale.shape == (M,) - # allocate output - dx = ( - torch.empty_like(x) - if x_dtype is None - else torch.empty(M, N, dtype=x_dtype, device=x.device) - ) - dresidual_in = ( - torch.empty_like(x) - if has_residual - and (dx.dtype != x.dtype or dropout_p > 0.0 or rowscale is not None or has_x1) - else None - ) - dx1 = torch.empty_like(dx) if (has_x1 and dropout_p > 0.0) else None - y = torch.empty(M, N, dtype=dy.dtype, device=dy.device) if recompute_output else None - if recompute_output: - assert weight1 is None, "recompute_output is not supported with parallel LayerNorm" - - # Less than 64KB per feature: enqueue fused kernel - MAX_FUSED_SIZE = 65536 // x.element_size() - BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) - if N > BLOCK_N: - raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") - # Increasing the multiple (e.g. 8) will allow more thread blocks to be launched and hide the - # latency of the gmem reads/writes, but will increase the time of summing up dw / db. - sm_count = torch.cuda.get_device_properties(x.device).multi_processor_count * 8 - _dw = torch.empty((sm_count, N), dtype=torch.float32, device=weight.device) - _db = ( - torch.empty((sm_count, N), dtype=torch.float32, device=bias.device) - if bias is not None - else None - ) - _dw1 = torch.empty_like(_dw) if weight1 is not None else None - _db1 = torch.empty_like(_db) if bias1 is not None else None - rows_per_program = math.ceil(M / sm_count) - grid = (sm_count,) - with torch.cuda.device(x.device.index): - torch.library.wrap_triton(_layer_norm_bwd_kernel)[grid]( - x, - weight, - bias, - y, - dy, - dx, - _dw, - _db, - dresidual, - weight1, - dy1, - dx1, - _dw1, - _db1, - dresidual_in, - rowscale, - seeds, - mean, - rstd, - x.stride(0), - 0 if not recompute_output else y.stride(0), - dy.stride(0), - dx.stride(0), - dresidual.stride(0) if dresidual is not None else 0, - dy1.stride(0) if dy1 is not None else 0, - dx1.stride(0) if dx1 is not None else 0, - dresidual_in.stride(0) if dresidual_in is not None else 0, - M, - N, - eps, - dropout_p, - # Passing bool make torch inductor very unhappy since it then tries to compare to int_max - int(zero_centered_weight), - rows_per_program, - is_rms_norm, - BLOCK_N, - dresidual is not None, - dresidual_in is not None, - bias is not None, - dropout_p > 0.0, - HAS_ROWSCALE=rowscale is not None, - HAS_DY1=dy1 is not None, - HAS_DX1=dx1 is not None, - HAS_B1=bias1 is not None, - RECOMPUTE_OUTPUT=y is not None, - ) - dw = _dw.sum(0).to(weight.dtype) - db = _db.sum(0).to(bias.dtype) if bias is not None else None - dw1 = _dw1.sum(0).to(weight1.dtype) if weight1 is not None else None - db1 = _db1.sum(0).to(bias1.dtype) if bias1 is not None else None - # dresidual_in and dx1 could be None, the wrapper will handle assigning them from dx - return dx, dw, db, dresidual_in, dx1, dw1, db1, y - - -class LayerNormFn(torch.autograd.Function): - - @staticmethod - def forward( - ctx, - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - residual_in_fp32=False, - zero_centered_weight=False, - is_rms_norm=False, - return_dropout_mask=False, - out_dtype=None, - out=None, - residual_out=None - ): - x_shape_og = x.shape - # reshape input data into 2D tensor - x = maybe_contiguous_lastdim(x.reshape(-1, x.shape[-1])) - if residual is not None: - assert residual.shape == x_shape_og - residual = maybe_contiguous_lastdim(residual.reshape(-1, residual.shape[-1])) - if x1 is not None: - assert x1.shape == x_shape_og - assert rowscale is None, "rowscale is not supported with parallel LayerNorm" - x1 = maybe_contiguous_lastdim(x1.reshape(-1, x1.shape[-1])) - weight = weight.contiguous() - bias = maybe_contiguous(bias) - weight1 = maybe_contiguous(weight1) - bias1 = maybe_contiguous(bias1) - if rowscale is not None: - rowscale = rowscale.reshape(-1).contiguous() - residual_dtype = ( - residual.dtype - if residual is not None - else (torch.float32 if residual_in_fp32 else None) - ) - if out is not None: - out = out.reshape(-1, out.shape[-1]) - if residual_out is not None: - residual_out = residual_out.reshape(-1, residual_out.shape[-1]) - y, y1, mean, rstd, residual_out, seeds, dropout_mask, dropout_mask1 = _layer_norm_fwd( - x, - weight, - bias, - eps, - residual, - x1, - weight1, - bias1, - dropout_p=dropout_p, - rowscale=rowscale, - out_dtype=out_dtype, - residual_dtype=residual_dtype, - zero_centered_weight=zero_centered_weight, - is_rms_norm=is_rms_norm, - return_dropout_mask=return_dropout_mask, - out=out, - residual_out=residual_out, - ) - ctx.save_for_backward( - residual_out, weight, bias, weight1, bias1, rowscale, seeds, mean, rstd - ) - ctx.x_shape_og = x_shape_og - ctx.eps = eps - ctx.dropout_p = dropout_p - ctx.is_rms_norm = is_rms_norm - ctx.has_residual = residual is not None - ctx.has_x1 = x1 is not None - ctx.prenorm = prenorm - ctx.x_dtype = x.dtype - ctx.zero_centered_weight = zero_centered_weight - y = y.reshape(x_shape_og) - y1 = y1.reshape(x_shape_og) if y1 is not None else None - residual_out = residual_out.reshape(x_shape_og) if residual_out is not None else None - dropout_mask = dropout_mask.reshape(x_shape_og) if dropout_mask is not None else None - dropout_mask1 = dropout_mask1.reshape(x_shape_og) if dropout_mask1 is not None else None - if not return_dropout_mask: - if weight1 is None: - return y if not prenorm else (y, residual_out) - else: - return (y, y1) if not prenorm else (y, y1, residual_out) - else: - if weight1 is None: - return ( - (y, dropout_mask, dropout_mask1) - if not prenorm - else (y, residual_out, dropout_mask, dropout_mask1) - ) - else: - return ( - (y, y1, dropout_mask, dropout_mask1) - if not prenorm - else (y, y1, residual_out, dropout_mask, dropout_mask1) - ) - - @staticmethod - def backward(ctx, dy, *args): - x, weight, bias, weight1, bias1, rowscale, seeds, mean, rstd = ctx.saved_tensors - dy = dy.reshape(-1, dy.shape[-1]) - if weight1 is not None: - dy1, args = args[0], args[1:] - dy1 = dy1.reshape(-1, dy1.shape[-1]) - assert dy1.shape == x.shape - else: - dy1 = None - if ctx.prenorm: - dresidual = args[0] - dresidual = dresidual.reshape(-1, dresidual.shape[-1]) - assert dresidual.shape == x.shape - else: - dresidual = None - dx, dw, db, dresidual_in, dx1, dw1, db1, _ = _layer_norm_bwd( - dy, - x, - weight, - bias, - ctx.eps, - mean, - rstd, - dresidual, - dy1, - weight1, - bias1, - seeds, - ctx.dropout_p, - rowscale, - ctx.has_residual, - ctx.has_x1, - ctx.zero_centered_weight, - ctx.is_rms_norm, - x_dtype=ctx.x_dtype, - recompute_output=False, - ) - return ( - dx.reshape(ctx.x_shape_og), - dw, - db, - dresidual_in.reshape(ctx.x_shape_og) if ctx.has_residual else None, - dx1.reshape(ctx.x_shape_og) if dx1 is not None else None, - dw1, - db1, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) - - -def layer_norm_fn( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - residual_in_fp32=False, - zero_centered_weight=False, - is_rms_norm=False, - return_dropout_mask=False, - out_dtype=None, - out=None, - residual_out=None -): - return LayerNormFn.apply( - x, - weight, - bias, - residual, - x1, - weight1, - bias1, - eps, - dropout_p, - rowscale, - prenorm, - residual_in_fp32, - zero_centered_weight, - is_rms_norm, - return_dropout_mask, - out_dtype, - out, - residual_out - ) - - -def rms_norm_fn( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - residual_in_fp32=False, - zero_centered_weight=False, - return_dropout_mask=False, - out_dtype=None, - out=None, - residual_out=None -): - return LayerNormFn.apply( - x, - weight, - bias, - residual, - x1, - weight1, - bias1, - eps, - dropout_p, - rowscale, - prenorm, - residual_in_fp32, - zero_centered_weight, - True, - return_dropout_mask, - out_dtype, - out, - residual_out - ) - - -class RMSNorm(torch.nn.Module): - - def __init__(self, hidden_size, eps=1e-5, dropout_p=0.0, zero_centered_weight=False, - device=None, dtype=None): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.eps = eps - if dropout_p > 0.0: - self.drop = torch.nn.Dropout(dropout_p) - else: - self.drop = None - self.zero_centered_weight = zero_centered_weight - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.register_parameter("bias", None) - self.reset_parameters() - - def reset_parameters(self): - if not self.zero_centered_weight: - torch.nn.init.ones_(self.weight) - else: - torch.nn.init.zeros_(self.weight) - - def forward(self, x, residual=None, prenorm=False, residual_in_fp32=False): - return rms_norm_fn( - x, - self.weight, - self.bias, - residual=residual, - eps=self.eps, - dropout_p=self.drop.p if self.drop is not None and self.training else 0.0, - prenorm=prenorm, - residual_in_fp32=residual_in_fp32, - zero_centered_weight=self.zero_centered_weight, - ) - - -class LayerNormLinearFn(torch.autograd.Function): - - @staticmethod - @custom_fwd - def forward( - ctx, - x, - norm_weight, - norm_bias, - linear_weight, - linear_bias, - residual=None, - eps=1e-6, - prenorm=False, - residual_in_fp32=False, - is_rms_norm=False, - ): - x_shape_og = x.shape - # reshape input data into 2D tensor - x = maybe_contiguous_lastdim(x.reshape(-1, x.shape[-1])) - if residual is not None: - assert residual.shape == x_shape_og - residual = maybe_contiguous_lastdim(residual.reshape(-1, residual.shape[-1])) - norm_weight = norm_weight.contiguous() - norm_bias = maybe_contiguous(norm_bias) - residual_dtype = ( - residual.dtype - if residual is not None - else (torch.float32 if residual_in_fp32 else None) - ) - y, _, mean, rstd, residual_out, *rest = _layer_norm_fwd( - x, - norm_weight, - norm_bias, - eps, - residual, - out_dtype=None if not torch.is_autocast_enabled() else torch.get_autocast_dtype("cuda"), - residual_dtype=residual_dtype, - is_rms_norm=is_rms_norm, - ) - y = y.reshape(x_shape_og) - dtype = torch.get_autocast_dtype("cuda") if torch.is_autocast_enabled() else y.dtype - linear_weight = linear_weight.to(dtype) - linear_bias = linear_bias.to(dtype) if linear_bias is not None else None - out = F.linear(y.to(linear_weight.dtype), linear_weight, linear_bias) - # We don't store y, will be recomputed in the backward pass to save memory - ctx.save_for_backward(residual_out, norm_weight, norm_bias, linear_weight, mean, rstd) - ctx.x_shape_og = x_shape_og - ctx.eps = eps - ctx.is_rms_norm = is_rms_norm - ctx.has_residual = residual is not None - ctx.prenorm = prenorm - ctx.x_dtype = x.dtype - ctx.linear_bias_is_none = linear_bias is None - return out if not prenorm else (out, residual_out.reshape(x_shape_og)) - - @staticmethod - @custom_bwd - def backward(ctx, dout, *args): - x, norm_weight, norm_bias, linear_weight, mean, rstd = ctx.saved_tensors - dout = dout.reshape(-1, dout.shape[-1]) - dy = F.linear(dout, linear_weight.t()) - dlinear_bias = None if ctx.linear_bias_is_none else dout.sum(0) - dy = maybe_contiguous_lastdim(dy) - assert dy.shape == x.shape - if ctx.prenorm: - dresidual = args[0] - dresidual = maybe_contiguous_lastdim(dresidual.reshape(-1, dresidual.shape[-1])) - assert dresidual.shape == x.shape - else: - dresidual = None - dx, dnorm_weight, dnorm_bias, dresidual_in, _, _, _, y = _layer_norm_bwd( - dy, - x, - norm_weight, - norm_bias, - ctx.eps, - mean, - rstd, - dresidual=dresidual, - has_residual=ctx.has_residual, - is_rms_norm=ctx.is_rms_norm, - x_dtype=ctx.x_dtype, - recompute_output=True, - ) - dlinear_weight = torch.einsum("bo,bi->oi", dout, y) - return ( - dx.reshape(ctx.x_shape_og), - dnorm_weight, - dnorm_bias, - dlinear_weight, - dlinear_bias, - dresidual_in.reshape(ctx.x_shape_og) if ctx.has_residual else None, - None, - None, - None, - None, - ) - - -def layer_norm_linear_fn( - x, - norm_weight, - norm_bias, - linear_weight, - linear_bias, - residual=None, - eps=1e-6, - prenorm=False, - residual_in_fp32=False, - is_rms_norm=False, -): - return LayerNormLinearFn.apply( - x, - norm_weight, - norm_bias, - linear_weight, - linear_bias, - residual, - eps, - prenorm, - residual_in_fp32, - is_rms_norm, - ) diff --git a/build/torch212-cxx11-cu130-aarch64-linux/ops/triton/linear.py b/build/torch212-cxx11-cu130-aarch64-linux/ops/triton/linear.py deleted file mode 100644 index a8966dbc345ab0e593df0124451ee7be3dae131a..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu130-aarch64-linux/ops/triton/linear.py +++ /dev/null @@ -1,594 +0,0 @@ -# Adapted from https://github.com/ELS-RD/kernl/blob/main/src/kernl/implementations/linear_layer.py -# and https://github.com/openai/triton/blob/master/python/triton/ops/matmul.py -from typing import Optional - -import torch -import triton -import triton.language as tl -from triton.ops.matmul_perf_model import early_config_prune, estimate_matmul_time - -from flash_attn.ops.triton.k_activations import ( - gelu, - gelu_approx, - gelu_approx_grad, - gelu_grad, - squared_relu, - squared_relu_grad, -) - -# CREDITS: Initially inspired by the Triton tutorial on matrix multiplications - - -def init_to_zero(name): - return lambda nargs: nargs[name].zero_() - - -def get_configs_io_bound(): - configs = [] - for num_stages in [2, 3, 4, 5, 6]: - for block_m in [16, 32]: - for block_k in [32, 64]: - for block_n in [32, 64, 128, 256]: - num_warps = 2 if block_n <= 64 else 4 - configs.append( - triton.Config( - { - "BLOCK_M": block_m, - "BLOCK_N": block_n, - "BLOCK_K": block_k, - "SPLIT_K": 1, - }, - num_stages=num_stages, - num_warps=num_warps, - ) - ) - # split_k not used - # for split_k in [2, 4, 8, 16]: - # configs.append(triton.Config( - # {'BLOCK_M': block_m, 'BLOCK_N': block_n, 'BLOCK_K': block_k, 'SPLIT_K': split_k}, - # num_stages=num_stages, num_warps=num_warps, pre_hook=init_to_zero('C'))) - return configs - - -@triton.autotune( - configs=[ - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - # good for int8 - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=4, - num_warps=4, - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - ] - + get_configs_io_bound(), - key=["CACHE_KEY_M", "CACHE_KEY_N", "CACHE_KEY_K"], - prune_configs_by={ - "early_config_prune": early_config_prune, - "perf_model": estimate_matmul_time, - "top_k": 10, - }, -) -@triton.heuristics( - { - "EVEN_K": lambda args: args["K"] % (args["BLOCK_K"] * args["SPLIT_K"]) == 0, - } -) -@triton.jit -def kernel_fwd( - C, # Pointers to matrices - ACT_INPUT, - A, - B, - bias, - # Matrix dimensions - M, - N, - K, - CACHE_KEY_M, - CACHE_KEY_N, - CACHE_KEY_K, - # The stride variables represent how much to increase the ptr by when moving by 1 - # element in a particular dimension. E.g. stride_am is how much to increase a_ptr - # by to get the element one row down (A has M rows) - stride_cm, - # stride_cn, # Assume that stride_cn == 1 - stride_am, - stride_ak, - stride_bn, - stride_bk, - # Meta-parameters - BLOCK_M: tl.constexpr, - GROUP_M: tl.constexpr, - BLOCK_N: tl.constexpr, - BLOCK_K: tl.constexpr, - # split k not used, not performant with activation, kept because early_config_prune is expecting it - SPLIT_K: tl.constexpr, - EVEN_K: tl.constexpr, - A_ROWMAJOR: tl.constexpr, - B_COLMAJOR: tl.constexpr, - BIAS: tl.constexpr, - SAVE_ACT_INPUT: tl.constexpr, - ACTIVATION: tl.constexpr, -): - - """ - Kernel for computing Out = activation(A x W + C) - - Input has shape (M, K) - - Weight has shape (K, N) - - Bias has shape (N,) - - Output has shape (M, N) - - ActInputs (optional) has shape (M, N) - 'ActInputs' optionally saves the A x W + C intermediate for backward computations - This kernel will consolidate over K - """ - - pid = tl.program_id(axis=0) - - grid_m = (M + BLOCK_M - 1) // BLOCK_M - grid_n = (N + BLOCK_N - 1) // BLOCK_N - # re-order program ID for better L2 performance - width = GROUP_M * grid_n - group_id = pid // width - group_size = min(grid_m - group_id * GROUP_M, GROUP_M) - pid_m = group_id * GROUP_M + (pid % group_size) - pid_n = (pid % width) // (group_size) - - # now compute the block that each program will go through - # rm (resp. rn) denotes a range of indices - # for rows (resp. col) of C - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - # trick to avoid masking on M and N axis - ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M) - rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N) - rk = tl.arange(0, BLOCK_K) - - if A_ROWMAJOR: - A = A + (ram[:, None] * stride_am + rk[None, :]) - else: - A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak) - if B_COLMAJOR: - B = B + (rk[:, None] + rbn[None, :] * stride_bn) - else: - B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn) - - acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) - - for k in range(K, 0, -BLOCK_K): - if EVEN_K: - a = tl.load(A) - b = tl.load(B) - else: - a = tl.load(A, mask=rk[None, :] < k, other=0.0) - b = tl.load(B, mask=rk[:, None] < k, other=0.0) - acc += tl.dot(a, b) - - if A_ROWMAJOR: - A += BLOCK_K - else: - A += BLOCK_K * stride_ak - if B_COLMAJOR: - B += BLOCK_K - else: - B += BLOCK_K * stride_bk - - # Putting bias after the matmul (instead of before) is faster, idk why - if BIAS: - bias = tl.load(bias + rn, mask=rn < N, other=0.0).to(tl.float32) - acc += bias[None, :] - - # optional: save the activation inputs - if SAVE_ACT_INPUT: - # act_in_ptrs = ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :] * stride_cn - act_in_ptrs = ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :] - tl.store(act_in_ptrs, acc) - - # optional: fused activation (while the data is in shared memory) - if ACTIVATION == "gelu": - acc = gelu(acc) - elif ACTIVATION == "gelu_approx": - acc = gelu_approx(acc) - elif ACTIVATION == "squared_relu": - acc = squared_relu(acc) - # rematerialize rm and rn to save registers - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - - # write back result - # C = C + rm[:, None] * stride_cm + rn[None, :] * stride_cn - C = C + rm[:, None] * stride_cm + rn[None, :] - mask = (rm < M)[:, None] & (rn < N)[None, :] - tl.store(C, acc) - - -def triton_linear_act( - x: torch.Tensor, - weight: torch.Tensor, - bias: Optional[torch.Tensor] = None, - activation: str = "id", - save_act_input: bool = False, -) -> torch.Tensor: - """ - Compute e = activation(x @ weight.T + bias). - This wrapper kicks the `kernel_fwd` Triton kernel - :param x: input tensor - :param weight: weight matrix - :param bias: an optional bias tensor - :param activation: Activation name. Needs to be a Triton kernel. - :param act_input: an optional tensor to save the activation inputs (for backward) - :return: result tensor - """ - # if torch.is_autocast_enabled(): - # dtype = torch.get_autocast_gpu_dtype() - # x, weight, bias = [a.to(dtype=dtype) for a in [x, weight, bias]] - - assert activation in ["id", "gelu", "gelu_approx", "squared_relu"] - - batch_shape, n = x.shape[:-1], x.shape[-1] - batch_dim = batch_shape.numel() - x_reshaped = x.reshape(batch_dim, n) - - if x_reshaped.stride(0) > 1 and x_reshaped.stride(1) > 1: - x_reshaped = x_reshaped.contiguous() - if weight.stride(0) > 1 and weight.stride(1) > 1: - weight = weight.contiguous() - bias = bias.contiguous() if bias is not None else None - - assert ( - x.dtype == weight.dtype - ), f"Input and weight must have the same dtype, got {x.dtype} and {weight.dtype}" - if bias is not None: - assert ( - x.dtype == bias.dtype - ), f"Input and bias must have the same dtype, got {x.dtype} and {bias.dtype}" - assert ( - x_reshaped.shape[1] == weight.shape[1] - ), f"Incompatible dimensions: {x_reshaped.shape} - {weight.shape}" - - assert ( - bias is None or bias.shape[0] == weight.shape[0] - ), "Incompatible dimensions in between weight and bias" - - M, K = x_reshaped.shape - N, K = weight.shape - - output = torch.empty((M, N), device=x.device, dtype=x.dtype) - act_input = torch.empty_like(output) if save_act_input else None - - # 1D launch kernel where each block gets its own program. - grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),) # noqa - - kernel_fwd[grid]( - output, - act_input, - x_reshaped, - weight, # data ptrs - bias if bias is not None else x, # auto skip bias if not present - M, # shapes - N, - K, - M // 32, # key for triton cache (limit number of compilations) - N // 32, - K // 32, - stride_cm=output.stride(0), # strides - # stride_cn=output.stride(1), - stride_am=x_reshaped.stride(0), - stride_ak=x_reshaped.stride(1), - stride_bk=weight.stride(1), - stride_bn=weight.stride(0), - BIAS=bias is not None, # optional fused bias - SAVE_ACT_INPUT=save_act_input, # optional save activation inputs - ACTIVATION=activation, # optional fused activation - A_ROWMAJOR=x_reshaped.stride(1) == 1, - B_COLMAJOR=weight.stride(1) == 1, - GROUP_M=8, # speed optimization: group the programs - ) - - if not save_act_input: - return output.reshape(*batch_shape, output.shape[-1]) - else: - return ( - output.reshape(*batch_shape, output.shape[-1]), - act_input.reshape(*batch_shape, act_input.shape[-1]), - ) - - -@triton.autotune( - configs=[ - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - # good for int8 - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=4, - num_warps=4, - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - ] - + get_configs_io_bound(), - key=["CACHE_KEY_M", "CACHE_KEY_N", "CACHE_KEY_K"], - prune_configs_by={ - "early_config_prune": early_config_prune, - "perf_model": estimate_matmul_time, - "top_k": 10, - }, -) -@triton.heuristics( - { - "EVEN_K": lambda args: args["K"] % (args["BLOCK_K"] * args["SPLIT_K"]) == 0, - } -) -@triton.jit -def kernel_bwd( - C, # Pointers to matrices - ACT_INPUT, - A, - B, - # Matrix dimensions - M, - N, - K, - CACHE_KEY_M, - CACHE_KEY_N, - CACHE_KEY_K, - # The stride variables represent how much to increase the ptr by when moving by 1 - # element in a particular dimension. E.g. stride_am is how much to increase a_ptr - # by to get the element one row down (A has M rows) - stride_cm, - # stride_cn, # Assume that stride_cn == 1 - stride_am, - stride_ak, - stride_bk, - stride_bn, - # Meta-parameters - BLOCK_M: tl.constexpr, - GROUP_M: tl.constexpr, - BLOCK_N: tl.constexpr, - BLOCK_K: tl.constexpr, - # split k not used, not performant with activation, kept because early_config_prune is expecting it - SPLIT_K: tl.constexpr, - EVEN_K: tl.constexpr, - ACTIVATION: tl.constexpr, -): - - """ - Kernel for computing Out = activation(A x W + C) - - Input has shape (M, K) - - Weight has shape (K, N) - - Output has shape (M, N) - - ActInputs (optional) has shape (M, N) - 'ActInputs' optionally saves the A x W + C intermediate for backward computations - This kernel will consolidate over K - """ - - pid = tl.program_id(axis=0) - - grid_m = (M + BLOCK_M - 1) // BLOCK_M - grid_n = (N + BLOCK_N - 1) // BLOCK_N - # re-order program ID for better L2 performance - width = GROUP_M * grid_n - group_id = pid // width - group_size = min(grid_m - group_id * GROUP_M, GROUP_M) - pid_m = group_id * GROUP_M + (pid % group_size) - pid_n = (pid % width) // (group_size) - - # now compute the block that each program will go through - # rm (resp. rn) denotes a range of indices - # for rows (resp. col) of C - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - # trick to avoid masking on M and N axis - ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M) - rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N) - rk = tl.arange(0, BLOCK_K) - - A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak) - B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn) - - acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) - - for k in range(K, 0, -BLOCK_K): - if EVEN_K: - a = tl.load(A) - b = tl.load(B) - else: - a = tl.load(A, mask=rk[None, :] < k, other=0.0) - b = tl.load(B, mask=rk[:, None] < k, other=0.0) - acc += tl.dot(a, b) - - A += BLOCK_K * stride_ak - B += BLOCK_K * stride_bk - - # optional: fused activation (while the data is in shared memory) - if ACTIVATION != "id": - act_in_ptrs = ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :] - act_input = tl.load(act_in_ptrs).to(acc.dtype) - if ACTIVATION == "gelu": - acc *= gelu_grad(act_input) - elif ACTIVATION == "gelu_approx": - acc *= gelu_approx_grad(act_input) - elif ACTIVATION == "squared_relu": - acc *= squared_relu_grad(act_input) - - # rematerialize rm and rn to save registers - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - - # write back result - C = C + rm[:, None] * stride_cm + rn[None, :] - mask = (rm < M)[:, None] & (rn < N)[None, :] - tl.store(C, acc, mask=mask) - - -def triton_dgrad_act( - grad_output: torch.Tensor, - weight: torch.Tensor, - activation: str = "id", - act_input: Optional[torch.Tensor] = None, -) -> torch.Tensor: - """ - Compute e = activation(grad_output @ weight + bias). - This wrapper kicks the `kernel_fwd` Triton kernel - :param grad_output: input tensor - :param weight: weight matrix - :param activation: Activation name. Needs to be a Triton kernel. - :param act_input: an optional tensor to save the activation inputs (for backward) - :return: result tensor - """ - assert activation in ["id", "gelu", "gelu_approx", "squared_relu"] - - batch_shape, n = grad_output.shape[:-1], grad_output.shape[-1] - batch_dim = batch_shape.numel() - grad_output_reshaped = grad_output.reshape(batch_dim, n) - - if grad_output_reshaped.stride(0) > 1 and grad_output_reshaped.stride(1) > 1: - grad_output_reshaped = grad_output_reshaped.contiguous() - if weight.stride(0) > 1 and weight.stride(1) > 1: - weight = weight.contiguous() - - assert ( - grad_output.dtype == weight.dtype - ), f"grad_output and weight must have the same dtype, got {grad_output.dtype} and {weight.dtype}" - assert ( - grad_output_reshaped.shape[1] == weight.shape[0] - ), f"Incompatible dimensions: {grad_output_reshaped.shape} - {weight.shape}" - if activation != "id": - assert act_input is not None, f"act_input is required for activation {activation}" - - # M, N, K in bwd are different from M, N, K in fwd - M, K = grad_output_reshaped.shape - K, N = weight.shape - - grad_input = torch.empty((M, N), device=grad_output.device, dtype=grad_output.dtype) - - # 1D launch kernel where each block gets its own program. - grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),) # noqa - - kernel_bwd[grid]( - grad_input, - act_input, - grad_output_reshaped, - weight, # data ptrs - M, # shapes - N, - K, - M // 32, # key for triton cache (limit number of compilations) - N // 32, - K // 32, - stride_cm=grad_input.stride(0), # strides - # stride_cn=grad_input.stride(1), - stride_am=grad_output_reshaped.stride(0), - stride_ak=grad_output_reshaped.stride(1), - stride_bk=weight.stride(0), - stride_bn=weight.stride(1), - ACTIVATION=activation, # optional fused activation - GROUP_M=8, # speed optimization: group the programs - ) - - return grad_input.reshape(*batch_shape, grad_input.shape[-1]) diff --git a/build/torch212-cxx11-cu130-aarch64-linux/ops/triton/mlp.py b/build/torch212-cxx11-cu130-aarch64-linux/ops/triton/mlp.py deleted file mode 100644 index 059f4f8a5e174c1f4824e43d313fca18eaa799b8..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu130-aarch64-linux/ops/triton/mlp.py +++ /dev/null @@ -1,149 +0,0 @@ -# The triton fused matmul + sqrelu is faster for fp16 but slower for bf16, compared -# to naive implementation. -import fused_dense_lib as fused_dense_cuda -import torch -import torch.nn as nn -import torch.nn.functional as F - -from flash_attn.utils.torch import custom_fwd, custom_bwd -from flash_attn.ops.activations import sqrelu_bwd, sqrelu_fwd -from flash_attn.ops.triton.linear import triton_dgrad_act, triton_linear_act - - -class FusedDenseSqreluDenseFunc(torch.autograd.Function): - @staticmethod - @custom_fwd - def forward(ctx, x, weight1, bias1, weight2, bias2, checkpoint_lvl=0): - """checkpoint_lvl: - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute act_input and gelu_out in the bwd - """ - if torch.is_autocast_enabled(): - dtype = torch.get_autocast_gpu_dtype() - x, weight1, bias1, weight2, bias2 = [ - a.to(dtype=dtype) for a in [x, weight1, bias1, weight2, bias2] - ] - is_bf16 = x.dtype == torch.bfloat16 - assert checkpoint_lvl in [0, 1, 2] - x = x.contiguous() - weight1 = weight1.contiguous() - bias1 = bias1.contiguous() - weight2 = weight2.contiguous() - bias2 = bias2.contiguous() - batch_shape, n = x.shape[:-1], x.shape[-1] - batch_dim = batch_shape.numel() - if is_bf16: - act_input = fused_dense_cuda.linear_bias_forward( - x.reshape(batch_dim, n), weight1, bias1 - ) - output1 = sqrelu_fwd(act_input) - else: - save_act_input = checkpoint_lvl != 2 - result = triton_linear_act( - x.reshape(batch_dim, n), - weight1, - bias1, - activation="squared_relu", - save_act_input=save_act_input, - ) - if save_act_input: - output1, act_input = result - else: - output1 = result - output2 = fused_dense_cuda.linear_bias_forward(output1, weight2, bias2) - ctx.checkpoint_lvl = checkpoint_lvl - if checkpoint_lvl == 0: - ctx.save_for_backward(x, weight1, bias1, weight2, act_input, output1) - elif checkpoint_lvl == 1: - ctx.save_for_backward(x, weight1, bias1, weight2, act_input) - elif checkpoint_lvl == 2: - ctx.save_for_backward(x, weight1, bias1, weight2) - return output2.reshape(*batch_shape, output2.shape[-1]) - - @staticmethod - @custom_bwd - def backward(ctx, grad_output): - grad_output = grad_output.contiguous() - checkpoint_lvl = ctx.checkpoint_lvl - x, weight1, bias1, weight2, *rest = ctx.saved_tensors - batch_shape, n = x.shape[:-1], x.shape[-1] - batch_dim = batch_shape.numel() - is_bf16 = x.dtype == torch.bfloat16 - if checkpoint_lvl == 0: - act_input, output1 = rest - elif checkpoint_lvl == 1: - (act_input,) = rest - output1 = sqrelu_fwd(act_input) - elif checkpoint_lvl == 2: - if is_bf16: - act_input = fused_dense_cuda.linear_bias_forward( - x.reshape(batch_dim, n), weight1, bias1 - ) - output1 = sqrelu_fwd(act_input) - else: - output1, act_input = triton_linear_act( - x.reshape(batch_dim, n), - weight1, - bias1, - activation="squared_relu", - save_act_input=True, - ) - - if is_bf16: - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad(output1, grad_output) - grad_output1 = grad_output @ weight2 - grad_act_input = sqrelu_bwd(grad_output1, act_input) - grad_input, grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_backward( - x.reshape(batch_dim, n), weight1, grad_act_input - ) - else: - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad(output1, grad_output) - grad_act_input = triton_dgrad_act( - grad_output, weight2, activation="squared_relu", act_input=act_input - ) - grad_input, grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_backward( - x.reshape(batch_dim, n), weight1, grad_act_input - ) - return grad_input.reshape_as(x), grad_weight1, grad_bias1, grad_weight2, grad_bias2, None - - -fused_dense_sqrelu_dense_function = FusedDenseSqreluDenseFunc.apply - - -class FusedDenseSqreluDense(nn.Module): - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - bias1=True, - bias2=True, - checkpoint_lvl=0, - device=None, - dtype=None, - ): - """ - checkpoint_lvl (increasing lvl means slower but more memory saving): - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute gelu_in and gelu_out in the bwd - """ - assert checkpoint_lvl in [0, 1, 2] - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features * 4 - assert bias1 == True, "DenseSqreluDense module without bias is currently not supported" - assert bias2 == True, "DenseSqreluDense module without bias is currently not supported" - self.checkpoint_lvl = checkpoint_lvl - self.fc1 = nn.Linear(in_features, hidden_features, bias=bias1, **factory_kwargs) - self.fc2 = nn.Linear(hidden_features, out_features, bias=bias2, **factory_kwargs) - - def forward(self, x): - assert x.is_cuda - return fused_dense_sqrelu_dense_function( - x, self.fc1.weight, self.fc1.bias, self.fc2.weight, self.fc2.bias, self.checkpoint_lvl - ) diff --git a/build/torch212-cxx11-cu130-aarch64-linux/ops/triton/rotary.py b/build/torch212-cxx11-cu130-aarch64-linux/ops/triton/rotary.py deleted file mode 100644 index 6ad3ab412a084270120167ad32b67ad1720f2c37..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu130-aarch64-linux/ops/triton/rotary.py +++ /dev/null @@ -1,186 +0,0 @@ -# Copyright (c) 2025, Tri Dao. -# As of 2025-04-23, we require triton >= 3.0 - -from typing import Optional, Union - -import torch - -import triton -import triton.language as tl - - -@triton.jit -def rotary_kernel( - OUT, # Pointers to matrices - X, - COS, - SIN, - CU_SEQLENS, - SEQLEN_OFFSETS, # this could be int or a pointer - # Matrix dimensions - seqlen, - nheads, - seqlen_ro, - # strides - stride_out_batch, - stride_out_seqlen, - stride_out_nheads, - stride_out_headdim, - stride_x_batch, - stride_x_seqlen, - stride_x_nheads, - stride_x_headdim, - # Meta-parameters - # We want ROTARY_DIM to be constexpr, otherwise the triton compiler doesn't know that - # the mask is constant every 8 elements, and it will generate LDG.16 instead of LDG.128 - ROTARY_DIM: tl.constexpr, - IS_SEQLEN_OFFSETS_TENSOR: tl.constexpr, - IS_VARLEN: tl.constexpr, - INTERLEAVED: tl.constexpr, - CONJUGATE: tl.constexpr, - BLOCK_H: tl.constexpr, - BLOCK_M: tl.constexpr, -): - BLOCK_K: tl.constexpr = triton.next_power_of_2(ROTARY_DIM) - ROTARY_DIM_HALF = ROTARY_DIM // 2 - pid_head = tl.program_id(axis=0) - pid_m = tl.program_id(axis=1) - pid_batch = tl.program_id(axis=2) - - if not IS_VARLEN: - X = X + pid_batch * stride_x_batch - OUT = OUT + pid_batch * stride_out_batch - else: - start_idx = tl.load(CU_SEQLENS + pid_batch) - seqlen = tl.load(CU_SEQLENS + pid_batch + 1) - start_idx - X = X + start_idx * stride_x_seqlen - OUT = OUT + start_idx * stride_out_seqlen - - if pid_m * BLOCK_M >= seqlen: - return - - rh = pid_head * BLOCK_H + tl.arange(0, BLOCK_H) - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - if not IS_SEQLEN_OFFSETS_TENSOR: - rm_cs = rm + SEQLEN_OFFSETS - else: - rm_cs = rm + tl.load(SEQLEN_OFFSETS + pid_batch) - - rk_half = tl.arange(0, BLOCK_K // 2) - COS = COS + (rm_cs[:, None] * ROTARY_DIM_HALF + rk_half[None, :]) - SIN = SIN + (rm_cs[:, None] * ROTARY_DIM_HALF + rk_half[None, :]) - mask_cs = (rm_cs[:, None] < seqlen_ro) & (rk_half[None, :] < ROTARY_DIM_HALF) - cos = tl.load(COS, mask=mask_cs, other=1.0).to(tl.float32) - sin = tl.load(SIN, mask=mask_cs, other=0.0).to(tl.float32) - if CONJUGATE: - sin = -sin - - if not INTERLEAVED: - # Load the 1st and 2nd halves of X, do calculation, then store to 1st and 2nd halves of OUT - X = X + (rh[:, None, None] * stride_x_nheads + rm[None, :, None] * stride_x_seqlen + rk_half[None, None, :] * stride_x_headdim) - OUT = OUT + (rh[:, None, None] * stride_out_nheads + rm[None, :, None] * stride_out_seqlen + rk_half[None, None, :] * stride_out_headdim) - mask = (rh[:, None, None] < nheads) & (rm[None, :, None] < seqlen) & (rk_half[None, None, :] < ROTARY_DIM_HALF) - x0 = tl.load(X, mask=mask, other=0.0).to(tl.float32) - x1 = tl.load(X + ROTARY_DIM_HALF * stride_x_headdim, mask=mask, other=0.0,).to(tl.float32) - o0 = x0 * cos - x1 * sin - o1 = x0 * sin + x1 * cos - tl.store(OUT, o0, mask=mask) - tl.store(OUT + ROTARY_DIM_HALF * stride_out_headdim, o1, mask=mask) - else: - rk = tl.arange(0, BLOCK_K) - X = X + (rh[:, None, None] * stride_x_nheads + rm[None, :, None] * stride_x_seqlen + rk[None, None, :] * stride_x_headdim) - OUT = OUT + (rh[:, None, None] * stride_out_nheads + rm[None, :, None] * stride_out_seqlen + rk[None, None, :] * stride_out_headdim) - mask = (rh[:, None, None] < nheads) & (rm[None, :, None] < seqlen) & (rk[None, None, :] < ROTARY_DIM) - x = tl.load(X, mask=mask, other=0.0).to(tl.float32) - x0, x1 = tl.split(tl.reshape(x, [BLOCK_H, BLOCK_M, BLOCK_K // 2, 2])) - o0 = x0 * cos - x1 * sin - o1 = x0 * sin + x1 * cos - o = tl.reshape(tl.join(o0, o1), [BLOCK_H, BLOCK_M, BLOCK_K]) - tl.store(OUT, o, mask=mask) - - -def apply_rotary( - x: torch.Tensor, - cos: torch.Tensor, - sin: torch.Tensor, - seqlen_offsets: Union[int, torch.Tensor] = 0, - cu_seqlens: Optional[torch.Tensor] = None, - max_seqlen: Optional[int] = None, - interleaved=False, - inplace=False, - conjugate=False, -) -> torch.Tensor: - """ - Arguments: - x: (batch, seqlen, nheads, headdim) if cu_seqlens is None - else (total_seqlen, nheads, headdim). - cos: (seqlen_ro, rotary_dim / 2) - sin: (seqlen_ro, rotary_dim / 2) - seqlen_offsets: integer or integer tensor of size (batch,) - cu_seqlens: (batch + 1,) or None - max_seqlen: int - Returns: - y: (batch, seqlen, nheads, headdim) - """ - is_varlen = cu_seqlens is not None - if not is_varlen: - batch, seqlen, nheads, headdim = x.shape - else: - assert max_seqlen is not None, "If cu_seqlens is passed in, then max_seqlen must be passed" - total_seqlen, nheads, headdim = x.shape - batch_p_1 = cu_seqlens.shape[0] - batch = batch_p_1 - 1 - seqlen = max_seqlen - seqlen_ro, rotary_dim = cos.shape - assert sin.shape == cos.shape - rotary_dim *= 2 - assert rotary_dim <= headdim, "rotary_dim must be <= headdim" - assert headdim <= 256, "Only support headdim <= 256" - assert seqlen_ro >= seqlen, "seqlen_ro must be >= seqlen" - - cos, sin = cos.contiguous(), sin.contiguous() - if isinstance(seqlen_offsets, torch.Tensor): - assert seqlen_offsets.shape == (batch,) - assert seqlen_offsets.dtype in [torch.int32, torch.int64] - seqlen_offsets = seqlen_offsets.contiguous() - else: - assert seqlen_offsets + seqlen <= seqlen_ro - - output = torch.empty_like(x) if not inplace else x - if rotary_dim < headdim and not inplace: - output[..., rotary_dim:].copy_(x[..., rotary_dim:]) - - grid = lambda META: (triton.cdiv(nheads, META["BLOCK_H"]), triton.cdiv(seqlen, META["BLOCK_M"]), batch) # noqa - BLOCK_M = 8 if rotary_dim <= 128 else 4 - - # Need this, otherwise Triton tries to launch from cuda:0 and we get - # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?) - device_ctx = torch.cuda.device(x.device.index) if x.device.type == 'cuda' else torch.xpu.device(x.device.index) - with device_ctx: - torch.library.wrap_triton(rotary_kernel)[grid]( - output, # data ptrs - x, - cos, - sin, - cu_seqlens, - seqlen_offsets, - seqlen, # shapes - nheads, - seqlen_ro, - output.stride(0) if not is_varlen else 0, # batch_strides if not varlen else 0 - output.stride(-3), # seqlen_stride or total_seqlen_stride - output.stride(-2), # nheads_stride - output.stride(-1), # headdim_stride - x.stride(0) if not is_varlen else 0, # batch_strides if not varlen else 0 - x.stride(-3), # seqlen stride or total_seqlen_stride - x.stride(-2), # nheads stride - x.stride(-1), # headdim stride - rotary_dim, - isinstance(seqlen_offsets, torch.Tensor), - is_varlen, - interleaved, - conjugate, - BLOCK_M=BLOCK_M, - BLOCK_H=2, - ) - return output diff --git a/build/torch212-cxx11-cu130-x86_64-linux/__init__.py b/build/torch212-cxx11-cu130-x86_64-linux/__init__.py deleted file mode 100644 index ecc2f9d896b6c93f90b0a1499856dc0612177422..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu130-x86_64-linux/__init__.py +++ /dev/null @@ -1,393 +0,0 @@ -from typing import Optional, List -import torch -from ._ops import ops as flash_attn_ops -from .flash_attn_interface import ( - flash_attn_func, - flash_attn_kvpacked_func, - flash_attn_qkvpacked_func, - flash_attn_varlen_func, - flash_attn_varlen_kvpacked_func, - flash_attn_varlen_qkvpacked_func, - flash_attn_with_kvcache, -) - - -def fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - return_softmax: bool = False, - gen: Optional[torch.Generator] = None, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention. - - Args: - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Optional output tensor, same shape as q - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - return_softmax: Whether to return softmax weights - gen: Optional random number generator - - Returns: - List of tensors: [output, softmax_lse, (softmax if return_softmax)] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.fwd( - q, - k, - v, - out, - alibi_slopes, - p_dropout, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - gen, - ) - - -def varlen_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - out: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - block_table: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - max_seqlen_q: int = 0, - max_seqlen_k: int = 0, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - zero_tensors: bool = False, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - return_softmax: bool = False, - gen: Optional[torch.Generator] = None, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention with variable sequence lengths. - - Args: - q: Query tensor of shape [total_q, num_heads, head_size] - k: Key tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - v: Value tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1] - cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1] - out: Optional output tensor of shape [total_q, num_heads, head_size] - seqused_k: Optional tensor specifying how many keys to use per batch element [batch_size] - leftpad_k: Optional left padding for keys of shape [batch_size] - block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq] - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - max_seqlen_q: Maximum sequence length for queries - max_seqlen_k: Maximum sequence length for keys - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - zero_tensors: Whether to zero tensors before computation - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - return_softmax: Whether to return softmax weights - gen: Optional random number generator - - Returns: - List of tensors: [output, softmax_lse, (softmax if return_softmax)] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.varlen_fwd( - q, - k, - v, - out, - cu_seqlens_q, - cu_seqlens_k, - seqused_k, - leftpad_k, - block_table, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - p_dropout, - softmax_scale, - zero_tensors, - is_causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - gen, - ) - - -def bwd( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - p_dropout: float = 0.0, - 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, - gen: Optional[torch.Generator] = None, - rng_state: Optional[torch.Tensor] = None, -) -> List[torch.Tensor]: - """ - Backward pass for multi-head attention. - - Args: - dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size] - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size] - softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q] - dq: Optional gradient tensor for queries, same shape as q - dk: Optional gradient tensor for keys, same shape as k - dv: Optional gradient tensor for values, same shape as v - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - deterministic: Whether to use deterministic algorithms - gen: Optional random number generator - rng_state: Optional RNG state from forward pass - - Returns: - List of tensors: [dq, dk, dv] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - alibi_slopes, - p_dropout, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - deterministic, - gen, - rng_state, - ) - - -def varlen_bwd( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - max_seqlen_q: int = 0, - max_seqlen_k: int = 0, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - zero_tensors: bool = False, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - deterministic: bool = False, - gen: Optional[torch.Generator] = None, - rng_state: Optional[torch.Tensor] = None, -) -> List[torch.Tensor]: - """ - Backward pass for multi-head attention with variable sequence lengths. - - Args: - dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size] - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size] - softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q] - cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1] - cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1] - dq: Optional gradient tensor for queries, same shape as q - dk: Optional gradient tensor for keys, same shape as k - dv: Optional gradient tensor for values, same shape as v - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - max_seqlen_q: Maximum sequence length for queries - max_seqlen_k: Maximum sequence length for keys - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - zero_tensors: Whether to zero tensors before computation - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - deterministic: Whether to use deterministic algorithms - gen: Optional random number generator - rng_state: Optional RNG state from forward pass - - Returns: - List of tensors: [dq, dk, dv] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.varlen_bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - p_dropout, - softmax_scale, - zero_tensors, - is_causal, - window_size_left, - window_size_right, - softcap, - deterministic, - gen, - rng_state, - ) - - -def fwd_kvcache( - q: torch.Tensor, - kcache: torch.Tensor, - vcache: torch.Tensor, - k: Optional[torch.Tensor] = None, - v: Optional[torch.Tensor] = None, - seqlens_k: Optional[torch.Tensor] = None, - rotary_cos: Optional[torch.Tensor] = None, - rotary_sin: Optional[torch.Tensor] = None, - cache_batch_idx: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - block_table: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - out: 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, - is_rotary_interleaved: bool = False, - num_splits: int = 1, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention with KV cache. - - Args: - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - kcache: Key cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - vcache: Value cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - k: Optional new keys tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size] - v: Optional new values tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size] - seqlens_k: Optional sequence lengths for keys of shape [batch_size] - rotary_cos: Optional rotary cosine tensor of shape [seqlen_ro, rotary_dim/2] - rotary_sin: Optional rotary sine tensor of shape [seqlen_ro, rotary_dim/2] - cache_batch_idx: Optional indices to index into the KV cache - leftpad_k: Optional left padding for keys of shape [batch_size] - block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq] - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - out: Optional output tensor, same shape as q - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - is_rotary_interleaved: Whether rotary embeddings are interleaved - num_splits: Number of splits for computation - - Returns: - List of tensors: [output, softmax_lse] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.fwd_kvcache( - q, - kcache, - vcache, - k, - v, - seqlens_k, - rotary_cos, - rotary_sin, - cache_batch_idx, - leftpad_k, - block_table, - alibi_slopes, - out, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - is_rotary_interleaved, - num_splits, - ) diff --git a/build/torch212-cxx11-cu130-x86_64-linux/_flash_attn2_cuda_f12afc9.abi3.so b/build/torch212-cxx11-cu130-x86_64-linux/_flash_attn2_cuda_f12afc9.abi3.so deleted file mode 100644 index e681b43ef1c5d1072c029b930bae40fdbebd1eb4..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu130-x86_64-linux/_flash_attn2_cuda_f12afc9.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:c364417d85bce469fe6d8e044191b42df1a83e269b62f718110b077551fc2d2d -size 1007151216 diff --git a/build/torch212-cxx11-cu130-x86_64-linux/_ops.py b/build/torch212-cxx11-cu130-x86_64-linux/_ops.py deleted file mode 100644 index af5b731cb04b5e6ce54e862db0501f3683c3163f..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu130-x86_64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _flash_attn2_cuda_f12afc9 -ops = torch.ops._flash_attn2_cuda_f12afc9 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_flash_attn2_cuda_f12afc9::{op_name}" diff --git a/build/torch212-cxx11-cu130-x86_64-linux/bert_padding.py b/build/torch212-cxx11-cu130-x86_64-linux/bert_padding.py deleted file mode 100644 index 3c2d35159a014a9d03aabead9e52e009168696ea..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu130-x86_64-linux/bert_padding.py +++ /dev/null @@ -1,218 +0,0 @@ -# Adapted from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/padding.py - -import torch -import torch.nn.functional as F -from einops import rearrange, repeat - - -class IndexFirstAxis(torch.autograd.Function): - @staticmethod - def forward(ctx, input, indices): - ctx.save_for_backward(indices) - assert input.ndim >= 2 - ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] - second_dim = other_shape.numel() - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - # return input[indices] - return torch.gather( - rearrange(input, "b ... -> b (...)"), 0, repeat(indices, "z -> z d", d=second_dim) - ).reshape(-1, *other_shape) - - @staticmethod - def backward(ctx, grad_output): - (indices,) = ctx.saved_tensors - assert grad_output.ndim >= 2 - other_shape = grad_output.shape[1:] - grad_output = rearrange(grad_output, "b ... -> b (...)") - grad_input = torch.zeros( - [ctx.first_axis_dim, grad_output.shape[1]], - device=grad_output.device, - dtype=grad_output.dtype, - ) - # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. - # grad_input[indices] = grad_output - grad_input.scatter_(0, repeat(indices, "z -> z d", d=grad_output.shape[1]), grad_output) - return grad_input.reshape(ctx.first_axis_dim, *other_shape), None - - -index_first_axis = IndexFirstAxis.apply - - -class IndexPutFirstAxis(torch.autograd.Function): - @staticmethod - def forward(ctx, values, indices, first_axis_dim): - ctx.save_for_backward(indices) - assert indices.ndim == 1 - assert values.ndim >= 2 - output = torch.zeros( - first_axis_dim, *values.shape[1:], device=values.device, dtype=values.dtype - ) - # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. - output[indices] = values - # output.scatter_(0, repeat(indices, 'z -> z d', d=values.shape[1]), values) - return output - - @staticmethod - def backward(ctx, grad_output): - (indices,) = ctx.saved_tensors - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - grad_values = grad_output[indices] - # grad_values = torch.gather(grad_output, 0, repeat(indices, 'z -> z d', d=grad_output.shape[1])) - return grad_values, None, None - - -index_put_first_axis = IndexPutFirstAxis.apply - - -class IndexFirstAxisResidual(torch.autograd.Function): - @staticmethod - def forward(ctx, input, indices): - ctx.save_for_backward(indices) - assert input.ndim >= 2 - ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] - second_dim = other_shape.numel() - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - output = input[indices] - # We don't want to reshape input (b ... -> b (...)) since it could change the channel_last - # memory format to channel_first. In other words, input might not be contiguous. - # If we don't detach, Pytorch complains about output being a view and is being modified inplace - return output, input.detach() - - @staticmethod - def backward(ctx, grad_output, grad_residual): - (indices,) = ctx.saved_tensors - assert grad_output.ndim >= 2 - other_shape = grad_output.shape[1:] - assert grad_residual.shape[1:] == other_shape - grad_input = grad_residual - # grad_input[indices] += grad_output - indices = indices.reshape(indices.shape[0], *((1,) * (grad_output.ndim - 1))) - indices = indices.expand_as(grad_output) - grad_input.scatter_add_(0, indices, grad_output) - return grad_input.reshape(ctx.first_axis_dim, *other_shape), None - - -index_first_axis_residual = IndexFirstAxisResidual.apply - - -def unpad_input(hidden_states, attention_mask, unused_mask=None): - """ - Arguments: - hidden_states: (batch, seqlen, ...) - attention_mask: (batch, seqlen), bool / int, 1 means valid and 0 means not valid. - unused_mask: (batch, seqlen), bool / int, 1 means the element is allocated but unused. - Return: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens selected in attention_mask + unused_mask. - indices: (total_nnz), the indices of masked tokens from the flattened input sequence. - cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. - max_seqlen_in_batch: int - seqused: (batch), returns the number of tokens selected in attention_mask + unused_mask. - """ - all_masks = (attention_mask + unused_mask) if unused_mask is not None else attention_mask - seqlens_in_batch = all_masks.sum(dim=-1, dtype=torch.int32) - used_seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) - indices = torch.nonzero(all_masks.flatten(), as_tuple=False).flatten() - max_seqlen_in_batch = seqlens_in_batch.max().item() - cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) - # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the - # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim - # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to - # index with integer indices. Moreover, torch's index is a bit slower than it needs to be, - # so we write custom forward and backward to make it a bit faster. - return ( - index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices), - indices, - cu_seqlens, - max_seqlen_in_batch, - used_seqlens_in_batch, - ) - - -def unpad_input_for_concatenated_sequences(hidden_states, attention_mask_in_length): - """ - Supports concatenating short samples in one sequence. The attention_mask_in_length is utilized to mask other short samples. It helps efficient training of variant lengths-based samples (e.g., the supervised fine-tuning task in large language model). - The motivation for this function is explained [here](https://github.com/Dao-AILab/flash-attention/issues/432#issuecomment-1668822286). - - For example, if batch = 3 and seqlen = 6, the attention_mask_in_length is: - ``` - [ - [2, 3, 0, 0, 0, 0], - [3, 2, 0, 0, 0, 0], - [6, 0, 0, 0, 0, 0] - ] - ``` - , which refers to the 3D-attention mask: - ``` - [ - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [0, 0, 1, 0, 0, 0], - [0, 0, 1, 1, 0, 0], - [0, 0, 1, 1, 1, 0], - [0, 0, 0, 0, 0, 1] - ], - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [0, 0, 0, 1, 0, 0], - [0, 0, 0, 1, 1, 0], - [0, 0, 0, 0, 0, 1] - ], - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [1, 1, 1, 1, 0, 0], - [1, 1, 1, 1, 1, 0], - [1, 1, 1, 1, 1, 1] - ] - ] - ```. - - Arguments: - hidden_states: (batch, seqlen, ...) - attention_mask_in_length: (batch, seqlen), int, a nonzero number (e.g., 1, 2, 3, etc.) means length of concatenated sequence in b-th batch, and 0 means none. - Return: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. - indices: (total_nnz), the indices of non-masked tokens from the flattened input sequence. - cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. - max_seqlen_in_batch: int - """ - length = attention_mask_in_length.sum(dim=-1) - seqlen = attention_mask_in_length.size(-1) - attention_mask_2d = torch.arange(seqlen, device=length.device, dtype=length.dtype).expand(len(length), seqlen) < length.unsqueeze(1) - real_indices_idx = torch.nonzero(attention_mask_in_length.flatten(), as_tuple=False).flatten() - seqlens_in_batch = attention_mask_in_length.flatten()[real_indices_idx] - indices = torch.nonzero(attention_mask_2d.flatten(), as_tuple=False).flatten() - max_seqlen_in_batch = seqlens_in_batch.max().item() - cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) - # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the - # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim - # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to - # index with integer indices. Moreover, torch's index is a bit slower than it needs to be, - # so we write custom forward and backward to make it a bit faster. - return ( - index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices), - indices, - cu_seqlens, - max_seqlen_in_batch, - ) - - -def pad_input(hidden_states, indices, batch, seqlen): - """ - Arguments: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. - indices: (total_nnz), the indices that represent the non-masked tokens of the original padded input sequence. - batch: int, batch size for the padded sequence. - seqlen: int, maximum sequence length for the padded sequence. - Return: - hidden_states: (batch, seqlen, ...) - """ - dim = hidden_states.shape[-1] - # output = torch.zeros((batch * seqlen), dim, device=hidden_states.device, dtype=hidden_states.dtype) - # output[indices] = hidden_states - output = index_put_first_axis(hidden_states, indices, batch * seqlen) - return rearrange(output, "(b s) ... -> b s ...", b=batch) diff --git a/build/torch212-cxx11-cu130-x86_64-linux/flash_attn2/__init__.py b/build/torch212-cxx11-cu130-x86_64-linux/flash_attn2/__init__.py deleted file mode 100644 index a9b2672c1cd85b74c1b3ded0fc0b2100e1aeac23..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu130-x86_64-linux/flash_attn2/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -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/torch212-cxx11-cu130-x86_64-linux/flash_attn_interface.py b/build/torch212-cxx11-cu130-x86_64-linux/flash_attn_interface.py deleted file mode 100644 index b014e569dc9812a9c840d1043f2cebe767b4d477..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu130-x86_64-linux/flash_attn_interface.py +++ /dev/null @@ -1,1631 +0,0 @@ -# Copyright (c) 2023, Tri Dao. - -from typing import Optional, Sequence, Tuple, Union - -import torch -import torch.nn as nn -import os - -# # isort: off -# # We need to import the CUDA kernels after importing torch -# USE_TRITON_ROCM = os.getenv("FLASH_ATTENTION_TRITON_AMD_ENABLE", "FALSE") == "TRUE" -# if USE_TRITON_ROCM: -# from .flash_attn_triton_amd import interface_fa as flash_attn -# else: -# import flash_attn_2_cuda as flash_attn - - -from . import _ops -from ._ops import ops as flash_attn -from ._ops import add_op_namespace_prefix - -# # isort: on - -def maybe_contiguous(x): - return x.contiguous() if x is not None and x.stride(-1) != 1 else x - - -def _get_device(): - if torch.xpu.is_available(): - return "xpu" - elif torch.cuda.is_available(): - return "cuda" - else: - return "cpu" - - -def _get_block_size_n(device, head_dim, is_dropout, is_causal): - assert head_dim <= 256 - - if device.type == "xpu": - if head_dim <= 96: - return 64 - elif head_dim <= 128: - return 32 - elif head_dim <= 256: - return 64 - else: - return 32 - - # This should match the block sizes in the CUDA kernel - major, minor = torch.cuda.get_device_capability(device) - is_sm8x = major == 8 and minor > 0 # Only include sm86 and sm89, exclude sm80 (A100) - is_sm80 = major == 8 and minor == 0 - is_sm90 = major == 9 and minor == 0 - if head_dim <= 32: - return 128 - if head_dim <= 64: - return 128 if not is_dropout else 64 - elif head_dim <= 96: - return 64 - elif head_dim <= 128: - if is_sm8x: - return 64 if (not is_dropout and is_causal) else 32 - else: - return 64 if not is_dropout else 32 - elif head_dim <= 192: - return 64 - elif head_dim <= 224: - return 64 - elif head_dim <= 256: - return 64 - - -def round_multiple(x, m): - return (x + m - 1) // m * m - - -# torch.compile() support is only enabled for pytorch >= 2.4 -# The reason for this is that we are using the new custom_op and register_fake -# APIs, which support inplace modification of inputs in the function itself -if torch.__version__ >= "2.4.0": - _torch_custom_op_wrapper = torch.library.custom_op - _torch_register_fake_wrapper = torch.library.register_fake -else: - def noop_custom_op_wrapper(name, fn=None, /, *, mutates_args, device_types=None, schema=None): - def wrap(func): - return func - if fn is None: - return wrap - return fn - def noop_register_fake_wrapper(op, fn=None, /, *, lib=None, _stacklevel=1): - def wrap(func): - return func - if fn is None: - return wrap - return fn - _torch_custom_op_wrapper = noop_custom_op_wrapper - _torch_register_fake_wrapper = noop_register_fake_wrapper - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_forward"), mutates_args=(), device_types=_get_device()) -def _flash_attn_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - return_softmax: bool -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - out, softmax_lse, S_dmask, rng_state = flash_attn.fwd( - q, - k, - v, - None, - alibi_slopes, - dropout_p, - softmax_scale, - causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - None, - ) - return out, softmax_lse, S_dmask, rng_state - - -@_torch_register_fake_wrapper(add_op_namespace_prefix("_flash_attn_forward")) -def _flash_attn_forward_fake( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - return_softmax: bool -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - batch_size, seqlen_q, num_heads, head_size = q.shape - seqlen_k = k.shape[1] - out = torch.empty_like(q) - softmax_lse = torch.empty((batch_size, num_heads, seqlen_q), dtype=torch.float32, device=q.device, layout=q.layout) - p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) - if return_softmax: - p = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128), round_multiple(seqlen_k, 128)), dtype=q.dtype, device=q.device, layout=q.layout) - rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) - - return out, softmax_lse, p, rng_state - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_forward = _ops.ops._flash_attn_forward -else: - _wrapped_flash_attn_forward = _flash_attn_forward - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_varlen_forward"), mutates_args=(), device_types=_get_device()) -def _flash_attn_varlen_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - alibi_slopes: Optional[torch.Tensor] = None, - return_softmax: bool = False, - block_table: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - out, softmax_lse, S_dmask, rng_state = flash_attn.varlen_fwd( - q, - k, - v, - None, - cu_seqlens_q, - cu_seqlens_k, - seqused_k, - leftpad_k, - block_table, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - zero_tensors, - causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - None, - ) - # if out.isnan().any() or softmax_lse.isnan().any(): - # breakpoint() - return out, softmax_lse, S_dmask, rng_state - - -@_torch_register_fake_wrapper(add_op_namespace_prefix("_flash_attn_varlen_forward")) -def _flash_attn_varlen_forward_fake( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - alibi_slopes: Optional[torch.Tensor] = None, - return_softmax: bool = False, - block_table: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - paged_kv = block_table is not None - batch_size = cu_seqlens_q.numel() - 1 - total_q, num_heads, _ = q.shape - - out = torch.empty_like(q) - softmax_lse = torch.empty((num_heads, total_q), dtype=torch.float32, device=q.device, layout=q.layout) - p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) - seqlen_q_rounded = round_multiple(max_seqlen_q, 128) - seqlen_k_rounded = round_multiple(max_seqlen_k, 128) - if return_softmax: - p = torch.empty((batch_size, num_heads, seqlen_q_rounded, seqlen_k_rounded), dtype=q.dtype, device=q.device, layout=q.layout) - rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) - return out, softmax_lse, p, rng_state - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_varlen_forward = _ops.ops._flash_attn_varlen_forward -else: - _wrapped_flash_attn_varlen_forward = _flash_attn_varlen_forward - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_backward"), mutates_args=("dq", "dk", "dv"), device_types=_get_device()) -def _flash_attn_backward( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, -) -> 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)] - ( - dq, - dk, - dv, - softmax_d, - ) = flash_attn.bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - alibi_slopes, - dropout_p, - softmax_scale, - causal, - window_size_left, - window_size_right, - softcap, - deterministic, - None, - rng_state, - ) - return softmax_d - - -@_torch_register_fake_wrapper(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, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, -) -> torch.Tensor: - dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] - if dq is None: - dq = torch.empty_like(q) - if dk is None: - dk = torch.empty_like(k) - if dv is None: - dv = torch.empty_like(v) - batch_size, seqlen_q, num_heads, _ = q.shape - softmax_d = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128)), device=q.device, dtype=torch.float32) - - return softmax_d - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_backward = _ops.ops._flash_attn_backward -else: - _wrapped_flash_attn_backward = _flash_attn_backward - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_varlen_backward"), mutates_args=("dq", "dk", "dv"), device_types=_get_device()) -def _flash_attn_varlen_backward( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> 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)] - ( - dq, - dk, - dv, - softmax_d, - ) = flash_attn.varlen_bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - zero_tensors, - causal, - window_size_left, - window_size_right, - softcap, - deterministic, - None, - rng_state, - ) - # if dk.isnan().any() or dk.isnan().any() or dv.isnan().any() or softmax_d.isnan().any(): - # breakpoint() - return softmax_d - - -@_torch_register_fake_wrapper(add_op_namespace_prefix("_flash_attn_varlen_backward")) -def _flash_attn_varlen_backward_fake( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> torch.Tensor: - dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] - batch_size = cu_seqlens_q.numel() - 1 - total_q, num_heads, _ = q.shape - - if dq is None: - dq = torch.empty_like(q) - if dk is None: - dk = torch.empty_like(k) - if dv is None: - dv = torch.empty_like(v) - softmax_d = torch.empty((num_heads, total_q + 128 * batch_size), device=q.device, dtype=torch.float32) - - return softmax_d - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_varlen_backward = _ops.ops._flash_attn_varlen_backward -else: - _wrapped_flash_attn_varlen_backward = _flash_attn_varlen_backward - - -class FlashAttnQKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and qkv.requires_grad - if softmax_scale is None: - softmax_scale = qkv.shape[-1] ** (-0.5) - q, k, v = qkv[:, :, 0].detach(), qkv[:, :, 1].detach(), qkv[:, :, 2].detach() - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) - dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dqkv[:, :, 0], - dqkv[:, :, 1], - dqkv[:, :, 2], - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dqkv = dqkv[..., : dout.shape[-1]] # We could have padded the head dimension - return dqkv, None, None, None, None, None, None, None, None, None - - -class FlashAttnVarlenQKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - cu_seqlens, - max_seqlen, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and qkv.requires_grad - if softmax_scale is None: - softmax_scale = qkv.shape[-1] ** (-0.5) - q, k, v = qkv[:, 0].detach(), qkv[:, 1].detach(), qkv[:, 2].detach() - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens, - cu_seqlens, - max_seqlen, - max_seqlen, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=None, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens, rng_state) - ctx.dropout_p = dropout_p - ctx.max_seqlen = max_seqlen - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens, rng_state = ctx.saved_tensors - qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) - dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dqkv[:, 0], - dqkv[:, 1], - dqkv[:, 2], - cu_seqlens, - cu_seqlens, - ctx.max_seqlen, - ctx.max_seqlen, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - 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 FlashAttnKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - kv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, kv] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - k, v = kv[:, :, 0].detach(), kv[:, :, 1].detach() - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - dq = torch.empty_like(q) - kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) - dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dkv[:, :, 0], - dkv[:, :, 1], - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dkv = dkv[..., : dout.shape[-1]] - return dq, dkv, None, None, None, None, None, None, None, None, None - - -class FlashAttnVarlenKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, kv] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - k, v = kv[:, 0].detach(), kv[:, 1].detach() - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=None, - ) - if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state - ) - ctx.dropout_p = dropout_p - 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.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors - dq = torch.empty_like(q) - kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) - dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dkv[:, 0], - dkv[:, 1], - cu_seqlens_q, - cu_seqlens_k, - ctx.max_seqlen_q, - ctx.max_seqlen_k, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dkv = dkv[..., : dout.shape[-1]] - return dq, dkv, None, None, None, None, None, None, None, None, None, None, None, None, None - - -class FlashAttnFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - k, - v, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, k, v] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dk = dk[..., : dout.shape[-1]] - dv = dv[..., : dout.shape[-1]] - return dq, dk, dv, 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, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - block_table, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, k, v] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=block_table, - ) - if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state - ) - ctx.dropout_p = dropout_p - 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.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors - dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - ctx.max_seqlen_q, - ctx.max_seqlen_k, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dk = dk[..., : dout.shape[-1]] - dv = dv[..., : dout.shape[-1]] - return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None - - -def flash_attn_qkvpacked_func( - qkv, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # <=0.0 means deactivate - alibi_slopes=None, - deterministic=False, - 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, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_kvpacked_func( - q, - kv, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, -): - """dropout_p should be set to 0.0 during evaluation - If 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 K, V. - 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) - kv: (batch_size, seqlen, 2, 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. - softcap: float. Anything > 0 activates softcapping 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). - 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 FlashAttnKVPackedFunc.apply( - q, - kv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_func( - q, - k, - v, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - 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). - 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 FlashAttnFunc.apply( - q, - k, - v, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_qkvpacked_func( - qkv, - cu_seqlens, - max_seqlen, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - 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_varlen_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_varlen_kvpacked_func and flash_attn_varlen_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: (total, 3, nheads, headdim), where total = total number of tokens in the batch. - cu_seqlens: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into qkv. - max_seqlen: int. Maximum sequence length in the batch. - 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenQKVPackedFunc.apply( - qkv, - cu_seqlens, - max_seqlen, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_kvpacked_func( - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, -): - """dropout_p should be set to 0.0 during evaluation - If 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 K, V. - 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: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch. - kv: (total_k, 2, nheads_k, headdim), where total_k = total number of key tokens in the batch. - cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into q. - cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into kv. - max_seqlen_q: int. Maximum query sequence length in the batch. - max_seqlen_k: int. Maximum key sequence length in the batch. - 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 + 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenKVPackedFunc.apply( - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_func( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, - block_table=None, -): - """dropout_p should be set to 0.0 during evaluation - Supports multi-query and grouped-query attention (MQA/GQA) by passing in K, V 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: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch. - k: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. - v: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. - cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into q. - cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into kv. - max_seqlen_q: int. Maximum query sequence length in the batch. - max_seqlen_k: int. Maximum key sequence length in the batch. - 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 + 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenFunc.apply( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - block_table, - False if q.device.type == "cpu" else torch.is_grad_enabled(), - ) - - -def flash_attn_with_kvcache( - q, - k_cache, - v_cache, - k=None, - v=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, - block_table: Optional[torch.Tensor] = None, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - rotary_interleaved=True, - alibi_slopes=None, - num_splits=0, - 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 block_table, - or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_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) if there's no block_table, - or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_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). Similar to k. - 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. - block_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). - 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. - 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" - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-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) - cache_batch_idx = maybe_contiguous(cache_batch_idx) - block_table = maybe_contiguous(block_table) - out, softmax_lse = flash_attn.fwd_kvcache( - q, - k_cache, - v_cache, - k, - v, - cache_seqlens, - rotary_cos, - rotary_sin, - cache_batch_idx, - cache_leftpad, - block_table, - alibi_slopes, - None, - softmax_scale, - causal, - window_size[0], - window_size[1], - softcap, - rotary_interleaved, - num_splits, - ) - return (out, softmax_lse) if return_softmax_lse else out diff --git a/build/torch212-cxx11-cu130-x86_64-linux/layers/__init__.py b/build/torch212-cxx11-cu130-x86_64-linux/layers/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/build/torch212-cxx11-cu130-x86_64-linux/layers/patch_embed.py b/build/torch212-cxx11-cu130-x86_64-linux/layers/patch_embed.py deleted file mode 100644 index 05562f8e8bcdb58e947c6f402a49eacd2d031871..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu130-x86_64-linux/layers/patch_embed.py +++ /dev/null @@ -1,67 +0,0 @@ -# We use the same API as https://github.com/rwightman/pytorch-image-models/blob/v0.6.11/timm/models/layers/patch_embed.py -# But we use nn.Linear instead of Conv2d and it's about 8x faster. - -from functools import partial - -import torch.nn as nn -from einops import rearrange -from torch import _assert -from torch.nn.modules.utils import _pair - -try: - from flash_attn.ops.fused_dense import FusedDense -except ImportError: - FusedDense = None - - -class PatchEmbed(nn.Module): - """2D Image to Patch Embedding""" - - def __init__( - self, - img_size=224, - patch_size=16, - in_chans=3, - embed_dim=768, - norm_layer=None, - flatten=True, - bias=True, - fused_bias_fc=False, - ): - super().__init__() - img_size = _pair(img_size) - patch_size = _pair(patch_size) - self.img_size = img_size - self.patch_size = patch_size - self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1]) - self.num_patches = self.grid_size[0] * self.grid_size[1] - self.flatten = flatten - if fused_bias_fc and FusedDense is None: - raise ImportError("fused_dense is not installed") - - linear_cls = nn.Linear if not fused_bias_fc or not bias else FusedDense - self.proj = linear_cls(in_chans * patch_size[0] * patch_size[1], embed_dim, bias=bias) - self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() - - def forward(self, x): - _, _, H, W = x.shape - _assert( - H == self.img_size[0], - f"Input image height ({H}) doesn't match model ({self.img_size[0]}).", - ) - _assert( - W == self.img_size[1], - f"Input image width ({W}) doesn't match model ({self.img_size[1]}).", - ) - x = self.proj( - rearrange( - x, - "b c (h p1) (w p2) -> b h w (c p1 p2)", - p1=self.patch_size[0], - p2=self.patch_size[1], - ) - ) - if self.flatten: - x = rearrange(x, "b h w c -> b (h w) c") - x = self.norm(x) - return x diff --git a/build/torch212-cxx11-cu130-x86_64-linux/layers/rotary.py b/build/torch212-cxx11-cu130-x86_64-linux/layers/rotary.py deleted file mode 100644 index d1bfc21fc7de1dd287e8f382847b194a48075981..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu130-x86_64-linux/layers/rotary.py +++ /dev/null @@ -1,483 +0,0 @@ -# Copyright (c) 2025, Tri Dao - -import math -from functools import partial -from typing import Optional, Tuple, Union - -import torch -from torch import Tensor - -from einops import rearrange, repeat -# from flash_attn.ops.triton.rotary import apply_rotary -from ..ops.triton.rotary import apply_rotary - - -def rotate_half(x, interleaved=False): - if not interleaved: - x1, x2 = x.chunk(2, dim=-1) - return torch.cat((-x2, x1), dim=-1) - else: - x1, x2 = x[..., ::2], x[..., 1::2] - return rearrange(torch.stack((-x2, x1), dim=-1), "... d two -> ... (d two)", two=2) - - -def apply_rotary_emb_torch(x, cos, sin, interleaved=False): - """ - x: (batch_size, seqlen, nheads, headdim) - cos, sin: (seqlen, rotary_dim / 2) or (batch_size, seqlen, rotary_dim / 2) - """ - ro_dim = cos.shape[-1] * 2 - assert ro_dim <= x.shape[-1] - cos = repeat(cos, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)") - sin = repeat(sin, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)") - return torch.cat( - [x[..., :ro_dim] * cos + rotate_half(x[..., :ro_dim], interleaved) * sin, x[..., ro_dim:]], - dim=-1, - ) - - -class ApplyRotaryEmb(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x, - cos, - sin, - interleaved=False, - inplace=False, - seqlen_offsets: Union[int, Tensor] = 0, - cu_seqlens: Optional[Tensor] = None, - max_seqlen: Optional[int] = None, - ): - out = apply_rotary( - x, - cos, - sin, - seqlen_offsets=seqlen_offsets, - cu_seqlens=cu_seqlens, - max_seqlen=max_seqlen, - interleaved=interleaved, - inplace=inplace, - ) - if isinstance(seqlen_offsets, int): - ctx.save_for_backward(cos, sin, cu_seqlens) # Can't save int with save_for_backward - ctx.seqlen_offsets = seqlen_offsets - else: - ctx.save_for_backward(cos, sin, cu_seqlens, seqlen_offsets) - ctx.seqlen_offsets = None - ctx.interleaved = interleaved - ctx.inplace = inplace - ctx.max_seqlen = max_seqlen - return out if not inplace else x - - @staticmethod - def backward(ctx, do): - seqlen_offsets = ctx.seqlen_offsets - if seqlen_offsets is None: - cos, sin, cu_seqlens, seqlen_offsets = ctx.saved_tensors - else: - cos, sin, cu_seqlens = ctx.saved_tensors - dx = apply_rotary( - do, - cos, - sin, - seqlen_offsets=seqlen_offsets, - cu_seqlens=cu_seqlens, - max_seqlen=ctx.max_seqlen, - interleaved=ctx.interleaved, - inplace=ctx.inplace, - conjugate=True, - ) - return dx, None, None, None, None, None, None, None - - -def apply_rotary_emb( - x, - cos, - sin, - interleaved=False, - inplace=False, - seqlen_offsets: Union[int, Tensor] = 0, - cu_seqlens: Optional[Tensor] = None, - max_seqlen: Optional[int] = None, -): - """ - Arguments: - x: (batch_size, seqlen, nheads, headdim) if cu_seqlens is None - else (total_seqlen, nheads, headdim) - cos, sin: (seqlen_rotary, rotary_dim / 2) - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead - of 1st half and 2nd half (GPT-NeoX style). - inplace: if True, apply rotary embedding in-place. - seqlen_offsets: (batch_size,) or int. Each sequence in x is shifted by this amount. - Most commonly used in inference when we have KV cache. - cu_seqlens: (batch + 1,) or None - max_seqlen: int - Return: - out: (batch_size, seqlen, nheads, headdim) if cu_seqlens is None - else (total_seqlen, nheads, headdim) - rotary_dim must be <= headdim - Apply rotary embedding to the first rotary_dim of x. - """ - return ApplyRotaryEmb.apply( - x, cos, sin, interleaved, inplace, seqlen_offsets, cu_seqlens, max_seqlen - ) - - -# For backward compatibility -apply_rotary_emb_func = apply_rotary_emb - - -def _apply_rotary_emb_qkv( - qkv, - cos, - sin, - cos_k=None, - sin_k=None, - interleaved=False, - inplace=False, - conjugate=False, - seqlen_offsets: Union[int, Tensor] = 0, - num_heads_q: Optional[int] = None, -): - apply_rotary_fn = partial( - apply_rotary, - interleaved=interleaved, - inplace=inplace, - conjugate=conjugate, - seqlen_offsets=seqlen_offsets - ) - if cos_k is None and sin_k is None and qkv.is_contiguous(): - # Call 1 kernel instead of 2 kernels - # We need qkv to be contiguous so that when we reshape to combine (3, nheads) - # dimensions, we get the same tensor - if qkv.dim() == 5: - batch, seqlen, three, nheads, headdim = qkv.shape - assert three == 3 - # qk = rearrange(qkv[:, :, :2], "b s t h d -> b s (t h) d") - qk = qkv[:, :, :2].reshape(batch, seqlen, -1, headdim) - qk = apply_rotary_fn(qk, cos, sin) - else: - assert qkv.dim() == 4 - assert num_heads_q is not None - num_heads_k = (qkv.shape[2] - num_heads_q) // 2 - assert qkv.shape[2] == num_heads_q + 2 * num_heads_k - qk = qkv[:, :, :num_heads_q + num_heads_k] - qk = apply_rotary_fn(qk, cos, sin) - if not inplace: - if qkv.dim() == 5: - qkv = torch.cat([rearrange(qk, "b s (t h) d -> b s t h d", t=2), qkv[:, :, 2:]], dim=2) - else: - qkv = torch.cat([qk, qkv[:, :, num_heads_q + num_heads_k :]], dim=2) - else: - cos_k = cos if cos_k is None else cos_k - sin_k = sin if sin_k is None else sin_k - if qkv.dim() == 5: - batch, seqlen, three, nheads, headdim = qkv.shape - assert three == 3 - q, k = qkv[:, :, 0], qkv[:, :, 1] - else: - assert qkv.dim() == 4 - assert num_heads_q is not None - num_heads_k = (qkv.shape[2] - num_heads_q) // 2 - assert qkv.shape[2] == num_heads_q + 2 * num_heads_k - q, k = qkv[:, :, :num_heads_q], qkv[:, :, num_heads_q : num_heads_q + num_heads_k] - q = apply_rotary_fn(q, cos, sin) - k = apply_rotary_fn(k, cos_k, sin_k) - if not inplace: - if qkv.dim() == 5: - qkv = torch.stack([q, k, qkv[:, :, 2]], dim=2) - else: - qkv = torch.cat([q, k, qkv[:, :, num_heads_q + num_heads_k:]], dim=2) - return qkv - - -class ApplyRotaryEmbQKV_(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - cos, - sin, - cos_k=None, - sin_k=None, - interleaved=False, - seqlen_offsets: Union[int, torch.Tensor] = 0, - num_heads_q: Optional[int] = None, - ): - # apply_rotary_emb_qkv_inplace( - qkv = _apply_rotary_emb_qkv( - qkv, cos, sin, cos_k, sin_k, interleaved=interleaved, inplace=True, - seqlen_offsets=seqlen_offsets, num_heads_q=num_heads_q, - ) - if isinstance(seqlen_offsets, int): - ctx.save_for_backward(cos, sin, cos_k, sin_k) - ctx.seqlen_offsets = seqlen_offsets - else: - ctx.save_for_backward(cos, sin, cos_k, sin_k, seqlen_offsets) - ctx.seqlen_offsets = None - ctx.interleaved = interleaved - ctx.num_heads_q = num_heads_q - return qkv - - @staticmethod - def backward(ctx, dqkv): - seqlen_offsets = ctx.seqlen_offsets - if seqlen_offsets is None: - cos, sin, cos_k, sin_k, seqlen_offsets = ctx.saved_tensors - else: - cos, sin, cos_k, sin_k = ctx.saved_tensors - dqkv = _apply_rotary_emb_qkv( - dqkv, cos, sin, cos_k, sin_k, interleaved=ctx.interleaved, inplace=True, - seqlen_offsets=seqlen_offsets, num_heads_q=ctx.num_heads_q, conjugate=True, - ) - return dqkv, None, None, None, None, None, None, None - - -def apply_rotary_emb_qkv_( - qkv, - cos, - sin, - cos_k=None, - sin_k=None, - interleaved=False, - seqlen_offsets: Union[int, torch.Tensor] = 0, - num_heads_q: Optional[int] = None, -): - """ - Arguments: - qkv: (batch_size, seqlen, 3, nheads, headdim) or (batch_size, seqlen, num_heads_q + 2 * num_heads_k, headdim). - If qkv has shape (batch_size, seqlen, num_heads_q + 2 * num_heads_k, headdim) (e.g. MQA / GQA), - then num_heads_q must be provided. - cos, sin: (seqlen, rotary_dim / 2) - cos_k, sin_k: (seqlen, rotary_dim / 2), optional - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead of - 1st half and 2nd half (GPT-NeoX style). - seqlen_offsets: (batch_size,) or int. Each sequence in Q and K is shifted by this amount. - Most commonly used in inference when we have KV cache. - Return: - qkv: (batch_size, seqlen, 3, nheads, headdim) or (batch_size, seqlen, num_heads_q + 2 * num_heads_k, headdim) - rotary_dim must be <= headdim - Apply rotary embedding *inplace* to the first rotary_dim of Q and K. - """ - return ApplyRotaryEmbQKV_.apply( - qkv, cos, sin, cos_k, sin_k, interleaved, seqlen_offsets, num_heads_q - ) - - -class ApplyRotaryEmbKV_(torch.autograd.Function): - - @staticmethod - def forward(ctx, kv, cos, sin, interleaved=False, seqlen_offsets: Union[int, torch.Tensor] = 0): - batch, seqlen, two, nheads, headdim = kv.shape - assert two == 2 - k = kv[:, :, 0] - apply_rotary( - k, cos, sin, seqlen_offsets=seqlen_offsets, interleaved=interleaved, inplace=True - ) - if isinstance(seqlen_offsets, int): - ctx.save_for_backward(cos, sin) # Can't save int with save_for_backward - ctx.seqlen_offsets = seqlen_offsets - else: - ctx.save_for_backward(cos, sin, seqlen_offsets) - ctx.seqlen_offsets = None - ctx.interleaved = interleaved - return kv - - @staticmethod - def backward(ctx, dkv): - seqlen_offsets = ctx.seqlen_offsets - if seqlen_offsets is None: - cos, sin, seqlen_offsets = ctx.saved_tensors - else: - cos, sin = ctx.saved_tensors - apply_rotary( - dkv[:, :, 0], - cos, - sin, - seqlen_offsets=seqlen_offsets, - interleaved=ctx.interleaved, - inplace=True, - conjugate=True, - ) - return dkv, None, None, None, None - - -apply_rotary_emb_kv_ = ApplyRotaryEmbKV_.apply - - -def apply_rotary_emb_kv_( - kv, - cos, - sin, - interleaved=False, - seqlen_offsets: Union[int, torch.Tensor] = 0, -): - """ - Arguments: - kv: (batch_size, seqlen, 2, nheads, headdim) - cos, sin: (seqlen, rotary_dim / 2) - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead of - 1st half and 2nd half (GPT-NeoX style). - seqlen_offsets: (batch_size,) or int. Each sequence in Q and K is shifted by this amount. - Most commonly used in inference when we have KV cache. - Return: - kv: (batch_size, seqlen, 2, nheads, headdim) - rotary_dim must be <= headdim - Apply rotary embedding *inplace* to the first rotary_dim of K. - """ - return ApplyRotaryEmbKV_.apply(kv, cos, sin, interleaved, seqlen_offsets) - - -class RotaryEmbedding(torch.nn.Module): - """ - The rotary position embeddings from RoFormer_ (Su et. al). - A crucial insight from the method is that the query and keys are - transformed by rotation matrices which depend on the relative positions. - - Other implementations are available in the Rotary Transformer repo_ and in - GPT-NeoX_, GPT-NeoX was an inspiration - - .. _RoFormer: https://arxiv.org/abs/2104.09864 - .. _repo: https://github.com/ZhuiyiTechnology/roformer - .. _GPT-NeoX: https://github.com/EleutherAI/gpt-neox - - If scale_base is not None, this implements XPos (Sun et al., https://arxiv.org/abs/2212.10554). - A recommended value for scale_base is 512: https://github.com/HazyResearch/flash-attention/issues/96 - Reference: https://github.com/sunyt32/torchscale/blob/main/torchscale/component/xpos_relative_position.py - """ - - def __init__( - self, - dim: int, - base=10000.0, - interleaved=False, - scale_base=None, - device=None, - ): - """ - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead - of 1st half and 2nd half (GPT-NeoX style). - """ - super().__init__() - self.dim = dim - self.base = float(base) - # Generate and save the inverse frequency buffer (non trainable) - inv_freq = self._compute_inv_freq(device) - self.register_buffer("inv_freq", inv_freq, persistent=False) - self.interleaved = interleaved - self.scale_base = scale_base - scale = ( - (torch.arange(0, dim, 2, device=device, dtype=torch.float32) + 0.4 * dim) / (1.4 * dim) - if scale_base is not None - else None - ) - self.register_buffer("scale", scale, persistent=False) - - self._seq_len_cached = 0 - self._cos_cached = None - self._sin_cached = None - self._cos_k_cached = None - self._sin_k_cached = None - - def _compute_inv_freq(self, device=None): - return 1.0 / ( - self.base - ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim) - ) - - def _update_cos_sin_cache(self, seqlen, device=None, dtype=None): - # Reset the tables if the sequence length has changed, - # if we're on a new device (possibly due to tracing for instance), - # or if we're switching from inference mode to training - if ( - seqlen > self._seq_len_cached - or self._cos_cached is None - or self._cos_cached.device != device - or self._cos_cached.dtype != dtype - or (self.training and self._cos_cached.is_inference()) - ): - self._seq_len_cached = seqlen - # We want fp32 here, not self.inv_freq.dtype, since the model could be loaded in bf16 - # And the output of arange can be quite large, so bf16 would lose a lot of precision. - t = torch.arange(seqlen, device=device, dtype=torch.float32) - # We want fp32 here as well since inv_freq will be multiplied with t, and the output - # will be large. Having it in bf16 will lose a lot of precision and cause the - # cos & sin output to change significantly. - # We want to recompute self.inv_freq if it was not loaded in fp32 - if self.inv_freq.dtype != torch.float32: - inv_freq = self._compute_inv_freq(device=device) - else: - inv_freq = self.inv_freq - # Don't do einsum, it converts fp32 to bf16 under AMP - # freqs = torch.einsum("i,j->ij", t, self.inv_freq) - freqs = torch.outer(t, inv_freq) - if self.scale is None: - self._cos_cached = torch.cos(freqs).to(dtype) - self._sin_cached = torch.sin(freqs).to(dtype) - else: - power = ( - torch.arange(seqlen, dtype=self.scale.dtype, device=self.scale.device) - - seqlen // 2 - ) / self.scale_base - scale = self.scale.to(device=power.device) ** rearrange(power, "s -> s 1") - # We want the multiplication by scale to happen in fp32 - self._cos_cached = (torch.cos(freqs) * scale).to(dtype) - self._sin_cached = (torch.sin(freqs) * scale).to(dtype) - self._cos_k_cached = (torch.cos(freqs) / scale).to(dtype) - self._sin_k_cached = (torch.sin(freqs) / scale).to(dtype) - - def forward( - self, - qkv: torch.Tensor, - kv: Optional[torch.Tensor] = None, - seqlen_offset: Union[int, torch.Tensor] = 0, - max_seqlen: Optional[int] = None, - num_heads_q: Optional[int] = None, - ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: - """ - qkv: (batch, seqlen, 3, nheads, headdim) or (batch, seqlen, num_heads_q + 2 * num_heads_k, headdim) - if kv is none, else it's just q of shape (batch, seqlen, nheads, headdim). - If qkv has shape (batch, seqlen, num_heads_q + 2 * num_heads_k, headdim) (e.g. MQA / GQA), - then num_heads_q must be provided. - kv: (batch, seqlen, 2, nheads, headdim) - seqlen_offset: (batch_size,) or int. Each sequence in x is shifted by this amount. - Most commonly used in inference when we have KV cache. - If it's a tensor of shape (batch_size,), then to update the cos / sin cache, one - should pass in max_seqlen, which will update the cos / sin cache up to that length. - Apply rotary embedding *inplace* to qkv and / or kv. - """ - seqlen = qkv.shape[1] - if max_seqlen is not None: - self._update_cos_sin_cache(max_seqlen, device=qkv.device, dtype=qkv.dtype) - elif isinstance(seqlen_offset, int): - self._update_cos_sin_cache(seqlen + seqlen_offset, device=qkv.device, dtype=qkv.dtype) - if kv is None: - return apply_rotary_emb_qkv_( - qkv, - self._cos_cached, - self._sin_cached, - self._cos_k_cached if self.scale is not None else None, - self._sin_k_cached if self.scale is not None else None, - interleaved=self.interleaved, - seqlen_offsets=seqlen_offset, - num_heads_q=num_heads_q, - ) - else: - q = qkv - q = apply_rotary_emb_func( - q, - self._cos_cached, - self._sin_cached, - interleaved=self.interleaved, - inplace=True, - seqlen_offsets=seqlen_offset, - ) - kv = apply_rotary_emb_kv_( - kv, - self._cos_cached if self.scale is None else self._cos_k_cached, - self._sin_cached if self.scale is None else self._sin_k_cached, - interleaved=self.interleaved, - seqlen_offsets=seqlen_offset, - ) - return q, kv diff --git a/build/torch212-cxx11-cu130-x86_64-linux/metadata.json b/build/torch212-cxx11-cu130-x86_64-linux/metadata.json deleted file mode 100644 index ccee10f045c9d05828407a033909dbe8803b3d77..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu130-x86_64-linux/metadata.json +++ /dev/null @@ -1,42 +0,0 @@ -{ - "name": "flash-attn2", - "id": "_flash_attn2_cuda_f12afc9", - "version": 1, - "license": "BSD-3-Clause", - "python-depends": [], - "backend": { - "type": "cuda", - "archs": [ - "10.0", - "12.0", - "8.0", - "9.0" - ] - }, - "digest": { - "algorithm": "sha256", - "files": { - "__init__.py": "Cd9MEdd/WUYnBmfiQpC4SyPqAxrTilTjvsvswNju4zc=", - "_flash_attn2_cuda_f12afc9.abi3.so": "w2RBfYW85Gn+bY4EQZG0LfGoPiabYvcYEQsHdVH8LS0=", - "_ops.py": "GQbmTTeSOU+7Q2nI68Hz+1RvEogyfEFccx1ubUYw+2s=", - "bert_padding.py": "gF1EmsdJ+HpQ86MRQ4VxDw+Sb/RVISdQALdNnoByHlw=", - "flash_attn2/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY=", - "flash_attn_interface.py": "W7a9bHR2ly6QL3aKqsBqxWA7M7ZWYcZPmtsPJu6rAfY=", - "layers/__init__.py": "47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU=", - "layers/patch_embed.py": "H58CgME/qSOPTZLOG08wFgrQS1j34pvNwMPrkTj3Ek4=", - "layers/rotary.py": "rFIdNrZQFNawwLD6d7t5LE2m+M93wwDZ8EuExw+gfD8=", - "ops/__init__.py": "47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU=", - "ops/activations.py": "t5lzNg1In8LP6bKeTnyeMizwqjv27JGbJ6ylPdGvZYg=", - "ops/fused_dense.py": "BqT8VTSCGdLJAjh6WBMkCRFYrdVodPeoElWBvH6AIzA=", - "ops/layer_norm.py": "zr7NXIm+2mtEynTp1CS0fbFGI2Mqdp41dY4AfDWF6EQ=", - "ops/rms_norm.py": "XEnihcj0a4aSz4LO55m5iKGVn4HKTeKN8TIyHjuDgxI=", - "ops/triton/__init__.py": "AbpHGcgLb+kRsJGnwFEktk7uzpZOCcBY74+YBdrKVGs=", - "ops/triton/cross_entropy.py": "y56nmuOtpSNxM7sDZcUSgXKIsIEayfuvrU8bAVog/nw=", - "ops/triton/k_activations.py": "+Z3vIyO4JkqBMipKsPvhzmxljtBdIhJCsl/M+/ESqBo=", - "ops/triton/layer_norm.py": "BBQdJiqopbIS12viUqzx3SvOQCNbDSAYLJXhpyMGyC8=", - "ops/triton/linear.py": "OtRvKz8xdpl+7v3q/ZTaS9fdBt9XrzMyapgRr50uBbM=", - "ops/triton/mlp.py": "uWA+D7HyuRUsOT+4nhi3vvCDNZnNEXwYcHMQ02UREfQ=", - "ops/triton/rotary.py": "0YSrlbep0lJmzLpTiANs0slTFeGcQQqR+el6uzbSLvI=" - } - } -} \ No newline at end of file diff --git a/build/torch212-cxx11-cu130-x86_64-linux/metadata.json.sigstore b/build/torch212-cxx11-cu130-x86_64-linux/metadata.json.sigstore deleted file mode 100644 index 9652830fdb8855977edac2c7fffbf77cf093761e..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu130-x86_64-linux/metadata.json.sigstore +++ /dev/null @@ -1 +0,0 @@ -{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json","verificationMaterial":{"certificate":{"rawBytes":"MIIHSTCCBtCgAwIBAgIUFDByNopW2mWM/XlH8WR535HA+WYwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwNjI0MDEwMTI0WhcNMjYwNjI0MDExMTI0WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEHx9Td9aqj40xx9bojpbhxKheY/ic/ts0iVKQOiun3fbr6xBLhGzYBUjLON7Gn6ValTivkDX4Fzjq/4Chx54B1aOCBe8wggXrMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUbnAdN//JkWSasDUtAQRVVGS470swHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoZjEyYWZjOTJjNDBiYmQ3NmQ5OTBkOTExZWJhYmYyOGU0OTBhN2JkZTATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoZjEyYWZjOTJjNDBiYmQ3NmQ5OTBkOTExZWJhYmYyOGU0OTBhN2JkZTAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoZjEyYWZjOTJjNDBiYmQ3NmQ5OTBkOTExZWJhYmYyOGU0OTBhN2JkZTAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKGYxMmFmYzkyYzQwYmJkNzZkOTkwZDkxMWViYWJmMjhlNDkwYTdiZGUwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMjgwMzYyNjg5MzUvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBiQYKKwYBBAHWeQIEAgR7BHkAdwB1AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABnvclZ10AAAQDAEYwRAIgavhHVv4Yr5JQpaECVgkElYUNjitRJzbpA+wsDRp04igCIATz+88+JMmwsPMOwAWAzE8+os2yrle+snJdEfnau7DnMAoGCCqGSM49BAMDA2cAMGQCMF7LhtAhU6gGQaANzI+G15pjV0ho4Evz1mmDPhGHGVuAaH+TqQtW7+DY07+oLWl/XgIwXnnXUVErb1DOGCWhXcVV5kS4TLbC8olR4acgHEl5yDkKg+/fMitWFgxcTfoTK0Su"},"tlogEntries":[{"logIndex":"1933772581","logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="},"kindVersion":{"kind":"hashedrekord","version":"0.0.1"},"integratedTime":"1782262884","inclusionPromise":{"signedEntryTimestamp":"MEUCIQDhUSktpg+YVulMd6nNiQ+mVMFyVk1+9FXBsjqcxHOz0wIgPf6a/7D1SMIiLTt49xvdPTZDnYBHfyhW0IueymFIQ0E="},"inclusionProof":{"logIndex":"1811868319","rootHash":"ADX2RGp0bl3/wnENEwStGD5LkGXrc2t8MMql70gYR+I=","treeSize":"1811868404","hashes":["t/sH6MoIc2eujbXl+p8UYmAsN7PA/b9WpGeRdYwJKh0=","AOBaqiQTY/N1Yk0j3XznYD60dlG0SMlyUIHtuhiwoDU=","9BCwNbHLjP5PjLTvUblkgH3cC+ykUUOxAm4iRHTtM6c=","263UvYN90pHlLVmxTGAgmtnkGY02kD2DfKwyQVf9sS0=","VV3fxED1GRui4b7EksyBLJGBRwnJW3NT1gx7fSBz50s=","xe5j3w7Cg58hpGqx6Vna0h+zKPl9Hr542q1ZR4ymnKY=","PhAYQupUy5fc0E30eQa/SFETkYgUoCFZ20LF5d0fgKM=","scsGY3fU98OnnafnUlncq7QETC3+OJNgtRzvA/QwwEc=","eEf0n4aXyn0P4Z/HNLjdj0hqsSjjouR8dwugBBfshiA=","fqtytk2if0R+M3kWRpiPUrBEhf2z8y5Li6MeOuWmiks=","mzp9m+2At7GZdgfx0m7RwtO9YzltCZ6UhZBy4p+lUdM=","ipWWLDkqay5JAYe6P4Nrqsck/1C0C5uwqT9QbFsiY1E=","COKvK4TO+dkpqUDGbTbSKewajXzAonx1B1YaRb/ymLw=","dQlR+DXGDKXm5V1fupxaMSOwNZ9F1MybXyyxB96yj80=","JKJdxSZwKPO6QFzLTxHCrUyI5sDU/CwY57zDlA9ni6Q=","OgfnQCh7xr8xUTErIw//YAKaKMDwRJWiThvYExQvLJM=","S8Oe9ZicoJqbVuL1I0dIdTgVpLc82Cm2CMOcV83sKVs=","E1O2ilWlpFRLWGdbmJUMDZ2qbGu/eDZ12m6YjQ9/JPA=","nyv2TDlfkz0XfRwjY98zsbUbOkBZ+Tgz7PML0c2uN14=","OMfDh4k9YD0E+t4sjqJsh+vg8D1ueo1cC07exyxmr6s=","+vHKm4wNwGJ46ceyZ8mePD7Vn1SediStNHCStRoTSWQ=","5HwWjuGNDisetbiWM0cWIXM65TIM3EJlqk13ntCQlcU=","NDIRlqxO0IaGr94q86MIU7VCAN3yEEgOHS4loya5lu4=","daxmZaajRpZV+JxHiOYZhJBiSKN5ucqjh2WnGbHhirw=","DOCeoSMovIvLExkhIvisow9AuNXgeWs4ECkyR6EcqYU="],"checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n1811868404\nADX2RGp0bl3/wnENEwStGD5LkGXrc2t8MMql70gYR+I=\n\n— rekor.sigstore.dev wNI9ajBFAiAlDLGmhchbVNTNIJy/FsLjTMhJxHLDRTMsvhGjXSB+zAIhAPkAcPqgaT65cxZ9A+ctIJZyVZ4lBLQukg64ZZYSQRFr\n"}},"canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiJiMTQ0NmVjOGYzNTBkOTJjYTIwNjBjMmFmMzEwMGZhMDFhYjE4MzZkYTU5OTNkODNhODRmOGNkYTViNWRlNzk5In19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJQ01XYVJpL0tTcG5HcnZVOGlwOWtuOVB2aUtXRUUyQWdxd0hISHArR0hlMUFpRUFwaEZFa0ZMTklycEJSMjZPWFJ3SExCOTBHOXNPSncrSmJXdzlXQzdDVENNPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRWRU5EUW5SRFowRjNTVUpCWjBsVlJrUkNlVTV2Y0ZjeWJWZE5MMWhzU0RoWFVqVXpOVWhCSzFkWmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDVxU1RCTlJFVjNUVlJKTUZkb1kwNU5hbGwzVG1wSk1FMUVSWGhOVkVrd1YycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZJZURsVVpEbGhjV28wTUhoNE9XSnZhbkJpYUhoTGFHVlpMMmxqTDNSek1HbFdTMUVLVDJsMWJqTm1ZbkkyZUVKTWFFZDZXVUpWYWt4UFRqZEhialpXWVd4VWFYWnJSRmcwUm5wcWNTODBRMmg0TlRSQ01XRlBRMEpsT0hkbloxaHlUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlZpYmtGa0NrNHZMMHByVjFOaGMwUlZkRUZSVWxaV1IxTTBOekJ6ZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOWFha1Y1V1ZkYWFrOVVTbXBPUkVKcFdXMVJNMDV0VVRWUFZFSnJDazlVUlhoYVYwcG9XVzFaZVU5SFZUQlBWRUpvVGpKS2ExcFVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMXBxUlhsWlYxcHFUMVJLYWs1RVFtbFpiVkV6VG0xUk5VOVVRbXRQVkVWNFdsZEthRmx0V1hrS1QwZFZNRTlVUW1oT01rcHJXbFJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMXBxUlhsWlYxcHFUMVJLYWs1RVFtbFpiVkV6VG0xUk5VOVVRbXNLVDFSRmVGcFhTbWhaYlZsNVQwZFZNRTlVUW1oT01rcHJXbFJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwZFplRTF0Um0wS1dYcHJlVmw2VVhkWmJVcHJUbnBhYTA5VWEzZGFSR3Q0VFZkV2FWbFhTbTFOYW1oc1RrUnJkMWxVWkdsYVIxVjNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOYW1kM1RYcFplVTVxWnpWTmVsVjJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwVVZsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTTjBKSWEwRUtaSGRDTVVGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtNTJZMnhhTVRCQlFVRlJSQXBCUlZsM1VrRkpaMkYyYUVoV2RqUlpjalZLVVhCaFJVTldaMnRGYkZsVlRtcHBkRkpLZW1Kd1FTdDNjMFJTY0RBMGFXZERTVUZVZWlzNE9DdEtUVzEzQ25OUVRVOTNRVmRCZWtVNEsyOXpNbmx5YkdVcmMyNUtaRVZtYm1GMU4wUnVUVUZ2UjBORGNVZFRUVFE1UWtGTlJFRXlZMEZOUjFGRFRVWTNUR2gwUVdnS1ZUWm5SMUZoUVU1NlNTdEhNVFZ3YWxZd2FHODBSWFo2TVcxdFJGQm9SMGhIVm5WQllVZ3JWSEZSZEZjM0swUlpNRGNyYjB4WGJDOVlaMGwzV0c1dVdBcFZWa1Z5WWpGRVQwZERWMmhZWTFaV05XdFRORlJNWWtNNGIyeFNOR0ZqWjBoRmJEVjVSR3RMWnlzdlprMXBkRmRHWjNoalZHWnZWRXN3VTNVS0xTMHRMUzFGVGtRZ1EwVlNWRWxHU1VOQlZFVXRMUzB0TFFvPSJ9fX19"}],"timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyTADAgEAMIICwAYJKoZIhvcNAQcCoIICsTCCAq0CAQMxDTALBglghkgBZQMEAgEwgbcGCyqGSIb3DQEJEAEEoIGnBIGkMIGhAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgODd/U6X/MxyB9QlRMxGHoiMWQaWB/aZdYFPP1pppl9UCFANF9uj22gAgErrTQcUpIWPItUiiGA8yMDI2MDYyNDAxMDEyNFowAwIBAaAypDAwLjEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MRUwEwYDVQQDEwxzaWdzdG9yZS10c2GgADGCAdswggHXAgEBMFEwOTEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MSAwHgYDVQQDExdzaWdzdG9yZS10c2Etc2VsZnNpZ25lZAIUOhNULwyQYe68wUMvy4qOiyojiwwwCwYJYIZIAWUDBAIBoIH8MBoGCSqGSIb3DQEJAzENBgsqhkiG9w0BCRABBDAcBgkqhkiG9w0BCQUxDxcNMjYwNjI0MDEwMTI0WjAvBgkqhkiG9w0BCQQxIgQgZToxx2m7P02p4FZ7UceQ8UgWYimCO/H5ZZvjfbG2MGMwgY4GCyqGSIb3DQEJEAIvMX8wfTB7MHkEIIX5J7wHq2LKw7RDVsEO/IGyxog/2nq55thw2dE6zQW3MFUwPaQ7MDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAoGCCqGSM49BAMCBGcwZQIwV6JFlcpT8gojSf7YR9CKQIDKmjqf0a3fenBkO9gSuqP22xfLsaAy5ZepkHPtj07QAjEAs+Ys6rGJPX3Wof1kRPN0wU6mvyZw5qTyj3FjSoMjCybT9a8CNUzqVEKUpQPESUNG"}]}},"messageSignature":{"messageDigest":{"algorithm":"SHA2_256","digest":"sURuyPNQ2SyiBgwq8xAPoBqxg22lmT2DqE+M2ltd55k="},"signature":"MEUCICMWaRi/KSpnGrvU8ip9kn9PviKWEE2AgqwHHHp+GHe1AiEAphFEkFLNIrpBR26OXRwHLB90G9sOJw+JbWw9WC7CTCM="}} \ No newline at end of file diff --git a/build/torch212-cxx11-cu130-x86_64-linux/ops/__init__.py b/build/torch212-cxx11-cu130-x86_64-linux/ops/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/build/torch212-cxx11-cu130-x86_64-linux/ops/activations.py b/build/torch212-cxx11-cu130-x86_64-linux/ops/activations.py deleted file mode 100644 index 7c09649fc41e12d5a360c5672825d8380bc7ec80..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu130-x86_64-linux/ops/activations.py +++ /dev/null @@ -1,135 +0,0 @@ -# Copied from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/model/layers/activations.py -import math - -import torch -import torch.nn as nn -import torch.nn.functional as F - -# 1/sqrt(2*pi)-> 0.3989423 -# 1/sqrt(2) -> 0.70710678 -# sqrt(2/pi) -> 0.79788456 - -# this function is tanh approximation of gelu -# actual gelu is: -# x * 0.5 * (1.0 + torch.erf(x * 0.70710678)) -@torch.jit.script -def bias_gelu(y, bias): - x = bias + y - return (x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))).to(dtype=y.dtype) - - -# gradient of tanh approximation of gelu -# gradient of actual gelu is: -# 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x) -@torch.jit.script -def bias_gelu_back(g, y, bias): - """Assume that y has shape (B, D) and bias has shape (D)""" - x = bias + y - tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)) - # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243 - ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( - 1 + tanh_out - ) - grad_y = ff * g - return grad_y.to(dtype=y.dtype), grad_y.sum(dim=(0), dtype=bias.dtype) - - -class GeLUFunction(torch.autograd.Function): - @staticmethod - # bias is an optional argument - def forward(ctx, input, bias): - ctx.save_for_backward(input, bias) - return bias_gelu(input, bias) - - @staticmethod - def backward(ctx, grad_output): - input, bias = ctx.saved_tensors - tmp = bias_gelu_back(grad_output, input, bias) - return tmp, tmp - - -bias_gelu_impl = GeLUFunction.apply - -# this function is tanh approximation of gelu -# actual gelu is: -# x * 0.5 * (1.0 + torch.erf(x * 0.70710678)) -@torch.jit.script -def gelu_fwd(x): - return (x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))).to(dtype=x.dtype) - - -# gradient of tanh approximation of gelu -# gradient of actual gelu is: -# 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x) -@torch.jit.script -def gelu_bwd(g, x): - tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)) - # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243 - ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( - 1 + tanh_out - ) - return (ff * g).to(dtype=x.dtype) - - -class FastGeLUFunction(torch.autograd.Function): - @staticmethod - # bias is an optional argument - def forward(ctx, input): - ctx.save_for_backward(input) - return gelu_fwd(input) - - @staticmethod - def backward(ctx, grad_output): - (input,) = ctx.saved_tensors - tmp = gelu_bwd(grad_output, input) - return tmp - - -fast_gelu_impl = FastGeLUFunction.apply - - -@torch.jit.script -def relu_bwd(g, x): - return torch.where(x >= 0, g, 0.0).to(dtype=x.dtype) - - -@torch.jit.script -def sqrelu_fwd(x): - r = F.relu(x) - return (r * r).to(dtype=x.dtype) - - -@torch.jit.script -def sqrelu_bwd(g, x): - return (2.0 * g * F.relu(x)).to(dtype=x.dtype) - - -swiglu_fwd_codestring = """ -template T swiglu_fwd(T x, T y) { - return float(x) * float(y) / (1.0f + ::exp(-float(x))); -} -""" -swiglu_bwd_codestring = """ -template void swiglu_bwd(T x, T y, T g, T& dx, T& dy) { - float x_sigmoid = 1.0f / (1.0f + ::exp(-float(x))); - dx = x_sigmoid * (1 + float(x) * (1.0f - x_sigmoid)) * float(g) * float(y); - dy = float(x) * x_sigmoid * float(g); -} -""" -swiglu_fwd = torch.cuda.jiterator._create_jit_fn(swiglu_fwd_codestring) -swiglu_bwd = torch.cuda.jiterator._create_multi_output_jit_fn(swiglu_bwd_codestring, num_outputs=2) - - -class SwiGLUFunction(torch.autograd.Function): - - @staticmethod - def forward(ctx, x, y): - ctx.save_for_backward(x, y) - return swiglu_fwd(x, y) - - @staticmethod - def backward(ctx, dout): - x, y = ctx.saved_tensors - return swiglu_bwd(x, y, dout) - -swiglu = SwiGLUFunction.apply diff --git a/build/torch212-cxx11-cu130-x86_64-linux/ops/fused_dense.py b/build/torch212-cxx11-cu130-x86_64-linux/ops/fused_dense.py deleted file mode 100644 index 6b4033d134e4093fe278f7b3f8c7d3128ce9f36d..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu130-x86_64-linux/ops/fused_dense.py +++ /dev/null @@ -1,688 +0,0 @@ -# Copyright (c) 2023, Tri Dao. -# Inspired by https://github.com/NVIDIA/apex/blob/master/apex/fused_dense/fused_dense.py -# We make it work with pytorch amp and with bfloat16. -# The TensorParallel linear modules are inspired by https://github.com/NVIDIA/apex/blob/master/apex/transformer/tensor_parallel/layers.py -from functools import partial -from typing import Optional - -# import fused_dense_cuda # from apex -import fused_dense_lib as fused_dense_cuda -import torch -import torch.nn as nn -import torch.nn.functional as F -from torch import Tensor -from torch.distributed import ProcessGroup - -from flash_attn.utils.torch import custom_fwd, custom_bwd -from flash_attn.ops.activations import gelu_bwd, relu_bwd, sqrelu_bwd, sqrelu_fwd -from flash_attn.utils.distributed import ( - all_gather_raw, - all_reduce, - all_reduce_raw, - reduce_scatter, - reduce_scatter_raw, -) - - -class FusedDenseFunc(torch.autograd.Function): - @staticmethod - @custom_fwd - def forward( - ctx, x, weight, bias, return_residual=False, process_group=None, sequence_parallel=True - ): - """ - If process_group is not None and sequence_parallel=True, we're doing Tensor Parallel - with sequence parallelism: we do an all_gather_raw of x before doing the matmul. - """ - ctx.compute_weight_gradient = weight.requires_grad - ctx.return_residual = return_residual - ctx.process_group = process_group - ctx.sequence_parallel = sequence_parallel - - if torch.is_autocast_enabled(): - x = x.to(dtype=torch.get_autocast_gpu_dtype()) - x = x.contiguous() - if process_group is not None and sequence_parallel: - # We want to kick off the all_gather early, before weight dtype conversion - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - else: - total_x = x - - if torch.is_autocast_enabled(): - weight = weight.to(dtype=torch.get_autocast_gpu_dtype()) - bias = bias.to(dtype=torch.get_autocast_gpu_dtype()) if bias is not None else None - weight = weight.contiguous() - if process_group is not None and sequence_parallel: - handle_x.wait() - batch_shape, n = total_x.shape[:-1], total_x.shape[-1] - batch_dim = batch_shape.numel() - # https://github.com/pytorch/pytorch/blob/5b51849b48a7dbccd297286cc0110def4706f9e7/aten/src/ATen/native/cuda/Blas.cpp#L174 - if min(batch_dim, n, *weight.shape) > 65535 * 32: - raise RuntimeError("fused_dense only supports matrix dims <= 2M") - output = F.linear(total_x, weight, bias) - if ctx.compute_weight_gradient: - ctx.save_for_backward(x, weight) - else: - ctx.save_for_backward(weight) - return output if not return_residual else (output, x) - - @staticmethod - @custom_bwd - def backward(ctx, grad_output, *args): - grad_output = grad_output.contiguous() - if ctx.return_residual: - (grad_input,) = args - grad_input = grad_input.contiguous() - process_group = ctx.process_group - sequence_parallel = ctx.sequence_parallel - if ctx.compute_weight_gradient: - x, weight = ctx.saved_tensors - if process_group is not None and sequence_parallel: - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - else: - total_x = x - else: - (weight,) = ctx.saved_tensors - total_x = None - batch_shape = grad_output.shape[:-1] - batch_dim = batch_shape.numel() - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - if ctx.needs_input_grad[0]: - if not ctx.return_residual: - grad_input = F.linear(grad_output, weight.t()) - else: - grad_input = torch.addmm( - grad_input.reshape(batch_dim, grad_input.shape[-1]), grad_output, weight - ) - grad_input = grad_input.reshape(*batch_shape, grad_input.shape[-1]) - if process_group is not None: - reduce_fn = reduce_scatter_raw if sequence_parallel else all_reduce_raw - grad_input, handle_grad_input = reduce_fn(grad_input, process_group, async_op=True) - else: - grad_input = None - if ctx.needs_input_grad[1]: - assert ctx.compute_weight_gradient - if process_group is not None and sequence_parallel: - handle_x.wait() - grad_weight, grad_bias = fused_dense_cuda.linear_bias_wgrad( - total_x.reshape(batch_dim, total_x.shape[-1]), grad_output, ctx.needs_input_grad[2] - ) - else: - grad_weight = None - grad_bias = grad_output if ctx.needs_input_grad[2] else None - if process_group is not None and ctx.needs_input_grad[0]: - handle_grad_input.wait() - return grad_input, grad_weight, grad_bias, None, None, None - - -def fused_dense_func( - x: Tensor, - weight: Tensor, - bias: Optional[Tensor] = None, - return_residual: bool = False, - process_group: Optional[ProcessGroup] = None, - sequence_parallel: bool = True, -): - dtype_eligible = x.dtype in [torch.float16, torch.bfloat16] or ( - x.dtype == torch.float32 and torch.is_autocast_enabled() - ) - if x.is_cuda and weight.is_cuda and (bias is None or bias.is_cuda) and dtype_eligible: - return FusedDenseFunc.apply( - x, weight, bias, return_residual, process_group, sequence_parallel - ) - else: - assert process_group is None - out = F.linear(x, weight, bias) - return out if not return_residual else (out, x) - - -class FusedDense(nn.Linear): - def __init__( - self, - in_features: int, - out_features: int, - bias: bool = True, - return_residual: bool = False, - device=None, - dtype=None, - ) -> None: - super().__init__(in_features, out_features, bias=bias, device=device, dtype=dtype) - self.return_residual = return_residual - - def forward(self, x, process_group=None): - """ - If process_group is not None, we're doing Tensor Parallel with sequence parallelism: - we do an all_gather of x before doing the matmul. - """ - return fused_dense_func( - x, - self.weight, - self.bias, - return_residual=self.return_residual, - process_group=process_group, - ) - - -class ColumnParallelLinear(nn.Linear): - def __init__( - self, - in_features: int, - out_features: int, - process_group: ProcessGroup, - bias: bool = True, - sequence_parallel=True, - multiple_of=1, - device=None, - dtype=None, - ) -> None: - world_size = torch.distributed.get_world_size(process_group) - if out_features % multiple_of: - raise ValueError(f"out_features ({out_features}) must be a multiple of {multiple_of}") - multiple = out_features // multiple_of - # We want to split @multiple across world_size, but it could be an uneven split - div = multiple // world_size - mod = multiple % world_size - # The first @mod ranks get @div + 1 copies, the rest get @div copies - local_multiple = div + int(torch.distributed.get_rank(process_group) < mod) - super().__init__( - in_features, local_multiple * multiple_of, bias=bias, device=device, dtype=dtype - ) - self.process_group = process_group - self.sequence_parallel = sequence_parallel - - def forward(self, x): - # If self.sequence_parallel is True, we're doing Tensor Parallel with sequence parallelism: - # we do an all_gather of x before doing the matmul. - # If not, then the input is already gathered. - return fused_dense_func( - x, - self.weight, - self.bias, - process_group=self.process_group, - sequence_parallel=self.sequence_parallel, - ) - - -class RowParallelLinear(nn.Linear): - def __init__( - self, - in_features: int, - out_features: int, - process_group: ProcessGroup, - bias: bool = True, - sequence_parallel=True, - multiple_of=1, - device=None, - dtype=None, - ) -> None: - world_size = torch.distributed.get_world_size(process_group) - rank = torch.distributed.get_rank(process_group) - if in_features % multiple_of: - raise ValueError(f"in_features ({in_features}) must be a multiple of {multiple_of}") - multiple = in_features // multiple_of - # We want to split @multiple across world_size, but it could be an uneven split - div = multiple // world_size - mod = multiple % world_size - # The first @mod ranks get @div + 1 copies, the rest get @div copies - local_multiple = div + int(torch.distributed.get_rank(process_group) < mod) - # Only rank 0 will have bias - super().__init__( - local_multiple * multiple_of, - out_features, - bias=bias and rank == 0, - device=device, - dtype=dtype, - ) - self.process_group = process_group - self.sequence_parallel = sequence_parallel - - def forward(self, x): - """ - We're doing Tensor Parallel with sequence parallelism: we do the matmul and then - a reduce_scatter of the result. - """ - out = fused_dense_func(x, self.weight, self.bias) - reduce_fn = reduce_scatter if self.sequence_parallel else all_reduce - return reduce_fn(out, self.process_group) - - -class FusedMLPFunc(torch.autograd.Function): - @staticmethod - @custom_fwd - def forward( - ctx, - x, - weight1, - bias1, - weight2, - bias2, - activation="gelu_approx", - save_pre_act=True, - return_residual=False, - checkpoint_lvl=0, - heuristic=0, - process_group=None, - sequence_parallel=True, - ): - """ - If process_group is not None and sequence_parallel=True, we're doing Tensor Parallel - with sequence parallelism: we do an all_gather of x before doing the matmul. - If sequence_parallel=False, then the input is already gathered. - - checkpoint_lvl: - 0: no recomputation in the bwd - 1: recompute gelu_out / relu_out in the bwd - 2: recompute pre_act and gelu_out / relu_out in the bwd - """ - assert -1 <= heuristic <= 4 - assert activation in ["gelu_approx", "relu", "sqrelu"] - if activation == "sqrelu": - assert heuristic == -1 - if not save_pre_act: - checkpoint_lvl = 2 - assert checkpoint_lvl in [0, 1, 2] - ctx.return_residual = return_residual - ctx.process_group = process_group - ctx.sequence_parallel = sequence_parallel - ctx.checkpoint_lvl = checkpoint_lvl - ctx.activation = activation - ctx.heuristic = heuristic - - if torch.is_autocast_enabled(): - x = x.to(dtype=torch.get_autocast_gpu_dtype()) - x = x.contiguous() - if process_group is not None and sequence_parallel: - # We want to kick off the all_gather early, before weight dtype conversion - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - else: - total_x = x - - if torch.is_autocast_enabled(): - dtype = torch.get_autocast_gpu_dtype() - weight1, weight2 = [a.to(dtype=dtype) for a in [weight1, weight2]] - bias1 = bias1.to(dtype=dtype) if bias1 is not None else None - bias2 = bias2.to(dtype=dtype) if bias2 is not None else None - weight1 = weight1.contiguous() - bias1 = bias1.contiguous() if bias1 is not None else None - weight2 = weight2.contiguous() - bias2 = bias2.contiguous() if bias2 is not None else None - if process_group is not None and sequence_parallel: - handle_x.wait() - batch_shape, n = total_x.shape[:-1], total_x.shape[-1] - batch_dim = batch_shape.numel() - # https://github.com/pytorch/pytorch/blob/5b51849b48a7dbccd297286cc0110def4706f9e7/aten/src/ATen/native/cuda/Blas.cpp#L174 - if min(batch_dim, n, *weight1.shape, *weight2.shape) > 65535 * 32: - raise RuntimeError("fused_dense only supports matrix dims <= 2M") - if heuristic == -1: - pre_act = F.linear(total_x, weight1, bias1) - activation_fn = ( - partial(F.gelu, approximate="tanh") - if activation == "gelu_approx" - else (sqrelu_fwd if activation == "sqrelu" else F.relu) - ) - with torch.jit.fuser("fuser2"): - output1 = activation_fn(pre_act) - # This is before adding bias1 - # pre_act = F.linear(total_x.reshape(batch_dim, n), weight1) - # with torch.jit.fuser('fuser2'): - # output1 = bias_gelu(pre_act, bias1) - else: - is_gelu = activation == "gelu_approx" - output1, *rest = fused_dense_cuda.linear_act_forward( - total_x.reshape(batch_dim, n), weight1, bias1, is_gelu, save_pre_act, heuristic - ) - if save_pre_act: - pre_act = rest[0] - output2 = F.linear(output1, weight2, bias2) - if checkpoint_lvl == 0 or (checkpoint_lvl == 1 and activation == "relu"): - # For RELU the pre_act is very small (just a bit-mask) so we just save it - ctx.save_for_backward(x, weight1, weight2, pre_act, output1) - elif checkpoint_lvl == 1: - ctx.save_for_backward(x, weight1, weight2, pre_act) - elif checkpoint_lvl == 2: - ctx.save_for_backward(x, weight1, weight2, bias1) - output2 = output2.reshape(*batch_shape, output2.shape[-1]) - return output2 if not return_residual else (output2, x) - - @staticmethod - @custom_bwd - def backward(ctx, grad_output, *args): - grad_output = grad_output.contiguous() - checkpoint_lvl = ctx.checkpoint_lvl - activation = ctx.activation - activation_fn = ( - partial(F.gelu, approximate="tanh") - if activation == "gelu_approx" - else (sqrelu_fwd if activation == "sqrelu" else F.relu) - ) - if ctx.return_residual: - (grad_input,) = args - grad_input = grad_input.contiguous() - process_group = ctx.process_group - sequence_parallel = ctx.sequence_parallel - x, weight1, weight2, *rest = ctx.saved_tensors - if process_group is None or not sequence_parallel: - total_x = x - batch_shape = grad_output.shape[:-1] - batch_dim = batch_shape.numel() - if checkpoint_lvl in [0, 1]: - if process_group is not None and sequence_parallel: - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - if checkpoint_lvl == 0 or (checkpoint_lvl == 1 and activation == "relu"): - pre_act, output1 = rest - elif checkpoint_lvl == 1: - (pre_act,) = rest - with torch.jit.fuser("fuser2"): - output1 = activation_fn(pre_act) - elif checkpoint_lvl == 2: - (bias1,) = rest - if process_group is not None and sequence_parallel: - total_x, _ = all_gather_raw(x, process_group) - if ctx.heuristic == -1: - pre_act = F.linear(total_x, weight1, bias1) - with torch.jit.fuser("fuser2"): - output1 = activation_fn(pre_act) - else: - output1, pre_act = fused_dense_cuda.linear_act_forward( - total_x.reshape(batch_dim, total_x.shape[-1]), - weight1, - bias1, - activation == "gelu_approx", - True, - ctx.heuristic, - ) - - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - output1 = output1.reshape(batch_dim, output1.shape[-1]) - pre_act = pre_act.reshape(batch_dim, pre_act.shape[-1]) - if ctx.needs_input_grad[3]: - grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad( - output1, grad_output, ctx.needs_input_grad[4] - ) - else: - grad_weight2 = None - grad_bias2 = grad_output if ctx.needs_input_grad[4] else None - if ctx.heuristic == -1: - # grad_pre_act = matmul_dgelu(grad_output, weight2, pre_act) - grad_output1 = F.linear(grad_output, weight2.t()) - activation_grad_fn = ( - gelu_bwd - if activation == "gelu_approx" - else (sqrelu_bwd if activation == "sqrelu" else relu_bwd) - ) - with torch.jit.fuser("fuser2"): - grad_pre_act = activation_grad_fn(grad_output1, pre_act) - else: - # The cublasLt epilogue has to compute both gelu/relu grad and bias grad, we can't - # just compute gelu/relu grad - grad_pre_act, grad_bias1 = fused_dense_cuda.bias_act_linear_dgrad_bgrad( - weight2, grad_output, pre_act, activation == "gelu_approx", ctx.heuristic - ) - if not ctx.needs_input_grad[2]: - grad_bias1 = None - if ctx.needs_input_grad[0]: - if not ctx.return_residual: - grad_input = F.linear(grad_pre_act, weight1.t()) - else: - grad_input = torch.addmm( - grad_input.reshape(batch_dim, grad_input.shape[-1]), grad_pre_act, weight1 - ) - grad_input = grad_input.reshape(*batch_shape, grad_input.shape[-1]) - if process_group is not None: - reduce_fn = reduce_scatter_raw if sequence_parallel else all_reduce_raw - grad_input, handle_grad_input = reduce_fn(grad_input, process_group, async_op=True) - else: - grad_input = None - if ctx.heuristic == -1: - if ctx.needs_input_grad[1]: - if process_group is not None and sequence_parallel and checkpoint_lvl != 2: - handle_x.wait() - grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_wgrad( - total_x.reshape(batch_dim, total_x.shape[-1]), - grad_pre_act, - ctx.needs_input_grad[2], - ) - else: - grad_weight1 = None - grad_bias1 = grad_pre_act if ctx.needs_input_grad[2] else None - else: - if ctx.needs_input_grad[1]: - if process_group is not None and sequence_parallel and checkpoint_lvl != 2: - handle_x.wait() - grad_weight1 = F.linear( - grad_pre_act.t(), total_x.reshape(batch_dim, total_x.shape[-1]).t() - ) - else: - grad_weight1 = None - if process_group is not None and ctx.needs_input_grad[0]: - handle_grad_input.wait() - return ( - grad_input, - grad_weight1, - grad_bias1, - grad_weight2, - grad_bias2, - None, - None, - None, - None, - None, - None, - None, - ) - - -def fused_mlp_func( - x: Tensor, - weight1: Tensor, - weight2: Tensor, - bias1: Optional[Tensor] = None, - bias2: Optional[Tensor] = None, - activation: str = "gelu_approx", - save_pre_act: bool = True, - return_residual: bool = False, - checkpoint_lvl: int = 0, - heuristic: int = 0, - process_group: Optional[ProcessGroup] = None, - sequence_parallel: bool = True, -): - assert activation in ["gelu_approx", "relu", "sqrelu"] - dtype_eligible = x.dtype in [torch.float16, torch.bfloat16] or ( - x.dtype == torch.float32 and torch.is_autocast_enabled() - ) - # If we save pre-activation, dimension must be divisible by 128 (relu) or 8 (gelu) - dim_eligible = not save_pre_act or (x.shape[-1] % (128 if activation == "relu" else 8) == 0) - if ( - x.is_cuda - and weight1.is_cuda - and weight2.is_cuda - and (bias1 is None or bias1.is_cuda) - and (bias2 is None or bias2.is_cuda) - and dtype_eligible - and dim_eligible - ): - return FusedMLPFunc.apply( - x, - weight1, - bias1, - weight2, - bias2, - activation, - save_pre_act, - return_residual, - checkpoint_lvl, - heuristic, - process_group, - sequence_parallel, - ) - else: - assert process_group is None - pre_act = F.linear(x, weight1, bias1) - activation_fn = ( - partial(F.gelu, approximate="tanh") - if activation == "gelu_approx" - else partial(F.relu, inplace=True) - ) - output1 = activation_fn(pre_act) - output2 = F.linear(output1, weight2, bias2) - return output2 if not return_residual else (output2, x) - - -class FusedMLP(nn.Module): - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - bias1=True, - bias2=True, - activation="gelu_approx", - return_residual=False, - checkpoint_lvl=0, - heuristic="auto", - device=None, - dtype=None, - ): - """ - If process_group is not None, we're doing Tensor Parallel with sequence parallelism: - we do an all_gather of x before doing the matmul, gelu, then matmul. - Finally we do a reduce_scatter of the output. - - checkpoint_lvl (increasing lvl means slower but more memory saving): - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute pre_act and gelu_out in the bwd - heuristic: - -1: don't fuse gemm + gelu (separate kernel) - 0..4: use this heuristic for the algo section in the fused gemm + gelu - 'auto': heuristic will be picked automatically: - For CUDA >= 11.8, we set heuristic=0 for both fp16 and bf16 for best perf. - For CUDA <= 11.7, we set heuristic=1 for fp16 and heuristic=-1 for bf16. - For H100, we set heuristic=-1 for both fp16 and bf16 as the fused cuBlasLt implementation - is slower than the unfused version. - return_residual: whether to return the input x along with the output. This is for - performance reason: for post-norm architecture, returning the input allows us - to fuse the backward of nn.Linear with the residual connection. - """ - assert checkpoint_lvl in [0, 1, 2] - assert activation in ["gelu_approx", "relu", "sqrelu"] - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features * 4 - self.activation = activation - self.return_residual = return_residual - self.checkpoint_lvl = checkpoint_lvl - self.heuristic = heuristic if activation != "sqrelu" else -1 - self.fc1 = nn.Linear(in_features, hidden_features, bias=bias1, **factory_kwargs) - self.fc2 = nn.Linear(hidden_features, out_features, bias=bias2, **factory_kwargs) - - def forward(self, x, process_group=None): - dtype = x.dtype if not torch.is_autocast_enabled() else torch.get_autocast_gpu_dtype() - if self.heuristic == "auto": - if self.activation == "gelu_approx": - if torch.cuda.get_device_capability("cuda") == (9, 0): - heuristic = -1 - else: - cuda_ver = tuple(map(int, torch.version.cuda.split("."))) - heuristic = 0 if cuda_ver >= (11, 8) else (1 if dtype == torch.float16 else -1) - else: - heuristic = 0 - else: - heuristic = self.heuristic - out = fused_mlp_func( - x, - self.fc1.weight, - self.fc2.weight, - self.fc1.bias, - self.fc2.bias, - activation=self.activation, - save_pre_act=self.training, - return_residual=self.return_residual, - checkpoint_lvl=self.checkpoint_lvl, - heuristic=heuristic, - process_group=process_group, - ) - if self.return_residual: - out, x = out - if process_group is not None: - out = reduce_scatter(out, process_group) - return out if not self.return_residual else (out, x) - - -class ParallelFusedMLP(nn.Module): - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - activation="gelu_approx", - process_group: ProcessGroup = None, - bias1=True, - bias2=True, - sequence_parallel=True, - checkpoint_lvl=0, - heuristic="auto", - device=None, - dtype=None, - ): - """ - process_group is required. We're doing Tensor Parallel with sequence parallelism: - we do an all_gather of x before doing the matmul, gelu, then matmul. - Finally we do a reduce_scatter of the output. - - checkpoint_lvl (increasing lvl means slower but more memory saving): - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute pre_act and gelu_out in the bwd - heuristic: - -1: don't fuse gemm + gelu (separate kernel) - 0..4: use this heuristic for the algo section in the fused gemm + gelu - 'auto': heuristic will be picked automatically: - For CUDA >= 11.8, we set heuristic=0 for both fp16 and bf16 for best perf. - For CUDA <= 11.7, we set heuristic=1 for fp16 and heuristic=-1 for bf16. - """ - assert checkpoint_lvl in [0, 1, 2] - assert activation in ["gelu_approx", "relu", "sqrelu"] - assert process_group is not None - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features * 4 - self.activation = activation - self.process_group = process_group - self.sequence_parallel = sequence_parallel - self.checkpoint_lvl = checkpoint_lvl - self.heuristic = heuristic if activation != "sqrelu" else -1 - self.fc1 = ColumnParallelLinear( - in_features, hidden_features, process_group, bias=bias1, **factory_kwargs - ) - self.fc2 = RowParallelLinear( - hidden_features, out_features, process_group, bias=bias2, **factory_kwargs - ) - - def forward(self, x): - dtype = x.dtype if not torch.is_autocast_enabled() else torch.get_autocast_gpu_dtype() - if self.heuristic == "auto": - if self.activation == "gelu_approx": - cuda_ver = tuple(map(int, torch.version.cuda.split("."))) - heuristic = 0 if cuda_ver >= (11, 8) else (1 if dtype == torch.float16 else -1) - else: - heuristic = 0 - else: - heuristic = self.heuristic - out = fused_mlp_func( - x, - self.fc1.weight, - self.fc2.weight, - self.fc1.bias, - self.fc2.bias, - activation=self.activation, - save_pre_act=self.training, - checkpoint_lvl=self.checkpoint_lvl, - heuristic=heuristic, - process_group=self.process_group, - sequence_parallel=self.sequence_parallel, - ) - reduce_fn = reduce_scatter if self.sequence_parallel else all_reduce - return reduce_fn(out, self.process_group) diff --git a/build/torch212-cxx11-cu130-x86_64-linux/ops/layer_norm.py b/build/torch212-cxx11-cu130-x86_64-linux/ops/layer_norm.py deleted file mode 100644 index 4b6cd798fd02844ef9cd3897f8ab95e490e638bf..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu130-x86_64-linux/ops/layer_norm.py +++ /dev/null @@ -1,800 +0,0 @@ -# Copyright (c) 2022, Tri Dao. -# Adapted from https://github.com/NVIDIA/apex/blob/master/apex/contrib/layer_norm/layer_norm.py - -import dropout_layer_norm -import torch -from torch.nn import init - - -def maybe_align(x, alignment_in_bytes=16): - """Assume that x already has last dim divisible by alignment_in_bytes""" - # TD [2023-07-04] I'm not 100% sure that clone will align the memory - # https://discuss.pytorch.org/t/how-to-ensure-that-tensor-data-ptr-is-aligned-to-16-bytes/183440 - return x if x.data_ptr() % alignment_in_bytes == 0 else x.clone() - - -def _dropout_add_layer_norm_forward( - x0, - residual, - gamma, - beta, - rowscale, - colscale, - dropout_p, - epsilon, - residual_in_fp32=False, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes""" - hidden_size = gamma.numel() - x0mat = x0.view((-1, hidden_size)) - residualmat = residual.view((-1, hidden_size)) if residual is not None else None - rowscale = rowscale.view(-1) if rowscale is not None else None - zmat, xmat, dmask, mu, rsigma = dropout_layer_norm.dropout_add_ln_fwd( - x0mat, - residualmat, - gamma, - beta, - rowscale, - colscale, - None, - None, - dropout_p, - epsilon, - 1.0, - 0, - None, - residual_in_fp32, - is_rms_norm, - ) - # dmask is None if dropout_p == 0.0 - # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype - return zmat, xmat if xmat is not None else x0mat, dmask, mu, rsigma - - -def _dropout_add_layer_norm_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - rowscale, - colscale, - dropout_p, - has_residual, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes - dx == None means that it was a post-norm architecture - (x = drop(x0) + residual was not returned in the fwd). - x0 must not be None if we have colscale. - """ - hidden_size = gamma.numel() - xmat = x.view((-1, hidden_size)) - dzmat = dz.view(xmat.shape) - dxmat = dx.view(xmat.shape) if dx is not None else None - x0mat = x0.view((-1, hidden_size)) if x0 is not None else None - rowscale = rowscale.view(-1) if rowscale is not None else None - if colscale is not None: - assert x0 is not None, "x0 is required to compute the gradient of colscale" - dx0mat, dresidualmat, dgamma, dbeta, _, _, *rest = dropout_layer_norm.dropout_add_ln_bwd( - dzmat, - dxmat, - xmat, - x0mat, - dmask, - mu, - rsigma, - gamma, - rowscale, - colscale, - None, - None, - dropout_p, - 1.0, - 0, - has_residual, - is_rms_norm, - ) - # dresidualmat is None if not has_residual - if colscale is None: - return dx0mat, dresidualmat, dgamma, dbeta - else: - dcolscale = rest[0] - return dx0mat, dresidualmat, dgamma, dbeta, dcolscale - - -def _dropout_add_layer_norm_subset_forward( - x0, - residual, - gamma, - beta, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32=False, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes""" - hidden_size = gamma.numel() - x0mat = x0.view((-1, hidden_size)) - residualmat = residual.view((-1, hidden_size)) if residual is not None else None - x0_subset = x0_subset.view(-1) if x0_subset is not None else None - out_subset = out_subset.view(-1) if out_subset is not None else None - zmat, xmat, dmask, mu, rsigma = dropout_layer_norm.dropout_add_ln_fwd( - x0mat, - residualmat, - gamma, - beta, - None, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - None, - residual_in_fp32, - is_rms_norm, - ) - # dmask is None if dropout_p == 0.0 - # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype - return zmat, xmat if xmat is not None else x0mat, dmask, mu, rsigma - - -def _dropout_add_layer_norm_subset_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - colscale, - x0_subset, - out_subset, - dropout_p, - rowscale_const, - x0_numrows, - has_residual, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes - dx == None means that it was a post-norm architecture - (x = drop(x0) + residual was not returned in the fwd). - x0 must not be None if we have colscale. - """ - hidden_size = gamma.numel() - xmat = x.view((-1, hidden_size)) - dzmat = dz.view(-1, hidden_size) - dxmat = dx.view(xmat.shape) if dx is not None else None - x0mat = x0.view((-1, hidden_size)) if x0 is not None else None - x0_subset = x0_subset.view(-1) if x0_subset is not None else None - out_subset = out_subset.view(-1) if out_subset is not None else None - if colscale is not None: - assert x0 is not None, "x0 is required to compute the gradient of colscale" - dx0mat, dresidualmat, dgamma, dbeta, _, _, *rest = dropout_layer_norm.dropout_add_ln_bwd( - dzmat, - dxmat, - xmat, - x0mat, - dmask, - mu, - rsigma, - gamma, - None, - colscale, - x0_subset, - out_subset, - dropout_p, - rowscale_const, - x0_numrows, - has_residual, - is_rms_norm, - ) - # dresidualmat is None if not has_residual - if colscale is None: - return dx0mat, dresidualmat, dgamma, dbeta - else: - dcolscale = rest[0] - return dx0mat, dresidualmat, dgamma, dbeta, dcolscale - - -def _dropout_add_layer_norm_parallel_residual_forward( - x0, - x1, - residual, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - residual_in_fp32=False, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes""" - hidden_size = gamma0.numel() - x0mat = x0.view((-1, hidden_size)) - x1mat = x1.view((-1, hidden_size)) if x1 is not None else None - residualmat = residual.view((-1, hidden_size)) if residual is not None else None - ( - z0mat, - z1mat, - xmat, - dmask0, - dmask1, - mu, - rsigma, - ) = dropout_layer_norm.dropout_add_ln_parallel_residual_fwd( - x0mat, - x1mat, - residualmat, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - None, - residual_in_fp32, - is_rms_norm, - ) - # dmask0 and dmask1 are None if dropout_p == 0.0 - # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype - return z0mat, z1mat, xmat if xmat is not None else x0mat, dmask0, dmask1, mu, rsigma - - -def _dropout_add_layer_norm_parallel_residual_backward( - dz0, - dz1, - dx, - x, - dmask0, - dmask1, - mu, - rsigma, - gamma0, - gamma1, - dropout_p, - has_x1, - has_residual, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes - dx == None means that it was a post-norm architecture - (x = drop(x0) + residual was not returned in the fwd). - """ - hidden_size = gamma0.numel() - xmat = x.view((-1, hidden_size)) - dz0mat = dz0.view(xmat.shape) - dz1mat = dz1.view(xmat.shape) if dz1 is not None else None - dxmat = dx.view(xmat.shape) if dx is not None else None - ( - dx0mat, - dx1mat, - dresidualmat, - dgamma0, - dbeta0, - dgamma1, - dbeta1, - *rest, - ) = dropout_layer_norm.dropout_add_ln_parallel_residual_bwd( - dz0mat, - dz1mat, - dxmat, - xmat, - dmask0, - dmask1, - mu, - rsigma, - gamma0, - gamma1, - dropout_p, - has_x1, - has_residual, - is_rms_norm, - ) - # dresidualmat is None if not has_residual - return dx0mat, dx1mat, dresidualmat, dgamma0, dbeta0, dgamma1, dbeta1 - - -class DropoutAddLayerNormFn(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x0, - residual, - gamma, - beta, - rowscale, - colscale, - dropout_p, - epsilon, - residual_in_fp32=False, - prenorm=False, - is_rms_norm=False, - return_dmask=False, - ): - x0 = maybe_align(x0.contiguous(), 16) - residual = maybe_align(residual.contiguous(), 16) if residual is not None else None - gamma = maybe_align(gamma.contiguous(), 16) - beta = maybe_align(beta.contiguous(), 16) if beta is not None else None - rowscale = maybe_align(rowscale.contiguous(), 16) if rowscale is not None else None - colscale = maybe_align(colscale.contiguous(), 16) if colscale is not None else None - zmat, xmat, dmask, mu, rsigma = _dropout_add_layer_norm_forward( - x0, - residual, - gamma, - beta, - rowscale, - colscale, - dropout_p, - epsilon, - residual_in_fp32, - is_rms_norm, - ) - # Only need to save x0 if we need to compute gradient wrt colscale - x0_saved = x0 if colscale is not None else None - ctx.save_for_backward( - xmat.view(x0.shape), x0_saved, dmask, gamma, mu, rsigma, rowscale, colscale - ) - ctx.prenorm = prenorm - ctx.dropout_p = dropout_p - ctx.has_residual = residual is not None - ctx.is_rms_norm = is_rms_norm - ctx.has_beta = beta is not None - if not return_dmask: - return ( - zmat.view(x0.shape) if not prenorm else (zmat.view(x0.shape), xmat.view(x0.shape)) - ) - else: - dmask = ( - dmask.view(x0.shape) - if dropout_p > 0.0 - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - ctx.mark_non_differentiable(dmask) - return ( - (zmat.view(x0.shape), dmask) - if not prenorm - else (zmat.view(x0.shape), xmat.view(x0.shape), dmask) - ) - - @staticmethod - def backward(ctx, dz, *args): - # assert dz.is_contiguous() - dz = maybe_align(dz.contiguous(), 16) # this happens! - dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None - x, x0, dmask, gamma, mu, rsigma, rowscale, colscale = ctx.saved_tensors - # x0 is None if colscale is None - dropout_p = ctx.dropout_p - has_residual = ctx.has_residual - dx0mat, dresidualmat, dgamma, dbeta, *rest = _dropout_add_layer_norm_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - rowscale, - colscale, - dropout_p, - has_residual, - ctx.is_rms_norm, - ) - dx0 = dx0mat.view(x.shape) - dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None - dcolscale = rest[0] if colscale is not None else None - return ( - dx0, - dresidual, - dgamma, - dbeta if ctx.has_beta else None, - None, - dcolscale, - None, - None, - None, - None, - None, - None, - ) - - -class DropoutAddLayerNormSubsetFn(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x0, - residual, - gamma, - beta, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32=False, - prenorm=False, - is_rms_norm=False, - return_dmask=False, - ): - x0 = maybe_align(x0.contiguous(), 16) - residual = maybe_align(residual.contiguous(), 16) if residual is not None else None - gamma = maybe_align(gamma.contiguous(), 16) - beta = maybe_align(beta.contiguous(), 16) if beta is not None else None - colscale = maybe_align(colscale.contiguous(), 16) if colscale is not None else None - zmat, xmat, dmask, mu, rsigma = _dropout_add_layer_norm_subset_forward( - x0, - residual, - gamma, - beta, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32, - is_rms_norm, - ) - # Only need to save x0 if we need to compute gradient wrt colscale - x0_saved = x0 if colscale is not None else None - x_shape = (-1, *x0.shape[1:]) - ctx.save_for_backward( - xmat.view(x_shape), x0_saved, dmask, gamma, mu, rsigma, colscale, x0_subset, out_subset - ) - ctx.prenorm = prenorm - ctx.dropout_p = dropout_p - ctx.rowscale_const = rowscale_const - ctx.x0_numrows = x0.shape[:-1].numel() - ctx.has_residual = residual is not None - ctx.is_rms_norm = is_rms_norm - ctx.has_beta = beta is not None - z_shape = (-1, *x0.shape[1:]) - if not return_dmask: - return zmat.view(z_shape) if not prenorm else (zmat.view(z_shape), xmat.view(x0.shape)) - else: - z = zmat.view(z_shape) - dmask = ( - dmask.view(x0.shape) - if dropout_p > 0.0 - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - ctx.mark_non_differentiable(dmask) - return (z, dmask) if not prenorm else (z, xmat.view(x_shape), dmask) - - @staticmethod - def backward(ctx, dz, *args): - # assert dz.is_contiguous() - dz = maybe_align(dz.contiguous(), 16) # this happens! - dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None - x, x0, dmask, gamma, mu, rsigma, colscale, x0_subset, out_subset = ctx.saved_tensors - # x0 is None if colscale is None - dropout_p = ctx.dropout_p - has_residual = ctx.has_residual - dx0mat, dresidualmat, dgamma, dbeta, *rest = _dropout_add_layer_norm_subset_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - colscale, - x0_subset, - out_subset, - dropout_p, - ctx.rowscale_const, - ctx.x0_numrows, - has_residual, - ctx.is_rms_norm, - ) - dx0 = dx0mat.view(-1, *x.shape[1:]) - dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None - dcolscale = rest[0] if colscale is not None else None - return ( - dx0, - dresidual, - dgamma, - dbeta if ctx.has_beta else None, - dcolscale, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) - - -class DropoutAddLayerNormParallelResidualFn(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x0, - x1, - residual, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - residual_in_fp32=False, - prenorm=False, - is_rms_norm=False, - return_dmask=False, - ): - x0 = maybe_align(x0.contiguous(), 16) - x1 = maybe_align(x1.contiguous(), 16) if x1 is not None else None - residual = maybe_align(residual.contiguous(), 16) if residual is not None else None - gamma0 = maybe_align(gamma0.contiguous(), 16) - beta0 = maybe_align(beta0.contiguous(), 16) if beta0 is not None else None - gamma1 = maybe_align(gamma1.contiguous(), 16) if gamma1 is not None else None - beta1 = maybe_align(beta1.contiguous(), 16) if beta1 is not None else None - ( - z0mat, - z1mat, - xmat, - dmask0, - dmask1, - mu, - rsigma, - ) = _dropout_add_layer_norm_parallel_residual_forward( - x0, - x1, - residual, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - residual_in_fp32, - is_rms_norm, - ) - ctx.save_for_backward(xmat.view(x0.shape), dmask0, dmask1, gamma0, gamma1, mu, rsigma) - ctx.prenorm = prenorm - ctx.dropout_p = dropout_p - ctx.has_x1 = x1 is not None - ctx.has_residual = residual is not None - ctx.is_rms_norm = is_rms_norm - ctx.has_beta = beta0 is not None - z = (z0mat.view(x0.shape), z1mat.view(x0.shape) if z1mat is not None else None) - if not return_dmask: - return z if not prenorm else (*z, xmat.view(x0.shape)) - else: - dmask0 = ( - dmask0.view(x0.shape) - if dropout_p > 0.0 - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - dmask1 = ( - dmask1.view(x0.shape) - if dropout_p > 0.0 and x1 is not None - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - ctx.mark_non_differentiable(dmask0) - ctx.mark_non_differentiable(dmask1) - return ( - (*z, dmask0, dmask1) if not prenorm else (*z, xmat.view(x0.shape), dmask0, dmask1) - ) - - @staticmethod - def backward(ctx, dz0, dz1, *args): - dz0 = maybe_align(dz0.contiguous(), 16) # this happens! - dz1 = maybe_align(dz1.contiguous(), 16) if dz1 is not None else None - dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None - x, dmask0, dmask1, gamma0, gamma1, mu, rsigma = ctx.saved_tensors - dropout_p = ctx.dropout_p - has_x1 = ctx.has_x1 - has_residual = ctx.has_residual - ( - dx0mat, - dx1mat, - dresidualmat, - dgamma0, - dbeta0, - dgamma1, - dbeta1, - ) = _dropout_add_layer_norm_parallel_residual_backward( - dz0, - dz1, - dx, - x, - dmask0, - dmask1, - mu, - rsigma, - gamma0, - gamma1, - dropout_p, - has_x1, - has_residual, - ctx.is_rms_norm, - ) - dx0 = dx0mat.view(x.shape) - dx1 = dx1mat.view(x.shape) if dx1mat is not None else None - dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None - return ( - dx0, - dx1, - dresidual, - dgamma0, - dbeta0 if ctx.has_beta else None, - dgamma1, - dbeta1 if ctx.has_beta else None, - None, - None, - None, - None, - None, - None, - ) - - -def layer_norm(x, weight, bias, epsilon): - return DropoutAddLayerNormFn.apply(x, None, weight, bias, None, None, 0.0, epsilon, False) - - -def dropout_add_layer_norm( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - rowscale=None, - layerscale=None, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormFn.apply( - x0, - residual, - weight, - bias, - rowscale, - layerscale, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - False, - return_dropout_mask, - ) - - -def dropout_add_layer_norm_subset( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - layerscale=None, - x0_subset=None, - out_subset=None, - rowscale_const=1.0, - out_numrows=0, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormSubsetFn.apply( - x0, - residual, - weight, - bias, - layerscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32, - prenorm, - False, - return_dropout_mask, - ) - - -def dropout_add_layer_norm_parallel_residual( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormParallelResidualFn.apply( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - False, - return_dropout_mask, - ) - - -class DropoutAddLayerNorm(torch.nn.Module): - def __init__( - self, - hidden_size, - prenorm=False, - p=0.0, - eps=1e-5, - residual_in_fp32=False, - device=None, - dtype=None, - ): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.prenorm = prenorm - self.p = p - self.eps = eps - self.residual_in_fp32 = residual_in_fp32 - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.bias = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.reset_parameters() - - def reset_parameters(self): - init.ones_(self.weight) - init.zeros_(self.bias) - - def forward(self, x0, residual=None): - return dropout_add_layer_norm( - x0, - residual, - self.weight, - self.bias, - self.p if self.training else 0.0, - self.eps, - prenorm=self.prenorm, - residual_in_fp32=self.residual_in_fp32, - ) diff --git a/build/torch212-cxx11-cu130-x86_64-linux/ops/rms_norm.py b/build/torch212-cxx11-cu130-x86_64-linux/ops/rms_norm.py deleted file mode 100644 index 068348d61290e3839dd082b540d898578ba1e8e2..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu130-x86_64-linux/ops/rms_norm.py +++ /dev/null @@ -1,174 +0,0 @@ -# Copyright (c) 2022, Tri Dao. -# Adapted from https://github.com/NVIDIA/apex/blob/master/apex/contrib/layer_norm/layer_norm.py - -import torch -from torch.nn import init - -from flash_attn.ops.layer_norm import ( - DropoutAddLayerNormFn, - DropoutAddLayerNormParallelResidualFn, - DropoutAddLayerNormSubsetFn, -) - - -def rms_norm(x, weight, epsilon): - return DropoutAddLayerNormFn.apply( - x, None, weight, None, None, None, 0.0, epsilon, False, False, True - ) - - -def dropout_add_rms_norm( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - rowscale=None, - layerscale=None, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormFn.apply( - x0, - residual, - weight, - bias, - rowscale, - layerscale, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - True, - return_dropout_mask, - ) - - -def dropout_add_rms_norm_subset( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - layerscale=None, - x0_subset=None, - out_subset=None, - rowscale_const=1.0, - out_numrows=0, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormSubsetFn.apply( - x0, - residual, - weight, - bias, - layerscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32, - prenorm, - True, - return_dropout_mask, - ) - - -def dropout_add_rms_norm_parallel_residual( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormParallelResidualFn.apply( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - True, - return_dropout_mask, - ) - - -class RMSNorm(torch.nn.Module): - def __init__(self, hidden_size, eps=1e-5, device=None, dtype=None): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.eps = eps - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.register_parameter("bias", None) - self.reset_parameters() - - def reset_parameters(self): - init.ones_(self.weight) - - def forward(self, x): - return rms_norm(x, self.weight, self.eps) - - -class DropoutAddRMSNorm(torch.nn.Module): - def __init__( - self, - hidden_size, - prenorm=False, - p=0.0, - eps=1e-5, - residual_in_fp32=False, - device=None, - dtype=None, - ): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.prenorm = prenorm - self.p = p - self.eps = eps - self.residual_in_fp32 = residual_in_fp32 - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.register_parameter("bias", None) - self.reset_parameters() - - def reset_parameters(self): - init.ones_(self.weight) - - def forward(self, x0, residual=None): - return dropout_add_rms_norm( - x0, - residual, - self.weight, - None, - self.p if self.training else 0.0, - self.eps, - prenorm=self.prenorm, - residual_in_fp32=self.residual_in_fp32, - ) diff --git a/build/torch212-cxx11-cu130-x86_64-linux/ops/triton/__init__.py b/build/torch212-cxx11-cu130-x86_64-linux/ops/triton/__init__.py deleted file mode 100644 index 8b137891791fe96927ad78e64b0aad7bded08bdc..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu130-x86_64-linux/ops/triton/__init__.py +++ /dev/null @@ -1 +0,0 @@ - diff --git a/build/torch212-cxx11-cu130-x86_64-linux/ops/triton/cross_entropy.py b/build/torch212-cxx11-cu130-x86_64-linux/ops/triton/cross_entropy.py deleted file mode 100644 index 1b5a415b73f236f3e05fb14b9141959559e18526..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu130-x86_64-linux/ops/triton/cross_entropy.py +++ /dev/null @@ -1,330 +0,0 @@ -# Copyright (c) 2023, Tri Dao. - -from typing import Tuple, Optional, Union - -import torch -import torch.nn.functional as F - -import triton -import triton.language as tl - -# `all_gather_into_tensor` and `reduce_scatter_tensor` are new placeholders for -# `_all_gather_base` and `_reduce_scatter_base`. They require the most recent -# version of PyTorch. The following 2 lines are for backward compatibility with -# older PyTorch. -if "all_gather_into_tensor" not in dir(torch.distributed): - torch.distributed.all_gather_into_tensor = torch.distributed._all_gather_base - - -@triton.heuristics( - { - "HAS_SMOOTHING": lambda args: args["smoothing"] > 0.0, - } -) -@triton.jit -def cross_entropy_fwd_kernel( - loss_ptr, # data ptrs - lse_ptr, - z_loss_ptr, - logits_ptr, - labels_ptr, - smoothing, - logit_scale, - lse_square_scale, - ignore_index, - total_classes, - class_start_idx, # Useful for tensor parallel when each rank only has a subset of classes - n_cols, # shapes - logits_row_stride, # strides - BLOCK_SIZE: tl.constexpr, - HAS_SMOOTHING: tl.constexpr, - # if SPLIT (e.g. tensor parallel), don't include the LSE in the loss since it's not the final LSE - SPLIT: tl.constexpr, - PRECOMPUTED_LSE: tl.constexpr, # If LSE is already computed (also no smoothing and logit_scale == 1.0) -): - row_idx = tl.program_id(0) - logits_ptr = logits_ptr + row_idx * logits_row_stride.to(tl.int64) - sum_logits = 0.0 # For smoothing - if not PRECOMPUTED_LSE: - # Statistics for online softmax - m_i = -float("inf") - l_i = 0.0 - for col_offset in range(0, n_cols, BLOCK_SIZE): - cols = col_offset + tl.arange(0, BLOCK_SIZE) - logits = tl.load(logits_ptr + cols, mask=cols < n_cols, other=-float("inf")).to( - tl.float32 - ) * logit_scale - if HAS_SMOOTHING: - sum_logits += tl.sum(tl.where(cols < n_cols, logits, 0.0)) - m_i_new = tl.maximum(m_i, tl.max(logits)) - l_i = tl.exp(m_i - m_i_new) * l_i + tl.sum(tl.exp(logits - m_i_new)) - m_i = m_i_new - lse = tl.log(l_i) + m_i - tl.store(lse_ptr + row_idx, lse) - else: - lse = tl.load(lse_ptr + row_idx) - label_idx = tl.load(labels_ptr + row_idx) - if label_idx == ignore_index: - loss = 0.0 - z_loss = 0.0 - else: - label_idx -= class_start_idx - if label_idx >= 0 and label_idx < n_cols: - logits_label = tl.load(logits_ptr + label_idx) * logit_scale - if HAS_SMOOTHING: - loss = ( - (lse if not SPLIT else 0.0) - - smoothing * sum_logits / total_classes - - (1 - smoothing) * logits_label - ) - else: - loss = (lse if not SPLIT else 0.0) - logits_label - else: - # If label is out of bounds, we set the CE loss to 0.0. But we still want the smoothing loss - if HAS_SMOOTHING: - loss = smoothing * ((lse if not SPLIT else 0.0) - sum_logits / total_classes) - else: - loss = 0.0 - if not SPLIT: - z_loss = lse_square_scale * lse * lse - loss += z_loss - else: - z_loss = 0.0 - tl.store(loss_ptr + row_idx, loss) - if not SPLIT: - tl.store(z_loss_ptr + row_idx, z_loss) - - -@triton.heuristics( - { - "HAS_SMOOTHING": lambda args: args["smoothing"] > 0.0, - } -) -@triton.jit -def cross_entropy_bwd_kernel( - dlogits_ptr, # data ptrs - dloss_ptr, - logits_ptr, - lse_ptr, - labels_ptr, - smoothing, - logit_scale, - lse_square_scale, - ignore_index, - total_classes, - class_start_idx, # Useful for tensor parallel when each rank only has a subset of classes - n_cols, # shapes - logits_row_stride, # strides - dlogits_row_stride, - dloss_row_stride, - BLOCK_SIZE: tl.constexpr, - HAS_SMOOTHING: tl.constexpr, -): - row_idx = tl.program_id(0) - col_block_idx = tl.program_id(1) - logits_ptr = logits_ptr + row_idx * logits_row_stride.to(tl.int64) - dlogits_ptr = dlogits_ptr + row_idx * dlogits_row_stride.to(tl.int64) - col_offsets = col_block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - label_idx = tl.load(labels_ptr + row_idx) - if label_idx != ignore_index: - dloss = tl.load(dloss_ptr + row_idx * dloss_row_stride) - else: - dloss = 0.0 - logits = tl.load(logits_ptr + col_offsets, mask=col_offsets < n_cols, other=-float("inf")).to( - tl.float32 - ) * logit_scale - lse = tl.load(lse_ptr + row_idx) - probs = tl.exp(logits - lse) - probs += 2.0 * lse_square_scale * lse * probs - label_idx -= class_start_idx - if HAS_SMOOTHING: - smooth_positive = 1.0 - smoothing - smooth_negative = smoothing / total_classes - probs = tl.where(col_offsets == label_idx, probs - smooth_positive, probs) - smooth_negative - else: - probs = tl.where(col_offsets == label_idx, probs - 1.0, probs) - tl.store(dlogits_ptr + col_offsets, (dloss * logit_scale) * probs, mask=col_offsets < n_cols) - - -class CrossEntropyLoss(torch.autograd.Function): - - @staticmethod - def forward( - ctx, - logits, - labels, - precomputed_lse=None, - smoothing=0.0, - logit_scale=1.0, - lse_square_scale=0.0, - ignore_index=-100, - inplace_backward=False, - process_group=None, - ): - # For some reason Triton generates wrong code when labels has dtype long and its address - # is not aligned to 16 bytes. The ld.global.b64 seems to load the wrong label index. - if labels.dtype == torch.long and labels.data_ptr() % 16 != 0: - labels = F.pad(labels, (0, 1))[..., :-1] - assert labels.data_ptr() % 16 == 0 - assert logit_scale > 0.0 - n_rows, n_cols = logits.shape - assert labels.shape == (n_rows,) - world_size = 1 if process_group is None else torch.distributed.get_world_size(process_group) - total_classes = world_size * n_cols - rank = 0 if process_group is None else torch.distributed.get_rank(process_group) - class_start_idx = rank * n_cols - use_precomputed_lse = precomputed_lse is not None and logit_scale == 1.0 and smoothing == 0.0 - - if logits.stride(-1) != 1: - logits = logits.contiguous() - MAX_BLOCK_SIZE = 16 * 1024 - BLOCK_SIZE = min(triton.next_power_of_2(n_cols), MAX_BLOCK_SIZE) - num_warps = ( - 4 - if BLOCK_SIZE < 2048 - else (8 if BLOCK_SIZE < 8192 else (16 if BLOCK_SIZE < 128 * 1024 else 32)) - ) - losses = torch.empty(n_rows, dtype=torch.float, device=logits.device) - if use_precomputed_lse: - assert precomputed_lse.shape == (n_rows,) - lse = precomputed_lse.contiguous() - else: - lse = torch.empty(n_rows, dtype=torch.float, device=logits.device) - z_losses = torch.empty(n_rows, dtype=torch.float, device=logits.device) - # Need this, otherwise Triton tries to launch from cuda:0 and we get - # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?) - with torch.cuda.device(logits.device.index): - cross_entropy_fwd_kernel[(n_rows,)]( - losses, # data ptrs - lse, - z_losses, - logits, - labels, - smoothing, - logit_scale, - lse_square_scale, - ignore_index, - total_classes, - class_start_idx, - n_cols, # shapes - logits.stride(0), # strides - BLOCK_SIZE=BLOCK_SIZE, # constants - SPLIT=world_size > 1, - PRECOMPUTED_LSE=use_precomputed_lse, - num_warps=num_warps, - ) - - if world_size > 1: - # If there's no smoothing, if labels are in the vocab of this partition, losses contains - # - predicted logit, and 0 otherwise. - # If there's smoothing=0.1, for labels in the vocab of this partition, losses contains - # -0.9 * predicted logit - 0.1 * sum logit / total_classes. - # For labels not in the vocab of this partition, losses contains - # -0.1 * sum logit / total_classes. - if world_size > 1: - lse_allgather = torch.empty(world_size, n_rows, dtype=lse.dtype, device=lse.device) - torch.distributed.all_gather_into_tensor(lse_allgather, lse, group=process_group) - handle_losses = torch.distributed.all_reduce( - losses, op=torch.distributed.ReduceOp.SUM, group=process_group, async_op=True - ) - lse = torch.logsumexp(lse_allgather, dim=0) - handle_losses.wait() - # After the allreduce, if there's no smoothing, the total losses are - predicted_logit, - # we just have to add the (global) lse. - # If there's smoothing=0.1, the total losses are - # -0.9 * predicted_logit - 0.1 * sum logit / total_classes. - # Again, we just have to add the (global) lse. - losses += lse - if lse_square_scale != 0.0: - z_losses = lse_square_scale * lse.square() - z_losses.masked_fill_(labels == ignore_index, 0.0) - losses += z_losses - else: - z_losses = torch.zeros_like(losses) - losses.masked_fill_(labels == ignore_index, 0.0) - - ctx.save_for_backward(logits, lse, labels) - ctx.mark_non_differentiable(z_losses) - ctx.smoothing = smoothing - ctx.logit_scale = logit_scale - ctx.lse_square_scale = lse_square_scale - ctx.ignore_index = ignore_index - ctx.total_classes = total_classes - ctx.class_start_idx = class_start_idx - ctx.inplace_backward = inplace_backward - return losses, z_losses - - @staticmethod - def backward(ctx, grad_losses, grad_z_losses): - del grad_z_losses # z_losses are only for logging. - - logits, lse, labels = ctx.saved_tensors - dlogits = logits if ctx.inplace_backward else torch.empty_like(logits) - n_rows, n_cols = logits.shape - BLOCK_SIZE = min(triton.next_power_of_2(n_cols), 4 * 1024) - num_warps = 4 if BLOCK_SIZE < 2048 else (8 if BLOCK_SIZE < 8192 else 16) - grid = lambda META: (n_rows, triton.cdiv(n_cols, META["BLOCK_SIZE"])) # noqa - # Need this, otherwise Triton tries to launch from cuda:0 and we get - # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?) - with torch.cuda.device(logits.device.index): - cross_entropy_bwd_kernel[grid]( - dlogits, # data ptrs - grad_losses, - logits, - lse, - labels, - ctx.smoothing, - ctx.logit_scale, - ctx.lse_square_scale, - ctx.ignore_index, - ctx.total_classes, - ctx.class_start_idx, - n_cols, # shapes - logits.stride(0), # strides - dlogits.stride(0), - grad_losses.stride(0), - BLOCK_SIZE=BLOCK_SIZE, # constants - num_warps=num_warps, - ) - return dlogits, None, None, None, None, None, None, None, None, None - - -def cross_entropy_loss( - logits: torch.Tensor, - labels: torch.Tensor, - precomputed_lse: Optional[torch.Tensor] = None, - label_smoothing: float = 0.0, - logit_scale: float = 1.0, - lse_square_scale: float = 0.0, - ignore_index=-100, - inplace_backward: bool = False, - process_group=None, -) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Arguments: - logits: (batch, vocab_size) - labels: (batch,) - label_smoothing: float - logit_scale: float. Multiply logits by this scale before calculating the loss. - lse_square_scale: float. If > 0, we add lse_square_scale * lse(logits) ^ 2 to the loss. - This is also referred to as "z-loss". - ignore_index: int. If labels == ignore_index, the loss is set to 0.0. - inplace_backward: bool. If True, we do the backward pass in-place by modifying the logits. - This saves memory. - process_group: if not None, we're doing Tensor Parallel: each process is responsible for - one part of the vocab. The loss will be aggregated across processes. - Returns: - losses: (batch,), float - z_losses: (batch,), float - """ - return CrossEntropyLoss.apply( - logits, - labels, - precomputed_lse, - label_smoothing, - logit_scale, - lse_square_scale, - ignore_index, - inplace_backward, - process_group, - ) diff --git a/build/torch212-cxx11-cu130-x86_64-linux/ops/triton/k_activations.py b/build/torch212-cxx11-cu130-x86_64-linux/ops/triton/k_activations.py deleted file mode 100644 index efb83c358eb4a85d069ee340a3c83f418f9a805b..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu130-x86_64-linux/ops/triton/k_activations.py +++ /dev/null @@ -1,162 +0,0 @@ -# Adapted from https://github.com/facebookresearch/xformers/blob/main/xformers/triton/k_activations.py -# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved. -# -# This source code is licensed under the BSD license found in the -# LICENSE file in the root directory of this source tree. - -import math -from enum import Enum -from typing import Optional - -import triton -import triton.language as tl - -_sqrt2pi = math.sqrt(2.0 / math.pi) -_sqrt1_2 = math.sqrt(1.0 / 2) -_gaussian_pdf_normalization = 1.0 / math.sqrt(2 * math.pi) - - -class Activation(str, Enum): - SquaredReLU = "squared_relu" - GeLU = "gelu" - GeLUApprox = "gelu_approx" - LeakyReLU = "leaky_relu" - ReLU = "relu" - - -def get_triton_activation_kernel(activation: Optional[Activation]): - return ( - { - Activation.ReLU: relu, - Activation.LeakyReLU: leaky_relu, - Activation.GeLU: gelu, - Activation.GeLUApprox: gelu_approx, - Activation.SquaredReLU: squared_relu, - }[activation] - if activation - else None - ) - - -def get_triton_activation_bwd_kernel(activation: Optional[Activation]): - return ( - { - Activation.ReLU: relu_grad, - Activation.LeakyReLU: leaky_relu_grad, - Activation.GeLU: gelu_grad, - Activation.GeLUApprox: gelu_approx_grad, - Activation.SquaredReLU: squared_relu_grad, - }[activation] - if activation - else None - ) - - -@triton.jit -def tanh(x): - # Tanh is just a scaled sigmoid - return 2 * tl.sigmoid(2 * x) - 1 - - -@triton.jit -def cosh(x): - exp_x = tl.exp(x) - return (exp_x + 1.0 / exp_x) * 0.5 - - -# a Triton implementation of the most used activations -# See for instance http://arxiv.org/abs/1606.08415 for an overview - -# ReLU -@triton.jit -def relu(x): - """ - ReLU_ activation function - - .. _ReLU: https://pytorch.org/docs/stable/generated/torch.nn.ReLU.html - """ - zero = 0.0 - return tl.where(x >= 0, x, zero.to(x.dtype)) - - -@triton.jit -def relu_grad(x): - # ReLU is different from other activations - # in that it does not require the input to retrospectively compute its gradient - # here the input is the downstream gradient, and we return the upstream gradient directly - zero = 0.0 - one = 1.0 - return tl.where(x >= 0, one.to(x.dtype), zero.to(x.dtype)) - - -@triton.jit -def squared_relu(x): - """ - Squared ReLU activation, as proposed in the Primer_ paper. - - .. _Primer: https://arxiv.org/abs/2109.08668 - """ - x_ = relu(x) - return (x_ * x_).to(x.dtype) - - -@triton.jit -def squared_relu_grad(x): - return tl.where(x >= 0, 2.0 * x, 0.0) - - -# Leaky ReLU -@triton.jit -def leaky_relu(x): - """ - LeakyReLU_ activation - - .. _LeakyReLU: https://pytorch.org/docs/stable/generated/torch.nn.LeakyReLU.html - """ - scale = 0.01 + 0.0 - scale = scale.to(x.dtype) - return tl.where(x >= 0, x, scale * x) - - -@triton.jit -def leaky_relu_grad(x): - min_grad = 0.01 - max_grad = 1 - - min_grad = min_grad.to(x.dtype) - max_grad = max_grad.to(x.dtype) - - return tl.where(x >= 0, max_grad, min_grad) - - -@triton.jit -def gelu(x): - """Gaussian Error Linear Unit (GELU)""" - return x * 0.5 * (1.0 + tl.libdevice.erf(x * _sqrt1_2)) - - -@triton.jit -def gelu_grad(x): - cdf = 0.5 * (1.0 + tl.libdevice.erf(x * _sqrt1_2)) - pdf = tl.exp(-0.5 * x * x) * _gaussian_pdf_normalization - return cdf + x * pdf - - -@triton.jit -def gelu_approx(x): - """ - GeLU_ activation - Gaussian error linear unit, with tanh approximation - - .. _GeLU: https://arxiv.org/pdf/1606.08415.pdf - """ - return 0.5 * x * (1.0 + tanh(_sqrt2pi * x * (1.0 + 0.044715 * x * x))) - - -@triton.jit -def gelu_approx_grad(x): - # CREDITS: Fast implementation proposed in - # https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/fused_bias_gelu.py#L30 - tanh_out = tanh(0.79788456 * x * (1 + 0.044715 * x * x)) - return 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( - 1 + tanh_out - ) diff --git a/build/torch212-cxx11-cu130-x86_64-linux/ops/triton/layer_norm.py b/build/torch212-cxx11-cu130-x86_64-linux/ops/triton/layer_norm.py deleted file mode 100644 index 5e6f0b3697b612e6a19031bbc55ac76ad233dc85..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu130-x86_64-linux/ops/triton/layer_norm.py +++ /dev/null @@ -1,1254 +0,0 @@ -# Copyright (c) 2024, Tri Dao. -# Implement dropout + residual + layer_norm / rms_norm. - -# Based on the Triton LayerNorm tutorial: https://triton-lang.org/main/getting-started/tutorials/05-layer-norm.html -# For the backward pass, we keep weight_grad and bias_grad in registers and accumulate. -# This is faster for dimensions up to 8k, but after that it's much slower due to register spilling. -# The models we train have hidden dim up to 8k anyway (e.g. Llama 70B), so this is fine. - -import math -from typing import Optional, List - -import torch -import torch.nn.functional as F -from torch import Tensor - -import triton -import triton.language as tl - -from flash_attn.utils.torch import custom_fwd, custom_bwd -from flash_attn.utils.library import triton_op - -from ..._ops import add_op_namespace_prefix - - -def maybe_contiguous_lastdim(x): - return x.contiguous() if x is not None and x.stride(-1) != 1 else x - - -def maybe_contiguous(x): - return x.contiguous() if x is not None else None - - -def triton_autotune_configs(): - # Return configs with a valid warp count for the current device - configs = [] - # Maximum threads per block is architecture-dependent in theory, but in reality all are 1024 - max_threads_per_block = 1024 - # Default to warp size 32 if not defined by device - warp_size = getattr(torch.cuda.get_device_properties(torch.cuda.current_device()), "warp_size", 32) - # Autotune for warp counts which are powers of 2 and do not exceed thread per block limit - return [triton.Config({}, num_warps=warp_count) for warp_count in [1, 2, 4, 8, 16, 32] - if warp_count * warp_size <= max_threads_per_block] - # return [triton.Config({}, num_warps=8)] - - -def layer_norm_ref( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - zero_centered_weight=False, - dropout_mask=None, - dropout_mask1=None, - upcast=False, -): - dtype = x.dtype - if upcast: - x = x.float() - weight = weight.float() - bias = bias.float() if bias is not None else None - residual = residual.float() if residual is not None else residual - x1 = x1.float() if x1 is not None else None - weight1 = weight1.float() if weight1 is not None else None - bias1 = bias1.float() if bias1 is not None else None - if zero_centered_weight: - weight = weight + 1.0 - if weight1 is not None: - weight1 = weight1 + 1.0 - if x1 is not None: - assert rowscale is None, "rowscale is not supported with parallel LayerNorm" - if rowscale is not None: - x = x * rowscale[..., None] - if dropout_p > 0.0: - if dropout_mask is not None: - x = x.masked_fill(~dropout_mask, 0.0) / (1.0 - dropout_p) - else: - x = F.dropout(x, p=dropout_p) - if x1 is not None: - if dropout_mask1 is not None: - x1 = x1.masked_fill(~dropout_mask1, 0.0) / (1.0 - dropout_p) - else: - x1 = F.dropout(x1, p=dropout_p) - if x1 is not None: - x = x + x1 - if residual is not None: - x = (x + residual).to(x.dtype) - out = F.layer_norm(x.to(weight.dtype), x.shape[-1:], weight=weight, bias=bias, eps=eps).to( - dtype - ) - if weight1 is None: - return out if not prenorm else (out, x) - else: - out1 = F.layer_norm( - x.to(weight1.dtype), x.shape[-1:], weight=weight1, bias=bias1, eps=eps - ).to(dtype) - return (out, out1) if not prenorm else (out, out1, x) - - -def rms_norm_ref( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - zero_centered_weight=False, - dropout_mask=None, - dropout_mask1=None, - upcast=False, -): - dtype = x.dtype - if upcast: - x = x.float() - weight = weight.float() - bias = bias.float() if bias is not None else None - residual = residual.float() if residual is not None else residual - x1 = x1.float() if x1 is not None else None - weight1 = weight1.float() if weight1 is not None else None - bias1 = bias1.float() if bias1 is not None else None - if zero_centered_weight: - weight = weight + 1.0 - if weight1 is not None: - weight1 = weight1 + 1.0 - if x1 is not None: - assert rowscale is None, "rowscale is not supported with parallel LayerNorm" - if rowscale is not None: - x = x * rowscale[..., None] - if dropout_p > 0.0: - if dropout_mask is not None: - x = x.masked_fill(~dropout_mask, 0.0) / (1.0 - dropout_p) - else: - x = F.dropout(x, p=dropout_p) - if x1 is not None: - if dropout_mask1 is not None: - x1 = x1.masked_fill(~dropout_mask1, 0.0) / (1.0 - dropout_p) - else: - x1 = F.dropout(x1, p=dropout_p) - if x1 is not None: - x = x + x1 - if residual is not None: - x = (x + residual).to(x.dtype) - rstd = 1 / torch.sqrt((x.square()).mean(dim=-1, keepdim=True) + eps) - out = ((x * rstd * weight) + bias if bias is not None else (x * rstd * weight)).to(dtype) - if weight1 is None: - return out if not prenorm else (out, x) - else: - out1 = ((x * rstd * weight1) + bias1 if bias1 is not None else (x * rstd * weight1)).to( - dtype - ) - return (out, out1) if not prenorm else (out, out1, x) - - -@triton.autotune( - configs=triton_autotune_configs(), - key=["N", "HAS_RESIDUAL", "STORE_RESIDUAL_OUT", "IS_RMS_NORM", "HAS_BIAS", "HAS_X1", "HAS_W1", "HAS_B1"], -) -# torch compile doesn't like triton.heuristics, so we set these manually when calling the kernel -# @triton.heuristics({"HAS_BIAS": lambda args: args["B"] is not None}) -# @triton.heuristics({"HAS_RESIDUAL": lambda args: args["RESIDUAL"] is not None}) -# @triton.heuristics({"HAS_X1": lambda args: args["X1"] is not None}) -# @triton.heuristics({"HAS_W1": lambda args: args["W1"] is not None}) -# @triton.heuristics({"HAS_B1": lambda args: args["B1"] is not None}) -@triton.jit -def _layer_norm_fwd_1pass_kernel( - X, # pointer to the input - Y, # pointer to the output - W, # pointer to the weights - B, # pointer to the biases - RESIDUAL, # pointer to the residual - X1, - W1, - B1, - Y1, - RESIDUAL_OUT, # pointer to the residual - ROWSCALE, - SEEDS, # Dropout seeds for each row - DROPOUT_MASK, - DROPOUT_MASK1, - Mean, # pointer to the mean - Rstd, # pointer to the 1/std - stride_x_row, # how much to increase the pointer when moving by 1 row - stride_y_row, - stride_res_row, - stride_res_out_row, - stride_x1_row, - stride_y1_row, - M, # number of rows in X - N, # number of columns in X - eps, # epsilon to avoid division by zero - dropout_p, # Dropout probability - zero_centered_weight, # If true, add 1.0 to the weight - IS_RMS_NORM: tl.constexpr, - BLOCK_N: tl.constexpr, - HAS_RESIDUAL: tl.constexpr, - STORE_RESIDUAL_OUT: tl.constexpr, - HAS_BIAS: tl.constexpr, - HAS_DROPOUT: tl.constexpr, - STORE_DROPOUT_MASK: tl.constexpr, - HAS_ROWSCALE: tl.constexpr, - HAS_X1: tl.constexpr, - HAS_W1: tl.constexpr, - HAS_B1: tl.constexpr, -): - # Map the program id to the row of X and Y it should compute. - row = tl.program_id(0) - X += row * stride_x_row - Y += row * stride_y_row - if HAS_RESIDUAL: - RESIDUAL += row * stride_res_row - if STORE_RESIDUAL_OUT: - RESIDUAL_OUT += row * stride_res_out_row - if HAS_X1: - X1 += row * stride_x1_row - if HAS_W1: - Y1 += row * stride_y1_row - # Compute mean and variance - cols = tl.arange(0, BLOCK_N) - x = tl.load(X + cols, mask=cols < N, other=0.0).to(tl.float32) - if HAS_ROWSCALE: - rowscale = tl.load(ROWSCALE + row).to(tl.float32) - x *= rowscale - if HAS_DROPOUT: - # Compute dropout mask - # 7 rounds is good enough, and reduces register pressure - keep_mask = tl.rand(tl.load(SEEDS + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - x = tl.where(keep_mask, x / (1.0 - dropout_p), 0.0) - if STORE_DROPOUT_MASK: - tl.store(DROPOUT_MASK + row * N + cols, keep_mask, mask=cols < N) - if HAS_X1: - x1 = tl.load(X1 + cols, mask=cols < N, other=0.0).to(tl.float32) - if HAS_ROWSCALE: - rowscale = tl.load(ROWSCALE + M + row).to(tl.float32) - x1 *= rowscale - if HAS_DROPOUT: - # Compute dropout mask - # 7 rounds is good enough, and reduces register pressure - keep_mask = ( - tl.rand(tl.load(SEEDS + M + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - ) - x1 = tl.where(keep_mask, x1 / (1.0 - dropout_p), 0.0) - if STORE_DROPOUT_MASK: - tl.store(DROPOUT_MASK1 + row * N + cols, keep_mask, mask=cols < N) - x += x1 - if HAS_RESIDUAL: - residual = tl.load(RESIDUAL + cols, mask=cols < N, other=0.0).to(tl.float32) - x += residual - if STORE_RESIDUAL_OUT: - tl.store(RESIDUAL_OUT + cols, x, mask=cols < N) - if not IS_RMS_NORM: - mean = tl.sum(x, axis=0) / N - tl.store(Mean + row, mean) - xbar = tl.where(cols < N, x - mean, 0.0) - var = tl.sum(xbar * xbar, axis=0) / N - else: - xbar = tl.where(cols < N, x, 0.0) - var = tl.sum(xbar * xbar, axis=0) / N - rstd = 1 / tl.sqrt(var + eps) - tl.store(Rstd + row, rstd) - # Normalize and apply linear transformation - mask = cols < N - w = tl.load(W + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w += 1.0 - if HAS_BIAS: - b = tl.load(B + cols, mask=mask).to(tl.float32) - x_hat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd - y = x_hat * w + b if HAS_BIAS else x_hat * w - # Write output - tl.store(Y + cols, y, mask=mask) - if HAS_W1: - w1 = tl.load(W1 + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w1 += 1.0 - if HAS_B1: - b1 = tl.load(B1 + cols, mask=mask).to(tl.float32) - y1 = x_hat * w1 + b1 if HAS_B1 else x_hat * w1 - tl.store(Y1 + cols, y1, mask=mask) - - -def _layer_norm_fwd( - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - residual: Optional[Tensor] = None, - x1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - out_dtype: Optional[torch.dtype] = None, - residual_dtype: Optional[torch.dtype] = None, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - return_dropout_mask: bool = False, - out: Optional[Tensor] = None, - residual_out: Optional[Tensor] = None -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - # Need to wrap to handle the case where residual_out is a alias of x, which makes torch.library - # and torch.compile unhappy. Also allocate memory for out and residual_out if they are None - # so that _layer_norm_fwd_impl doesn't have to return them. - if out is None: - out = torch.empty_like(x, dtype=x.dtype if out_dtype is None else out_dtype) - if residual is not None: - residual_dtype = residual.dtype - if residual_out is None and ( - residual is not None - or (residual_dtype is not None and residual_dtype != x.dtype) - or dropout_p > 0.0 - or rowscale is not None - or x1 is not None - ): - residual_out = torch.empty_like( - x, dtype=residual_dtype if residual_dtype is not None else x.dtype - ) - else: - residual_out = None - y1, mean, rstd, seeds, dropout_mask, dropout_mask1 = _layer_norm_fwd_impl( - x, - weight, - bias, - eps, - out, - residual=residual, - x1=x1, - weight1=weight1, - bias1=bias1, - dropout_p=dropout_p, - rowscale=rowscale, - zero_centered_weight=zero_centered_weight, - is_rms_norm=is_rms_norm, - return_dropout_mask=return_dropout_mask, - residual_out=residual_out, - ) - # residual_out is None if residual is None and residual_dtype == input_dtype and dropout_p == 0.0 - if residual_out is None: - residual_out = x - return out, y1, mean, rstd, residual_out, seeds, dropout_mask, dropout_mask1 - - -# [2025-04-28] torch.library.triton_op ignores the schema argument, but here we need the schema -# since we're returning a tuple of tensors -@triton_op(add_op_namespace_prefix("layer_norm_fwd_impl"), mutates_args={"out", "residual_out"}, - schema="(Tensor x, Tensor weight, Tensor bias, float eps, Tensor(a!) out, Tensor? residual, Tensor? x1, Tensor? weight1, Tensor? bias1, float dropout_p, Tensor? rowscale, bool zero_centered_weight, bool is_rms_norm, bool return_dropout_mask, Tensor(a!)? residual_out) -> (Tensor y1, Tensor mean, Tensor rstd, Tensor seeds, Tensor dropout_mask, Tensor dropout_mask1)") -def _layer_norm_fwd_impl( - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - out: Tensor, - residual: Optional[Tensor] = None, - x1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - return_dropout_mask: bool = False, - residual_out: Optional[Tensor] = None -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - M, N = x.shape - assert x.stride(-1) == 1 - if residual is not None: - assert residual.stride(-1) == 1 - assert residual.shape == (M, N) - assert weight.shape == (N,) - assert weight.stride(-1) == 1 - if bias is not None: - assert bias.stride(-1) == 1 - assert bias.shape == (N,) - if x1 is not None: - assert x1.shape == x.shape - assert rowscale is None - assert x1.stride(-1) == 1 - if weight1 is not None: - assert weight1.shape == (N,) - assert weight1.stride(-1) == 1 - if bias1 is not None: - assert bias1.shape == (N,) - assert bias1.stride(-1) == 1 - if rowscale is not None: - assert rowscale.is_contiguous() - assert rowscale.shape == (M,) - assert out.shape == x.shape - assert out.stride(-1) == 1 - if residual_out is not None: - assert residual_out.shape == x.shape - assert residual_out.stride(-1) == 1 - if weight1 is not None: - y1 = torch.empty_like(out) - assert y1.stride(-1) == 1 - else: - y1 = None - mean = torch.empty((M,), dtype=torch.float32, device=x.device) if not is_rms_norm else None - rstd = torch.empty((M,), dtype=torch.float32, device=x.device) - if dropout_p > 0.0: - seeds = torch.randint( - 2**32, (M if x1 is None else 2 * M,), device=x.device, dtype=torch.int64 - ) - else: - seeds = None - if return_dropout_mask and dropout_p > 0.0: - dropout_mask = torch.empty(M, N, device=x.device, dtype=torch.bool) - if x1 is not None: - dropout_mask1 = torch.empty(M, N, device=x.device, dtype=torch.bool) - else: - dropout_mask1 = None - else: - dropout_mask, dropout_mask1 = None, None - # Less than 64KB per feature: enqueue fused kernel - MAX_FUSED_SIZE = 65536 // x.element_size() - BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) - if N > BLOCK_N: - raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") - with torch.cuda.device(x.device.index): - torch.library.wrap_triton(_layer_norm_fwd_1pass_kernel)[(M,)]( - x, - out, - weight, - bias, - residual, - x1, - weight1, - bias1, - y1, - residual_out, - rowscale, - seeds, - dropout_mask, - dropout_mask1, - mean, - rstd, - x.stride(0), - out.stride(0), - residual.stride(0) if residual is not None else 0, - residual_out.stride(0) if residual_out is not None else 0, - x1.stride(0) if x1 is not None else 0, - y1.stride(0) if y1 is not None else 0, - M, - N, - eps, - dropout_p, - # Passing bool make torch inductor very unhappy since it then tries to compare to int_max - int(zero_centered_weight), - is_rms_norm, - BLOCK_N, - residual is not None, - residual_out is not None, - bias is not None, - dropout_p > 0.0, - dropout_mask is not None, - rowscale is not None, - HAS_X1=x1 is not None, - HAS_W1=weight1 is not None, - HAS_B1=bias1 is not None, - ) - return y1, mean, rstd, seeds, dropout_mask, dropout_mask1 - - -@triton.autotune( - configs=triton_autotune_configs(), - key=["N", "HAS_DRESIDUAL", "STORE_DRESIDUAL", "IS_RMS_NORM", "HAS_BIAS", "HAS_DROPOUT"], -) -# torch compile doesn't like triton.heuristics, so we set these manually when calling the kernel -# @triton.heuristics({"HAS_BIAS": lambda args: args["B"] is not None}) -# @triton.heuristics({"HAS_DRESIDUAL": lambda args: args["DRESIDUAL"] is not None}) -# @triton.heuristics({"STORE_DRESIDUAL": lambda args: args["DRESIDUAL_IN"] is not None}) -# @triton.heuristics({"HAS_ROWSCALE": lambda args: args["ROWSCALE"] is not None}) -# @triton.heuristics({"HAS_DY1": lambda args: args["DY1"] is not None}) -# @triton.heuristics({"HAS_DX1": lambda args: args["DX1"] is not None}) -# @triton.heuristics({"HAS_B1": lambda args: args["DB1"] is not None}) -# @triton.heuristics({"RECOMPUTE_OUTPUT": lambda args: args["Y"] is not None}) -@triton.jit -def _layer_norm_bwd_kernel( - X, # pointer to the input - W, # pointer to the weights - B, # pointer to the biases - Y, # pointer to the output to be recomputed - DY, # pointer to the output gradient - DX, # pointer to the input gradient - DW, # pointer to the partial sum of weights gradient - DB, # pointer to the partial sum of biases gradient - DRESIDUAL, - W1, - DY1, - DX1, - DW1, - DB1, - DRESIDUAL_IN, - ROWSCALE, - SEEDS, - Mean, # pointer to the mean - Rstd, # pointer to the 1/std - stride_x_row, # how much to increase the pointer when moving by 1 row - stride_y_row, - stride_dy_row, - stride_dx_row, - stride_dres_row, - stride_dy1_row, - stride_dx1_row, - stride_dres_in_row, - M, # number of rows in X - N, # number of columns in X - eps, # epsilon to avoid division by zero - dropout_p, - zero_centered_weight, - rows_per_program, - IS_RMS_NORM: tl.constexpr, - BLOCK_N: tl.constexpr, - HAS_DRESIDUAL: tl.constexpr, - STORE_DRESIDUAL: tl.constexpr, - HAS_BIAS: tl.constexpr, - HAS_DROPOUT: tl.constexpr, - HAS_ROWSCALE: tl.constexpr, - HAS_DY1: tl.constexpr, - HAS_DX1: tl.constexpr, - HAS_B1: tl.constexpr, - RECOMPUTE_OUTPUT: tl.constexpr, -): - # Map the program id to the elements of X, DX, and DY it should compute. - row_block_id = tl.program_id(0) - row_start = row_block_id * rows_per_program - # Do not early exit if row_start >= M, because we need to write DW and DB - cols = tl.arange(0, BLOCK_N) - mask = cols < N - X += row_start * stride_x_row - if HAS_DRESIDUAL: - DRESIDUAL += row_start * stride_dres_row - if STORE_DRESIDUAL: - DRESIDUAL_IN += row_start * stride_dres_in_row - DY += row_start * stride_dy_row - DX += row_start * stride_dx_row - if HAS_DY1: - DY1 += row_start * stride_dy1_row - if HAS_DX1: - DX1 += row_start * stride_dx1_row - if RECOMPUTE_OUTPUT: - Y += row_start * stride_y_row - w = tl.load(W + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w += 1.0 - if RECOMPUTE_OUTPUT and HAS_BIAS: - b = tl.load(B + cols, mask=mask, other=0.0).to(tl.float32) - if HAS_DY1: - w1 = tl.load(W1 + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w1 += 1.0 - dw = tl.zeros((BLOCK_N,), dtype=tl.float32) - if HAS_BIAS: - db = tl.zeros((BLOCK_N,), dtype=tl.float32) - if HAS_DY1: - dw1 = tl.zeros((BLOCK_N,), dtype=tl.float32) - if HAS_B1: - db1 = tl.zeros((BLOCK_N,), dtype=tl.float32) - row_end = min((row_block_id + 1) * rows_per_program, M) - for row in range(row_start, row_end): - # Load data to SRAM - x = tl.load(X + cols, mask=mask, other=0).to(tl.float32) - dy = tl.load(DY + cols, mask=mask, other=0).to(tl.float32) - if HAS_DY1: - dy1 = tl.load(DY1 + cols, mask=mask, other=0).to(tl.float32) - if not IS_RMS_NORM: - mean = tl.load(Mean + row) - rstd = tl.load(Rstd + row) - # Compute dx - xhat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd - xhat = tl.where(mask, xhat, 0.0) - if RECOMPUTE_OUTPUT: - y = xhat * w + b if HAS_BIAS else xhat * w - tl.store(Y + cols, y, mask=mask) - wdy = w * dy - dw += dy * xhat - if HAS_BIAS: - db += dy - if HAS_DY1: - wdy += w1 * dy1 - dw1 += dy1 * xhat - if HAS_B1: - db1 += dy1 - if not IS_RMS_NORM: - c1 = tl.sum(xhat * wdy, axis=0) / N - c2 = tl.sum(wdy, axis=0) / N - dx = (wdy - (xhat * c1 + c2)) * rstd - else: - c1 = tl.sum(xhat * wdy, axis=0) / N - dx = (wdy - xhat * c1) * rstd - if HAS_DRESIDUAL: - dres = tl.load(DRESIDUAL + cols, mask=mask, other=0).to(tl.float32) - dx += dres - # Write dx - if STORE_DRESIDUAL: - tl.store(DRESIDUAL_IN + cols, dx, mask=mask) - if HAS_DX1: - if HAS_DROPOUT: - keep_mask = ( - tl.rand(tl.load(SEEDS + M + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - ) - dx1 = tl.where(keep_mask, dx / (1.0 - dropout_p), 0.0) - else: - dx1 = dx - tl.store(DX1 + cols, dx1, mask=mask) - if HAS_DROPOUT: - keep_mask = tl.rand(tl.load(SEEDS + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - dx = tl.where(keep_mask, dx / (1.0 - dropout_p), 0.0) - if HAS_ROWSCALE: - rowscale = tl.load(ROWSCALE + row).to(tl.float32) - dx *= rowscale - tl.store(DX + cols, dx, mask=mask) - - X += stride_x_row - if HAS_DRESIDUAL: - DRESIDUAL += stride_dres_row - if STORE_DRESIDUAL: - DRESIDUAL_IN += stride_dres_in_row - if RECOMPUTE_OUTPUT: - Y += stride_y_row - DY += stride_dy_row - DX += stride_dx_row - if HAS_DY1: - DY1 += stride_dy1_row - if HAS_DX1: - DX1 += stride_dx1_row - tl.store(DW + row_block_id * N + cols, dw, mask=mask) - if HAS_BIAS: - tl.store(DB + row_block_id * N + cols, db, mask=mask) - if HAS_DY1: - tl.store(DW1 + row_block_id * N + cols, dw1, mask=mask) - if HAS_B1: - tl.store(DB1 + row_block_id * N + cols, db1, mask=mask) - - -def _layer_norm_bwd( - dy: Tensor, - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - mean: Tensor, - rstd: Tensor, - dresidual: Optional[Tensor] = None, - dy1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - seeds: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - has_residual: bool = False, - has_x1: bool = False, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - x_dtype: Optional[torch.dtype] = None, - recompute_output: bool = False, -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - # Need to wrap to handle the case where dresidual_in or dx1 are aliases of x, - # which makes torch.library unhappy - dx, dw, db, dresidual_in, dx1, dw1, db1, y = _layer_norm_bwd_impl( - dy, - x, - weight, - bias, - eps, - mean, - rstd, - dresidual, - dy1, - weight1, - bias1, - seeds, - dropout_p, - rowscale, - has_residual, - has_x1, - zero_centered_weight, - is_rms_norm, - x_dtype=x_dtype, - recompute_output=recompute_output, - ) - # Don't need to compute dresidual_in separately in this case - if has_residual and dx.dtype == x.dtype and dropout_p == 0.0 and rowscale is None: - dresidual_in = dx - if has_x1 and dropout_p == 0.0: - dx1 = dx - return dx, dw, db, dresidual_in, dx1, dw1, db1, y - - - -@triton_op(add_op_namespace_prefix("layer_norm_bwd_impl"), mutates_args={}, - schema="(Tensor dy, Tensor x, Tensor weight, Tensor bias, float eps, Tensor mean, Tensor rstd, Tensor? dresidual, Tensor? dy1, Tensor? weight1, Tensor? bias1, Tensor? seeds, float dropout_p, Tensor? rowscale, bool has_residual, bool has_x1, bool zero_centered_weight, bool is_rms_norm, ScalarType? x_dtype, bool recompute_output) -> (Tensor dx, Tensor dw, Tensor db, Tensor dresidual_in, Tensor dx1, Tensor dw1, Tensor db1, Tensor y)", - allow_decomposition=False, # Don't let torch.compile trace inside - ) -def _layer_norm_bwd_impl( - dy: Tensor, - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - mean: Tensor, - rstd: Tensor, - dresidual: Optional[Tensor] = None, - dy1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - seeds: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - has_residual: bool = False, - has_x1: bool = False, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - x_dtype: Optional[torch.dtype] = None, - recompute_output: bool = False, -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - M, N = x.shape - assert x.stride(-1) == 1 - dy = maybe_contiguous_lastdim(dy) - assert dy.stride(-1) == 1 - assert dy.shape == (M, N) - if dresidual is not None: - dresidual = maybe_contiguous_lastdim(dresidual) - assert dresidual.stride(-1) == 1 - assert dresidual.shape == (M, N) - assert weight.shape == (N,) - assert weight.stride(-1) == 1 - if bias is not None: - assert bias.stride(-1) == 1 - assert bias.shape == (N,) - if dy1 is not None: - dy1 = maybe_contiguous_lastdim(dy1) - assert weight1 is not None - assert dy1.shape == dy.shape - assert dy1.stride(-1) == 1 - if weight1 is not None: - assert weight1.shape == (N,) - assert weight1.stride(-1) == 1 - if bias1 is not None: - assert bias1.shape == (N,) - assert bias1.stride(-1) == 1 - if seeds is not None: - assert seeds.is_contiguous() - assert seeds.shape == (M if not has_x1 else M * 2,) - if rowscale is not None: - assert rowscale.is_contiguous() - assert rowscale.shape == (M,) - # allocate output - dx = ( - torch.empty_like(x) - if x_dtype is None - else torch.empty(M, N, dtype=x_dtype, device=x.device) - ) - dresidual_in = ( - torch.empty_like(x) - if has_residual - and (dx.dtype != x.dtype or dropout_p > 0.0 or rowscale is not None or has_x1) - else None - ) - dx1 = torch.empty_like(dx) if (has_x1 and dropout_p > 0.0) else None - y = torch.empty(M, N, dtype=dy.dtype, device=dy.device) if recompute_output else None - if recompute_output: - assert weight1 is None, "recompute_output is not supported with parallel LayerNorm" - - # Less than 64KB per feature: enqueue fused kernel - MAX_FUSED_SIZE = 65536 // x.element_size() - BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) - if N > BLOCK_N: - raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") - # Increasing the multiple (e.g. 8) will allow more thread blocks to be launched and hide the - # latency of the gmem reads/writes, but will increase the time of summing up dw / db. - sm_count = torch.cuda.get_device_properties(x.device).multi_processor_count * 8 - _dw = torch.empty((sm_count, N), dtype=torch.float32, device=weight.device) - _db = ( - torch.empty((sm_count, N), dtype=torch.float32, device=bias.device) - if bias is not None - else None - ) - _dw1 = torch.empty_like(_dw) if weight1 is not None else None - _db1 = torch.empty_like(_db) if bias1 is not None else None - rows_per_program = math.ceil(M / sm_count) - grid = (sm_count,) - with torch.cuda.device(x.device.index): - torch.library.wrap_triton(_layer_norm_bwd_kernel)[grid]( - x, - weight, - bias, - y, - dy, - dx, - _dw, - _db, - dresidual, - weight1, - dy1, - dx1, - _dw1, - _db1, - dresidual_in, - rowscale, - seeds, - mean, - rstd, - x.stride(0), - 0 if not recompute_output else y.stride(0), - dy.stride(0), - dx.stride(0), - dresidual.stride(0) if dresidual is not None else 0, - dy1.stride(0) if dy1 is not None else 0, - dx1.stride(0) if dx1 is not None else 0, - dresidual_in.stride(0) if dresidual_in is not None else 0, - M, - N, - eps, - dropout_p, - # Passing bool make torch inductor very unhappy since it then tries to compare to int_max - int(zero_centered_weight), - rows_per_program, - is_rms_norm, - BLOCK_N, - dresidual is not None, - dresidual_in is not None, - bias is not None, - dropout_p > 0.0, - HAS_ROWSCALE=rowscale is not None, - HAS_DY1=dy1 is not None, - HAS_DX1=dx1 is not None, - HAS_B1=bias1 is not None, - RECOMPUTE_OUTPUT=y is not None, - ) - dw = _dw.sum(0).to(weight.dtype) - db = _db.sum(0).to(bias.dtype) if bias is not None else None - dw1 = _dw1.sum(0).to(weight1.dtype) if weight1 is not None else None - db1 = _db1.sum(0).to(bias1.dtype) if bias1 is not None else None - # dresidual_in and dx1 could be None, the wrapper will handle assigning them from dx - return dx, dw, db, dresidual_in, dx1, dw1, db1, y - - -class LayerNormFn(torch.autograd.Function): - - @staticmethod - def forward( - ctx, - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - residual_in_fp32=False, - zero_centered_weight=False, - is_rms_norm=False, - return_dropout_mask=False, - out_dtype=None, - out=None, - residual_out=None - ): - x_shape_og = x.shape - # reshape input data into 2D tensor - x = maybe_contiguous_lastdim(x.reshape(-1, x.shape[-1])) - if residual is not None: - assert residual.shape == x_shape_og - residual = maybe_contiguous_lastdim(residual.reshape(-1, residual.shape[-1])) - if x1 is not None: - assert x1.shape == x_shape_og - assert rowscale is None, "rowscale is not supported with parallel LayerNorm" - x1 = maybe_contiguous_lastdim(x1.reshape(-1, x1.shape[-1])) - weight = weight.contiguous() - bias = maybe_contiguous(bias) - weight1 = maybe_contiguous(weight1) - bias1 = maybe_contiguous(bias1) - if rowscale is not None: - rowscale = rowscale.reshape(-1).contiguous() - residual_dtype = ( - residual.dtype - if residual is not None - else (torch.float32 if residual_in_fp32 else None) - ) - if out is not None: - out = out.reshape(-1, out.shape[-1]) - if residual_out is not None: - residual_out = residual_out.reshape(-1, residual_out.shape[-1]) - y, y1, mean, rstd, residual_out, seeds, dropout_mask, dropout_mask1 = _layer_norm_fwd( - x, - weight, - bias, - eps, - residual, - x1, - weight1, - bias1, - dropout_p=dropout_p, - rowscale=rowscale, - out_dtype=out_dtype, - residual_dtype=residual_dtype, - zero_centered_weight=zero_centered_weight, - is_rms_norm=is_rms_norm, - return_dropout_mask=return_dropout_mask, - out=out, - residual_out=residual_out, - ) - ctx.save_for_backward( - residual_out, weight, bias, weight1, bias1, rowscale, seeds, mean, rstd - ) - ctx.x_shape_og = x_shape_og - ctx.eps = eps - ctx.dropout_p = dropout_p - ctx.is_rms_norm = is_rms_norm - ctx.has_residual = residual is not None - ctx.has_x1 = x1 is not None - ctx.prenorm = prenorm - ctx.x_dtype = x.dtype - ctx.zero_centered_weight = zero_centered_weight - y = y.reshape(x_shape_og) - y1 = y1.reshape(x_shape_og) if y1 is not None else None - residual_out = residual_out.reshape(x_shape_og) if residual_out is not None else None - dropout_mask = dropout_mask.reshape(x_shape_og) if dropout_mask is not None else None - dropout_mask1 = dropout_mask1.reshape(x_shape_og) if dropout_mask1 is not None else None - if not return_dropout_mask: - if weight1 is None: - return y if not prenorm else (y, residual_out) - else: - return (y, y1) if not prenorm else (y, y1, residual_out) - else: - if weight1 is None: - return ( - (y, dropout_mask, dropout_mask1) - if not prenorm - else (y, residual_out, dropout_mask, dropout_mask1) - ) - else: - return ( - (y, y1, dropout_mask, dropout_mask1) - if not prenorm - else (y, y1, residual_out, dropout_mask, dropout_mask1) - ) - - @staticmethod - def backward(ctx, dy, *args): - x, weight, bias, weight1, bias1, rowscale, seeds, mean, rstd = ctx.saved_tensors - dy = dy.reshape(-1, dy.shape[-1]) - if weight1 is not None: - dy1, args = args[0], args[1:] - dy1 = dy1.reshape(-1, dy1.shape[-1]) - assert dy1.shape == x.shape - else: - dy1 = None - if ctx.prenorm: - dresidual = args[0] - dresidual = dresidual.reshape(-1, dresidual.shape[-1]) - assert dresidual.shape == x.shape - else: - dresidual = None - dx, dw, db, dresidual_in, dx1, dw1, db1, _ = _layer_norm_bwd( - dy, - x, - weight, - bias, - ctx.eps, - mean, - rstd, - dresidual, - dy1, - weight1, - bias1, - seeds, - ctx.dropout_p, - rowscale, - ctx.has_residual, - ctx.has_x1, - ctx.zero_centered_weight, - ctx.is_rms_norm, - x_dtype=ctx.x_dtype, - recompute_output=False, - ) - return ( - dx.reshape(ctx.x_shape_og), - dw, - db, - dresidual_in.reshape(ctx.x_shape_og) if ctx.has_residual else None, - dx1.reshape(ctx.x_shape_og) if dx1 is not None else None, - dw1, - db1, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) - - -def layer_norm_fn( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - residual_in_fp32=False, - zero_centered_weight=False, - is_rms_norm=False, - return_dropout_mask=False, - out_dtype=None, - out=None, - residual_out=None -): - return LayerNormFn.apply( - x, - weight, - bias, - residual, - x1, - weight1, - bias1, - eps, - dropout_p, - rowscale, - prenorm, - residual_in_fp32, - zero_centered_weight, - is_rms_norm, - return_dropout_mask, - out_dtype, - out, - residual_out - ) - - -def rms_norm_fn( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - residual_in_fp32=False, - zero_centered_weight=False, - return_dropout_mask=False, - out_dtype=None, - out=None, - residual_out=None -): - return LayerNormFn.apply( - x, - weight, - bias, - residual, - x1, - weight1, - bias1, - eps, - dropout_p, - rowscale, - prenorm, - residual_in_fp32, - zero_centered_weight, - True, - return_dropout_mask, - out_dtype, - out, - residual_out - ) - - -class RMSNorm(torch.nn.Module): - - def __init__(self, hidden_size, eps=1e-5, dropout_p=0.0, zero_centered_weight=False, - device=None, dtype=None): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.eps = eps - if dropout_p > 0.0: - self.drop = torch.nn.Dropout(dropout_p) - else: - self.drop = None - self.zero_centered_weight = zero_centered_weight - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.register_parameter("bias", None) - self.reset_parameters() - - def reset_parameters(self): - if not self.zero_centered_weight: - torch.nn.init.ones_(self.weight) - else: - torch.nn.init.zeros_(self.weight) - - def forward(self, x, residual=None, prenorm=False, residual_in_fp32=False): - return rms_norm_fn( - x, - self.weight, - self.bias, - residual=residual, - eps=self.eps, - dropout_p=self.drop.p if self.drop is not None and self.training else 0.0, - prenorm=prenorm, - residual_in_fp32=residual_in_fp32, - zero_centered_weight=self.zero_centered_weight, - ) - - -class LayerNormLinearFn(torch.autograd.Function): - - @staticmethod - @custom_fwd - def forward( - ctx, - x, - norm_weight, - norm_bias, - linear_weight, - linear_bias, - residual=None, - eps=1e-6, - prenorm=False, - residual_in_fp32=False, - is_rms_norm=False, - ): - x_shape_og = x.shape - # reshape input data into 2D tensor - x = maybe_contiguous_lastdim(x.reshape(-1, x.shape[-1])) - if residual is not None: - assert residual.shape == x_shape_og - residual = maybe_contiguous_lastdim(residual.reshape(-1, residual.shape[-1])) - norm_weight = norm_weight.contiguous() - norm_bias = maybe_contiguous(norm_bias) - residual_dtype = ( - residual.dtype - if residual is not None - else (torch.float32 if residual_in_fp32 else None) - ) - y, _, mean, rstd, residual_out, *rest = _layer_norm_fwd( - x, - norm_weight, - norm_bias, - eps, - residual, - out_dtype=None if not torch.is_autocast_enabled() else torch.get_autocast_dtype("cuda"), - residual_dtype=residual_dtype, - is_rms_norm=is_rms_norm, - ) - y = y.reshape(x_shape_og) - dtype = torch.get_autocast_dtype("cuda") if torch.is_autocast_enabled() else y.dtype - linear_weight = linear_weight.to(dtype) - linear_bias = linear_bias.to(dtype) if linear_bias is not None else None - out = F.linear(y.to(linear_weight.dtype), linear_weight, linear_bias) - # We don't store y, will be recomputed in the backward pass to save memory - ctx.save_for_backward(residual_out, norm_weight, norm_bias, linear_weight, mean, rstd) - ctx.x_shape_og = x_shape_og - ctx.eps = eps - ctx.is_rms_norm = is_rms_norm - ctx.has_residual = residual is not None - ctx.prenorm = prenorm - ctx.x_dtype = x.dtype - ctx.linear_bias_is_none = linear_bias is None - return out if not prenorm else (out, residual_out.reshape(x_shape_og)) - - @staticmethod - @custom_bwd - def backward(ctx, dout, *args): - x, norm_weight, norm_bias, linear_weight, mean, rstd = ctx.saved_tensors - dout = dout.reshape(-1, dout.shape[-1]) - dy = F.linear(dout, linear_weight.t()) - dlinear_bias = None if ctx.linear_bias_is_none else dout.sum(0) - dy = maybe_contiguous_lastdim(dy) - assert dy.shape == x.shape - if ctx.prenorm: - dresidual = args[0] - dresidual = maybe_contiguous_lastdim(dresidual.reshape(-1, dresidual.shape[-1])) - assert dresidual.shape == x.shape - else: - dresidual = None - dx, dnorm_weight, dnorm_bias, dresidual_in, _, _, _, y = _layer_norm_bwd( - dy, - x, - norm_weight, - norm_bias, - ctx.eps, - mean, - rstd, - dresidual=dresidual, - has_residual=ctx.has_residual, - is_rms_norm=ctx.is_rms_norm, - x_dtype=ctx.x_dtype, - recompute_output=True, - ) - dlinear_weight = torch.einsum("bo,bi->oi", dout, y) - return ( - dx.reshape(ctx.x_shape_og), - dnorm_weight, - dnorm_bias, - dlinear_weight, - dlinear_bias, - dresidual_in.reshape(ctx.x_shape_og) if ctx.has_residual else None, - None, - None, - None, - None, - ) - - -def layer_norm_linear_fn( - x, - norm_weight, - norm_bias, - linear_weight, - linear_bias, - residual=None, - eps=1e-6, - prenorm=False, - residual_in_fp32=False, - is_rms_norm=False, -): - return LayerNormLinearFn.apply( - x, - norm_weight, - norm_bias, - linear_weight, - linear_bias, - residual, - eps, - prenorm, - residual_in_fp32, - is_rms_norm, - ) diff --git a/build/torch212-cxx11-cu130-x86_64-linux/ops/triton/linear.py b/build/torch212-cxx11-cu130-x86_64-linux/ops/triton/linear.py deleted file mode 100644 index a8966dbc345ab0e593df0124451ee7be3dae131a..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu130-x86_64-linux/ops/triton/linear.py +++ /dev/null @@ -1,594 +0,0 @@ -# Adapted from https://github.com/ELS-RD/kernl/blob/main/src/kernl/implementations/linear_layer.py -# and https://github.com/openai/triton/blob/master/python/triton/ops/matmul.py -from typing import Optional - -import torch -import triton -import triton.language as tl -from triton.ops.matmul_perf_model import early_config_prune, estimate_matmul_time - -from flash_attn.ops.triton.k_activations import ( - gelu, - gelu_approx, - gelu_approx_grad, - gelu_grad, - squared_relu, - squared_relu_grad, -) - -# CREDITS: Initially inspired by the Triton tutorial on matrix multiplications - - -def init_to_zero(name): - return lambda nargs: nargs[name].zero_() - - -def get_configs_io_bound(): - configs = [] - for num_stages in [2, 3, 4, 5, 6]: - for block_m in [16, 32]: - for block_k in [32, 64]: - for block_n in [32, 64, 128, 256]: - num_warps = 2 if block_n <= 64 else 4 - configs.append( - triton.Config( - { - "BLOCK_M": block_m, - "BLOCK_N": block_n, - "BLOCK_K": block_k, - "SPLIT_K": 1, - }, - num_stages=num_stages, - num_warps=num_warps, - ) - ) - # split_k not used - # for split_k in [2, 4, 8, 16]: - # configs.append(triton.Config( - # {'BLOCK_M': block_m, 'BLOCK_N': block_n, 'BLOCK_K': block_k, 'SPLIT_K': split_k}, - # num_stages=num_stages, num_warps=num_warps, pre_hook=init_to_zero('C'))) - return configs - - -@triton.autotune( - configs=[ - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - # good for int8 - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=4, - num_warps=4, - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - ] - + get_configs_io_bound(), - key=["CACHE_KEY_M", "CACHE_KEY_N", "CACHE_KEY_K"], - prune_configs_by={ - "early_config_prune": early_config_prune, - "perf_model": estimate_matmul_time, - "top_k": 10, - }, -) -@triton.heuristics( - { - "EVEN_K": lambda args: args["K"] % (args["BLOCK_K"] * args["SPLIT_K"]) == 0, - } -) -@triton.jit -def kernel_fwd( - C, # Pointers to matrices - ACT_INPUT, - A, - B, - bias, - # Matrix dimensions - M, - N, - K, - CACHE_KEY_M, - CACHE_KEY_N, - CACHE_KEY_K, - # The stride variables represent how much to increase the ptr by when moving by 1 - # element in a particular dimension. E.g. stride_am is how much to increase a_ptr - # by to get the element one row down (A has M rows) - stride_cm, - # stride_cn, # Assume that stride_cn == 1 - stride_am, - stride_ak, - stride_bn, - stride_bk, - # Meta-parameters - BLOCK_M: tl.constexpr, - GROUP_M: tl.constexpr, - BLOCK_N: tl.constexpr, - BLOCK_K: tl.constexpr, - # split k not used, not performant with activation, kept because early_config_prune is expecting it - SPLIT_K: tl.constexpr, - EVEN_K: tl.constexpr, - A_ROWMAJOR: tl.constexpr, - B_COLMAJOR: tl.constexpr, - BIAS: tl.constexpr, - SAVE_ACT_INPUT: tl.constexpr, - ACTIVATION: tl.constexpr, -): - - """ - Kernel for computing Out = activation(A x W + C) - - Input has shape (M, K) - - Weight has shape (K, N) - - Bias has shape (N,) - - Output has shape (M, N) - - ActInputs (optional) has shape (M, N) - 'ActInputs' optionally saves the A x W + C intermediate for backward computations - This kernel will consolidate over K - """ - - pid = tl.program_id(axis=0) - - grid_m = (M + BLOCK_M - 1) // BLOCK_M - grid_n = (N + BLOCK_N - 1) // BLOCK_N - # re-order program ID for better L2 performance - width = GROUP_M * grid_n - group_id = pid // width - group_size = min(grid_m - group_id * GROUP_M, GROUP_M) - pid_m = group_id * GROUP_M + (pid % group_size) - pid_n = (pid % width) // (group_size) - - # now compute the block that each program will go through - # rm (resp. rn) denotes a range of indices - # for rows (resp. col) of C - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - # trick to avoid masking on M and N axis - ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M) - rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N) - rk = tl.arange(0, BLOCK_K) - - if A_ROWMAJOR: - A = A + (ram[:, None] * stride_am + rk[None, :]) - else: - A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak) - if B_COLMAJOR: - B = B + (rk[:, None] + rbn[None, :] * stride_bn) - else: - B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn) - - acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) - - for k in range(K, 0, -BLOCK_K): - if EVEN_K: - a = tl.load(A) - b = tl.load(B) - else: - a = tl.load(A, mask=rk[None, :] < k, other=0.0) - b = tl.load(B, mask=rk[:, None] < k, other=0.0) - acc += tl.dot(a, b) - - if A_ROWMAJOR: - A += BLOCK_K - else: - A += BLOCK_K * stride_ak - if B_COLMAJOR: - B += BLOCK_K - else: - B += BLOCK_K * stride_bk - - # Putting bias after the matmul (instead of before) is faster, idk why - if BIAS: - bias = tl.load(bias + rn, mask=rn < N, other=0.0).to(tl.float32) - acc += bias[None, :] - - # optional: save the activation inputs - if SAVE_ACT_INPUT: - # act_in_ptrs = ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :] * stride_cn - act_in_ptrs = ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :] - tl.store(act_in_ptrs, acc) - - # optional: fused activation (while the data is in shared memory) - if ACTIVATION == "gelu": - acc = gelu(acc) - elif ACTIVATION == "gelu_approx": - acc = gelu_approx(acc) - elif ACTIVATION == "squared_relu": - acc = squared_relu(acc) - # rematerialize rm and rn to save registers - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - - # write back result - # C = C + rm[:, None] * stride_cm + rn[None, :] * stride_cn - C = C + rm[:, None] * stride_cm + rn[None, :] - mask = (rm < M)[:, None] & (rn < N)[None, :] - tl.store(C, acc) - - -def triton_linear_act( - x: torch.Tensor, - weight: torch.Tensor, - bias: Optional[torch.Tensor] = None, - activation: str = "id", - save_act_input: bool = False, -) -> torch.Tensor: - """ - Compute e = activation(x @ weight.T + bias). - This wrapper kicks the `kernel_fwd` Triton kernel - :param x: input tensor - :param weight: weight matrix - :param bias: an optional bias tensor - :param activation: Activation name. Needs to be a Triton kernel. - :param act_input: an optional tensor to save the activation inputs (for backward) - :return: result tensor - """ - # if torch.is_autocast_enabled(): - # dtype = torch.get_autocast_gpu_dtype() - # x, weight, bias = [a.to(dtype=dtype) for a in [x, weight, bias]] - - assert activation in ["id", "gelu", "gelu_approx", "squared_relu"] - - batch_shape, n = x.shape[:-1], x.shape[-1] - batch_dim = batch_shape.numel() - x_reshaped = x.reshape(batch_dim, n) - - if x_reshaped.stride(0) > 1 and x_reshaped.stride(1) > 1: - x_reshaped = x_reshaped.contiguous() - if weight.stride(0) > 1 and weight.stride(1) > 1: - weight = weight.contiguous() - bias = bias.contiguous() if bias is not None else None - - assert ( - x.dtype == weight.dtype - ), f"Input and weight must have the same dtype, got {x.dtype} and {weight.dtype}" - if bias is not None: - assert ( - x.dtype == bias.dtype - ), f"Input and bias must have the same dtype, got {x.dtype} and {bias.dtype}" - assert ( - x_reshaped.shape[1] == weight.shape[1] - ), f"Incompatible dimensions: {x_reshaped.shape} - {weight.shape}" - - assert ( - bias is None or bias.shape[0] == weight.shape[0] - ), "Incompatible dimensions in between weight and bias" - - M, K = x_reshaped.shape - N, K = weight.shape - - output = torch.empty((M, N), device=x.device, dtype=x.dtype) - act_input = torch.empty_like(output) if save_act_input else None - - # 1D launch kernel where each block gets its own program. - grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),) # noqa - - kernel_fwd[grid]( - output, - act_input, - x_reshaped, - weight, # data ptrs - bias if bias is not None else x, # auto skip bias if not present - M, # shapes - N, - K, - M // 32, # key for triton cache (limit number of compilations) - N // 32, - K // 32, - stride_cm=output.stride(0), # strides - # stride_cn=output.stride(1), - stride_am=x_reshaped.stride(0), - stride_ak=x_reshaped.stride(1), - stride_bk=weight.stride(1), - stride_bn=weight.stride(0), - BIAS=bias is not None, # optional fused bias - SAVE_ACT_INPUT=save_act_input, # optional save activation inputs - ACTIVATION=activation, # optional fused activation - A_ROWMAJOR=x_reshaped.stride(1) == 1, - B_COLMAJOR=weight.stride(1) == 1, - GROUP_M=8, # speed optimization: group the programs - ) - - if not save_act_input: - return output.reshape(*batch_shape, output.shape[-1]) - else: - return ( - output.reshape(*batch_shape, output.shape[-1]), - act_input.reshape(*batch_shape, act_input.shape[-1]), - ) - - -@triton.autotune( - configs=[ - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - # good for int8 - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=4, - num_warps=4, - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - ] - + get_configs_io_bound(), - key=["CACHE_KEY_M", "CACHE_KEY_N", "CACHE_KEY_K"], - prune_configs_by={ - "early_config_prune": early_config_prune, - "perf_model": estimate_matmul_time, - "top_k": 10, - }, -) -@triton.heuristics( - { - "EVEN_K": lambda args: args["K"] % (args["BLOCK_K"] * args["SPLIT_K"]) == 0, - } -) -@triton.jit -def kernel_bwd( - C, # Pointers to matrices - ACT_INPUT, - A, - B, - # Matrix dimensions - M, - N, - K, - CACHE_KEY_M, - CACHE_KEY_N, - CACHE_KEY_K, - # The stride variables represent how much to increase the ptr by when moving by 1 - # element in a particular dimension. E.g. stride_am is how much to increase a_ptr - # by to get the element one row down (A has M rows) - stride_cm, - # stride_cn, # Assume that stride_cn == 1 - stride_am, - stride_ak, - stride_bk, - stride_bn, - # Meta-parameters - BLOCK_M: tl.constexpr, - GROUP_M: tl.constexpr, - BLOCK_N: tl.constexpr, - BLOCK_K: tl.constexpr, - # split k not used, not performant with activation, kept because early_config_prune is expecting it - SPLIT_K: tl.constexpr, - EVEN_K: tl.constexpr, - ACTIVATION: tl.constexpr, -): - - """ - Kernel for computing Out = activation(A x W + C) - - Input has shape (M, K) - - Weight has shape (K, N) - - Output has shape (M, N) - - ActInputs (optional) has shape (M, N) - 'ActInputs' optionally saves the A x W + C intermediate for backward computations - This kernel will consolidate over K - """ - - pid = tl.program_id(axis=0) - - grid_m = (M + BLOCK_M - 1) // BLOCK_M - grid_n = (N + BLOCK_N - 1) // BLOCK_N - # re-order program ID for better L2 performance - width = GROUP_M * grid_n - group_id = pid // width - group_size = min(grid_m - group_id * GROUP_M, GROUP_M) - pid_m = group_id * GROUP_M + (pid % group_size) - pid_n = (pid % width) // (group_size) - - # now compute the block that each program will go through - # rm (resp. rn) denotes a range of indices - # for rows (resp. col) of C - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - # trick to avoid masking on M and N axis - ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M) - rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N) - rk = tl.arange(0, BLOCK_K) - - A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak) - B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn) - - acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) - - for k in range(K, 0, -BLOCK_K): - if EVEN_K: - a = tl.load(A) - b = tl.load(B) - else: - a = tl.load(A, mask=rk[None, :] < k, other=0.0) - b = tl.load(B, mask=rk[:, None] < k, other=0.0) - acc += tl.dot(a, b) - - A += BLOCK_K * stride_ak - B += BLOCK_K * stride_bk - - # optional: fused activation (while the data is in shared memory) - if ACTIVATION != "id": - act_in_ptrs = ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :] - act_input = tl.load(act_in_ptrs).to(acc.dtype) - if ACTIVATION == "gelu": - acc *= gelu_grad(act_input) - elif ACTIVATION == "gelu_approx": - acc *= gelu_approx_grad(act_input) - elif ACTIVATION == "squared_relu": - acc *= squared_relu_grad(act_input) - - # rematerialize rm and rn to save registers - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - - # write back result - C = C + rm[:, None] * stride_cm + rn[None, :] - mask = (rm < M)[:, None] & (rn < N)[None, :] - tl.store(C, acc, mask=mask) - - -def triton_dgrad_act( - grad_output: torch.Tensor, - weight: torch.Tensor, - activation: str = "id", - act_input: Optional[torch.Tensor] = None, -) -> torch.Tensor: - """ - Compute e = activation(grad_output @ weight + bias). - This wrapper kicks the `kernel_fwd` Triton kernel - :param grad_output: input tensor - :param weight: weight matrix - :param activation: Activation name. Needs to be a Triton kernel. - :param act_input: an optional tensor to save the activation inputs (for backward) - :return: result tensor - """ - assert activation in ["id", "gelu", "gelu_approx", "squared_relu"] - - batch_shape, n = grad_output.shape[:-1], grad_output.shape[-1] - batch_dim = batch_shape.numel() - grad_output_reshaped = grad_output.reshape(batch_dim, n) - - if grad_output_reshaped.stride(0) > 1 and grad_output_reshaped.stride(1) > 1: - grad_output_reshaped = grad_output_reshaped.contiguous() - if weight.stride(0) > 1 and weight.stride(1) > 1: - weight = weight.contiguous() - - assert ( - grad_output.dtype == weight.dtype - ), f"grad_output and weight must have the same dtype, got {grad_output.dtype} and {weight.dtype}" - assert ( - grad_output_reshaped.shape[1] == weight.shape[0] - ), f"Incompatible dimensions: {grad_output_reshaped.shape} - {weight.shape}" - if activation != "id": - assert act_input is not None, f"act_input is required for activation {activation}" - - # M, N, K in bwd are different from M, N, K in fwd - M, K = grad_output_reshaped.shape - K, N = weight.shape - - grad_input = torch.empty((M, N), device=grad_output.device, dtype=grad_output.dtype) - - # 1D launch kernel where each block gets its own program. - grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),) # noqa - - kernel_bwd[grid]( - grad_input, - act_input, - grad_output_reshaped, - weight, # data ptrs - M, # shapes - N, - K, - M // 32, # key for triton cache (limit number of compilations) - N // 32, - K // 32, - stride_cm=grad_input.stride(0), # strides - # stride_cn=grad_input.stride(1), - stride_am=grad_output_reshaped.stride(0), - stride_ak=grad_output_reshaped.stride(1), - stride_bk=weight.stride(0), - stride_bn=weight.stride(1), - ACTIVATION=activation, # optional fused activation - GROUP_M=8, # speed optimization: group the programs - ) - - return grad_input.reshape(*batch_shape, grad_input.shape[-1]) diff --git a/build/torch212-cxx11-cu130-x86_64-linux/ops/triton/mlp.py b/build/torch212-cxx11-cu130-x86_64-linux/ops/triton/mlp.py deleted file mode 100644 index 059f4f8a5e174c1f4824e43d313fca18eaa799b8..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu130-x86_64-linux/ops/triton/mlp.py +++ /dev/null @@ -1,149 +0,0 @@ -# The triton fused matmul + sqrelu is faster for fp16 but slower for bf16, compared -# to naive implementation. -import fused_dense_lib as fused_dense_cuda -import torch -import torch.nn as nn -import torch.nn.functional as F - -from flash_attn.utils.torch import custom_fwd, custom_bwd -from flash_attn.ops.activations import sqrelu_bwd, sqrelu_fwd -from flash_attn.ops.triton.linear import triton_dgrad_act, triton_linear_act - - -class FusedDenseSqreluDenseFunc(torch.autograd.Function): - @staticmethod - @custom_fwd - def forward(ctx, x, weight1, bias1, weight2, bias2, checkpoint_lvl=0): - """checkpoint_lvl: - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute act_input and gelu_out in the bwd - """ - if torch.is_autocast_enabled(): - dtype = torch.get_autocast_gpu_dtype() - x, weight1, bias1, weight2, bias2 = [ - a.to(dtype=dtype) for a in [x, weight1, bias1, weight2, bias2] - ] - is_bf16 = x.dtype == torch.bfloat16 - assert checkpoint_lvl in [0, 1, 2] - x = x.contiguous() - weight1 = weight1.contiguous() - bias1 = bias1.contiguous() - weight2 = weight2.contiguous() - bias2 = bias2.contiguous() - batch_shape, n = x.shape[:-1], x.shape[-1] - batch_dim = batch_shape.numel() - if is_bf16: - act_input = fused_dense_cuda.linear_bias_forward( - x.reshape(batch_dim, n), weight1, bias1 - ) - output1 = sqrelu_fwd(act_input) - else: - save_act_input = checkpoint_lvl != 2 - result = triton_linear_act( - x.reshape(batch_dim, n), - weight1, - bias1, - activation="squared_relu", - save_act_input=save_act_input, - ) - if save_act_input: - output1, act_input = result - else: - output1 = result - output2 = fused_dense_cuda.linear_bias_forward(output1, weight2, bias2) - ctx.checkpoint_lvl = checkpoint_lvl - if checkpoint_lvl == 0: - ctx.save_for_backward(x, weight1, bias1, weight2, act_input, output1) - elif checkpoint_lvl == 1: - ctx.save_for_backward(x, weight1, bias1, weight2, act_input) - elif checkpoint_lvl == 2: - ctx.save_for_backward(x, weight1, bias1, weight2) - return output2.reshape(*batch_shape, output2.shape[-1]) - - @staticmethod - @custom_bwd - def backward(ctx, grad_output): - grad_output = grad_output.contiguous() - checkpoint_lvl = ctx.checkpoint_lvl - x, weight1, bias1, weight2, *rest = ctx.saved_tensors - batch_shape, n = x.shape[:-1], x.shape[-1] - batch_dim = batch_shape.numel() - is_bf16 = x.dtype == torch.bfloat16 - if checkpoint_lvl == 0: - act_input, output1 = rest - elif checkpoint_lvl == 1: - (act_input,) = rest - output1 = sqrelu_fwd(act_input) - elif checkpoint_lvl == 2: - if is_bf16: - act_input = fused_dense_cuda.linear_bias_forward( - x.reshape(batch_dim, n), weight1, bias1 - ) - output1 = sqrelu_fwd(act_input) - else: - output1, act_input = triton_linear_act( - x.reshape(batch_dim, n), - weight1, - bias1, - activation="squared_relu", - save_act_input=True, - ) - - if is_bf16: - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad(output1, grad_output) - grad_output1 = grad_output @ weight2 - grad_act_input = sqrelu_bwd(grad_output1, act_input) - grad_input, grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_backward( - x.reshape(batch_dim, n), weight1, grad_act_input - ) - else: - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad(output1, grad_output) - grad_act_input = triton_dgrad_act( - grad_output, weight2, activation="squared_relu", act_input=act_input - ) - grad_input, grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_backward( - x.reshape(batch_dim, n), weight1, grad_act_input - ) - return grad_input.reshape_as(x), grad_weight1, grad_bias1, grad_weight2, grad_bias2, None - - -fused_dense_sqrelu_dense_function = FusedDenseSqreluDenseFunc.apply - - -class FusedDenseSqreluDense(nn.Module): - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - bias1=True, - bias2=True, - checkpoint_lvl=0, - device=None, - dtype=None, - ): - """ - checkpoint_lvl (increasing lvl means slower but more memory saving): - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute gelu_in and gelu_out in the bwd - """ - assert checkpoint_lvl in [0, 1, 2] - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features * 4 - assert bias1 == True, "DenseSqreluDense module without bias is currently not supported" - assert bias2 == True, "DenseSqreluDense module without bias is currently not supported" - self.checkpoint_lvl = checkpoint_lvl - self.fc1 = nn.Linear(in_features, hidden_features, bias=bias1, **factory_kwargs) - self.fc2 = nn.Linear(hidden_features, out_features, bias=bias2, **factory_kwargs) - - def forward(self, x): - assert x.is_cuda - return fused_dense_sqrelu_dense_function( - x, self.fc1.weight, self.fc1.bias, self.fc2.weight, self.fc2.bias, self.checkpoint_lvl - ) diff --git a/build/torch212-cxx11-cu130-x86_64-linux/ops/triton/rotary.py b/build/torch212-cxx11-cu130-x86_64-linux/ops/triton/rotary.py deleted file mode 100644 index 6ad3ab412a084270120167ad32b67ad1720f2c37..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu130-x86_64-linux/ops/triton/rotary.py +++ /dev/null @@ -1,186 +0,0 @@ -# Copyright (c) 2025, Tri Dao. -# As of 2025-04-23, we require triton >= 3.0 - -from typing import Optional, Union - -import torch - -import triton -import triton.language as tl - - -@triton.jit -def rotary_kernel( - OUT, # Pointers to matrices - X, - COS, - SIN, - CU_SEQLENS, - SEQLEN_OFFSETS, # this could be int or a pointer - # Matrix dimensions - seqlen, - nheads, - seqlen_ro, - # strides - stride_out_batch, - stride_out_seqlen, - stride_out_nheads, - stride_out_headdim, - stride_x_batch, - stride_x_seqlen, - stride_x_nheads, - stride_x_headdim, - # Meta-parameters - # We want ROTARY_DIM to be constexpr, otherwise the triton compiler doesn't know that - # the mask is constant every 8 elements, and it will generate LDG.16 instead of LDG.128 - ROTARY_DIM: tl.constexpr, - IS_SEQLEN_OFFSETS_TENSOR: tl.constexpr, - IS_VARLEN: tl.constexpr, - INTERLEAVED: tl.constexpr, - CONJUGATE: tl.constexpr, - BLOCK_H: tl.constexpr, - BLOCK_M: tl.constexpr, -): - BLOCK_K: tl.constexpr = triton.next_power_of_2(ROTARY_DIM) - ROTARY_DIM_HALF = ROTARY_DIM // 2 - pid_head = tl.program_id(axis=0) - pid_m = tl.program_id(axis=1) - pid_batch = tl.program_id(axis=2) - - if not IS_VARLEN: - X = X + pid_batch * stride_x_batch - OUT = OUT + pid_batch * stride_out_batch - else: - start_idx = tl.load(CU_SEQLENS + pid_batch) - seqlen = tl.load(CU_SEQLENS + pid_batch + 1) - start_idx - X = X + start_idx * stride_x_seqlen - OUT = OUT + start_idx * stride_out_seqlen - - if pid_m * BLOCK_M >= seqlen: - return - - rh = pid_head * BLOCK_H + tl.arange(0, BLOCK_H) - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - if not IS_SEQLEN_OFFSETS_TENSOR: - rm_cs = rm + SEQLEN_OFFSETS - else: - rm_cs = rm + tl.load(SEQLEN_OFFSETS + pid_batch) - - rk_half = tl.arange(0, BLOCK_K // 2) - COS = COS + (rm_cs[:, None] * ROTARY_DIM_HALF + rk_half[None, :]) - SIN = SIN + (rm_cs[:, None] * ROTARY_DIM_HALF + rk_half[None, :]) - mask_cs = (rm_cs[:, None] < seqlen_ro) & (rk_half[None, :] < ROTARY_DIM_HALF) - cos = tl.load(COS, mask=mask_cs, other=1.0).to(tl.float32) - sin = tl.load(SIN, mask=mask_cs, other=0.0).to(tl.float32) - if CONJUGATE: - sin = -sin - - if not INTERLEAVED: - # Load the 1st and 2nd halves of X, do calculation, then store to 1st and 2nd halves of OUT - X = X + (rh[:, None, None] * stride_x_nheads + rm[None, :, None] * stride_x_seqlen + rk_half[None, None, :] * stride_x_headdim) - OUT = OUT + (rh[:, None, None] * stride_out_nheads + rm[None, :, None] * stride_out_seqlen + rk_half[None, None, :] * stride_out_headdim) - mask = (rh[:, None, None] < nheads) & (rm[None, :, None] < seqlen) & (rk_half[None, None, :] < ROTARY_DIM_HALF) - x0 = tl.load(X, mask=mask, other=0.0).to(tl.float32) - x1 = tl.load(X + ROTARY_DIM_HALF * stride_x_headdim, mask=mask, other=0.0,).to(tl.float32) - o0 = x0 * cos - x1 * sin - o1 = x0 * sin + x1 * cos - tl.store(OUT, o0, mask=mask) - tl.store(OUT + ROTARY_DIM_HALF * stride_out_headdim, o1, mask=mask) - else: - rk = tl.arange(0, BLOCK_K) - X = X + (rh[:, None, None] * stride_x_nheads + rm[None, :, None] * stride_x_seqlen + rk[None, None, :] * stride_x_headdim) - OUT = OUT + (rh[:, None, None] * stride_out_nheads + rm[None, :, None] * stride_out_seqlen + rk[None, None, :] * stride_out_headdim) - mask = (rh[:, None, None] < nheads) & (rm[None, :, None] < seqlen) & (rk[None, None, :] < ROTARY_DIM) - x = tl.load(X, mask=mask, other=0.0).to(tl.float32) - x0, x1 = tl.split(tl.reshape(x, [BLOCK_H, BLOCK_M, BLOCK_K // 2, 2])) - o0 = x0 * cos - x1 * sin - o1 = x0 * sin + x1 * cos - o = tl.reshape(tl.join(o0, o1), [BLOCK_H, BLOCK_M, BLOCK_K]) - tl.store(OUT, o, mask=mask) - - -def apply_rotary( - x: torch.Tensor, - cos: torch.Tensor, - sin: torch.Tensor, - seqlen_offsets: Union[int, torch.Tensor] = 0, - cu_seqlens: Optional[torch.Tensor] = None, - max_seqlen: Optional[int] = None, - interleaved=False, - inplace=False, - conjugate=False, -) -> torch.Tensor: - """ - Arguments: - x: (batch, seqlen, nheads, headdim) if cu_seqlens is None - else (total_seqlen, nheads, headdim). - cos: (seqlen_ro, rotary_dim / 2) - sin: (seqlen_ro, rotary_dim / 2) - seqlen_offsets: integer or integer tensor of size (batch,) - cu_seqlens: (batch + 1,) or None - max_seqlen: int - Returns: - y: (batch, seqlen, nheads, headdim) - """ - is_varlen = cu_seqlens is not None - if not is_varlen: - batch, seqlen, nheads, headdim = x.shape - else: - assert max_seqlen is not None, "If cu_seqlens is passed in, then max_seqlen must be passed" - total_seqlen, nheads, headdim = x.shape - batch_p_1 = cu_seqlens.shape[0] - batch = batch_p_1 - 1 - seqlen = max_seqlen - seqlen_ro, rotary_dim = cos.shape - assert sin.shape == cos.shape - rotary_dim *= 2 - assert rotary_dim <= headdim, "rotary_dim must be <= headdim" - assert headdim <= 256, "Only support headdim <= 256" - assert seqlen_ro >= seqlen, "seqlen_ro must be >= seqlen" - - cos, sin = cos.contiguous(), sin.contiguous() - if isinstance(seqlen_offsets, torch.Tensor): - assert seqlen_offsets.shape == (batch,) - assert seqlen_offsets.dtype in [torch.int32, torch.int64] - seqlen_offsets = seqlen_offsets.contiguous() - else: - assert seqlen_offsets + seqlen <= seqlen_ro - - output = torch.empty_like(x) if not inplace else x - if rotary_dim < headdim and not inplace: - output[..., rotary_dim:].copy_(x[..., rotary_dim:]) - - grid = lambda META: (triton.cdiv(nheads, META["BLOCK_H"]), triton.cdiv(seqlen, META["BLOCK_M"]), batch) # noqa - BLOCK_M = 8 if rotary_dim <= 128 else 4 - - # Need this, otherwise Triton tries to launch from cuda:0 and we get - # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?) - device_ctx = torch.cuda.device(x.device.index) if x.device.type == 'cuda' else torch.xpu.device(x.device.index) - with device_ctx: - torch.library.wrap_triton(rotary_kernel)[grid]( - output, # data ptrs - x, - cos, - sin, - cu_seqlens, - seqlen_offsets, - seqlen, # shapes - nheads, - seqlen_ro, - output.stride(0) if not is_varlen else 0, # batch_strides if not varlen else 0 - output.stride(-3), # seqlen_stride or total_seqlen_stride - output.stride(-2), # nheads_stride - output.stride(-1), # headdim_stride - x.stride(0) if not is_varlen else 0, # batch_strides if not varlen else 0 - x.stride(-3), # seqlen stride or total_seqlen_stride - x.stride(-2), # nheads stride - x.stride(-1), # headdim stride - rotary_dim, - isinstance(seqlen_offsets, torch.Tensor), - is_varlen, - interleaved, - conjugate, - BLOCK_M=BLOCK_M, - BLOCK_H=2, - ) - return output diff --git a/build/torch212-cxx11-cu132-aarch64-linux/__init__.py b/build/torch212-cxx11-cu132-aarch64-linux/__init__.py deleted file mode 100644 index ecc2f9d896b6c93f90b0a1499856dc0612177422..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu132-aarch64-linux/__init__.py +++ /dev/null @@ -1,393 +0,0 @@ -from typing import Optional, List -import torch -from ._ops import ops as flash_attn_ops -from .flash_attn_interface import ( - flash_attn_func, - flash_attn_kvpacked_func, - flash_attn_qkvpacked_func, - flash_attn_varlen_func, - flash_attn_varlen_kvpacked_func, - flash_attn_varlen_qkvpacked_func, - flash_attn_with_kvcache, -) - - -def fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - return_softmax: bool = False, - gen: Optional[torch.Generator] = None, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention. - - Args: - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Optional output tensor, same shape as q - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - return_softmax: Whether to return softmax weights - gen: Optional random number generator - - Returns: - List of tensors: [output, softmax_lse, (softmax if return_softmax)] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.fwd( - q, - k, - v, - out, - alibi_slopes, - p_dropout, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - gen, - ) - - -def varlen_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - out: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - block_table: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - max_seqlen_q: int = 0, - max_seqlen_k: int = 0, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - zero_tensors: bool = False, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - return_softmax: bool = False, - gen: Optional[torch.Generator] = None, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention with variable sequence lengths. - - Args: - q: Query tensor of shape [total_q, num_heads, head_size] - k: Key tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - v: Value tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1] - cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1] - out: Optional output tensor of shape [total_q, num_heads, head_size] - seqused_k: Optional tensor specifying how many keys to use per batch element [batch_size] - leftpad_k: Optional left padding for keys of shape [batch_size] - block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq] - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - max_seqlen_q: Maximum sequence length for queries - max_seqlen_k: Maximum sequence length for keys - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - zero_tensors: Whether to zero tensors before computation - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - return_softmax: Whether to return softmax weights - gen: Optional random number generator - - Returns: - List of tensors: [output, softmax_lse, (softmax if return_softmax)] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.varlen_fwd( - q, - k, - v, - out, - cu_seqlens_q, - cu_seqlens_k, - seqused_k, - leftpad_k, - block_table, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - p_dropout, - softmax_scale, - zero_tensors, - is_causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - gen, - ) - - -def bwd( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - p_dropout: float = 0.0, - 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, - gen: Optional[torch.Generator] = None, - rng_state: Optional[torch.Tensor] = None, -) -> List[torch.Tensor]: - """ - Backward pass for multi-head attention. - - Args: - dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size] - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size] - softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q] - dq: Optional gradient tensor for queries, same shape as q - dk: Optional gradient tensor for keys, same shape as k - dv: Optional gradient tensor for values, same shape as v - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - deterministic: Whether to use deterministic algorithms - gen: Optional random number generator - rng_state: Optional RNG state from forward pass - - Returns: - List of tensors: [dq, dk, dv] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - alibi_slopes, - p_dropout, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - deterministic, - gen, - rng_state, - ) - - -def varlen_bwd( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - max_seqlen_q: int = 0, - max_seqlen_k: int = 0, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - zero_tensors: bool = False, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - deterministic: bool = False, - gen: Optional[torch.Generator] = None, - rng_state: Optional[torch.Tensor] = None, -) -> List[torch.Tensor]: - """ - Backward pass for multi-head attention with variable sequence lengths. - - Args: - dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size] - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size] - softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q] - cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1] - cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1] - dq: Optional gradient tensor for queries, same shape as q - dk: Optional gradient tensor for keys, same shape as k - dv: Optional gradient tensor for values, same shape as v - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - max_seqlen_q: Maximum sequence length for queries - max_seqlen_k: Maximum sequence length for keys - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - zero_tensors: Whether to zero tensors before computation - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - deterministic: Whether to use deterministic algorithms - gen: Optional random number generator - rng_state: Optional RNG state from forward pass - - Returns: - List of tensors: [dq, dk, dv] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.varlen_bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - p_dropout, - softmax_scale, - zero_tensors, - is_causal, - window_size_left, - window_size_right, - softcap, - deterministic, - gen, - rng_state, - ) - - -def fwd_kvcache( - q: torch.Tensor, - kcache: torch.Tensor, - vcache: torch.Tensor, - k: Optional[torch.Tensor] = None, - v: Optional[torch.Tensor] = None, - seqlens_k: Optional[torch.Tensor] = None, - rotary_cos: Optional[torch.Tensor] = None, - rotary_sin: Optional[torch.Tensor] = None, - cache_batch_idx: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - block_table: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - out: 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, - is_rotary_interleaved: bool = False, - num_splits: int = 1, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention with KV cache. - - Args: - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - kcache: Key cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - vcache: Value cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - k: Optional new keys tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size] - v: Optional new values tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size] - seqlens_k: Optional sequence lengths for keys of shape [batch_size] - rotary_cos: Optional rotary cosine tensor of shape [seqlen_ro, rotary_dim/2] - rotary_sin: Optional rotary sine tensor of shape [seqlen_ro, rotary_dim/2] - cache_batch_idx: Optional indices to index into the KV cache - leftpad_k: Optional left padding for keys of shape [batch_size] - block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq] - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - out: Optional output tensor, same shape as q - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - is_rotary_interleaved: Whether rotary embeddings are interleaved - num_splits: Number of splits for computation - - Returns: - List of tensors: [output, softmax_lse] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.fwd_kvcache( - q, - kcache, - vcache, - k, - v, - seqlens_k, - rotary_cos, - rotary_sin, - cache_batch_idx, - leftpad_k, - block_table, - alibi_slopes, - out, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - is_rotary_interleaved, - num_splits, - ) diff --git a/build/torch212-cxx11-cu132-aarch64-linux/_flash_attn2_cuda_f12afc9.abi3.so b/build/torch212-cxx11-cu132-aarch64-linux/_flash_attn2_cuda_f12afc9.abi3.so deleted file mode 100644 index f62c259f00a5259340c58a6b2a90eaf3dd1859b8..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu132-aarch64-linux/_flash_attn2_cuda_f12afc9.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:275e8292d1e114ef224ff7b76e068dd70e05c7e7be08fc47786bba38ffa37513 -size 1024074040 diff --git a/build/torch212-cxx11-cu132-aarch64-linux/_ops.py b/build/torch212-cxx11-cu132-aarch64-linux/_ops.py deleted file mode 100644 index af5b731cb04b5e6ce54e862db0501f3683c3163f..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu132-aarch64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _flash_attn2_cuda_f12afc9 -ops = torch.ops._flash_attn2_cuda_f12afc9 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_flash_attn2_cuda_f12afc9::{op_name}" diff --git a/build/torch212-cxx11-cu132-aarch64-linux/bert_padding.py b/build/torch212-cxx11-cu132-aarch64-linux/bert_padding.py deleted file mode 100644 index 3c2d35159a014a9d03aabead9e52e009168696ea..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu132-aarch64-linux/bert_padding.py +++ /dev/null @@ -1,218 +0,0 @@ -# Adapted from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/padding.py - -import torch -import torch.nn.functional as F -from einops import rearrange, repeat - - -class IndexFirstAxis(torch.autograd.Function): - @staticmethod - def forward(ctx, input, indices): - ctx.save_for_backward(indices) - assert input.ndim >= 2 - ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] - second_dim = other_shape.numel() - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - # return input[indices] - return torch.gather( - rearrange(input, "b ... -> b (...)"), 0, repeat(indices, "z -> z d", d=second_dim) - ).reshape(-1, *other_shape) - - @staticmethod - def backward(ctx, grad_output): - (indices,) = ctx.saved_tensors - assert grad_output.ndim >= 2 - other_shape = grad_output.shape[1:] - grad_output = rearrange(grad_output, "b ... -> b (...)") - grad_input = torch.zeros( - [ctx.first_axis_dim, grad_output.shape[1]], - device=grad_output.device, - dtype=grad_output.dtype, - ) - # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. - # grad_input[indices] = grad_output - grad_input.scatter_(0, repeat(indices, "z -> z d", d=grad_output.shape[1]), grad_output) - return grad_input.reshape(ctx.first_axis_dim, *other_shape), None - - -index_first_axis = IndexFirstAxis.apply - - -class IndexPutFirstAxis(torch.autograd.Function): - @staticmethod - def forward(ctx, values, indices, first_axis_dim): - ctx.save_for_backward(indices) - assert indices.ndim == 1 - assert values.ndim >= 2 - output = torch.zeros( - first_axis_dim, *values.shape[1:], device=values.device, dtype=values.dtype - ) - # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. - output[indices] = values - # output.scatter_(0, repeat(indices, 'z -> z d', d=values.shape[1]), values) - return output - - @staticmethod - def backward(ctx, grad_output): - (indices,) = ctx.saved_tensors - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - grad_values = grad_output[indices] - # grad_values = torch.gather(grad_output, 0, repeat(indices, 'z -> z d', d=grad_output.shape[1])) - return grad_values, None, None - - -index_put_first_axis = IndexPutFirstAxis.apply - - -class IndexFirstAxisResidual(torch.autograd.Function): - @staticmethod - def forward(ctx, input, indices): - ctx.save_for_backward(indices) - assert input.ndim >= 2 - ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] - second_dim = other_shape.numel() - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - output = input[indices] - # We don't want to reshape input (b ... -> b (...)) since it could change the channel_last - # memory format to channel_first. In other words, input might not be contiguous. - # If we don't detach, Pytorch complains about output being a view and is being modified inplace - return output, input.detach() - - @staticmethod - def backward(ctx, grad_output, grad_residual): - (indices,) = ctx.saved_tensors - assert grad_output.ndim >= 2 - other_shape = grad_output.shape[1:] - assert grad_residual.shape[1:] == other_shape - grad_input = grad_residual - # grad_input[indices] += grad_output - indices = indices.reshape(indices.shape[0], *((1,) * (grad_output.ndim - 1))) - indices = indices.expand_as(grad_output) - grad_input.scatter_add_(0, indices, grad_output) - return grad_input.reshape(ctx.first_axis_dim, *other_shape), None - - -index_first_axis_residual = IndexFirstAxisResidual.apply - - -def unpad_input(hidden_states, attention_mask, unused_mask=None): - """ - Arguments: - hidden_states: (batch, seqlen, ...) - attention_mask: (batch, seqlen), bool / int, 1 means valid and 0 means not valid. - unused_mask: (batch, seqlen), bool / int, 1 means the element is allocated but unused. - Return: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens selected in attention_mask + unused_mask. - indices: (total_nnz), the indices of masked tokens from the flattened input sequence. - cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. - max_seqlen_in_batch: int - seqused: (batch), returns the number of tokens selected in attention_mask + unused_mask. - """ - all_masks = (attention_mask + unused_mask) if unused_mask is not None else attention_mask - seqlens_in_batch = all_masks.sum(dim=-1, dtype=torch.int32) - used_seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) - indices = torch.nonzero(all_masks.flatten(), as_tuple=False).flatten() - max_seqlen_in_batch = seqlens_in_batch.max().item() - cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) - # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the - # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim - # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to - # index with integer indices. Moreover, torch's index is a bit slower than it needs to be, - # so we write custom forward and backward to make it a bit faster. - return ( - index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices), - indices, - cu_seqlens, - max_seqlen_in_batch, - used_seqlens_in_batch, - ) - - -def unpad_input_for_concatenated_sequences(hidden_states, attention_mask_in_length): - """ - Supports concatenating short samples in one sequence. The attention_mask_in_length is utilized to mask other short samples. It helps efficient training of variant lengths-based samples (e.g., the supervised fine-tuning task in large language model). - The motivation for this function is explained [here](https://github.com/Dao-AILab/flash-attention/issues/432#issuecomment-1668822286). - - For example, if batch = 3 and seqlen = 6, the attention_mask_in_length is: - ``` - [ - [2, 3, 0, 0, 0, 0], - [3, 2, 0, 0, 0, 0], - [6, 0, 0, 0, 0, 0] - ] - ``` - , which refers to the 3D-attention mask: - ``` - [ - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [0, 0, 1, 0, 0, 0], - [0, 0, 1, 1, 0, 0], - [0, 0, 1, 1, 1, 0], - [0, 0, 0, 0, 0, 1] - ], - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [0, 0, 0, 1, 0, 0], - [0, 0, 0, 1, 1, 0], - [0, 0, 0, 0, 0, 1] - ], - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [1, 1, 1, 1, 0, 0], - [1, 1, 1, 1, 1, 0], - [1, 1, 1, 1, 1, 1] - ] - ] - ```. - - Arguments: - hidden_states: (batch, seqlen, ...) - attention_mask_in_length: (batch, seqlen), int, a nonzero number (e.g., 1, 2, 3, etc.) means length of concatenated sequence in b-th batch, and 0 means none. - Return: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. - indices: (total_nnz), the indices of non-masked tokens from the flattened input sequence. - cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. - max_seqlen_in_batch: int - """ - length = attention_mask_in_length.sum(dim=-1) - seqlen = attention_mask_in_length.size(-1) - attention_mask_2d = torch.arange(seqlen, device=length.device, dtype=length.dtype).expand(len(length), seqlen) < length.unsqueeze(1) - real_indices_idx = torch.nonzero(attention_mask_in_length.flatten(), as_tuple=False).flatten() - seqlens_in_batch = attention_mask_in_length.flatten()[real_indices_idx] - indices = torch.nonzero(attention_mask_2d.flatten(), as_tuple=False).flatten() - max_seqlen_in_batch = seqlens_in_batch.max().item() - cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) - # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the - # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim - # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to - # index with integer indices. Moreover, torch's index is a bit slower than it needs to be, - # so we write custom forward and backward to make it a bit faster. - return ( - index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices), - indices, - cu_seqlens, - max_seqlen_in_batch, - ) - - -def pad_input(hidden_states, indices, batch, seqlen): - """ - Arguments: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. - indices: (total_nnz), the indices that represent the non-masked tokens of the original padded input sequence. - batch: int, batch size for the padded sequence. - seqlen: int, maximum sequence length for the padded sequence. - Return: - hidden_states: (batch, seqlen, ...) - """ - dim = hidden_states.shape[-1] - # output = torch.zeros((batch * seqlen), dim, device=hidden_states.device, dtype=hidden_states.dtype) - # output[indices] = hidden_states - output = index_put_first_axis(hidden_states, indices, batch * seqlen) - return rearrange(output, "(b s) ... -> b s ...", b=batch) diff --git a/build/torch212-cxx11-cu132-aarch64-linux/flash_attn2/__init__.py b/build/torch212-cxx11-cu132-aarch64-linux/flash_attn2/__init__.py deleted file mode 100644 index a9b2672c1cd85b74c1b3ded0fc0b2100e1aeac23..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu132-aarch64-linux/flash_attn2/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -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/torch212-cxx11-cu132-aarch64-linux/flash_attn_interface.py b/build/torch212-cxx11-cu132-aarch64-linux/flash_attn_interface.py deleted file mode 100644 index b014e569dc9812a9c840d1043f2cebe767b4d477..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu132-aarch64-linux/flash_attn_interface.py +++ /dev/null @@ -1,1631 +0,0 @@ -# Copyright (c) 2023, Tri Dao. - -from typing import Optional, Sequence, Tuple, Union - -import torch -import torch.nn as nn -import os - -# # isort: off -# # We need to import the CUDA kernels after importing torch -# USE_TRITON_ROCM = os.getenv("FLASH_ATTENTION_TRITON_AMD_ENABLE", "FALSE") == "TRUE" -# if USE_TRITON_ROCM: -# from .flash_attn_triton_amd import interface_fa as flash_attn -# else: -# import flash_attn_2_cuda as flash_attn - - -from . import _ops -from ._ops import ops as flash_attn -from ._ops import add_op_namespace_prefix - -# # isort: on - -def maybe_contiguous(x): - return x.contiguous() if x is not None and x.stride(-1) != 1 else x - - -def _get_device(): - if torch.xpu.is_available(): - return "xpu" - elif torch.cuda.is_available(): - return "cuda" - else: - return "cpu" - - -def _get_block_size_n(device, head_dim, is_dropout, is_causal): - assert head_dim <= 256 - - if device.type == "xpu": - if head_dim <= 96: - return 64 - elif head_dim <= 128: - return 32 - elif head_dim <= 256: - return 64 - else: - return 32 - - # This should match the block sizes in the CUDA kernel - major, minor = torch.cuda.get_device_capability(device) - is_sm8x = major == 8 and minor > 0 # Only include sm86 and sm89, exclude sm80 (A100) - is_sm80 = major == 8 and minor == 0 - is_sm90 = major == 9 and minor == 0 - if head_dim <= 32: - return 128 - if head_dim <= 64: - return 128 if not is_dropout else 64 - elif head_dim <= 96: - return 64 - elif head_dim <= 128: - if is_sm8x: - return 64 if (not is_dropout and is_causal) else 32 - else: - return 64 if not is_dropout else 32 - elif head_dim <= 192: - return 64 - elif head_dim <= 224: - return 64 - elif head_dim <= 256: - return 64 - - -def round_multiple(x, m): - return (x + m - 1) // m * m - - -# torch.compile() support is only enabled for pytorch >= 2.4 -# The reason for this is that we are using the new custom_op and register_fake -# APIs, which support inplace modification of inputs in the function itself -if torch.__version__ >= "2.4.0": - _torch_custom_op_wrapper = torch.library.custom_op - _torch_register_fake_wrapper = torch.library.register_fake -else: - def noop_custom_op_wrapper(name, fn=None, /, *, mutates_args, device_types=None, schema=None): - def wrap(func): - return func - if fn is None: - return wrap - return fn - def noop_register_fake_wrapper(op, fn=None, /, *, lib=None, _stacklevel=1): - def wrap(func): - return func - if fn is None: - return wrap - return fn - _torch_custom_op_wrapper = noop_custom_op_wrapper - _torch_register_fake_wrapper = noop_register_fake_wrapper - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_forward"), mutates_args=(), device_types=_get_device()) -def _flash_attn_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - return_softmax: bool -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - out, softmax_lse, S_dmask, rng_state = flash_attn.fwd( - q, - k, - v, - None, - alibi_slopes, - dropout_p, - softmax_scale, - causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - None, - ) - return out, softmax_lse, S_dmask, rng_state - - -@_torch_register_fake_wrapper(add_op_namespace_prefix("_flash_attn_forward")) -def _flash_attn_forward_fake( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - return_softmax: bool -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - batch_size, seqlen_q, num_heads, head_size = q.shape - seqlen_k = k.shape[1] - out = torch.empty_like(q) - softmax_lse = torch.empty((batch_size, num_heads, seqlen_q), dtype=torch.float32, device=q.device, layout=q.layout) - p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) - if return_softmax: - p = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128), round_multiple(seqlen_k, 128)), dtype=q.dtype, device=q.device, layout=q.layout) - rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) - - return out, softmax_lse, p, rng_state - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_forward = _ops.ops._flash_attn_forward -else: - _wrapped_flash_attn_forward = _flash_attn_forward - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_varlen_forward"), mutates_args=(), device_types=_get_device()) -def _flash_attn_varlen_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - alibi_slopes: Optional[torch.Tensor] = None, - return_softmax: bool = False, - block_table: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - out, softmax_lse, S_dmask, rng_state = flash_attn.varlen_fwd( - q, - k, - v, - None, - cu_seqlens_q, - cu_seqlens_k, - seqused_k, - leftpad_k, - block_table, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - zero_tensors, - causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - None, - ) - # if out.isnan().any() or softmax_lse.isnan().any(): - # breakpoint() - return out, softmax_lse, S_dmask, rng_state - - -@_torch_register_fake_wrapper(add_op_namespace_prefix("_flash_attn_varlen_forward")) -def _flash_attn_varlen_forward_fake( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - alibi_slopes: Optional[torch.Tensor] = None, - return_softmax: bool = False, - block_table: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - paged_kv = block_table is not None - batch_size = cu_seqlens_q.numel() - 1 - total_q, num_heads, _ = q.shape - - out = torch.empty_like(q) - softmax_lse = torch.empty((num_heads, total_q), dtype=torch.float32, device=q.device, layout=q.layout) - p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) - seqlen_q_rounded = round_multiple(max_seqlen_q, 128) - seqlen_k_rounded = round_multiple(max_seqlen_k, 128) - if return_softmax: - p = torch.empty((batch_size, num_heads, seqlen_q_rounded, seqlen_k_rounded), dtype=q.dtype, device=q.device, layout=q.layout) - rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) - return out, softmax_lse, p, rng_state - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_varlen_forward = _ops.ops._flash_attn_varlen_forward -else: - _wrapped_flash_attn_varlen_forward = _flash_attn_varlen_forward - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_backward"), mutates_args=("dq", "dk", "dv"), device_types=_get_device()) -def _flash_attn_backward( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, -) -> 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)] - ( - dq, - dk, - dv, - softmax_d, - ) = flash_attn.bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - alibi_slopes, - dropout_p, - softmax_scale, - causal, - window_size_left, - window_size_right, - softcap, - deterministic, - None, - rng_state, - ) - return softmax_d - - -@_torch_register_fake_wrapper(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, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, -) -> torch.Tensor: - dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] - if dq is None: - dq = torch.empty_like(q) - if dk is None: - dk = torch.empty_like(k) - if dv is None: - dv = torch.empty_like(v) - batch_size, seqlen_q, num_heads, _ = q.shape - softmax_d = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128)), device=q.device, dtype=torch.float32) - - return softmax_d - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_backward = _ops.ops._flash_attn_backward -else: - _wrapped_flash_attn_backward = _flash_attn_backward - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_varlen_backward"), mutates_args=("dq", "dk", "dv"), device_types=_get_device()) -def _flash_attn_varlen_backward( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> 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)] - ( - dq, - dk, - dv, - softmax_d, - ) = flash_attn.varlen_bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - zero_tensors, - causal, - window_size_left, - window_size_right, - softcap, - deterministic, - None, - rng_state, - ) - # if dk.isnan().any() or dk.isnan().any() or dv.isnan().any() or softmax_d.isnan().any(): - # breakpoint() - return softmax_d - - -@_torch_register_fake_wrapper(add_op_namespace_prefix("_flash_attn_varlen_backward")) -def _flash_attn_varlen_backward_fake( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> torch.Tensor: - dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] - batch_size = cu_seqlens_q.numel() - 1 - total_q, num_heads, _ = q.shape - - if dq is None: - dq = torch.empty_like(q) - if dk is None: - dk = torch.empty_like(k) - if dv is None: - dv = torch.empty_like(v) - softmax_d = torch.empty((num_heads, total_q + 128 * batch_size), device=q.device, dtype=torch.float32) - - return softmax_d - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_varlen_backward = _ops.ops._flash_attn_varlen_backward -else: - _wrapped_flash_attn_varlen_backward = _flash_attn_varlen_backward - - -class FlashAttnQKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and qkv.requires_grad - if softmax_scale is None: - softmax_scale = qkv.shape[-1] ** (-0.5) - q, k, v = qkv[:, :, 0].detach(), qkv[:, :, 1].detach(), qkv[:, :, 2].detach() - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) - dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dqkv[:, :, 0], - dqkv[:, :, 1], - dqkv[:, :, 2], - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dqkv = dqkv[..., : dout.shape[-1]] # We could have padded the head dimension - return dqkv, None, None, None, None, None, None, None, None, None - - -class FlashAttnVarlenQKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - cu_seqlens, - max_seqlen, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and qkv.requires_grad - if softmax_scale is None: - softmax_scale = qkv.shape[-1] ** (-0.5) - q, k, v = qkv[:, 0].detach(), qkv[:, 1].detach(), qkv[:, 2].detach() - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens, - cu_seqlens, - max_seqlen, - max_seqlen, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=None, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens, rng_state) - ctx.dropout_p = dropout_p - ctx.max_seqlen = max_seqlen - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens, rng_state = ctx.saved_tensors - qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) - dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dqkv[:, 0], - dqkv[:, 1], - dqkv[:, 2], - cu_seqlens, - cu_seqlens, - ctx.max_seqlen, - ctx.max_seqlen, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - 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 FlashAttnKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - kv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, kv] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - k, v = kv[:, :, 0].detach(), kv[:, :, 1].detach() - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - dq = torch.empty_like(q) - kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) - dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dkv[:, :, 0], - dkv[:, :, 1], - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dkv = dkv[..., : dout.shape[-1]] - return dq, dkv, None, None, None, None, None, None, None, None, None - - -class FlashAttnVarlenKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, kv] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - k, v = kv[:, 0].detach(), kv[:, 1].detach() - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=None, - ) - if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state - ) - ctx.dropout_p = dropout_p - 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.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors - dq = torch.empty_like(q) - kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) - dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dkv[:, 0], - dkv[:, 1], - cu_seqlens_q, - cu_seqlens_k, - ctx.max_seqlen_q, - ctx.max_seqlen_k, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dkv = dkv[..., : dout.shape[-1]] - return dq, dkv, None, None, None, None, None, None, None, None, None, None, None, None, None - - -class FlashAttnFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - k, - v, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, k, v] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dk = dk[..., : dout.shape[-1]] - dv = dv[..., : dout.shape[-1]] - return dq, dk, dv, 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, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - block_table, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, k, v] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=block_table, - ) - if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state - ) - ctx.dropout_p = dropout_p - 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.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors - dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - ctx.max_seqlen_q, - ctx.max_seqlen_k, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dk = dk[..., : dout.shape[-1]] - dv = dv[..., : dout.shape[-1]] - return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None - - -def flash_attn_qkvpacked_func( - qkv, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # <=0.0 means deactivate - alibi_slopes=None, - deterministic=False, - 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, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_kvpacked_func( - q, - kv, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, -): - """dropout_p should be set to 0.0 during evaluation - If 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 K, V. - 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) - kv: (batch_size, seqlen, 2, 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. - softcap: float. Anything > 0 activates softcapping 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). - 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 FlashAttnKVPackedFunc.apply( - q, - kv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_func( - q, - k, - v, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - 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). - 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 FlashAttnFunc.apply( - q, - k, - v, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_qkvpacked_func( - qkv, - cu_seqlens, - max_seqlen, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - 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_varlen_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_varlen_kvpacked_func and flash_attn_varlen_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: (total, 3, nheads, headdim), where total = total number of tokens in the batch. - cu_seqlens: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into qkv. - max_seqlen: int. Maximum sequence length in the batch. - 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenQKVPackedFunc.apply( - qkv, - cu_seqlens, - max_seqlen, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_kvpacked_func( - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, -): - """dropout_p should be set to 0.0 during evaluation - If 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 K, V. - 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: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch. - kv: (total_k, 2, nheads_k, headdim), where total_k = total number of key tokens in the batch. - cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into q. - cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into kv. - max_seqlen_q: int. Maximum query sequence length in the batch. - max_seqlen_k: int. Maximum key sequence length in the batch. - 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 + 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenKVPackedFunc.apply( - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_func( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, - block_table=None, -): - """dropout_p should be set to 0.0 during evaluation - Supports multi-query and grouped-query attention (MQA/GQA) by passing in K, V 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: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch. - k: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. - v: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. - cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into q. - cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into kv. - max_seqlen_q: int. Maximum query sequence length in the batch. - max_seqlen_k: int. Maximum key sequence length in the batch. - 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 + 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenFunc.apply( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - block_table, - False if q.device.type == "cpu" else torch.is_grad_enabled(), - ) - - -def flash_attn_with_kvcache( - q, - k_cache, - v_cache, - k=None, - v=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, - block_table: Optional[torch.Tensor] = None, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - rotary_interleaved=True, - alibi_slopes=None, - num_splits=0, - 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 block_table, - or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_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) if there's no block_table, - or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_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). Similar to k. - 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. - block_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). - 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. - 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" - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-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) - cache_batch_idx = maybe_contiguous(cache_batch_idx) - block_table = maybe_contiguous(block_table) - out, softmax_lse = flash_attn.fwd_kvcache( - q, - k_cache, - v_cache, - k, - v, - cache_seqlens, - rotary_cos, - rotary_sin, - cache_batch_idx, - cache_leftpad, - block_table, - alibi_slopes, - None, - softmax_scale, - causal, - window_size[0], - window_size[1], - softcap, - rotary_interleaved, - num_splits, - ) - return (out, softmax_lse) if return_softmax_lse else out diff --git a/build/torch212-cxx11-cu132-aarch64-linux/layers/__init__.py b/build/torch212-cxx11-cu132-aarch64-linux/layers/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/build/torch212-cxx11-cu132-aarch64-linux/layers/patch_embed.py b/build/torch212-cxx11-cu132-aarch64-linux/layers/patch_embed.py deleted file mode 100644 index 05562f8e8bcdb58e947c6f402a49eacd2d031871..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu132-aarch64-linux/layers/patch_embed.py +++ /dev/null @@ -1,67 +0,0 @@ -# We use the same API as https://github.com/rwightman/pytorch-image-models/blob/v0.6.11/timm/models/layers/patch_embed.py -# But we use nn.Linear instead of Conv2d and it's about 8x faster. - -from functools import partial - -import torch.nn as nn -from einops import rearrange -from torch import _assert -from torch.nn.modules.utils import _pair - -try: - from flash_attn.ops.fused_dense import FusedDense -except ImportError: - FusedDense = None - - -class PatchEmbed(nn.Module): - """2D Image to Patch Embedding""" - - def __init__( - self, - img_size=224, - patch_size=16, - in_chans=3, - embed_dim=768, - norm_layer=None, - flatten=True, - bias=True, - fused_bias_fc=False, - ): - super().__init__() - img_size = _pair(img_size) - patch_size = _pair(patch_size) - self.img_size = img_size - self.patch_size = patch_size - self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1]) - self.num_patches = self.grid_size[0] * self.grid_size[1] - self.flatten = flatten - if fused_bias_fc and FusedDense is None: - raise ImportError("fused_dense is not installed") - - linear_cls = nn.Linear if not fused_bias_fc or not bias else FusedDense - self.proj = linear_cls(in_chans * patch_size[0] * patch_size[1], embed_dim, bias=bias) - self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() - - def forward(self, x): - _, _, H, W = x.shape - _assert( - H == self.img_size[0], - f"Input image height ({H}) doesn't match model ({self.img_size[0]}).", - ) - _assert( - W == self.img_size[1], - f"Input image width ({W}) doesn't match model ({self.img_size[1]}).", - ) - x = self.proj( - rearrange( - x, - "b c (h p1) (w p2) -> b h w (c p1 p2)", - p1=self.patch_size[0], - p2=self.patch_size[1], - ) - ) - if self.flatten: - x = rearrange(x, "b h w c -> b (h w) c") - x = self.norm(x) - return x diff --git a/build/torch212-cxx11-cu132-aarch64-linux/layers/rotary.py b/build/torch212-cxx11-cu132-aarch64-linux/layers/rotary.py deleted file mode 100644 index d1bfc21fc7de1dd287e8f382847b194a48075981..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu132-aarch64-linux/layers/rotary.py +++ /dev/null @@ -1,483 +0,0 @@ -# Copyright (c) 2025, Tri Dao - -import math -from functools import partial -from typing import Optional, Tuple, Union - -import torch -from torch import Tensor - -from einops import rearrange, repeat -# from flash_attn.ops.triton.rotary import apply_rotary -from ..ops.triton.rotary import apply_rotary - - -def rotate_half(x, interleaved=False): - if not interleaved: - x1, x2 = x.chunk(2, dim=-1) - return torch.cat((-x2, x1), dim=-1) - else: - x1, x2 = x[..., ::2], x[..., 1::2] - return rearrange(torch.stack((-x2, x1), dim=-1), "... d two -> ... (d two)", two=2) - - -def apply_rotary_emb_torch(x, cos, sin, interleaved=False): - """ - x: (batch_size, seqlen, nheads, headdim) - cos, sin: (seqlen, rotary_dim / 2) or (batch_size, seqlen, rotary_dim / 2) - """ - ro_dim = cos.shape[-1] * 2 - assert ro_dim <= x.shape[-1] - cos = repeat(cos, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)") - sin = repeat(sin, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)") - return torch.cat( - [x[..., :ro_dim] * cos + rotate_half(x[..., :ro_dim], interleaved) * sin, x[..., ro_dim:]], - dim=-1, - ) - - -class ApplyRotaryEmb(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x, - cos, - sin, - interleaved=False, - inplace=False, - seqlen_offsets: Union[int, Tensor] = 0, - cu_seqlens: Optional[Tensor] = None, - max_seqlen: Optional[int] = None, - ): - out = apply_rotary( - x, - cos, - sin, - seqlen_offsets=seqlen_offsets, - cu_seqlens=cu_seqlens, - max_seqlen=max_seqlen, - interleaved=interleaved, - inplace=inplace, - ) - if isinstance(seqlen_offsets, int): - ctx.save_for_backward(cos, sin, cu_seqlens) # Can't save int with save_for_backward - ctx.seqlen_offsets = seqlen_offsets - else: - ctx.save_for_backward(cos, sin, cu_seqlens, seqlen_offsets) - ctx.seqlen_offsets = None - ctx.interleaved = interleaved - ctx.inplace = inplace - ctx.max_seqlen = max_seqlen - return out if not inplace else x - - @staticmethod - def backward(ctx, do): - seqlen_offsets = ctx.seqlen_offsets - if seqlen_offsets is None: - cos, sin, cu_seqlens, seqlen_offsets = ctx.saved_tensors - else: - cos, sin, cu_seqlens = ctx.saved_tensors - dx = apply_rotary( - do, - cos, - sin, - seqlen_offsets=seqlen_offsets, - cu_seqlens=cu_seqlens, - max_seqlen=ctx.max_seqlen, - interleaved=ctx.interleaved, - inplace=ctx.inplace, - conjugate=True, - ) - return dx, None, None, None, None, None, None, None - - -def apply_rotary_emb( - x, - cos, - sin, - interleaved=False, - inplace=False, - seqlen_offsets: Union[int, Tensor] = 0, - cu_seqlens: Optional[Tensor] = None, - max_seqlen: Optional[int] = None, -): - """ - Arguments: - x: (batch_size, seqlen, nheads, headdim) if cu_seqlens is None - else (total_seqlen, nheads, headdim) - cos, sin: (seqlen_rotary, rotary_dim / 2) - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead - of 1st half and 2nd half (GPT-NeoX style). - inplace: if True, apply rotary embedding in-place. - seqlen_offsets: (batch_size,) or int. Each sequence in x is shifted by this amount. - Most commonly used in inference when we have KV cache. - cu_seqlens: (batch + 1,) or None - max_seqlen: int - Return: - out: (batch_size, seqlen, nheads, headdim) if cu_seqlens is None - else (total_seqlen, nheads, headdim) - rotary_dim must be <= headdim - Apply rotary embedding to the first rotary_dim of x. - """ - return ApplyRotaryEmb.apply( - x, cos, sin, interleaved, inplace, seqlen_offsets, cu_seqlens, max_seqlen - ) - - -# For backward compatibility -apply_rotary_emb_func = apply_rotary_emb - - -def _apply_rotary_emb_qkv( - qkv, - cos, - sin, - cos_k=None, - sin_k=None, - interleaved=False, - inplace=False, - conjugate=False, - seqlen_offsets: Union[int, Tensor] = 0, - num_heads_q: Optional[int] = None, -): - apply_rotary_fn = partial( - apply_rotary, - interleaved=interleaved, - inplace=inplace, - conjugate=conjugate, - seqlen_offsets=seqlen_offsets - ) - if cos_k is None and sin_k is None and qkv.is_contiguous(): - # Call 1 kernel instead of 2 kernels - # We need qkv to be contiguous so that when we reshape to combine (3, nheads) - # dimensions, we get the same tensor - if qkv.dim() == 5: - batch, seqlen, three, nheads, headdim = qkv.shape - assert three == 3 - # qk = rearrange(qkv[:, :, :2], "b s t h d -> b s (t h) d") - qk = qkv[:, :, :2].reshape(batch, seqlen, -1, headdim) - qk = apply_rotary_fn(qk, cos, sin) - else: - assert qkv.dim() == 4 - assert num_heads_q is not None - num_heads_k = (qkv.shape[2] - num_heads_q) // 2 - assert qkv.shape[2] == num_heads_q + 2 * num_heads_k - qk = qkv[:, :, :num_heads_q + num_heads_k] - qk = apply_rotary_fn(qk, cos, sin) - if not inplace: - if qkv.dim() == 5: - qkv = torch.cat([rearrange(qk, "b s (t h) d -> b s t h d", t=2), qkv[:, :, 2:]], dim=2) - else: - qkv = torch.cat([qk, qkv[:, :, num_heads_q + num_heads_k :]], dim=2) - else: - cos_k = cos if cos_k is None else cos_k - sin_k = sin if sin_k is None else sin_k - if qkv.dim() == 5: - batch, seqlen, three, nheads, headdim = qkv.shape - assert three == 3 - q, k = qkv[:, :, 0], qkv[:, :, 1] - else: - assert qkv.dim() == 4 - assert num_heads_q is not None - num_heads_k = (qkv.shape[2] - num_heads_q) // 2 - assert qkv.shape[2] == num_heads_q + 2 * num_heads_k - q, k = qkv[:, :, :num_heads_q], qkv[:, :, num_heads_q : num_heads_q + num_heads_k] - q = apply_rotary_fn(q, cos, sin) - k = apply_rotary_fn(k, cos_k, sin_k) - if not inplace: - if qkv.dim() == 5: - qkv = torch.stack([q, k, qkv[:, :, 2]], dim=2) - else: - qkv = torch.cat([q, k, qkv[:, :, num_heads_q + num_heads_k:]], dim=2) - return qkv - - -class ApplyRotaryEmbQKV_(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - cos, - sin, - cos_k=None, - sin_k=None, - interleaved=False, - seqlen_offsets: Union[int, torch.Tensor] = 0, - num_heads_q: Optional[int] = None, - ): - # apply_rotary_emb_qkv_inplace( - qkv = _apply_rotary_emb_qkv( - qkv, cos, sin, cos_k, sin_k, interleaved=interleaved, inplace=True, - seqlen_offsets=seqlen_offsets, num_heads_q=num_heads_q, - ) - if isinstance(seqlen_offsets, int): - ctx.save_for_backward(cos, sin, cos_k, sin_k) - ctx.seqlen_offsets = seqlen_offsets - else: - ctx.save_for_backward(cos, sin, cos_k, sin_k, seqlen_offsets) - ctx.seqlen_offsets = None - ctx.interleaved = interleaved - ctx.num_heads_q = num_heads_q - return qkv - - @staticmethod - def backward(ctx, dqkv): - seqlen_offsets = ctx.seqlen_offsets - if seqlen_offsets is None: - cos, sin, cos_k, sin_k, seqlen_offsets = ctx.saved_tensors - else: - cos, sin, cos_k, sin_k = ctx.saved_tensors - dqkv = _apply_rotary_emb_qkv( - dqkv, cos, sin, cos_k, sin_k, interleaved=ctx.interleaved, inplace=True, - seqlen_offsets=seqlen_offsets, num_heads_q=ctx.num_heads_q, conjugate=True, - ) - return dqkv, None, None, None, None, None, None, None - - -def apply_rotary_emb_qkv_( - qkv, - cos, - sin, - cos_k=None, - sin_k=None, - interleaved=False, - seqlen_offsets: Union[int, torch.Tensor] = 0, - num_heads_q: Optional[int] = None, -): - """ - Arguments: - qkv: (batch_size, seqlen, 3, nheads, headdim) or (batch_size, seqlen, num_heads_q + 2 * num_heads_k, headdim). - If qkv has shape (batch_size, seqlen, num_heads_q + 2 * num_heads_k, headdim) (e.g. MQA / GQA), - then num_heads_q must be provided. - cos, sin: (seqlen, rotary_dim / 2) - cos_k, sin_k: (seqlen, rotary_dim / 2), optional - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead of - 1st half and 2nd half (GPT-NeoX style). - seqlen_offsets: (batch_size,) or int. Each sequence in Q and K is shifted by this amount. - Most commonly used in inference when we have KV cache. - Return: - qkv: (batch_size, seqlen, 3, nheads, headdim) or (batch_size, seqlen, num_heads_q + 2 * num_heads_k, headdim) - rotary_dim must be <= headdim - Apply rotary embedding *inplace* to the first rotary_dim of Q and K. - """ - return ApplyRotaryEmbQKV_.apply( - qkv, cos, sin, cos_k, sin_k, interleaved, seqlen_offsets, num_heads_q - ) - - -class ApplyRotaryEmbKV_(torch.autograd.Function): - - @staticmethod - def forward(ctx, kv, cos, sin, interleaved=False, seqlen_offsets: Union[int, torch.Tensor] = 0): - batch, seqlen, two, nheads, headdim = kv.shape - assert two == 2 - k = kv[:, :, 0] - apply_rotary( - k, cos, sin, seqlen_offsets=seqlen_offsets, interleaved=interleaved, inplace=True - ) - if isinstance(seqlen_offsets, int): - ctx.save_for_backward(cos, sin) # Can't save int with save_for_backward - ctx.seqlen_offsets = seqlen_offsets - else: - ctx.save_for_backward(cos, sin, seqlen_offsets) - ctx.seqlen_offsets = None - ctx.interleaved = interleaved - return kv - - @staticmethod - def backward(ctx, dkv): - seqlen_offsets = ctx.seqlen_offsets - if seqlen_offsets is None: - cos, sin, seqlen_offsets = ctx.saved_tensors - else: - cos, sin = ctx.saved_tensors - apply_rotary( - dkv[:, :, 0], - cos, - sin, - seqlen_offsets=seqlen_offsets, - interleaved=ctx.interleaved, - inplace=True, - conjugate=True, - ) - return dkv, None, None, None, None - - -apply_rotary_emb_kv_ = ApplyRotaryEmbKV_.apply - - -def apply_rotary_emb_kv_( - kv, - cos, - sin, - interleaved=False, - seqlen_offsets: Union[int, torch.Tensor] = 0, -): - """ - Arguments: - kv: (batch_size, seqlen, 2, nheads, headdim) - cos, sin: (seqlen, rotary_dim / 2) - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead of - 1st half and 2nd half (GPT-NeoX style). - seqlen_offsets: (batch_size,) or int. Each sequence in Q and K is shifted by this amount. - Most commonly used in inference when we have KV cache. - Return: - kv: (batch_size, seqlen, 2, nheads, headdim) - rotary_dim must be <= headdim - Apply rotary embedding *inplace* to the first rotary_dim of K. - """ - return ApplyRotaryEmbKV_.apply(kv, cos, sin, interleaved, seqlen_offsets) - - -class RotaryEmbedding(torch.nn.Module): - """ - The rotary position embeddings from RoFormer_ (Su et. al). - A crucial insight from the method is that the query and keys are - transformed by rotation matrices which depend on the relative positions. - - Other implementations are available in the Rotary Transformer repo_ and in - GPT-NeoX_, GPT-NeoX was an inspiration - - .. _RoFormer: https://arxiv.org/abs/2104.09864 - .. _repo: https://github.com/ZhuiyiTechnology/roformer - .. _GPT-NeoX: https://github.com/EleutherAI/gpt-neox - - If scale_base is not None, this implements XPos (Sun et al., https://arxiv.org/abs/2212.10554). - A recommended value for scale_base is 512: https://github.com/HazyResearch/flash-attention/issues/96 - Reference: https://github.com/sunyt32/torchscale/blob/main/torchscale/component/xpos_relative_position.py - """ - - def __init__( - self, - dim: int, - base=10000.0, - interleaved=False, - scale_base=None, - device=None, - ): - """ - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead - of 1st half and 2nd half (GPT-NeoX style). - """ - super().__init__() - self.dim = dim - self.base = float(base) - # Generate and save the inverse frequency buffer (non trainable) - inv_freq = self._compute_inv_freq(device) - self.register_buffer("inv_freq", inv_freq, persistent=False) - self.interleaved = interleaved - self.scale_base = scale_base - scale = ( - (torch.arange(0, dim, 2, device=device, dtype=torch.float32) + 0.4 * dim) / (1.4 * dim) - if scale_base is not None - else None - ) - self.register_buffer("scale", scale, persistent=False) - - self._seq_len_cached = 0 - self._cos_cached = None - self._sin_cached = None - self._cos_k_cached = None - self._sin_k_cached = None - - def _compute_inv_freq(self, device=None): - return 1.0 / ( - self.base - ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim) - ) - - def _update_cos_sin_cache(self, seqlen, device=None, dtype=None): - # Reset the tables if the sequence length has changed, - # if we're on a new device (possibly due to tracing for instance), - # or if we're switching from inference mode to training - if ( - seqlen > self._seq_len_cached - or self._cos_cached is None - or self._cos_cached.device != device - or self._cos_cached.dtype != dtype - or (self.training and self._cos_cached.is_inference()) - ): - self._seq_len_cached = seqlen - # We want fp32 here, not self.inv_freq.dtype, since the model could be loaded in bf16 - # And the output of arange can be quite large, so bf16 would lose a lot of precision. - t = torch.arange(seqlen, device=device, dtype=torch.float32) - # We want fp32 here as well since inv_freq will be multiplied with t, and the output - # will be large. Having it in bf16 will lose a lot of precision and cause the - # cos & sin output to change significantly. - # We want to recompute self.inv_freq if it was not loaded in fp32 - if self.inv_freq.dtype != torch.float32: - inv_freq = self._compute_inv_freq(device=device) - else: - inv_freq = self.inv_freq - # Don't do einsum, it converts fp32 to bf16 under AMP - # freqs = torch.einsum("i,j->ij", t, self.inv_freq) - freqs = torch.outer(t, inv_freq) - if self.scale is None: - self._cos_cached = torch.cos(freqs).to(dtype) - self._sin_cached = torch.sin(freqs).to(dtype) - else: - power = ( - torch.arange(seqlen, dtype=self.scale.dtype, device=self.scale.device) - - seqlen // 2 - ) / self.scale_base - scale = self.scale.to(device=power.device) ** rearrange(power, "s -> s 1") - # We want the multiplication by scale to happen in fp32 - self._cos_cached = (torch.cos(freqs) * scale).to(dtype) - self._sin_cached = (torch.sin(freqs) * scale).to(dtype) - self._cos_k_cached = (torch.cos(freqs) / scale).to(dtype) - self._sin_k_cached = (torch.sin(freqs) / scale).to(dtype) - - def forward( - self, - qkv: torch.Tensor, - kv: Optional[torch.Tensor] = None, - seqlen_offset: Union[int, torch.Tensor] = 0, - max_seqlen: Optional[int] = None, - num_heads_q: Optional[int] = None, - ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: - """ - qkv: (batch, seqlen, 3, nheads, headdim) or (batch, seqlen, num_heads_q + 2 * num_heads_k, headdim) - if kv is none, else it's just q of shape (batch, seqlen, nheads, headdim). - If qkv has shape (batch, seqlen, num_heads_q + 2 * num_heads_k, headdim) (e.g. MQA / GQA), - then num_heads_q must be provided. - kv: (batch, seqlen, 2, nheads, headdim) - seqlen_offset: (batch_size,) or int. Each sequence in x is shifted by this amount. - Most commonly used in inference when we have KV cache. - If it's a tensor of shape (batch_size,), then to update the cos / sin cache, one - should pass in max_seqlen, which will update the cos / sin cache up to that length. - Apply rotary embedding *inplace* to qkv and / or kv. - """ - seqlen = qkv.shape[1] - if max_seqlen is not None: - self._update_cos_sin_cache(max_seqlen, device=qkv.device, dtype=qkv.dtype) - elif isinstance(seqlen_offset, int): - self._update_cos_sin_cache(seqlen + seqlen_offset, device=qkv.device, dtype=qkv.dtype) - if kv is None: - return apply_rotary_emb_qkv_( - qkv, - self._cos_cached, - self._sin_cached, - self._cos_k_cached if self.scale is not None else None, - self._sin_k_cached if self.scale is not None else None, - interleaved=self.interleaved, - seqlen_offsets=seqlen_offset, - num_heads_q=num_heads_q, - ) - else: - q = qkv - q = apply_rotary_emb_func( - q, - self._cos_cached, - self._sin_cached, - interleaved=self.interleaved, - inplace=True, - seqlen_offsets=seqlen_offset, - ) - kv = apply_rotary_emb_kv_( - kv, - self._cos_cached if self.scale is None else self._cos_k_cached, - self._sin_cached if self.scale is None else self._sin_k_cached, - interleaved=self.interleaved, - seqlen_offsets=seqlen_offset, - ) - return q, kv diff --git a/build/torch212-cxx11-cu132-aarch64-linux/metadata.json b/build/torch212-cxx11-cu132-aarch64-linux/metadata.json deleted file mode 100644 index 439a0d1a68fbe6e3b33ef98719956a029d569573..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu132-aarch64-linux/metadata.json +++ /dev/null @@ -1,42 +0,0 @@ -{ - "name": "flash-attn2", - "id": "_flash_attn2_cuda_f12afc9", - "version": 1, - "license": "BSD-3-Clause", - "python-depends": [], - "backend": { - "type": "cuda", - "archs": [ - "10.0", - "12.0", - "8.0", - "9.0" - ] - }, - "digest": { - "algorithm": "sha256", - "files": { - "__init__.py": "Cd9MEdd/WUYnBmfiQpC4SyPqAxrTilTjvsvswNju4zc=", - "_flash_attn2_cuda_f12afc9.abi3.so": "J16CktHhFO8iT/e3bgaN1w4Fx+e+CPxHeGu6OP+jdRM=", - "_ops.py": "GQbmTTeSOU+7Q2nI68Hz+1RvEogyfEFccx1ubUYw+2s=", - "bert_padding.py": "gF1EmsdJ+HpQ86MRQ4VxDw+Sb/RVISdQALdNnoByHlw=", - "flash_attn2/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY=", - "flash_attn_interface.py": "W7a9bHR2ly6QL3aKqsBqxWA7M7ZWYcZPmtsPJu6rAfY=", - "layers/__init__.py": "47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU=", - "layers/patch_embed.py": "H58CgME/qSOPTZLOG08wFgrQS1j34pvNwMPrkTj3Ek4=", - "layers/rotary.py": "rFIdNrZQFNawwLD6d7t5LE2m+M93wwDZ8EuExw+gfD8=", - "ops/__init__.py": "47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU=", - "ops/activations.py": "t5lzNg1In8LP6bKeTnyeMizwqjv27JGbJ6ylPdGvZYg=", - "ops/fused_dense.py": "BqT8VTSCGdLJAjh6WBMkCRFYrdVodPeoElWBvH6AIzA=", - "ops/layer_norm.py": "zr7NXIm+2mtEynTp1CS0fbFGI2Mqdp41dY4AfDWF6EQ=", - "ops/rms_norm.py": "XEnihcj0a4aSz4LO55m5iKGVn4HKTeKN8TIyHjuDgxI=", - "ops/triton/__init__.py": "AbpHGcgLb+kRsJGnwFEktk7uzpZOCcBY74+YBdrKVGs=", - "ops/triton/cross_entropy.py": "y56nmuOtpSNxM7sDZcUSgXKIsIEayfuvrU8bAVog/nw=", - "ops/triton/k_activations.py": "+Z3vIyO4JkqBMipKsPvhzmxljtBdIhJCsl/M+/ESqBo=", - "ops/triton/layer_norm.py": "BBQdJiqopbIS12viUqzx3SvOQCNbDSAYLJXhpyMGyC8=", - "ops/triton/linear.py": "OtRvKz8xdpl+7v3q/ZTaS9fdBt9XrzMyapgRr50uBbM=", - "ops/triton/mlp.py": "uWA+D7HyuRUsOT+4nhi3vvCDNZnNEXwYcHMQ02UREfQ=", - "ops/triton/rotary.py": "0YSrlbep0lJmzLpTiANs0slTFeGcQQqR+el6uzbSLvI=" - } - } -} \ No newline at end of file diff --git a/build/torch212-cxx11-cu132-aarch64-linux/metadata.json.sigstore b/build/torch212-cxx11-cu132-aarch64-linux/metadata.json.sigstore deleted file mode 100644 index da29c222a5da455f9240fe424b763e69abf26f71..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu132-aarch64-linux/metadata.json.sigstore +++ /dev/null @@ -1 +0,0 @@ -{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHSzCCBtGgAwIBAgIUE3M6LhUCf0y6Y9O3pM9Mpc9497gwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwNjIzMjM0MDQ4WhcNMjYwNjIzMjM1MDQ4WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEYcfagHJbwo+02+yxRKDsW5kchzQoO8RwXeg6oAPOQjfivd6VofOIQ8SjiqtoIMyxOP41/xHE2i61Ok92m/6gBKOCBfAwggXsMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUDc46ij7cpKhAlWq4a8rWABOBOI8wHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoZjEyYWZjOTJjNDBiYmQ3NmQ5OTBkOTExZWJhYmYyOGU0OTBhN2JkZTATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoZjEyYWZjOTJjNDBiYmQ3NmQ5OTBkOTExZWJhYmYyOGU0OTBhN2JkZTAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoZjEyYWZjOTJjNDBiYmQ3NmQ5OTBkOTExZWJhYmYyOGU0OTBhN2JkZTAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKGYxMmFmYzkyYzQwYmJkNzZkOTkwZDkxMWViYWJmMjhlNDkwYTdiZGUwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMjgwMzYyNjg5MzUvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBigYKKwYBBAHWeQIEAgR8BHoAeAB2AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABnvbbnGAAAAQDAEcwRQIgG6745Q97j8c9hRHsGoVV9xkK+g78Fpcr7dEFDXp4HtECIQDSBHJ/T0K9LtdhOxfQWqd3WWHz6PXr5K7UjJfnqo8NozAKBggqhkjOPQQDAwNoADBlAjAW88vhBSIHLS6/HEjOGjI50BsVxodcDg2sqTw6co4I2U8tA/o44GT8wPFkyu3jUuwCMQCs9X2Vn6+/fAmJkouDpKSJ+yeDd8H040O+hU/YhpaekiATeFl4mJcu7M2dVX32pJk="}, "tlogEntries":[{"logIndex":"1933040488", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1782258048", "inclusionPromise":{"signedEntryTimestamp":"MEYCIQCy/RDTt8ME2yUxeLIVIWKW77M+y0fSq3JrY7OekbFBWAIhAOEL/u5zDBZzrfHAxs7bYnRUC/HTaVyOHMIvKSLXI+Gj"}, "inclusionProof":{"logIndex":"1811136226", "rootHash":"bjPkooQ0q3z0TuL3KhaKK20zhWb0rDeX0ApA1G0Ealg=", "treeSize":"1811136233", "hashes":["Lak+VyQEJdb1t3I7L5uQuz2wjSqDxIfNhHYrVmiPZUo=", "0wrvGYnNL9iTyp9vGZzgNPkWv12ByxcnrK/JTIKnD3s=", "Gv9/x8Lq9wnqI26T0QHB5szF3ivXj6axO4AN4LfTguE=", "H8p3CUBUJ5YAbCk5rYNYvgNdQbNRrA7RhtQA+T3kRbs=", "RyWX9+xB2qW/neWuTG1hD4sHW3hZmT2m8Y8bXsGu22k=", "C1fJxkfGKbCR3nQMFiBL5cyjyPxAxy4SwggV3YZHC+0=", "u1tOt8/ve6uKDIXXagbHz46wM+a03h7MlndDg1pLW0g=", "Q3HMPicFFutQ+mkf7xA/p6HtpWMkWvG66IB/xPYs3Vw=", "LuDs4vs9YQGnXlArOp7boDhzWABrNJ1pOn55UTTP5B8=", "WD8KwmQq5R8K69koYj5k4K7EKo2VF3LqJcx2xoXcd7o=", "99JTcUz0aJeZZocenAppVOLl8v91qMn0v7CmEjq9lws=", "OSDrt6jTKV6LYAkaG63p2BFz4SmbNmGb3uerhF5pOpQ=", "yehxDZkBWOGIzNDb93qxmQa6VBVo7kacVSoHyXMm3ZI=", "oUh+N71k3jWcTmlDDB6Esa5uqlCnbmzxTmDdut2A16M=", "2BsqkdK2omLvu+CpsFjpwqGi6m22H/7zcwCVkPwslM8=", "S8Oe9ZicoJqbVuL1I0dIdTgVpLc82Cm2CMOcV83sKVs=", "E1O2ilWlpFRLWGdbmJUMDZ2qbGu/eDZ12m6YjQ9/JPA=", "nyv2TDlfkz0XfRwjY98zsbUbOkBZ+Tgz7PML0c2uN14=", "OMfDh4k9YD0E+t4sjqJsh+vg8D1ueo1cC07exyxmr6s=", "+vHKm4wNwGJ46ceyZ8mePD7Vn1SediStNHCStRoTSWQ=", "5HwWjuGNDisetbiWM0cWIXM65TIM3EJlqk13ntCQlcU=", "NDIRlqxO0IaGr94q86MIU7VCAN3yEEgOHS4loya5lu4=", "daxmZaajRpZV+JxHiOYZhJBiSKN5ucqjh2WnGbHhirw=", "DOCeoSMovIvLExkhIvisow9AuNXgeWs4ECkyR6EcqYU="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n1811136233\nbjPkooQ0q3z0TuL3KhaKK20zhWb0rDeX0ApA1G0Ealg=\n\n— rekor.sigstore.dev wNI9ajBFAiA2BnOVOP7kAUHAazcvaKDzBVgE0SPHBvE+WFageSia6gIhAJnxQnAgi3+OJuFSpRg6C7OrGiiQvwSh7B43p39yXfQ6\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiJiNmExZWY2YzlhY2UwMDI5ZDFkOTg4OThhMjRhNWQwZmM2YjdjNjM1MDZlN2RhYjg2MTY1ZGUyZTBjYmJhNTU5In19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJUUQvRlhsZUxjZStNa2MzdnhVcVROeVI2TEhtVWtaM2dsUUdzQ2l3ZmdNUjhBSWdKcFpYS2xUZ2dlT1ZtQW5OcFRiclU2ck0xK1RwSGJuazUrdHZWN0ZpL1NZPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRla05EUW5SSFowRjNTVUpCWjBsVlJUTk5Oa3hvVlVObU1IazJXVGxQTTNCTk9VMXdZemswT1RkbmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDVxU1hwTmFrMHdUVVJSTkZkb1kwNU5hbGwzVG1wSmVrMXFUVEZOUkZFMFYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZaWTJaaFowaEtZbmR2S3pBeUszbDRVa3RFYzFjMWEyTm9lbEZ2VHpoU2QxaGxaellLYjBGUVQxRnFabWwyWkRaV2IyWlBTVkU0VTJwcGNYUnZTVTE1ZUU5UU5ERXZlRWhGTW1rMk1VOXJPVEp0THpablFrdFBRMEptUVhkbloxaHpUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlZFWXpRMkNtbHFOMk53UzJoQmJGZHhOR0U0Y2xkQlFrOUNUMGs0ZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOWFha1Y1V1ZkYWFrOVVTbXBPUkVKcFdXMVJNMDV0VVRWUFZFSnJDazlVUlhoYVYwcG9XVzFaZVU5SFZUQlBWRUpvVGpKS2ExcFVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMXBxUlhsWlYxcHFUMVJLYWs1RVFtbFpiVkV6VG0xUk5VOVVRbXRQVkVWNFdsZEthRmx0V1hrS1QwZFZNRTlVUW1oT01rcHJXbFJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMXBxUlhsWlYxcHFUMVJLYWs1RVFtbFpiVkV6VG0xUk5VOVVRbXNLVDFSRmVGcFhTbWhaYlZsNVQwZFZNRTlVUW1oT01rcHJXbFJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwZFplRTF0Um0wS1dYcHJlVmw2VVhkWmJVcHJUbnBhYTA5VWEzZGFSR3Q0VFZkV2FWbFhTbTFOYW1oc1RrUnJkMWxVWkdsYVIxVjNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOYW1kM1RYcFplVTVxWnpWTmVsVjJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwWjFsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTT0VKSWIwRUtaVUZDTWtGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtNTJZbUp1UjBGQlFVRlJSQXBCUldOM1VsRkpaMGMyTnpRMVVUazNhamhqT1doU1NITkhiMVpXT1hoclN5dG5OemhHY0dOeU4yUkZSa1JZY0RSSWRFVkRTVkZFVTBKSVNpOVVNRXM1Q2t4MFpHaFBlR1pSVjNGa00xZFhTSG8yVUZoeU5VczNWV3BLWm01eGJ6aE9iM3BCUzBKblozRm9hMnBQVUZGUlJFRjNUbTlCUkVKc1FXcEJWemc0ZG1nS1FsTkpTRXhUTmk5SVJXcFBSMnBKTlRCQ2MxWjRiMlJqUkdjeWMzRlVkelpqYnpSSk1sVTRkRUV2YnpRMFIxUTRkMUJHYTNsMU0ycFZkWGREVFZGRGN3bzVXREpXYmpZckwyWkJiVXByYjNWRWNFdFRTaXQ1WlVSa09FZ3dOREJQSzJoVkwxbG9jR0ZsYTJsQlZHVkdiRFJ0U21OMU4wMHlaRlpZTXpKd1NtczlDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyTADAgEAMIICwAYJKoZIhvcNAQcCoIICsTCCAq0CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgjpnx7xmjrk+jHgpaPc9cH5J35ECG5mRg9lzd0pUD+/QCFQDOmxP+jSNl9qPD0Mzxcis0vndPUBgPMjAyNjA2MjMyMzQwNDhaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHaMIIB1gIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MDYyMzIzNDA0OFowLwYJKoZIhvcNAQkEMSIEIMYhFE6wjj9owIgRLLAW+eujJfjNn+f+0eQge9JBUEFxMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRmMGQCMCCde2yIe0Nr0yqSmtDUwIhzw6vkhpjuPo1qxDvLBZfjVONYlyrLVx9fMizI0i/XOQIwdGqdoriop0tnWvtE2/r8slWpDlcyZ3Wz8vemyIhkkxvjWoWu02UKYirbaBAhF6qp"}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"tqHvbJrOACnR2YiYokpdD8a3xjUG59q4YWXeLgy7pVk="}, "signature":"MEUCIQD/FXleLce+Mkc3vxUqTNyR6LHmUkZ3glQGsCiwfgMR8AIgJpZXKlTggeOVmAnNpTbrU6rM1+TpHbnk5+tvV7Fi/SY="}} \ No newline at end of file diff --git a/build/torch212-cxx11-cu132-aarch64-linux/ops/__init__.py b/build/torch212-cxx11-cu132-aarch64-linux/ops/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/build/torch212-cxx11-cu132-aarch64-linux/ops/activations.py b/build/torch212-cxx11-cu132-aarch64-linux/ops/activations.py deleted file mode 100644 index 7c09649fc41e12d5a360c5672825d8380bc7ec80..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu132-aarch64-linux/ops/activations.py +++ /dev/null @@ -1,135 +0,0 @@ -# Copied from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/model/layers/activations.py -import math - -import torch -import torch.nn as nn -import torch.nn.functional as F - -# 1/sqrt(2*pi)-> 0.3989423 -# 1/sqrt(2) -> 0.70710678 -# sqrt(2/pi) -> 0.79788456 - -# this function is tanh approximation of gelu -# actual gelu is: -# x * 0.5 * (1.0 + torch.erf(x * 0.70710678)) -@torch.jit.script -def bias_gelu(y, bias): - x = bias + y - return (x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))).to(dtype=y.dtype) - - -# gradient of tanh approximation of gelu -# gradient of actual gelu is: -# 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x) -@torch.jit.script -def bias_gelu_back(g, y, bias): - """Assume that y has shape (B, D) and bias has shape (D)""" - x = bias + y - tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)) - # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243 - ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( - 1 + tanh_out - ) - grad_y = ff * g - return grad_y.to(dtype=y.dtype), grad_y.sum(dim=(0), dtype=bias.dtype) - - -class GeLUFunction(torch.autograd.Function): - @staticmethod - # bias is an optional argument - def forward(ctx, input, bias): - ctx.save_for_backward(input, bias) - return bias_gelu(input, bias) - - @staticmethod - def backward(ctx, grad_output): - input, bias = ctx.saved_tensors - tmp = bias_gelu_back(grad_output, input, bias) - return tmp, tmp - - -bias_gelu_impl = GeLUFunction.apply - -# this function is tanh approximation of gelu -# actual gelu is: -# x * 0.5 * (1.0 + torch.erf(x * 0.70710678)) -@torch.jit.script -def gelu_fwd(x): - return (x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))).to(dtype=x.dtype) - - -# gradient of tanh approximation of gelu -# gradient of actual gelu is: -# 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x) -@torch.jit.script -def gelu_bwd(g, x): - tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)) - # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243 - ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( - 1 + tanh_out - ) - return (ff * g).to(dtype=x.dtype) - - -class FastGeLUFunction(torch.autograd.Function): - @staticmethod - # bias is an optional argument - def forward(ctx, input): - ctx.save_for_backward(input) - return gelu_fwd(input) - - @staticmethod - def backward(ctx, grad_output): - (input,) = ctx.saved_tensors - tmp = gelu_bwd(grad_output, input) - return tmp - - -fast_gelu_impl = FastGeLUFunction.apply - - -@torch.jit.script -def relu_bwd(g, x): - return torch.where(x >= 0, g, 0.0).to(dtype=x.dtype) - - -@torch.jit.script -def sqrelu_fwd(x): - r = F.relu(x) - return (r * r).to(dtype=x.dtype) - - -@torch.jit.script -def sqrelu_bwd(g, x): - return (2.0 * g * F.relu(x)).to(dtype=x.dtype) - - -swiglu_fwd_codestring = """ -template T swiglu_fwd(T x, T y) { - return float(x) * float(y) / (1.0f + ::exp(-float(x))); -} -""" -swiglu_bwd_codestring = """ -template void swiglu_bwd(T x, T y, T g, T& dx, T& dy) { - float x_sigmoid = 1.0f / (1.0f + ::exp(-float(x))); - dx = x_sigmoid * (1 + float(x) * (1.0f - x_sigmoid)) * float(g) * float(y); - dy = float(x) * x_sigmoid * float(g); -} -""" -swiglu_fwd = torch.cuda.jiterator._create_jit_fn(swiglu_fwd_codestring) -swiglu_bwd = torch.cuda.jiterator._create_multi_output_jit_fn(swiglu_bwd_codestring, num_outputs=2) - - -class SwiGLUFunction(torch.autograd.Function): - - @staticmethod - def forward(ctx, x, y): - ctx.save_for_backward(x, y) - return swiglu_fwd(x, y) - - @staticmethod - def backward(ctx, dout): - x, y = ctx.saved_tensors - return swiglu_bwd(x, y, dout) - -swiglu = SwiGLUFunction.apply diff --git a/build/torch212-cxx11-cu132-aarch64-linux/ops/fused_dense.py b/build/torch212-cxx11-cu132-aarch64-linux/ops/fused_dense.py deleted file mode 100644 index 6b4033d134e4093fe278f7b3f8c7d3128ce9f36d..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu132-aarch64-linux/ops/fused_dense.py +++ /dev/null @@ -1,688 +0,0 @@ -# Copyright (c) 2023, Tri Dao. -# Inspired by https://github.com/NVIDIA/apex/blob/master/apex/fused_dense/fused_dense.py -# We make it work with pytorch amp and with bfloat16. -# The TensorParallel linear modules are inspired by https://github.com/NVIDIA/apex/blob/master/apex/transformer/tensor_parallel/layers.py -from functools import partial -from typing import Optional - -# import fused_dense_cuda # from apex -import fused_dense_lib as fused_dense_cuda -import torch -import torch.nn as nn -import torch.nn.functional as F -from torch import Tensor -from torch.distributed import ProcessGroup - -from flash_attn.utils.torch import custom_fwd, custom_bwd -from flash_attn.ops.activations import gelu_bwd, relu_bwd, sqrelu_bwd, sqrelu_fwd -from flash_attn.utils.distributed import ( - all_gather_raw, - all_reduce, - all_reduce_raw, - reduce_scatter, - reduce_scatter_raw, -) - - -class FusedDenseFunc(torch.autograd.Function): - @staticmethod - @custom_fwd - def forward( - ctx, x, weight, bias, return_residual=False, process_group=None, sequence_parallel=True - ): - """ - If process_group is not None and sequence_parallel=True, we're doing Tensor Parallel - with sequence parallelism: we do an all_gather_raw of x before doing the matmul. - """ - ctx.compute_weight_gradient = weight.requires_grad - ctx.return_residual = return_residual - ctx.process_group = process_group - ctx.sequence_parallel = sequence_parallel - - if torch.is_autocast_enabled(): - x = x.to(dtype=torch.get_autocast_gpu_dtype()) - x = x.contiguous() - if process_group is not None and sequence_parallel: - # We want to kick off the all_gather early, before weight dtype conversion - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - else: - total_x = x - - if torch.is_autocast_enabled(): - weight = weight.to(dtype=torch.get_autocast_gpu_dtype()) - bias = bias.to(dtype=torch.get_autocast_gpu_dtype()) if bias is not None else None - weight = weight.contiguous() - if process_group is not None and sequence_parallel: - handle_x.wait() - batch_shape, n = total_x.shape[:-1], total_x.shape[-1] - batch_dim = batch_shape.numel() - # https://github.com/pytorch/pytorch/blob/5b51849b48a7dbccd297286cc0110def4706f9e7/aten/src/ATen/native/cuda/Blas.cpp#L174 - if min(batch_dim, n, *weight.shape) > 65535 * 32: - raise RuntimeError("fused_dense only supports matrix dims <= 2M") - output = F.linear(total_x, weight, bias) - if ctx.compute_weight_gradient: - ctx.save_for_backward(x, weight) - else: - ctx.save_for_backward(weight) - return output if not return_residual else (output, x) - - @staticmethod - @custom_bwd - def backward(ctx, grad_output, *args): - grad_output = grad_output.contiguous() - if ctx.return_residual: - (grad_input,) = args - grad_input = grad_input.contiguous() - process_group = ctx.process_group - sequence_parallel = ctx.sequence_parallel - if ctx.compute_weight_gradient: - x, weight = ctx.saved_tensors - if process_group is not None and sequence_parallel: - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - else: - total_x = x - else: - (weight,) = ctx.saved_tensors - total_x = None - batch_shape = grad_output.shape[:-1] - batch_dim = batch_shape.numel() - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - if ctx.needs_input_grad[0]: - if not ctx.return_residual: - grad_input = F.linear(grad_output, weight.t()) - else: - grad_input = torch.addmm( - grad_input.reshape(batch_dim, grad_input.shape[-1]), grad_output, weight - ) - grad_input = grad_input.reshape(*batch_shape, grad_input.shape[-1]) - if process_group is not None: - reduce_fn = reduce_scatter_raw if sequence_parallel else all_reduce_raw - grad_input, handle_grad_input = reduce_fn(grad_input, process_group, async_op=True) - else: - grad_input = None - if ctx.needs_input_grad[1]: - assert ctx.compute_weight_gradient - if process_group is not None and sequence_parallel: - handle_x.wait() - grad_weight, grad_bias = fused_dense_cuda.linear_bias_wgrad( - total_x.reshape(batch_dim, total_x.shape[-1]), grad_output, ctx.needs_input_grad[2] - ) - else: - grad_weight = None - grad_bias = grad_output if ctx.needs_input_grad[2] else None - if process_group is not None and ctx.needs_input_grad[0]: - handle_grad_input.wait() - return grad_input, grad_weight, grad_bias, None, None, None - - -def fused_dense_func( - x: Tensor, - weight: Tensor, - bias: Optional[Tensor] = None, - return_residual: bool = False, - process_group: Optional[ProcessGroup] = None, - sequence_parallel: bool = True, -): - dtype_eligible = x.dtype in [torch.float16, torch.bfloat16] or ( - x.dtype == torch.float32 and torch.is_autocast_enabled() - ) - if x.is_cuda and weight.is_cuda and (bias is None or bias.is_cuda) and dtype_eligible: - return FusedDenseFunc.apply( - x, weight, bias, return_residual, process_group, sequence_parallel - ) - else: - assert process_group is None - out = F.linear(x, weight, bias) - return out if not return_residual else (out, x) - - -class FusedDense(nn.Linear): - def __init__( - self, - in_features: int, - out_features: int, - bias: bool = True, - return_residual: bool = False, - device=None, - dtype=None, - ) -> None: - super().__init__(in_features, out_features, bias=bias, device=device, dtype=dtype) - self.return_residual = return_residual - - def forward(self, x, process_group=None): - """ - If process_group is not None, we're doing Tensor Parallel with sequence parallelism: - we do an all_gather of x before doing the matmul. - """ - return fused_dense_func( - x, - self.weight, - self.bias, - return_residual=self.return_residual, - process_group=process_group, - ) - - -class ColumnParallelLinear(nn.Linear): - def __init__( - self, - in_features: int, - out_features: int, - process_group: ProcessGroup, - bias: bool = True, - sequence_parallel=True, - multiple_of=1, - device=None, - dtype=None, - ) -> None: - world_size = torch.distributed.get_world_size(process_group) - if out_features % multiple_of: - raise ValueError(f"out_features ({out_features}) must be a multiple of {multiple_of}") - multiple = out_features // multiple_of - # We want to split @multiple across world_size, but it could be an uneven split - div = multiple // world_size - mod = multiple % world_size - # The first @mod ranks get @div + 1 copies, the rest get @div copies - local_multiple = div + int(torch.distributed.get_rank(process_group) < mod) - super().__init__( - in_features, local_multiple * multiple_of, bias=bias, device=device, dtype=dtype - ) - self.process_group = process_group - self.sequence_parallel = sequence_parallel - - def forward(self, x): - # If self.sequence_parallel is True, we're doing Tensor Parallel with sequence parallelism: - # we do an all_gather of x before doing the matmul. - # If not, then the input is already gathered. - return fused_dense_func( - x, - self.weight, - self.bias, - process_group=self.process_group, - sequence_parallel=self.sequence_parallel, - ) - - -class RowParallelLinear(nn.Linear): - def __init__( - self, - in_features: int, - out_features: int, - process_group: ProcessGroup, - bias: bool = True, - sequence_parallel=True, - multiple_of=1, - device=None, - dtype=None, - ) -> None: - world_size = torch.distributed.get_world_size(process_group) - rank = torch.distributed.get_rank(process_group) - if in_features % multiple_of: - raise ValueError(f"in_features ({in_features}) must be a multiple of {multiple_of}") - multiple = in_features // multiple_of - # We want to split @multiple across world_size, but it could be an uneven split - div = multiple // world_size - mod = multiple % world_size - # The first @mod ranks get @div + 1 copies, the rest get @div copies - local_multiple = div + int(torch.distributed.get_rank(process_group) < mod) - # Only rank 0 will have bias - super().__init__( - local_multiple * multiple_of, - out_features, - bias=bias and rank == 0, - device=device, - dtype=dtype, - ) - self.process_group = process_group - self.sequence_parallel = sequence_parallel - - def forward(self, x): - """ - We're doing Tensor Parallel with sequence parallelism: we do the matmul and then - a reduce_scatter of the result. - """ - out = fused_dense_func(x, self.weight, self.bias) - reduce_fn = reduce_scatter if self.sequence_parallel else all_reduce - return reduce_fn(out, self.process_group) - - -class FusedMLPFunc(torch.autograd.Function): - @staticmethod - @custom_fwd - def forward( - ctx, - x, - weight1, - bias1, - weight2, - bias2, - activation="gelu_approx", - save_pre_act=True, - return_residual=False, - checkpoint_lvl=0, - heuristic=0, - process_group=None, - sequence_parallel=True, - ): - """ - If process_group is not None and sequence_parallel=True, we're doing Tensor Parallel - with sequence parallelism: we do an all_gather of x before doing the matmul. - If sequence_parallel=False, then the input is already gathered. - - checkpoint_lvl: - 0: no recomputation in the bwd - 1: recompute gelu_out / relu_out in the bwd - 2: recompute pre_act and gelu_out / relu_out in the bwd - """ - assert -1 <= heuristic <= 4 - assert activation in ["gelu_approx", "relu", "sqrelu"] - if activation == "sqrelu": - assert heuristic == -1 - if not save_pre_act: - checkpoint_lvl = 2 - assert checkpoint_lvl in [0, 1, 2] - ctx.return_residual = return_residual - ctx.process_group = process_group - ctx.sequence_parallel = sequence_parallel - ctx.checkpoint_lvl = checkpoint_lvl - ctx.activation = activation - ctx.heuristic = heuristic - - if torch.is_autocast_enabled(): - x = x.to(dtype=torch.get_autocast_gpu_dtype()) - x = x.contiguous() - if process_group is not None and sequence_parallel: - # We want to kick off the all_gather early, before weight dtype conversion - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - else: - total_x = x - - if torch.is_autocast_enabled(): - dtype = torch.get_autocast_gpu_dtype() - weight1, weight2 = [a.to(dtype=dtype) for a in [weight1, weight2]] - bias1 = bias1.to(dtype=dtype) if bias1 is not None else None - bias2 = bias2.to(dtype=dtype) if bias2 is not None else None - weight1 = weight1.contiguous() - bias1 = bias1.contiguous() if bias1 is not None else None - weight2 = weight2.contiguous() - bias2 = bias2.contiguous() if bias2 is not None else None - if process_group is not None and sequence_parallel: - handle_x.wait() - batch_shape, n = total_x.shape[:-1], total_x.shape[-1] - batch_dim = batch_shape.numel() - # https://github.com/pytorch/pytorch/blob/5b51849b48a7dbccd297286cc0110def4706f9e7/aten/src/ATen/native/cuda/Blas.cpp#L174 - if min(batch_dim, n, *weight1.shape, *weight2.shape) > 65535 * 32: - raise RuntimeError("fused_dense only supports matrix dims <= 2M") - if heuristic == -1: - pre_act = F.linear(total_x, weight1, bias1) - activation_fn = ( - partial(F.gelu, approximate="tanh") - if activation == "gelu_approx" - else (sqrelu_fwd if activation == "sqrelu" else F.relu) - ) - with torch.jit.fuser("fuser2"): - output1 = activation_fn(pre_act) - # This is before adding bias1 - # pre_act = F.linear(total_x.reshape(batch_dim, n), weight1) - # with torch.jit.fuser('fuser2'): - # output1 = bias_gelu(pre_act, bias1) - else: - is_gelu = activation == "gelu_approx" - output1, *rest = fused_dense_cuda.linear_act_forward( - total_x.reshape(batch_dim, n), weight1, bias1, is_gelu, save_pre_act, heuristic - ) - if save_pre_act: - pre_act = rest[0] - output2 = F.linear(output1, weight2, bias2) - if checkpoint_lvl == 0 or (checkpoint_lvl == 1 and activation == "relu"): - # For RELU the pre_act is very small (just a bit-mask) so we just save it - ctx.save_for_backward(x, weight1, weight2, pre_act, output1) - elif checkpoint_lvl == 1: - ctx.save_for_backward(x, weight1, weight2, pre_act) - elif checkpoint_lvl == 2: - ctx.save_for_backward(x, weight1, weight2, bias1) - output2 = output2.reshape(*batch_shape, output2.shape[-1]) - return output2 if not return_residual else (output2, x) - - @staticmethod - @custom_bwd - def backward(ctx, grad_output, *args): - grad_output = grad_output.contiguous() - checkpoint_lvl = ctx.checkpoint_lvl - activation = ctx.activation - activation_fn = ( - partial(F.gelu, approximate="tanh") - if activation == "gelu_approx" - else (sqrelu_fwd if activation == "sqrelu" else F.relu) - ) - if ctx.return_residual: - (grad_input,) = args - grad_input = grad_input.contiguous() - process_group = ctx.process_group - sequence_parallel = ctx.sequence_parallel - x, weight1, weight2, *rest = ctx.saved_tensors - if process_group is None or not sequence_parallel: - total_x = x - batch_shape = grad_output.shape[:-1] - batch_dim = batch_shape.numel() - if checkpoint_lvl in [0, 1]: - if process_group is not None and sequence_parallel: - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - if checkpoint_lvl == 0 or (checkpoint_lvl == 1 and activation == "relu"): - pre_act, output1 = rest - elif checkpoint_lvl == 1: - (pre_act,) = rest - with torch.jit.fuser("fuser2"): - output1 = activation_fn(pre_act) - elif checkpoint_lvl == 2: - (bias1,) = rest - if process_group is not None and sequence_parallel: - total_x, _ = all_gather_raw(x, process_group) - if ctx.heuristic == -1: - pre_act = F.linear(total_x, weight1, bias1) - with torch.jit.fuser("fuser2"): - output1 = activation_fn(pre_act) - else: - output1, pre_act = fused_dense_cuda.linear_act_forward( - total_x.reshape(batch_dim, total_x.shape[-1]), - weight1, - bias1, - activation == "gelu_approx", - True, - ctx.heuristic, - ) - - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - output1 = output1.reshape(batch_dim, output1.shape[-1]) - pre_act = pre_act.reshape(batch_dim, pre_act.shape[-1]) - if ctx.needs_input_grad[3]: - grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad( - output1, grad_output, ctx.needs_input_grad[4] - ) - else: - grad_weight2 = None - grad_bias2 = grad_output if ctx.needs_input_grad[4] else None - if ctx.heuristic == -1: - # grad_pre_act = matmul_dgelu(grad_output, weight2, pre_act) - grad_output1 = F.linear(grad_output, weight2.t()) - activation_grad_fn = ( - gelu_bwd - if activation == "gelu_approx" - else (sqrelu_bwd if activation == "sqrelu" else relu_bwd) - ) - with torch.jit.fuser("fuser2"): - grad_pre_act = activation_grad_fn(grad_output1, pre_act) - else: - # The cublasLt epilogue has to compute both gelu/relu grad and bias grad, we can't - # just compute gelu/relu grad - grad_pre_act, grad_bias1 = fused_dense_cuda.bias_act_linear_dgrad_bgrad( - weight2, grad_output, pre_act, activation == "gelu_approx", ctx.heuristic - ) - if not ctx.needs_input_grad[2]: - grad_bias1 = None - if ctx.needs_input_grad[0]: - if not ctx.return_residual: - grad_input = F.linear(grad_pre_act, weight1.t()) - else: - grad_input = torch.addmm( - grad_input.reshape(batch_dim, grad_input.shape[-1]), grad_pre_act, weight1 - ) - grad_input = grad_input.reshape(*batch_shape, grad_input.shape[-1]) - if process_group is not None: - reduce_fn = reduce_scatter_raw if sequence_parallel else all_reduce_raw - grad_input, handle_grad_input = reduce_fn(grad_input, process_group, async_op=True) - else: - grad_input = None - if ctx.heuristic == -1: - if ctx.needs_input_grad[1]: - if process_group is not None and sequence_parallel and checkpoint_lvl != 2: - handle_x.wait() - grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_wgrad( - total_x.reshape(batch_dim, total_x.shape[-1]), - grad_pre_act, - ctx.needs_input_grad[2], - ) - else: - grad_weight1 = None - grad_bias1 = grad_pre_act if ctx.needs_input_grad[2] else None - else: - if ctx.needs_input_grad[1]: - if process_group is not None and sequence_parallel and checkpoint_lvl != 2: - handle_x.wait() - grad_weight1 = F.linear( - grad_pre_act.t(), total_x.reshape(batch_dim, total_x.shape[-1]).t() - ) - else: - grad_weight1 = None - if process_group is not None and ctx.needs_input_grad[0]: - handle_grad_input.wait() - return ( - grad_input, - grad_weight1, - grad_bias1, - grad_weight2, - grad_bias2, - None, - None, - None, - None, - None, - None, - None, - ) - - -def fused_mlp_func( - x: Tensor, - weight1: Tensor, - weight2: Tensor, - bias1: Optional[Tensor] = None, - bias2: Optional[Tensor] = None, - activation: str = "gelu_approx", - save_pre_act: bool = True, - return_residual: bool = False, - checkpoint_lvl: int = 0, - heuristic: int = 0, - process_group: Optional[ProcessGroup] = None, - sequence_parallel: bool = True, -): - assert activation in ["gelu_approx", "relu", "sqrelu"] - dtype_eligible = x.dtype in [torch.float16, torch.bfloat16] or ( - x.dtype == torch.float32 and torch.is_autocast_enabled() - ) - # If we save pre-activation, dimension must be divisible by 128 (relu) or 8 (gelu) - dim_eligible = not save_pre_act or (x.shape[-1] % (128 if activation == "relu" else 8) == 0) - if ( - x.is_cuda - and weight1.is_cuda - and weight2.is_cuda - and (bias1 is None or bias1.is_cuda) - and (bias2 is None or bias2.is_cuda) - and dtype_eligible - and dim_eligible - ): - return FusedMLPFunc.apply( - x, - weight1, - bias1, - weight2, - bias2, - activation, - save_pre_act, - return_residual, - checkpoint_lvl, - heuristic, - process_group, - sequence_parallel, - ) - else: - assert process_group is None - pre_act = F.linear(x, weight1, bias1) - activation_fn = ( - partial(F.gelu, approximate="tanh") - if activation == "gelu_approx" - else partial(F.relu, inplace=True) - ) - output1 = activation_fn(pre_act) - output2 = F.linear(output1, weight2, bias2) - return output2 if not return_residual else (output2, x) - - -class FusedMLP(nn.Module): - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - bias1=True, - bias2=True, - activation="gelu_approx", - return_residual=False, - checkpoint_lvl=0, - heuristic="auto", - device=None, - dtype=None, - ): - """ - If process_group is not None, we're doing Tensor Parallel with sequence parallelism: - we do an all_gather of x before doing the matmul, gelu, then matmul. - Finally we do a reduce_scatter of the output. - - checkpoint_lvl (increasing lvl means slower but more memory saving): - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute pre_act and gelu_out in the bwd - heuristic: - -1: don't fuse gemm + gelu (separate kernel) - 0..4: use this heuristic for the algo section in the fused gemm + gelu - 'auto': heuristic will be picked automatically: - For CUDA >= 11.8, we set heuristic=0 for both fp16 and bf16 for best perf. - For CUDA <= 11.7, we set heuristic=1 for fp16 and heuristic=-1 for bf16. - For H100, we set heuristic=-1 for both fp16 and bf16 as the fused cuBlasLt implementation - is slower than the unfused version. - return_residual: whether to return the input x along with the output. This is for - performance reason: for post-norm architecture, returning the input allows us - to fuse the backward of nn.Linear with the residual connection. - """ - assert checkpoint_lvl in [0, 1, 2] - assert activation in ["gelu_approx", "relu", "sqrelu"] - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features * 4 - self.activation = activation - self.return_residual = return_residual - self.checkpoint_lvl = checkpoint_lvl - self.heuristic = heuristic if activation != "sqrelu" else -1 - self.fc1 = nn.Linear(in_features, hidden_features, bias=bias1, **factory_kwargs) - self.fc2 = nn.Linear(hidden_features, out_features, bias=bias2, **factory_kwargs) - - def forward(self, x, process_group=None): - dtype = x.dtype if not torch.is_autocast_enabled() else torch.get_autocast_gpu_dtype() - if self.heuristic == "auto": - if self.activation == "gelu_approx": - if torch.cuda.get_device_capability("cuda") == (9, 0): - heuristic = -1 - else: - cuda_ver = tuple(map(int, torch.version.cuda.split("."))) - heuristic = 0 if cuda_ver >= (11, 8) else (1 if dtype == torch.float16 else -1) - else: - heuristic = 0 - else: - heuristic = self.heuristic - out = fused_mlp_func( - x, - self.fc1.weight, - self.fc2.weight, - self.fc1.bias, - self.fc2.bias, - activation=self.activation, - save_pre_act=self.training, - return_residual=self.return_residual, - checkpoint_lvl=self.checkpoint_lvl, - heuristic=heuristic, - process_group=process_group, - ) - if self.return_residual: - out, x = out - if process_group is not None: - out = reduce_scatter(out, process_group) - return out if not self.return_residual else (out, x) - - -class ParallelFusedMLP(nn.Module): - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - activation="gelu_approx", - process_group: ProcessGroup = None, - bias1=True, - bias2=True, - sequence_parallel=True, - checkpoint_lvl=0, - heuristic="auto", - device=None, - dtype=None, - ): - """ - process_group is required. We're doing Tensor Parallel with sequence parallelism: - we do an all_gather of x before doing the matmul, gelu, then matmul. - Finally we do a reduce_scatter of the output. - - checkpoint_lvl (increasing lvl means slower but more memory saving): - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute pre_act and gelu_out in the bwd - heuristic: - -1: don't fuse gemm + gelu (separate kernel) - 0..4: use this heuristic for the algo section in the fused gemm + gelu - 'auto': heuristic will be picked automatically: - For CUDA >= 11.8, we set heuristic=0 for both fp16 and bf16 for best perf. - For CUDA <= 11.7, we set heuristic=1 for fp16 and heuristic=-1 for bf16. - """ - assert checkpoint_lvl in [0, 1, 2] - assert activation in ["gelu_approx", "relu", "sqrelu"] - assert process_group is not None - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features * 4 - self.activation = activation - self.process_group = process_group - self.sequence_parallel = sequence_parallel - self.checkpoint_lvl = checkpoint_lvl - self.heuristic = heuristic if activation != "sqrelu" else -1 - self.fc1 = ColumnParallelLinear( - in_features, hidden_features, process_group, bias=bias1, **factory_kwargs - ) - self.fc2 = RowParallelLinear( - hidden_features, out_features, process_group, bias=bias2, **factory_kwargs - ) - - def forward(self, x): - dtype = x.dtype if not torch.is_autocast_enabled() else torch.get_autocast_gpu_dtype() - if self.heuristic == "auto": - if self.activation == "gelu_approx": - cuda_ver = tuple(map(int, torch.version.cuda.split("."))) - heuristic = 0 if cuda_ver >= (11, 8) else (1 if dtype == torch.float16 else -1) - else: - heuristic = 0 - else: - heuristic = self.heuristic - out = fused_mlp_func( - x, - self.fc1.weight, - self.fc2.weight, - self.fc1.bias, - self.fc2.bias, - activation=self.activation, - save_pre_act=self.training, - checkpoint_lvl=self.checkpoint_lvl, - heuristic=heuristic, - process_group=self.process_group, - sequence_parallel=self.sequence_parallel, - ) - reduce_fn = reduce_scatter if self.sequence_parallel else all_reduce - return reduce_fn(out, self.process_group) diff --git a/build/torch212-cxx11-cu132-aarch64-linux/ops/layer_norm.py b/build/torch212-cxx11-cu132-aarch64-linux/ops/layer_norm.py deleted file mode 100644 index 4b6cd798fd02844ef9cd3897f8ab95e490e638bf..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu132-aarch64-linux/ops/layer_norm.py +++ /dev/null @@ -1,800 +0,0 @@ -# Copyright (c) 2022, Tri Dao. -# Adapted from https://github.com/NVIDIA/apex/blob/master/apex/contrib/layer_norm/layer_norm.py - -import dropout_layer_norm -import torch -from torch.nn import init - - -def maybe_align(x, alignment_in_bytes=16): - """Assume that x already has last dim divisible by alignment_in_bytes""" - # TD [2023-07-04] I'm not 100% sure that clone will align the memory - # https://discuss.pytorch.org/t/how-to-ensure-that-tensor-data-ptr-is-aligned-to-16-bytes/183440 - return x if x.data_ptr() % alignment_in_bytes == 0 else x.clone() - - -def _dropout_add_layer_norm_forward( - x0, - residual, - gamma, - beta, - rowscale, - colscale, - dropout_p, - epsilon, - residual_in_fp32=False, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes""" - hidden_size = gamma.numel() - x0mat = x0.view((-1, hidden_size)) - residualmat = residual.view((-1, hidden_size)) if residual is not None else None - rowscale = rowscale.view(-1) if rowscale is not None else None - zmat, xmat, dmask, mu, rsigma = dropout_layer_norm.dropout_add_ln_fwd( - x0mat, - residualmat, - gamma, - beta, - rowscale, - colscale, - None, - None, - dropout_p, - epsilon, - 1.0, - 0, - None, - residual_in_fp32, - is_rms_norm, - ) - # dmask is None if dropout_p == 0.0 - # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype - return zmat, xmat if xmat is not None else x0mat, dmask, mu, rsigma - - -def _dropout_add_layer_norm_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - rowscale, - colscale, - dropout_p, - has_residual, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes - dx == None means that it was a post-norm architecture - (x = drop(x0) + residual was not returned in the fwd). - x0 must not be None if we have colscale. - """ - hidden_size = gamma.numel() - xmat = x.view((-1, hidden_size)) - dzmat = dz.view(xmat.shape) - dxmat = dx.view(xmat.shape) if dx is not None else None - x0mat = x0.view((-1, hidden_size)) if x0 is not None else None - rowscale = rowscale.view(-1) if rowscale is not None else None - if colscale is not None: - assert x0 is not None, "x0 is required to compute the gradient of colscale" - dx0mat, dresidualmat, dgamma, dbeta, _, _, *rest = dropout_layer_norm.dropout_add_ln_bwd( - dzmat, - dxmat, - xmat, - x0mat, - dmask, - mu, - rsigma, - gamma, - rowscale, - colscale, - None, - None, - dropout_p, - 1.0, - 0, - has_residual, - is_rms_norm, - ) - # dresidualmat is None if not has_residual - if colscale is None: - return dx0mat, dresidualmat, dgamma, dbeta - else: - dcolscale = rest[0] - return dx0mat, dresidualmat, dgamma, dbeta, dcolscale - - -def _dropout_add_layer_norm_subset_forward( - x0, - residual, - gamma, - beta, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32=False, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes""" - hidden_size = gamma.numel() - x0mat = x0.view((-1, hidden_size)) - residualmat = residual.view((-1, hidden_size)) if residual is not None else None - x0_subset = x0_subset.view(-1) if x0_subset is not None else None - out_subset = out_subset.view(-1) if out_subset is not None else None - zmat, xmat, dmask, mu, rsigma = dropout_layer_norm.dropout_add_ln_fwd( - x0mat, - residualmat, - gamma, - beta, - None, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - None, - residual_in_fp32, - is_rms_norm, - ) - # dmask is None if dropout_p == 0.0 - # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype - return zmat, xmat if xmat is not None else x0mat, dmask, mu, rsigma - - -def _dropout_add_layer_norm_subset_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - colscale, - x0_subset, - out_subset, - dropout_p, - rowscale_const, - x0_numrows, - has_residual, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes - dx == None means that it was a post-norm architecture - (x = drop(x0) + residual was not returned in the fwd). - x0 must not be None if we have colscale. - """ - hidden_size = gamma.numel() - xmat = x.view((-1, hidden_size)) - dzmat = dz.view(-1, hidden_size) - dxmat = dx.view(xmat.shape) if dx is not None else None - x0mat = x0.view((-1, hidden_size)) if x0 is not None else None - x0_subset = x0_subset.view(-1) if x0_subset is not None else None - out_subset = out_subset.view(-1) if out_subset is not None else None - if colscale is not None: - assert x0 is not None, "x0 is required to compute the gradient of colscale" - dx0mat, dresidualmat, dgamma, dbeta, _, _, *rest = dropout_layer_norm.dropout_add_ln_bwd( - dzmat, - dxmat, - xmat, - x0mat, - dmask, - mu, - rsigma, - gamma, - None, - colscale, - x0_subset, - out_subset, - dropout_p, - rowscale_const, - x0_numrows, - has_residual, - is_rms_norm, - ) - # dresidualmat is None if not has_residual - if colscale is None: - return dx0mat, dresidualmat, dgamma, dbeta - else: - dcolscale = rest[0] - return dx0mat, dresidualmat, dgamma, dbeta, dcolscale - - -def _dropout_add_layer_norm_parallel_residual_forward( - x0, - x1, - residual, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - residual_in_fp32=False, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes""" - hidden_size = gamma0.numel() - x0mat = x0.view((-1, hidden_size)) - x1mat = x1.view((-1, hidden_size)) if x1 is not None else None - residualmat = residual.view((-1, hidden_size)) if residual is not None else None - ( - z0mat, - z1mat, - xmat, - dmask0, - dmask1, - mu, - rsigma, - ) = dropout_layer_norm.dropout_add_ln_parallel_residual_fwd( - x0mat, - x1mat, - residualmat, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - None, - residual_in_fp32, - is_rms_norm, - ) - # dmask0 and dmask1 are None if dropout_p == 0.0 - # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype - return z0mat, z1mat, xmat if xmat is not None else x0mat, dmask0, dmask1, mu, rsigma - - -def _dropout_add_layer_norm_parallel_residual_backward( - dz0, - dz1, - dx, - x, - dmask0, - dmask1, - mu, - rsigma, - gamma0, - gamma1, - dropout_p, - has_x1, - has_residual, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes - dx == None means that it was a post-norm architecture - (x = drop(x0) + residual was not returned in the fwd). - """ - hidden_size = gamma0.numel() - xmat = x.view((-1, hidden_size)) - dz0mat = dz0.view(xmat.shape) - dz1mat = dz1.view(xmat.shape) if dz1 is not None else None - dxmat = dx.view(xmat.shape) if dx is not None else None - ( - dx0mat, - dx1mat, - dresidualmat, - dgamma0, - dbeta0, - dgamma1, - dbeta1, - *rest, - ) = dropout_layer_norm.dropout_add_ln_parallel_residual_bwd( - dz0mat, - dz1mat, - dxmat, - xmat, - dmask0, - dmask1, - mu, - rsigma, - gamma0, - gamma1, - dropout_p, - has_x1, - has_residual, - is_rms_norm, - ) - # dresidualmat is None if not has_residual - return dx0mat, dx1mat, dresidualmat, dgamma0, dbeta0, dgamma1, dbeta1 - - -class DropoutAddLayerNormFn(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x0, - residual, - gamma, - beta, - rowscale, - colscale, - dropout_p, - epsilon, - residual_in_fp32=False, - prenorm=False, - is_rms_norm=False, - return_dmask=False, - ): - x0 = maybe_align(x0.contiguous(), 16) - residual = maybe_align(residual.contiguous(), 16) if residual is not None else None - gamma = maybe_align(gamma.contiguous(), 16) - beta = maybe_align(beta.contiguous(), 16) if beta is not None else None - rowscale = maybe_align(rowscale.contiguous(), 16) if rowscale is not None else None - colscale = maybe_align(colscale.contiguous(), 16) if colscale is not None else None - zmat, xmat, dmask, mu, rsigma = _dropout_add_layer_norm_forward( - x0, - residual, - gamma, - beta, - rowscale, - colscale, - dropout_p, - epsilon, - residual_in_fp32, - is_rms_norm, - ) - # Only need to save x0 if we need to compute gradient wrt colscale - x0_saved = x0 if colscale is not None else None - ctx.save_for_backward( - xmat.view(x0.shape), x0_saved, dmask, gamma, mu, rsigma, rowscale, colscale - ) - ctx.prenorm = prenorm - ctx.dropout_p = dropout_p - ctx.has_residual = residual is not None - ctx.is_rms_norm = is_rms_norm - ctx.has_beta = beta is not None - if not return_dmask: - return ( - zmat.view(x0.shape) if not prenorm else (zmat.view(x0.shape), xmat.view(x0.shape)) - ) - else: - dmask = ( - dmask.view(x0.shape) - if dropout_p > 0.0 - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - ctx.mark_non_differentiable(dmask) - return ( - (zmat.view(x0.shape), dmask) - if not prenorm - else (zmat.view(x0.shape), xmat.view(x0.shape), dmask) - ) - - @staticmethod - def backward(ctx, dz, *args): - # assert dz.is_contiguous() - dz = maybe_align(dz.contiguous(), 16) # this happens! - dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None - x, x0, dmask, gamma, mu, rsigma, rowscale, colscale = ctx.saved_tensors - # x0 is None if colscale is None - dropout_p = ctx.dropout_p - has_residual = ctx.has_residual - dx0mat, dresidualmat, dgamma, dbeta, *rest = _dropout_add_layer_norm_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - rowscale, - colscale, - dropout_p, - has_residual, - ctx.is_rms_norm, - ) - dx0 = dx0mat.view(x.shape) - dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None - dcolscale = rest[0] if colscale is not None else None - return ( - dx0, - dresidual, - dgamma, - dbeta if ctx.has_beta else None, - None, - dcolscale, - None, - None, - None, - None, - None, - None, - ) - - -class DropoutAddLayerNormSubsetFn(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x0, - residual, - gamma, - beta, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32=False, - prenorm=False, - is_rms_norm=False, - return_dmask=False, - ): - x0 = maybe_align(x0.contiguous(), 16) - residual = maybe_align(residual.contiguous(), 16) if residual is not None else None - gamma = maybe_align(gamma.contiguous(), 16) - beta = maybe_align(beta.contiguous(), 16) if beta is not None else None - colscale = maybe_align(colscale.contiguous(), 16) if colscale is not None else None - zmat, xmat, dmask, mu, rsigma = _dropout_add_layer_norm_subset_forward( - x0, - residual, - gamma, - beta, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32, - is_rms_norm, - ) - # Only need to save x0 if we need to compute gradient wrt colscale - x0_saved = x0 if colscale is not None else None - x_shape = (-1, *x0.shape[1:]) - ctx.save_for_backward( - xmat.view(x_shape), x0_saved, dmask, gamma, mu, rsigma, colscale, x0_subset, out_subset - ) - ctx.prenorm = prenorm - ctx.dropout_p = dropout_p - ctx.rowscale_const = rowscale_const - ctx.x0_numrows = x0.shape[:-1].numel() - ctx.has_residual = residual is not None - ctx.is_rms_norm = is_rms_norm - ctx.has_beta = beta is not None - z_shape = (-1, *x0.shape[1:]) - if not return_dmask: - return zmat.view(z_shape) if not prenorm else (zmat.view(z_shape), xmat.view(x0.shape)) - else: - z = zmat.view(z_shape) - dmask = ( - dmask.view(x0.shape) - if dropout_p > 0.0 - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - ctx.mark_non_differentiable(dmask) - return (z, dmask) if not prenorm else (z, xmat.view(x_shape), dmask) - - @staticmethod - def backward(ctx, dz, *args): - # assert dz.is_contiguous() - dz = maybe_align(dz.contiguous(), 16) # this happens! - dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None - x, x0, dmask, gamma, mu, rsigma, colscale, x0_subset, out_subset = ctx.saved_tensors - # x0 is None if colscale is None - dropout_p = ctx.dropout_p - has_residual = ctx.has_residual - dx0mat, dresidualmat, dgamma, dbeta, *rest = _dropout_add_layer_norm_subset_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - colscale, - x0_subset, - out_subset, - dropout_p, - ctx.rowscale_const, - ctx.x0_numrows, - has_residual, - ctx.is_rms_norm, - ) - dx0 = dx0mat.view(-1, *x.shape[1:]) - dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None - dcolscale = rest[0] if colscale is not None else None - return ( - dx0, - dresidual, - dgamma, - dbeta if ctx.has_beta else None, - dcolscale, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) - - -class DropoutAddLayerNormParallelResidualFn(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x0, - x1, - residual, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - residual_in_fp32=False, - prenorm=False, - is_rms_norm=False, - return_dmask=False, - ): - x0 = maybe_align(x0.contiguous(), 16) - x1 = maybe_align(x1.contiguous(), 16) if x1 is not None else None - residual = maybe_align(residual.contiguous(), 16) if residual is not None else None - gamma0 = maybe_align(gamma0.contiguous(), 16) - beta0 = maybe_align(beta0.contiguous(), 16) if beta0 is not None else None - gamma1 = maybe_align(gamma1.contiguous(), 16) if gamma1 is not None else None - beta1 = maybe_align(beta1.contiguous(), 16) if beta1 is not None else None - ( - z0mat, - z1mat, - xmat, - dmask0, - dmask1, - mu, - rsigma, - ) = _dropout_add_layer_norm_parallel_residual_forward( - x0, - x1, - residual, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - residual_in_fp32, - is_rms_norm, - ) - ctx.save_for_backward(xmat.view(x0.shape), dmask0, dmask1, gamma0, gamma1, mu, rsigma) - ctx.prenorm = prenorm - ctx.dropout_p = dropout_p - ctx.has_x1 = x1 is not None - ctx.has_residual = residual is not None - ctx.is_rms_norm = is_rms_norm - ctx.has_beta = beta0 is not None - z = (z0mat.view(x0.shape), z1mat.view(x0.shape) if z1mat is not None else None) - if not return_dmask: - return z if not prenorm else (*z, xmat.view(x0.shape)) - else: - dmask0 = ( - dmask0.view(x0.shape) - if dropout_p > 0.0 - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - dmask1 = ( - dmask1.view(x0.shape) - if dropout_p > 0.0 and x1 is not None - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - ctx.mark_non_differentiable(dmask0) - ctx.mark_non_differentiable(dmask1) - return ( - (*z, dmask0, dmask1) if not prenorm else (*z, xmat.view(x0.shape), dmask0, dmask1) - ) - - @staticmethod - def backward(ctx, dz0, dz1, *args): - dz0 = maybe_align(dz0.contiguous(), 16) # this happens! - dz1 = maybe_align(dz1.contiguous(), 16) if dz1 is not None else None - dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None - x, dmask0, dmask1, gamma0, gamma1, mu, rsigma = ctx.saved_tensors - dropout_p = ctx.dropout_p - has_x1 = ctx.has_x1 - has_residual = ctx.has_residual - ( - dx0mat, - dx1mat, - dresidualmat, - dgamma0, - dbeta0, - dgamma1, - dbeta1, - ) = _dropout_add_layer_norm_parallel_residual_backward( - dz0, - dz1, - dx, - x, - dmask0, - dmask1, - mu, - rsigma, - gamma0, - gamma1, - dropout_p, - has_x1, - has_residual, - ctx.is_rms_norm, - ) - dx0 = dx0mat.view(x.shape) - dx1 = dx1mat.view(x.shape) if dx1mat is not None else None - dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None - return ( - dx0, - dx1, - dresidual, - dgamma0, - dbeta0 if ctx.has_beta else None, - dgamma1, - dbeta1 if ctx.has_beta else None, - None, - None, - None, - None, - None, - None, - ) - - -def layer_norm(x, weight, bias, epsilon): - return DropoutAddLayerNormFn.apply(x, None, weight, bias, None, None, 0.0, epsilon, False) - - -def dropout_add_layer_norm( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - rowscale=None, - layerscale=None, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormFn.apply( - x0, - residual, - weight, - bias, - rowscale, - layerscale, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - False, - return_dropout_mask, - ) - - -def dropout_add_layer_norm_subset( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - layerscale=None, - x0_subset=None, - out_subset=None, - rowscale_const=1.0, - out_numrows=0, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormSubsetFn.apply( - x0, - residual, - weight, - bias, - layerscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32, - prenorm, - False, - return_dropout_mask, - ) - - -def dropout_add_layer_norm_parallel_residual( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormParallelResidualFn.apply( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - False, - return_dropout_mask, - ) - - -class DropoutAddLayerNorm(torch.nn.Module): - def __init__( - self, - hidden_size, - prenorm=False, - p=0.0, - eps=1e-5, - residual_in_fp32=False, - device=None, - dtype=None, - ): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.prenorm = prenorm - self.p = p - self.eps = eps - self.residual_in_fp32 = residual_in_fp32 - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.bias = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.reset_parameters() - - def reset_parameters(self): - init.ones_(self.weight) - init.zeros_(self.bias) - - def forward(self, x0, residual=None): - return dropout_add_layer_norm( - x0, - residual, - self.weight, - self.bias, - self.p if self.training else 0.0, - self.eps, - prenorm=self.prenorm, - residual_in_fp32=self.residual_in_fp32, - ) diff --git a/build/torch212-cxx11-cu132-aarch64-linux/ops/rms_norm.py b/build/torch212-cxx11-cu132-aarch64-linux/ops/rms_norm.py deleted file mode 100644 index 068348d61290e3839dd082b540d898578ba1e8e2..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu132-aarch64-linux/ops/rms_norm.py +++ /dev/null @@ -1,174 +0,0 @@ -# Copyright (c) 2022, Tri Dao. -# Adapted from https://github.com/NVIDIA/apex/blob/master/apex/contrib/layer_norm/layer_norm.py - -import torch -from torch.nn import init - -from flash_attn.ops.layer_norm import ( - DropoutAddLayerNormFn, - DropoutAddLayerNormParallelResidualFn, - DropoutAddLayerNormSubsetFn, -) - - -def rms_norm(x, weight, epsilon): - return DropoutAddLayerNormFn.apply( - x, None, weight, None, None, None, 0.0, epsilon, False, False, True - ) - - -def dropout_add_rms_norm( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - rowscale=None, - layerscale=None, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormFn.apply( - x0, - residual, - weight, - bias, - rowscale, - layerscale, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - True, - return_dropout_mask, - ) - - -def dropout_add_rms_norm_subset( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - layerscale=None, - x0_subset=None, - out_subset=None, - rowscale_const=1.0, - out_numrows=0, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormSubsetFn.apply( - x0, - residual, - weight, - bias, - layerscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32, - prenorm, - True, - return_dropout_mask, - ) - - -def dropout_add_rms_norm_parallel_residual( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormParallelResidualFn.apply( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - True, - return_dropout_mask, - ) - - -class RMSNorm(torch.nn.Module): - def __init__(self, hidden_size, eps=1e-5, device=None, dtype=None): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.eps = eps - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.register_parameter("bias", None) - self.reset_parameters() - - def reset_parameters(self): - init.ones_(self.weight) - - def forward(self, x): - return rms_norm(x, self.weight, self.eps) - - -class DropoutAddRMSNorm(torch.nn.Module): - def __init__( - self, - hidden_size, - prenorm=False, - p=0.0, - eps=1e-5, - residual_in_fp32=False, - device=None, - dtype=None, - ): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.prenorm = prenorm - self.p = p - self.eps = eps - self.residual_in_fp32 = residual_in_fp32 - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.register_parameter("bias", None) - self.reset_parameters() - - def reset_parameters(self): - init.ones_(self.weight) - - def forward(self, x0, residual=None): - return dropout_add_rms_norm( - x0, - residual, - self.weight, - None, - self.p if self.training else 0.0, - self.eps, - prenorm=self.prenorm, - residual_in_fp32=self.residual_in_fp32, - ) diff --git a/build/torch212-cxx11-cu132-aarch64-linux/ops/triton/__init__.py b/build/torch212-cxx11-cu132-aarch64-linux/ops/triton/__init__.py deleted file mode 100644 index 8b137891791fe96927ad78e64b0aad7bded08bdc..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu132-aarch64-linux/ops/triton/__init__.py +++ /dev/null @@ -1 +0,0 @@ - diff --git a/build/torch212-cxx11-cu132-aarch64-linux/ops/triton/cross_entropy.py b/build/torch212-cxx11-cu132-aarch64-linux/ops/triton/cross_entropy.py deleted file mode 100644 index 1b5a415b73f236f3e05fb14b9141959559e18526..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu132-aarch64-linux/ops/triton/cross_entropy.py +++ /dev/null @@ -1,330 +0,0 @@ -# Copyright (c) 2023, Tri Dao. - -from typing import Tuple, Optional, Union - -import torch -import torch.nn.functional as F - -import triton -import triton.language as tl - -# `all_gather_into_tensor` and `reduce_scatter_tensor` are new placeholders for -# `_all_gather_base` and `_reduce_scatter_base`. They require the most recent -# version of PyTorch. The following 2 lines are for backward compatibility with -# older PyTorch. -if "all_gather_into_tensor" not in dir(torch.distributed): - torch.distributed.all_gather_into_tensor = torch.distributed._all_gather_base - - -@triton.heuristics( - { - "HAS_SMOOTHING": lambda args: args["smoothing"] > 0.0, - } -) -@triton.jit -def cross_entropy_fwd_kernel( - loss_ptr, # data ptrs - lse_ptr, - z_loss_ptr, - logits_ptr, - labels_ptr, - smoothing, - logit_scale, - lse_square_scale, - ignore_index, - total_classes, - class_start_idx, # Useful for tensor parallel when each rank only has a subset of classes - n_cols, # shapes - logits_row_stride, # strides - BLOCK_SIZE: tl.constexpr, - HAS_SMOOTHING: tl.constexpr, - # if SPLIT (e.g. tensor parallel), don't include the LSE in the loss since it's not the final LSE - SPLIT: tl.constexpr, - PRECOMPUTED_LSE: tl.constexpr, # If LSE is already computed (also no smoothing and logit_scale == 1.0) -): - row_idx = tl.program_id(0) - logits_ptr = logits_ptr + row_idx * logits_row_stride.to(tl.int64) - sum_logits = 0.0 # For smoothing - if not PRECOMPUTED_LSE: - # Statistics for online softmax - m_i = -float("inf") - l_i = 0.0 - for col_offset in range(0, n_cols, BLOCK_SIZE): - cols = col_offset + tl.arange(0, BLOCK_SIZE) - logits = tl.load(logits_ptr + cols, mask=cols < n_cols, other=-float("inf")).to( - tl.float32 - ) * logit_scale - if HAS_SMOOTHING: - sum_logits += tl.sum(tl.where(cols < n_cols, logits, 0.0)) - m_i_new = tl.maximum(m_i, tl.max(logits)) - l_i = tl.exp(m_i - m_i_new) * l_i + tl.sum(tl.exp(logits - m_i_new)) - m_i = m_i_new - lse = tl.log(l_i) + m_i - tl.store(lse_ptr + row_idx, lse) - else: - lse = tl.load(lse_ptr + row_idx) - label_idx = tl.load(labels_ptr + row_idx) - if label_idx == ignore_index: - loss = 0.0 - z_loss = 0.0 - else: - label_idx -= class_start_idx - if label_idx >= 0 and label_idx < n_cols: - logits_label = tl.load(logits_ptr + label_idx) * logit_scale - if HAS_SMOOTHING: - loss = ( - (lse if not SPLIT else 0.0) - - smoothing * sum_logits / total_classes - - (1 - smoothing) * logits_label - ) - else: - loss = (lse if not SPLIT else 0.0) - logits_label - else: - # If label is out of bounds, we set the CE loss to 0.0. But we still want the smoothing loss - if HAS_SMOOTHING: - loss = smoothing * ((lse if not SPLIT else 0.0) - sum_logits / total_classes) - else: - loss = 0.0 - if not SPLIT: - z_loss = lse_square_scale * lse * lse - loss += z_loss - else: - z_loss = 0.0 - tl.store(loss_ptr + row_idx, loss) - if not SPLIT: - tl.store(z_loss_ptr + row_idx, z_loss) - - -@triton.heuristics( - { - "HAS_SMOOTHING": lambda args: args["smoothing"] > 0.0, - } -) -@triton.jit -def cross_entropy_bwd_kernel( - dlogits_ptr, # data ptrs - dloss_ptr, - logits_ptr, - lse_ptr, - labels_ptr, - smoothing, - logit_scale, - lse_square_scale, - ignore_index, - total_classes, - class_start_idx, # Useful for tensor parallel when each rank only has a subset of classes - n_cols, # shapes - logits_row_stride, # strides - dlogits_row_stride, - dloss_row_stride, - BLOCK_SIZE: tl.constexpr, - HAS_SMOOTHING: tl.constexpr, -): - row_idx = tl.program_id(0) - col_block_idx = tl.program_id(1) - logits_ptr = logits_ptr + row_idx * logits_row_stride.to(tl.int64) - dlogits_ptr = dlogits_ptr + row_idx * dlogits_row_stride.to(tl.int64) - col_offsets = col_block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - label_idx = tl.load(labels_ptr + row_idx) - if label_idx != ignore_index: - dloss = tl.load(dloss_ptr + row_idx * dloss_row_stride) - else: - dloss = 0.0 - logits = tl.load(logits_ptr + col_offsets, mask=col_offsets < n_cols, other=-float("inf")).to( - tl.float32 - ) * logit_scale - lse = tl.load(lse_ptr + row_idx) - probs = tl.exp(logits - lse) - probs += 2.0 * lse_square_scale * lse * probs - label_idx -= class_start_idx - if HAS_SMOOTHING: - smooth_positive = 1.0 - smoothing - smooth_negative = smoothing / total_classes - probs = tl.where(col_offsets == label_idx, probs - smooth_positive, probs) - smooth_negative - else: - probs = tl.where(col_offsets == label_idx, probs - 1.0, probs) - tl.store(dlogits_ptr + col_offsets, (dloss * logit_scale) * probs, mask=col_offsets < n_cols) - - -class CrossEntropyLoss(torch.autograd.Function): - - @staticmethod - def forward( - ctx, - logits, - labels, - precomputed_lse=None, - smoothing=0.0, - logit_scale=1.0, - lse_square_scale=0.0, - ignore_index=-100, - inplace_backward=False, - process_group=None, - ): - # For some reason Triton generates wrong code when labels has dtype long and its address - # is not aligned to 16 bytes. The ld.global.b64 seems to load the wrong label index. - if labels.dtype == torch.long and labels.data_ptr() % 16 != 0: - labels = F.pad(labels, (0, 1))[..., :-1] - assert labels.data_ptr() % 16 == 0 - assert logit_scale > 0.0 - n_rows, n_cols = logits.shape - assert labels.shape == (n_rows,) - world_size = 1 if process_group is None else torch.distributed.get_world_size(process_group) - total_classes = world_size * n_cols - rank = 0 if process_group is None else torch.distributed.get_rank(process_group) - class_start_idx = rank * n_cols - use_precomputed_lse = precomputed_lse is not None and logit_scale == 1.0 and smoothing == 0.0 - - if logits.stride(-1) != 1: - logits = logits.contiguous() - MAX_BLOCK_SIZE = 16 * 1024 - BLOCK_SIZE = min(triton.next_power_of_2(n_cols), MAX_BLOCK_SIZE) - num_warps = ( - 4 - if BLOCK_SIZE < 2048 - else (8 if BLOCK_SIZE < 8192 else (16 if BLOCK_SIZE < 128 * 1024 else 32)) - ) - losses = torch.empty(n_rows, dtype=torch.float, device=logits.device) - if use_precomputed_lse: - assert precomputed_lse.shape == (n_rows,) - lse = precomputed_lse.contiguous() - else: - lse = torch.empty(n_rows, dtype=torch.float, device=logits.device) - z_losses = torch.empty(n_rows, dtype=torch.float, device=logits.device) - # Need this, otherwise Triton tries to launch from cuda:0 and we get - # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?) - with torch.cuda.device(logits.device.index): - cross_entropy_fwd_kernel[(n_rows,)]( - losses, # data ptrs - lse, - z_losses, - logits, - labels, - smoothing, - logit_scale, - lse_square_scale, - ignore_index, - total_classes, - class_start_idx, - n_cols, # shapes - logits.stride(0), # strides - BLOCK_SIZE=BLOCK_SIZE, # constants - SPLIT=world_size > 1, - PRECOMPUTED_LSE=use_precomputed_lse, - num_warps=num_warps, - ) - - if world_size > 1: - # If there's no smoothing, if labels are in the vocab of this partition, losses contains - # - predicted logit, and 0 otherwise. - # If there's smoothing=0.1, for labels in the vocab of this partition, losses contains - # -0.9 * predicted logit - 0.1 * sum logit / total_classes. - # For labels not in the vocab of this partition, losses contains - # -0.1 * sum logit / total_classes. - if world_size > 1: - lse_allgather = torch.empty(world_size, n_rows, dtype=lse.dtype, device=lse.device) - torch.distributed.all_gather_into_tensor(lse_allgather, lse, group=process_group) - handle_losses = torch.distributed.all_reduce( - losses, op=torch.distributed.ReduceOp.SUM, group=process_group, async_op=True - ) - lse = torch.logsumexp(lse_allgather, dim=0) - handle_losses.wait() - # After the allreduce, if there's no smoothing, the total losses are - predicted_logit, - # we just have to add the (global) lse. - # If there's smoothing=0.1, the total losses are - # -0.9 * predicted_logit - 0.1 * sum logit / total_classes. - # Again, we just have to add the (global) lse. - losses += lse - if lse_square_scale != 0.0: - z_losses = lse_square_scale * lse.square() - z_losses.masked_fill_(labels == ignore_index, 0.0) - losses += z_losses - else: - z_losses = torch.zeros_like(losses) - losses.masked_fill_(labels == ignore_index, 0.0) - - ctx.save_for_backward(logits, lse, labels) - ctx.mark_non_differentiable(z_losses) - ctx.smoothing = smoothing - ctx.logit_scale = logit_scale - ctx.lse_square_scale = lse_square_scale - ctx.ignore_index = ignore_index - ctx.total_classes = total_classes - ctx.class_start_idx = class_start_idx - ctx.inplace_backward = inplace_backward - return losses, z_losses - - @staticmethod - def backward(ctx, grad_losses, grad_z_losses): - del grad_z_losses # z_losses are only for logging. - - logits, lse, labels = ctx.saved_tensors - dlogits = logits if ctx.inplace_backward else torch.empty_like(logits) - n_rows, n_cols = logits.shape - BLOCK_SIZE = min(triton.next_power_of_2(n_cols), 4 * 1024) - num_warps = 4 if BLOCK_SIZE < 2048 else (8 if BLOCK_SIZE < 8192 else 16) - grid = lambda META: (n_rows, triton.cdiv(n_cols, META["BLOCK_SIZE"])) # noqa - # Need this, otherwise Triton tries to launch from cuda:0 and we get - # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?) - with torch.cuda.device(logits.device.index): - cross_entropy_bwd_kernel[grid]( - dlogits, # data ptrs - grad_losses, - logits, - lse, - labels, - ctx.smoothing, - ctx.logit_scale, - ctx.lse_square_scale, - ctx.ignore_index, - ctx.total_classes, - ctx.class_start_idx, - n_cols, # shapes - logits.stride(0), # strides - dlogits.stride(0), - grad_losses.stride(0), - BLOCK_SIZE=BLOCK_SIZE, # constants - num_warps=num_warps, - ) - return dlogits, None, None, None, None, None, None, None, None, None - - -def cross_entropy_loss( - logits: torch.Tensor, - labels: torch.Tensor, - precomputed_lse: Optional[torch.Tensor] = None, - label_smoothing: float = 0.0, - logit_scale: float = 1.0, - lse_square_scale: float = 0.0, - ignore_index=-100, - inplace_backward: bool = False, - process_group=None, -) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Arguments: - logits: (batch, vocab_size) - labels: (batch,) - label_smoothing: float - logit_scale: float. Multiply logits by this scale before calculating the loss. - lse_square_scale: float. If > 0, we add lse_square_scale * lse(logits) ^ 2 to the loss. - This is also referred to as "z-loss". - ignore_index: int. If labels == ignore_index, the loss is set to 0.0. - inplace_backward: bool. If True, we do the backward pass in-place by modifying the logits. - This saves memory. - process_group: if not None, we're doing Tensor Parallel: each process is responsible for - one part of the vocab. The loss will be aggregated across processes. - Returns: - losses: (batch,), float - z_losses: (batch,), float - """ - return CrossEntropyLoss.apply( - logits, - labels, - precomputed_lse, - label_smoothing, - logit_scale, - lse_square_scale, - ignore_index, - inplace_backward, - process_group, - ) diff --git a/build/torch212-cxx11-cu132-aarch64-linux/ops/triton/k_activations.py b/build/torch212-cxx11-cu132-aarch64-linux/ops/triton/k_activations.py deleted file mode 100644 index efb83c358eb4a85d069ee340a3c83f418f9a805b..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu132-aarch64-linux/ops/triton/k_activations.py +++ /dev/null @@ -1,162 +0,0 @@ -# Adapted from https://github.com/facebookresearch/xformers/blob/main/xformers/triton/k_activations.py -# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved. -# -# This source code is licensed under the BSD license found in the -# LICENSE file in the root directory of this source tree. - -import math -from enum import Enum -from typing import Optional - -import triton -import triton.language as tl - -_sqrt2pi = math.sqrt(2.0 / math.pi) -_sqrt1_2 = math.sqrt(1.0 / 2) -_gaussian_pdf_normalization = 1.0 / math.sqrt(2 * math.pi) - - -class Activation(str, Enum): - SquaredReLU = "squared_relu" - GeLU = "gelu" - GeLUApprox = "gelu_approx" - LeakyReLU = "leaky_relu" - ReLU = "relu" - - -def get_triton_activation_kernel(activation: Optional[Activation]): - return ( - { - Activation.ReLU: relu, - Activation.LeakyReLU: leaky_relu, - Activation.GeLU: gelu, - Activation.GeLUApprox: gelu_approx, - Activation.SquaredReLU: squared_relu, - }[activation] - if activation - else None - ) - - -def get_triton_activation_bwd_kernel(activation: Optional[Activation]): - return ( - { - Activation.ReLU: relu_grad, - Activation.LeakyReLU: leaky_relu_grad, - Activation.GeLU: gelu_grad, - Activation.GeLUApprox: gelu_approx_grad, - Activation.SquaredReLU: squared_relu_grad, - }[activation] - if activation - else None - ) - - -@triton.jit -def tanh(x): - # Tanh is just a scaled sigmoid - return 2 * tl.sigmoid(2 * x) - 1 - - -@triton.jit -def cosh(x): - exp_x = tl.exp(x) - return (exp_x + 1.0 / exp_x) * 0.5 - - -# a Triton implementation of the most used activations -# See for instance http://arxiv.org/abs/1606.08415 for an overview - -# ReLU -@triton.jit -def relu(x): - """ - ReLU_ activation function - - .. _ReLU: https://pytorch.org/docs/stable/generated/torch.nn.ReLU.html - """ - zero = 0.0 - return tl.where(x >= 0, x, zero.to(x.dtype)) - - -@triton.jit -def relu_grad(x): - # ReLU is different from other activations - # in that it does not require the input to retrospectively compute its gradient - # here the input is the downstream gradient, and we return the upstream gradient directly - zero = 0.0 - one = 1.0 - return tl.where(x >= 0, one.to(x.dtype), zero.to(x.dtype)) - - -@triton.jit -def squared_relu(x): - """ - Squared ReLU activation, as proposed in the Primer_ paper. - - .. _Primer: https://arxiv.org/abs/2109.08668 - """ - x_ = relu(x) - return (x_ * x_).to(x.dtype) - - -@triton.jit -def squared_relu_grad(x): - return tl.where(x >= 0, 2.0 * x, 0.0) - - -# Leaky ReLU -@triton.jit -def leaky_relu(x): - """ - LeakyReLU_ activation - - .. _LeakyReLU: https://pytorch.org/docs/stable/generated/torch.nn.LeakyReLU.html - """ - scale = 0.01 + 0.0 - scale = scale.to(x.dtype) - return tl.where(x >= 0, x, scale * x) - - -@triton.jit -def leaky_relu_grad(x): - min_grad = 0.01 - max_grad = 1 - - min_grad = min_grad.to(x.dtype) - max_grad = max_grad.to(x.dtype) - - return tl.where(x >= 0, max_grad, min_grad) - - -@triton.jit -def gelu(x): - """Gaussian Error Linear Unit (GELU)""" - return x * 0.5 * (1.0 + tl.libdevice.erf(x * _sqrt1_2)) - - -@triton.jit -def gelu_grad(x): - cdf = 0.5 * (1.0 + tl.libdevice.erf(x * _sqrt1_2)) - pdf = tl.exp(-0.5 * x * x) * _gaussian_pdf_normalization - return cdf + x * pdf - - -@triton.jit -def gelu_approx(x): - """ - GeLU_ activation - Gaussian error linear unit, with tanh approximation - - .. _GeLU: https://arxiv.org/pdf/1606.08415.pdf - """ - return 0.5 * x * (1.0 + tanh(_sqrt2pi * x * (1.0 + 0.044715 * x * x))) - - -@triton.jit -def gelu_approx_grad(x): - # CREDITS: Fast implementation proposed in - # https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/fused_bias_gelu.py#L30 - tanh_out = tanh(0.79788456 * x * (1 + 0.044715 * x * x)) - return 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( - 1 + tanh_out - ) diff --git a/build/torch212-cxx11-cu132-aarch64-linux/ops/triton/layer_norm.py b/build/torch212-cxx11-cu132-aarch64-linux/ops/triton/layer_norm.py deleted file mode 100644 index 5e6f0b3697b612e6a19031bbc55ac76ad233dc85..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu132-aarch64-linux/ops/triton/layer_norm.py +++ /dev/null @@ -1,1254 +0,0 @@ -# Copyright (c) 2024, Tri Dao. -# Implement dropout + residual + layer_norm / rms_norm. - -# Based on the Triton LayerNorm tutorial: https://triton-lang.org/main/getting-started/tutorials/05-layer-norm.html -# For the backward pass, we keep weight_grad and bias_grad in registers and accumulate. -# This is faster for dimensions up to 8k, but after that it's much slower due to register spilling. -# The models we train have hidden dim up to 8k anyway (e.g. Llama 70B), so this is fine. - -import math -from typing import Optional, List - -import torch -import torch.nn.functional as F -from torch import Tensor - -import triton -import triton.language as tl - -from flash_attn.utils.torch import custom_fwd, custom_bwd -from flash_attn.utils.library import triton_op - -from ..._ops import add_op_namespace_prefix - - -def maybe_contiguous_lastdim(x): - return x.contiguous() if x is not None and x.stride(-1) != 1 else x - - -def maybe_contiguous(x): - return x.contiguous() if x is not None else None - - -def triton_autotune_configs(): - # Return configs with a valid warp count for the current device - configs = [] - # Maximum threads per block is architecture-dependent in theory, but in reality all are 1024 - max_threads_per_block = 1024 - # Default to warp size 32 if not defined by device - warp_size = getattr(torch.cuda.get_device_properties(torch.cuda.current_device()), "warp_size", 32) - # Autotune for warp counts which are powers of 2 and do not exceed thread per block limit - return [triton.Config({}, num_warps=warp_count) for warp_count in [1, 2, 4, 8, 16, 32] - if warp_count * warp_size <= max_threads_per_block] - # return [triton.Config({}, num_warps=8)] - - -def layer_norm_ref( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - zero_centered_weight=False, - dropout_mask=None, - dropout_mask1=None, - upcast=False, -): - dtype = x.dtype - if upcast: - x = x.float() - weight = weight.float() - bias = bias.float() if bias is not None else None - residual = residual.float() if residual is not None else residual - x1 = x1.float() if x1 is not None else None - weight1 = weight1.float() if weight1 is not None else None - bias1 = bias1.float() if bias1 is not None else None - if zero_centered_weight: - weight = weight + 1.0 - if weight1 is not None: - weight1 = weight1 + 1.0 - if x1 is not None: - assert rowscale is None, "rowscale is not supported with parallel LayerNorm" - if rowscale is not None: - x = x * rowscale[..., None] - if dropout_p > 0.0: - if dropout_mask is not None: - x = x.masked_fill(~dropout_mask, 0.0) / (1.0 - dropout_p) - else: - x = F.dropout(x, p=dropout_p) - if x1 is not None: - if dropout_mask1 is not None: - x1 = x1.masked_fill(~dropout_mask1, 0.0) / (1.0 - dropout_p) - else: - x1 = F.dropout(x1, p=dropout_p) - if x1 is not None: - x = x + x1 - if residual is not None: - x = (x + residual).to(x.dtype) - out = F.layer_norm(x.to(weight.dtype), x.shape[-1:], weight=weight, bias=bias, eps=eps).to( - dtype - ) - if weight1 is None: - return out if not prenorm else (out, x) - else: - out1 = F.layer_norm( - x.to(weight1.dtype), x.shape[-1:], weight=weight1, bias=bias1, eps=eps - ).to(dtype) - return (out, out1) if not prenorm else (out, out1, x) - - -def rms_norm_ref( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - zero_centered_weight=False, - dropout_mask=None, - dropout_mask1=None, - upcast=False, -): - dtype = x.dtype - if upcast: - x = x.float() - weight = weight.float() - bias = bias.float() if bias is not None else None - residual = residual.float() if residual is not None else residual - x1 = x1.float() if x1 is not None else None - weight1 = weight1.float() if weight1 is not None else None - bias1 = bias1.float() if bias1 is not None else None - if zero_centered_weight: - weight = weight + 1.0 - if weight1 is not None: - weight1 = weight1 + 1.0 - if x1 is not None: - assert rowscale is None, "rowscale is not supported with parallel LayerNorm" - if rowscale is not None: - x = x * rowscale[..., None] - if dropout_p > 0.0: - if dropout_mask is not None: - x = x.masked_fill(~dropout_mask, 0.0) / (1.0 - dropout_p) - else: - x = F.dropout(x, p=dropout_p) - if x1 is not None: - if dropout_mask1 is not None: - x1 = x1.masked_fill(~dropout_mask1, 0.0) / (1.0 - dropout_p) - else: - x1 = F.dropout(x1, p=dropout_p) - if x1 is not None: - x = x + x1 - if residual is not None: - x = (x + residual).to(x.dtype) - rstd = 1 / torch.sqrt((x.square()).mean(dim=-1, keepdim=True) + eps) - out = ((x * rstd * weight) + bias if bias is not None else (x * rstd * weight)).to(dtype) - if weight1 is None: - return out if not prenorm else (out, x) - else: - out1 = ((x * rstd * weight1) + bias1 if bias1 is not None else (x * rstd * weight1)).to( - dtype - ) - return (out, out1) if not prenorm else (out, out1, x) - - -@triton.autotune( - configs=triton_autotune_configs(), - key=["N", "HAS_RESIDUAL", "STORE_RESIDUAL_OUT", "IS_RMS_NORM", "HAS_BIAS", "HAS_X1", "HAS_W1", "HAS_B1"], -) -# torch compile doesn't like triton.heuristics, so we set these manually when calling the kernel -# @triton.heuristics({"HAS_BIAS": lambda args: args["B"] is not None}) -# @triton.heuristics({"HAS_RESIDUAL": lambda args: args["RESIDUAL"] is not None}) -# @triton.heuristics({"HAS_X1": lambda args: args["X1"] is not None}) -# @triton.heuristics({"HAS_W1": lambda args: args["W1"] is not None}) -# @triton.heuristics({"HAS_B1": lambda args: args["B1"] is not None}) -@triton.jit -def _layer_norm_fwd_1pass_kernel( - X, # pointer to the input - Y, # pointer to the output - W, # pointer to the weights - B, # pointer to the biases - RESIDUAL, # pointer to the residual - X1, - W1, - B1, - Y1, - RESIDUAL_OUT, # pointer to the residual - ROWSCALE, - SEEDS, # Dropout seeds for each row - DROPOUT_MASK, - DROPOUT_MASK1, - Mean, # pointer to the mean - Rstd, # pointer to the 1/std - stride_x_row, # how much to increase the pointer when moving by 1 row - stride_y_row, - stride_res_row, - stride_res_out_row, - stride_x1_row, - stride_y1_row, - M, # number of rows in X - N, # number of columns in X - eps, # epsilon to avoid division by zero - dropout_p, # Dropout probability - zero_centered_weight, # If true, add 1.0 to the weight - IS_RMS_NORM: tl.constexpr, - BLOCK_N: tl.constexpr, - HAS_RESIDUAL: tl.constexpr, - STORE_RESIDUAL_OUT: tl.constexpr, - HAS_BIAS: tl.constexpr, - HAS_DROPOUT: tl.constexpr, - STORE_DROPOUT_MASK: tl.constexpr, - HAS_ROWSCALE: tl.constexpr, - HAS_X1: tl.constexpr, - HAS_W1: tl.constexpr, - HAS_B1: tl.constexpr, -): - # Map the program id to the row of X and Y it should compute. - row = tl.program_id(0) - X += row * stride_x_row - Y += row * stride_y_row - if HAS_RESIDUAL: - RESIDUAL += row * stride_res_row - if STORE_RESIDUAL_OUT: - RESIDUAL_OUT += row * stride_res_out_row - if HAS_X1: - X1 += row * stride_x1_row - if HAS_W1: - Y1 += row * stride_y1_row - # Compute mean and variance - cols = tl.arange(0, BLOCK_N) - x = tl.load(X + cols, mask=cols < N, other=0.0).to(tl.float32) - if HAS_ROWSCALE: - rowscale = tl.load(ROWSCALE + row).to(tl.float32) - x *= rowscale - if HAS_DROPOUT: - # Compute dropout mask - # 7 rounds is good enough, and reduces register pressure - keep_mask = tl.rand(tl.load(SEEDS + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - x = tl.where(keep_mask, x / (1.0 - dropout_p), 0.0) - if STORE_DROPOUT_MASK: - tl.store(DROPOUT_MASK + row * N + cols, keep_mask, mask=cols < N) - if HAS_X1: - x1 = tl.load(X1 + cols, mask=cols < N, other=0.0).to(tl.float32) - if HAS_ROWSCALE: - rowscale = tl.load(ROWSCALE + M + row).to(tl.float32) - x1 *= rowscale - if HAS_DROPOUT: - # Compute dropout mask - # 7 rounds is good enough, and reduces register pressure - keep_mask = ( - tl.rand(tl.load(SEEDS + M + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - ) - x1 = tl.where(keep_mask, x1 / (1.0 - dropout_p), 0.0) - if STORE_DROPOUT_MASK: - tl.store(DROPOUT_MASK1 + row * N + cols, keep_mask, mask=cols < N) - x += x1 - if HAS_RESIDUAL: - residual = tl.load(RESIDUAL + cols, mask=cols < N, other=0.0).to(tl.float32) - x += residual - if STORE_RESIDUAL_OUT: - tl.store(RESIDUAL_OUT + cols, x, mask=cols < N) - if not IS_RMS_NORM: - mean = tl.sum(x, axis=0) / N - tl.store(Mean + row, mean) - xbar = tl.where(cols < N, x - mean, 0.0) - var = tl.sum(xbar * xbar, axis=0) / N - else: - xbar = tl.where(cols < N, x, 0.0) - var = tl.sum(xbar * xbar, axis=0) / N - rstd = 1 / tl.sqrt(var + eps) - tl.store(Rstd + row, rstd) - # Normalize and apply linear transformation - mask = cols < N - w = tl.load(W + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w += 1.0 - if HAS_BIAS: - b = tl.load(B + cols, mask=mask).to(tl.float32) - x_hat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd - y = x_hat * w + b if HAS_BIAS else x_hat * w - # Write output - tl.store(Y + cols, y, mask=mask) - if HAS_W1: - w1 = tl.load(W1 + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w1 += 1.0 - if HAS_B1: - b1 = tl.load(B1 + cols, mask=mask).to(tl.float32) - y1 = x_hat * w1 + b1 if HAS_B1 else x_hat * w1 - tl.store(Y1 + cols, y1, mask=mask) - - -def _layer_norm_fwd( - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - residual: Optional[Tensor] = None, - x1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - out_dtype: Optional[torch.dtype] = None, - residual_dtype: Optional[torch.dtype] = None, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - return_dropout_mask: bool = False, - out: Optional[Tensor] = None, - residual_out: Optional[Tensor] = None -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - # Need to wrap to handle the case where residual_out is a alias of x, which makes torch.library - # and torch.compile unhappy. Also allocate memory for out and residual_out if they are None - # so that _layer_norm_fwd_impl doesn't have to return them. - if out is None: - out = torch.empty_like(x, dtype=x.dtype if out_dtype is None else out_dtype) - if residual is not None: - residual_dtype = residual.dtype - if residual_out is None and ( - residual is not None - or (residual_dtype is not None and residual_dtype != x.dtype) - or dropout_p > 0.0 - or rowscale is not None - or x1 is not None - ): - residual_out = torch.empty_like( - x, dtype=residual_dtype if residual_dtype is not None else x.dtype - ) - else: - residual_out = None - y1, mean, rstd, seeds, dropout_mask, dropout_mask1 = _layer_norm_fwd_impl( - x, - weight, - bias, - eps, - out, - residual=residual, - x1=x1, - weight1=weight1, - bias1=bias1, - dropout_p=dropout_p, - rowscale=rowscale, - zero_centered_weight=zero_centered_weight, - is_rms_norm=is_rms_norm, - return_dropout_mask=return_dropout_mask, - residual_out=residual_out, - ) - # residual_out is None if residual is None and residual_dtype == input_dtype and dropout_p == 0.0 - if residual_out is None: - residual_out = x - return out, y1, mean, rstd, residual_out, seeds, dropout_mask, dropout_mask1 - - -# [2025-04-28] torch.library.triton_op ignores the schema argument, but here we need the schema -# since we're returning a tuple of tensors -@triton_op(add_op_namespace_prefix("layer_norm_fwd_impl"), mutates_args={"out", "residual_out"}, - schema="(Tensor x, Tensor weight, Tensor bias, float eps, Tensor(a!) out, Tensor? residual, Tensor? x1, Tensor? weight1, Tensor? bias1, float dropout_p, Tensor? rowscale, bool zero_centered_weight, bool is_rms_norm, bool return_dropout_mask, Tensor(a!)? residual_out) -> (Tensor y1, Tensor mean, Tensor rstd, Tensor seeds, Tensor dropout_mask, Tensor dropout_mask1)") -def _layer_norm_fwd_impl( - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - out: Tensor, - residual: Optional[Tensor] = None, - x1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - return_dropout_mask: bool = False, - residual_out: Optional[Tensor] = None -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - M, N = x.shape - assert x.stride(-1) == 1 - if residual is not None: - assert residual.stride(-1) == 1 - assert residual.shape == (M, N) - assert weight.shape == (N,) - assert weight.stride(-1) == 1 - if bias is not None: - assert bias.stride(-1) == 1 - assert bias.shape == (N,) - if x1 is not None: - assert x1.shape == x.shape - assert rowscale is None - assert x1.stride(-1) == 1 - if weight1 is not None: - assert weight1.shape == (N,) - assert weight1.stride(-1) == 1 - if bias1 is not None: - assert bias1.shape == (N,) - assert bias1.stride(-1) == 1 - if rowscale is not None: - assert rowscale.is_contiguous() - assert rowscale.shape == (M,) - assert out.shape == x.shape - assert out.stride(-1) == 1 - if residual_out is not None: - assert residual_out.shape == x.shape - assert residual_out.stride(-1) == 1 - if weight1 is not None: - y1 = torch.empty_like(out) - assert y1.stride(-1) == 1 - else: - y1 = None - mean = torch.empty((M,), dtype=torch.float32, device=x.device) if not is_rms_norm else None - rstd = torch.empty((M,), dtype=torch.float32, device=x.device) - if dropout_p > 0.0: - seeds = torch.randint( - 2**32, (M if x1 is None else 2 * M,), device=x.device, dtype=torch.int64 - ) - else: - seeds = None - if return_dropout_mask and dropout_p > 0.0: - dropout_mask = torch.empty(M, N, device=x.device, dtype=torch.bool) - if x1 is not None: - dropout_mask1 = torch.empty(M, N, device=x.device, dtype=torch.bool) - else: - dropout_mask1 = None - else: - dropout_mask, dropout_mask1 = None, None - # Less than 64KB per feature: enqueue fused kernel - MAX_FUSED_SIZE = 65536 // x.element_size() - BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) - if N > BLOCK_N: - raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") - with torch.cuda.device(x.device.index): - torch.library.wrap_triton(_layer_norm_fwd_1pass_kernel)[(M,)]( - x, - out, - weight, - bias, - residual, - x1, - weight1, - bias1, - y1, - residual_out, - rowscale, - seeds, - dropout_mask, - dropout_mask1, - mean, - rstd, - x.stride(0), - out.stride(0), - residual.stride(0) if residual is not None else 0, - residual_out.stride(0) if residual_out is not None else 0, - x1.stride(0) if x1 is not None else 0, - y1.stride(0) if y1 is not None else 0, - M, - N, - eps, - dropout_p, - # Passing bool make torch inductor very unhappy since it then tries to compare to int_max - int(zero_centered_weight), - is_rms_norm, - BLOCK_N, - residual is not None, - residual_out is not None, - bias is not None, - dropout_p > 0.0, - dropout_mask is not None, - rowscale is not None, - HAS_X1=x1 is not None, - HAS_W1=weight1 is not None, - HAS_B1=bias1 is not None, - ) - return y1, mean, rstd, seeds, dropout_mask, dropout_mask1 - - -@triton.autotune( - configs=triton_autotune_configs(), - key=["N", "HAS_DRESIDUAL", "STORE_DRESIDUAL", "IS_RMS_NORM", "HAS_BIAS", "HAS_DROPOUT"], -) -# torch compile doesn't like triton.heuristics, so we set these manually when calling the kernel -# @triton.heuristics({"HAS_BIAS": lambda args: args["B"] is not None}) -# @triton.heuristics({"HAS_DRESIDUAL": lambda args: args["DRESIDUAL"] is not None}) -# @triton.heuristics({"STORE_DRESIDUAL": lambda args: args["DRESIDUAL_IN"] is not None}) -# @triton.heuristics({"HAS_ROWSCALE": lambda args: args["ROWSCALE"] is not None}) -# @triton.heuristics({"HAS_DY1": lambda args: args["DY1"] is not None}) -# @triton.heuristics({"HAS_DX1": lambda args: args["DX1"] is not None}) -# @triton.heuristics({"HAS_B1": lambda args: args["DB1"] is not None}) -# @triton.heuristics({"RECOMPUTE_OUTPUT": lambda args: args["Y"] is not None}) -@triton.jit -def _layer_norm_bwd_kernel( - X, # pointer to the input - W, # pointer to the weights - B, # pointer to the biases - Y, # pointer to the output to be recomputed - DY, # pointer to the output gradient - DX, # pointer to the input gradient - DW, # pointer to the partial sum of weights gradient - DB, # pointer to the partial sum of biases gradient - DRESIDUAL, - W1, - DY1, - DX1, - DW1, - DB1, - DRESIDUAL_IN, - ROWSCALE, - SEEDS, - Mean, # pointer to the mean - Rstd, # pointer to the 1/std - stride_x_row, # how much to increase the pointer when moving by 1 row - stride_y_row, - stride_dy_row, - stride_dx_row, - stride_dres_row, - stride_dy1_row, - stride_dx1_row, - stride_dres_in_row, - M, # number of rows in X - N, # number of columns in X - eps, # epsilon to avoid division by zero - dropout_p, - zero_centered_weight, - rows_per_program, - IS_RMS_NORM: tl.constexpr, - BLOCK_N: tl.constexpr, - HAS_DRESIDUAL: tl.constexpr, - STORE_DRESIDUAL: tl.constexpr, - HAS_BIAS: tl.constexpr, - HAS_DROPOUT: tl.constexpr, - HAS_ROWSCALE: tl.constexpr, - HAS_DY1: tl.constexpr, - HAS_DX1: tl.constexpr, - HAS_B1: tl.constexpr, - RECOMPUTE_OUTPUT: tl.constexpr, -): - # Map the program id to the elements of X, DX, and DY it should compute. - row_block_id = tl.program_id(0) - row_start = row_block_id * rows_per_program - # Do not early exit if row_start >= M, because we need to write DW and DB - cols = tl.arange(0, BLOCK_N) - mask = cols < N - X += row_start * stride_x_row - if HAS_DRESIDUAL: - DRESIDUAL += row_start * stride_dres_row - if STORE_DRESIDUAL: - DRESIDUAL_IN += row_start * stride_dres_in_row - DY += row_start * stride_dy_row - DX += row_start * stride_dx_row - if HAS_DY1: - DY1 += row_start * stride_dy1_row - if HAS_DX1: - DX1 += row_start * stride_dx1_row - if RECOMPUTE_OUTPUT: - Y += row_start * stride_y_row - w = tl.load(W + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w += 1.0 - if RECOMPUTE_OUTPUT and HAS_BIAS: - b = tl.load(B + cols, mask=mask, other=0.0).to(tl.float32) - if HAS_DY1: - w1 = tl.load(W1 + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w1 += 1.0 - dw = tl.zeros((BLOCK_N,), dtype=tl.float32) - if HAS_BIAS: - db = tl.zeros((BLOCK_N,), dtype=tl.float32) - if HAS_DY1: - dw1 = tl.zeros((BLOCK_N,), dtype=tl.float32) - if HAS_B1: - db1 = tl.zeros((BLOCK_N,), dtype=tl.float32) - row_end = min((row_block_id + 1) * rows_per_program, M) - for row in range(row_start, row_end): - # Load data to SRAM - x = tl.load(X + cols, mask=mask, other=0).to(tl.float32) - dy = tl.load(DY + cols, mask=mask, other=0).to(tl.float32) - if HAS_DY1: - dy1 = tl.load(DY1 + cols, mask=mask, other=0).to(tl.float32) - if not IS_RMS_NORM: - mean = tl.load(Mean + row) - rstd = tl.load(Rstd + row) - # Compute dx - xhat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd - xhat = tl.where(mask, xhat, 0.0) - if RECOMPUTE_OUTPUT: - y = xhat * w + b if HAS_BIAS else xhat * w - tl.store(Y + cols, y, mask=mask) - wdy = w * dy - dw += dy * xhat - if HAS_BIAS: - db += dy - if HAS_DY1: - wdy += w1 * dy1 - dw1 += dy1 * xhat - if HAS_B1: - db1 += dy1 - if not IS_RMS_NORM: - c1 = tl.sum(xhat * wdy, axis=0) / N - c2 = tl.sum(wdy, axis=0) / N - dx = (wdy - (xhat * c1 + c2)) * rstd - else: - c1 = tl.sum(xhat * wdy, axis=0) / N - dx = (wdy - xhat * c1) * rstd - if HAS_DRESIDUAL: - dres = tl.load(DRESIDUAL + cols, mask=mask, other=0).to(tl.float32) - dx += dres - # Write dx - if STORE_DRESIDUAL: - tl.store(DRESIDUAL_IN + cols, dx, mask=mask) - if HAS_DX1: - if HAS_DROPOUT: - keep_mask = ( - tl.rand(tl.load(SEEDS + M + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - ) - dx1 = tl.where(keep_mask, dx / (1.0 - dropout_p), 0.0) - else: - dx1 = dx - tl.store(DX1 + cols, dx1, mask=mask) - if HAS_DROPOUT: - keep_mask = tl.rand(tl.load(SEEDS + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - dx = tl.where(keep_mask, dx / (1.0 - dropout_p), 0.0) - if HAS_ROWSCALE: - rowscale = tl.load(ROWSCALE + row).to(tl.float32) - dx *= rowscale - tl.store(DX + cols, dx, mask=mask) - - X += stride_x_row - if HAS_DRESIDUAL: - DRESIDUAL += stride_dres_row - if STORE_DRESIDUAL: - DRESIDUAL_IN += stride_dres_in_row - if RECOMPUTE_OUTPUT: - Y += stride_y_row - DY += stride_dy_row - DX += stride_dx_row - if HAS_DY1: - DY1 += stride_dy1_row - if HAS_DX1: - DX1 += stride_dx1_row - tl.store(DW + row_block_id * N + cols, dw, mask=mask) - if HAS_BIAS: - tl.store(DB + row_block_id * N + cols, db, mask=mask) - if HAS_DY1: - tl.store(DW1 + row_block_id * N + cols, dw1, mask=mask) - if HAS_B1: - tl.store(DB1 + row_block_id * N + cols, db1, mask=mask) - - -def _layer_norm_bwd( - dy: Tensor, - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - mean: Tensor, - rstd: Tensor, - dresidual: Optional[Tensor] = None, - dy1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - seeds: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - has_residual: bool = False, - has_x1: bool = False, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - x_dtype: Optional[torch.dtype] = None, - recompute_output: bool = False, -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - # Need to wrap to handle the case where dresidual_in or dx1 are aliases of x, - # which makes torch.library unhappy - dx, dw, db, dresidual_in, dx1, dw1, db1, y = _layer_norm_bwd_impl( - dy, - x, - weight, - bias, - eps, - mean, - rstd, - dresidual, - dy1, - weight1, - bias1, - seeds, - dropout_p, - rowscale, - has_residual, - has_x1, - zero_centered_weight, - is_rms_norm, - x_dtype=x_dtype, - recompute_output=recompute_output, - ) - # Don't need to compute dresidual_in separately in this case - if has_residual and dx.dtype == x.dtype and dropout_p == 0.0 and rowscale is None: - dresidual_in = dx - if has_x1 and dropout_p == 0.0: - dx1 = dx - return dx, dw, db, dresidual_in, dx1, dw1, db1, y - - - -@triton_op(add_op_namespace_prefix("layer_norm_bwd_impl"), mutates_args={}, - schema="(Tensor dy, Tensor x, Tensor weight, Tensor bias, float eps, Tensor mean, Tensor rstd, Tensor? dresidual, Tensor? dy1, Tensor? weight1, Tensor? bias1, Tensor? seeds, float dropout_p, Tensor? rowscale, bool has_residual, bool has_x1, bool zero_centered_weight, bool is_rms_norm, ScalarType? x_dtype, bool recompute_output) -> (Tensor dx, Tensor dw, Tensor db, Tensor dresidual_in, Tensor dx1, Tensor dw1, Tensor db1, Tensor y)", - allow_decomposition=False, # Don't let torch.compile trace inside - ) -def _layer_norm_bwd_impl( - dy: Tensor, - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - mean: Tensor, - rstd: Tensor, - dresidual: Optional[Tensor] = None, - dy1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - seeds: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - has_residual: bool = False, - has_x1: bool = False, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - x_dtype: Optional[torch.dtype] = None, - recompute_output: bool = False, -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - M, N = x.shape - assert x.stride(-1) == 1 - dy = maybe_contiguous_lastdim(dy) - assert dy.stride(-1) == 1 - assert dy.shape == (M, N) - if dresidual is not None: - dresidual = maybe_contiguous_lastdim(dresidual) - assert dresidual.stride(-1) == 1 - assert dresidual.shape == (M, N) - assert weight.shape == (N,) - assert weight.stride(-1) == 1 - if bias is not None: - assert bias.stride(-1) == 1 - assert bias.shape == (N,) - if dy1 is not None: - dy1 = maybe_contiguous_lastdim(dy1) - assert weight1 is not None - assert dy1.shape == dy.shape - assert dy1.stride(-1) == 1 - if weight1 is not None: - assert weight1.shape == (N,) - assert weight1.stride(-1) == 1 - if bias1 is not None: - assert bias1.shape == (N,) - assert bias1.stride(-1) == 1 - if seeds is not None: - assert seeds.is_contiguous() - assert seeds.shape == (M if not has_x1 else M * 2,) - if rowscale is not None: - assert rowscale.is_contiguous() - assert rowscale.shape == (M,) - # allocate output - dx = ( - torch.empty_like(x) - if x_dtype is None - else torch.empty(M, N, dtype=x_dtype, device=x.device) - ) - dresidual_in = ( - torch.empty_like(x) - if has_residual - and (dx.dtype != x.dtype or dropout_p > 0.0 or rowscale is not None or has_x1) - else None - ) - dx1 = torch.empty_like(dx) if (has_x1 and dropout_p > 0.0) else None - y = torch.empty(M, N, dtype=dy.dtype, device=dy.device) if recompute_output else None - if recompute_output: - assert weight1 is None, "recompute_output is not supported with parallel LayerNorm" - - # Less than 64KB per feature: enqueue fused kernel - MAX_FUSED_SIZE = 65536 // x.element_size() - BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) - if N > BLOCK_N: - raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") - # Increasing the multiple (e.g. 8) will allow more thread blocks to be launched and hide the - # latency of the gmem reads/writes, but will increase the time of summing up dw / db. - sm_count = torch.cuda.get_device_properties(x.device).multi_processor_count * 8 - _dw = torch.empty((sm_count, N), dtype=torch.float32, device=weight.device) - _db = ( - torch.empty((sm_count, N), dtype=torch.float32, device=bias.device) - if bias is not None - else None - ) - _dw1 = torch.empty_like(_dw) if weight1 is not None else None - _db1 = torch.empty_like(_db) if bias1 is not None else None - rows_per_program = math.ceil(M / sm_count) - grid = (sm_count,) - with torch.cuda.device(x.device.index): - torch.library.wrap_triton(_layer_norm_bwd_kernel)[grid]( - x, - weight, - bias, - y, - dy, - dx, - _dw, - _db, - dresidual, - weight1, - dy1, - dx1, - _dw1, - _db1, - dresidual_in, - rowscale, - seeds, - mean, - rstd, - x.stride(0), - 0 if not recompute_output else y.stride(0), - dy.stride(0), - dx.stride(0), - dresidual.stride(0) if dresidual is not None else 0, - dy1.stride(0) if dy1 is not None else 0, - dx1.stride(0) if dx1 is not None else 0, - dresidual_in.stride(0) if dresidual_in is not None else 0, - M, - N, - eps, - dropout_p, - # Passing bool make torch inductor very unhappy since it then tries to compare to int_max - int(zero_centered_weight), - rows_per_program, - is_rms_norm, - BLOCK_N, - dresidual is not None, - dresidual_in is not None, - bias is not None, - dropout_p > 0.0, - HAS_ROWSCALE=rowscale is not None, - HAS_DY1=dy1 is not None, - HAS_DX1=dx1 is not None, - HAS_B1=bias1 is not None, - RECOMPUTE_OUTPUT=y is not None, - ) - dw = _dw.sum(0).to(weight.dtype) - db = _db.sum(0).to(bias.dtype) if bias is not None else None - dw1 = _dw1.sum(0).to(weight1.dtype) if weight1 is not None else None - db1 = _db1.sum(0).to(bias1.dtype) if bias1 is not None else None - # dresidual_in and dx1 could be None, the wrapper will handle assigning them from dx - return dx, dw, db, dresidual_in, dx1, dw1, db1, y - - -class LayerNormFn(torch.autograd.Function): - - @staticmethod - def forward( - ctx, - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - residual_in_fp32=False, - zero_centered_weight=False, - is_rms_norm=False, - return_dropout_mask=False, - out_dtype=None, - out=None, - residual_out=None - ): - x_shape_og = x.shape - # reshape input data into 2D tensor - x = maybe_contiguous_lastdim(x.reshape(-1, x.shape[-1])) - if residual is not None: - assert residual.shape == x_shape_og - residual = maybe_contiguous_lastdim(residual.reshape(-1, residual.shape[-1])) - if x1 is not None: - assert x1.shape == x_shape_og - assert rowscale is None, "rowscale is not supported with parallel LayerNorm" - x1 = maybe_contiguous_lastdim(x1.reshape(-1, x1.shape[-1])) - weight = weight.contiguous() - bias = maybe_contiguous(bias) - weight1 = maybe_contiguous(weight1) - bias1 = maybe_contiguous(bias1) - if rowscale is not None: - rowscale = rowscale.reshape(-1).contiguous() - residual_dtype = ( - residual.dtype - if residual is not None - else (torch.float32 if residual_in_fp32 else None) - ) - if out is not None: - out = out.reshape(-1, out.shape[-1]) - if residual_out is not None: - residual_out = residual_out.reshape(-1, residual_out.shape[-1]) - y, y1, mean, rstd, residual_out, seeds, dropout_mask, dropout_mask1 = _layer_norm_fwd( - x, - weight, - bias, - eps, - residual, - x1, - weight1, - bias1, - dropout_p=dropout_p, - rowscale=rowscale, - out_dtype=out_dtype, - residual_dtype=residual_dtype, - zero_centered_weight=zero_centered_weight, - is_rms_norm=is_rms_norm, - return_dropout_mask=return_dropout_mask, - out=out, - residual_out=residual_out, - ) - ctx.save_for_backward( - residual_out, weight, bias, weight1, bias1, rowscale, seeds, mean, rstd - ) - ctx.x_shape_og = x_shape_og - ctx.eps = eps - ctx.dropout_p = dropout_p - ctx.is_rms_norm = is_rms_norm - ctx.has_residual = residual is not None - ctx.has_x1 = x1 is not None - ctx.prenorm = prenorm - ctx.x_dtype = x.dtype - ctx.zero_centered_weight = zero_centered_weight - y = y.reshape(x_shape_og) - y1 = y1.reshape(x_shape_og) if y1 is not None else None - residual_out = residual_out.reshape(x_shape_og) if residual_out is not None else None - dropout_mask = dropout_mask.reshape(x_shape_og) if dropout_mask is not None else None - dropout_mask1 = dropout_mask1.reshape(x_shape_og) if dropout_mask1 is not None else None - if not return_dropout_mask: - if weight1 is None: - return y if not prenorm else (y, residual_out) - else: - return (y, y1) if not prenorm else (y, y1, residual_out) - else: - if weight1 is None: - return ( - (y, dropout_mask, dropout_mask1) - if not prenorm - else (y, residual_out, dropout_mask, dropout_mask1) - ) - else: - return ( - (y, y1, dropout_mask, dropout_mask1) - if not prenorm - else (y, y1, residual_out, dropout_mask, dropout_mask1) - ) - - @staticmethod - def backward(ctx, dy, *args): - x, weight, bias, weight1, bias1, rowscale, seeds, mean, rstd = ctx.saved_tensors - dy = dy.reshape(-1, dy.shape[-1]) - if weight1 is not None: - dy1, args = args[0], args[1:] - dy1 = dy1.reshape(-1, dy1.shape[-1]) - assert dy1.shape == x.shape - else: - dy1 = None - if ctx.prenorm: - dresidual = args[0] - dresidual = dresidual.reshape(-1, dresidual.shape[-1]) - assert dresidual.shape == x.shape - else: - dresidual = None - dx, dw, db, dresidual_in, dx1, dw1, db1, _ = _layer_norm_bwd( - dy, - x, - weight, - bias, - ctx.eps, - mean, - rstd, - dresidual, - dy1, - weight1, - bias1, - seeds, - ctx.dropout_p, - rowscale, - ctx.has_residual, - ctx.has_x1, - ctx.zero_centered_weight, - ctx.is_rms_norm, - x_dtype=ctx.x_dtype, - recompute_output=False, - ) - return ( - dx.reshape(ctx.x_shape_og), - dw, - db, - dresidual_in.reshape(ctx.x_shape_og) if ctx.has_residual else None, - dx1.reshape(ctx.x_shape_og) if dx1 is not None else None, - dw1, - db1, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) - - -def layer_norm_fn( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - residual_in_fp32=False, - zero_centered_weight=False, - is_rms_norm=False, - return_dropout_mask=False, - out_dtype=None, - out=None, - residual_out=None -): - return LayerNormFn.apply( - x, - weight, - bias, - residual, - x1, - weight1, - bias1, - eps, - dropout_p, - rowscale, - prenorm, - residual_in_fp32, - zero_centered_weight, - is_rms_norm, - return_dropout_mask, - out_dtype, - out, - residual_out - ) - - -def rms_norm_fn( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - residual_in_fp32=False, - zero_centered_weight=False, - return_dropout_mask=False, - out_dtype=None, - out=None, - residual_out=None -): - return LayerNormFn.apply( - x, - weight, - bias, - residual, - x1, - weight1, - bias1, - eps, - dropout_p, - rowscale, - prenorm, - residual_in_fp32, - zero_centered_weight, - True, - return_dropout_mask, - out_dtype, - out, - residual_out - ) - - -class RMSNorm(torch.nn.Module): - - def __init__(self, hidden_size, eps=1e-5, dropout_p=0.0, zero_centered_weight=False, - device=None, dtype=None): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.eps = eps - if dropout_p > 0.0: - self.drop = torch.nn.Dropout(dropout_p) - else: - self.drop = None - self.zero_centered_weight = zero_centered_weight - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.register_parameter("bias", None) - self.reset_parameters() - - def reset_parameters(self): - if not self.zero_centered_weight: - torch.nn.init.ones_(self.weight) - else: - torch.nn.init.zeros_(self.weight) - - def forward(self, x, residual=None, prenorm=False, residual_in_fp32=False): - return rms_norm_fn( - x, - self.weight, - self.bias, - residual=residual, - eps=self.eps, - dropout_p=self.drop.p if self.drop is not None and self.training else 0.0, - prenorm=prenorm, - residual_in_fp32=residual_in_fp32, - zero_centered_weight=self.zero_centered_weight, - ) - - -class LayerNormLinearFn(torch.autograd.Function): - - @staticmethod - @custom_fwd - def forward( - ctx, - x, - norm_weight, - norm_bias, - linear_weight, - linear_bias, - residual=None, - eps=1e-6, - prenorm=False, - residual_in_fp32=False, - is_rms_norm=False, - ): - x_shape_og = x.shape - # reshape input data into 2D tensor - x = maybe_contiguous_lastdim(x.reshape(-1, x.shape[-1])) - if residual is not None: - assert residual.shape == x_shape_og - residual = maybe_contiguous_lastdim(residual.reshape(-1, residual.shape[-1])) - norm_weight = norm_weight.contiguous() - norm_bias = maybe_contiguous(norm_bias) - residual_dtype = ( - residual.dtype - if residual is not None - else (torch.float32 if residual_in_fp32 else None) - ) - y, _, mean, rstd, residual_out, *rest = _layer_norm_fwd( - x, - norm_weight, - norm_bias, - eps, - residual, - out_dtype=None if not torch.is_autocast_enabled() else torch.get_autocast_dtype("cuda"), - residual_dtype=residual_dtype, - is_rms_norm=is_rms_norm, - ) - y = y.reshape(x_shape_og) - dtype = torch.get_autocast_dtype("cuda") if torch.is_autocast_enabled() else y.dtype - linear_weight = linear_weight.to(dtype) - linear_bias = linear_bias.to(dtype) if linear_bias is not None else None - out = F.linear(y.to(linear_weight.dtype), linear_weight, linear_bias) - # We don't store y, will be recomputed in the backward pass to save memory - ctx.save_for_backward(residual_out, norm_weight, norm_bias, linear_weight, mean, rstd) - ctx.x_shape_og = x_shape_og - ctx.eps = eps - ctx.is_rms_norm = is_rms_norm - ctx.has_residual = residual is not None - ctx.prenorm = prenorm - ctx.x_dtype = x.dtype - ctx.linear_bias_is_none = linear_bias is None - return out if not prenorm else (out, residual_out.reshape(x_shape_og)) - - @staticmethod - @custom_bwd - def backward(ctx, dout, *args): - x, norm_weight, norm_bias, linear_weight, mean, rstd = ctx.saved_tensors - dout = dout.reshape(-1, dout.shape[-1]) - dy = F.linear(dout, linear_weight.t()) - dlinear_bias = None if ctx.linear_bias_is_none else dout.sum(0) - dy = maybe_contiguous_lastdim(dy) - assert dy.shape == x.shape - if ctx.prenorm: - dresidual = args[0] - dresidual = maybe_contiguous_lastdim(dresidual.reshape(-1, dresidual.shape[-1])) - assert dresidual.shape == x.shape - else: - dresidual = None - dx, dnorm_weight, dnorm_bias, dresidual_in, _, _, _, y = _layer_norm_bwd( - dy, - x, - norm_weight, - norm_bias, - ctx.eps, - mean, - rstd, - dresidual=dresidual, - has_residual=ctx.has_residual, - is_rms_norm=ctx.is_rms_norm, - x_dtype=ctx.x_dtype, - recompute_output=True, - ) - dlinear_weight = torch.einsum("bo,bi->oi", dout, y) - return ( - dx.reshape(ctx.x_shape_og), - dnorm_weight, - dnorm_bias, - dlinear_weight, - dlinear_bias, - dresidual_in.reshape(ctx.x_shape_og) if ctx.has_residual else None, - None, - None, - None, - None, - ) - - -def layer_norm_linear_fn( - x, - norm_weight, - norm_bias, - linear_weight, - linear_bias, - residual=None, - eps=1e-6, - prenorm=False, - residual_in_fp32=False, - is_rms_norm=False, -): - return LayerNormLinearFn.apply( - x, - norm_weight, - norm_bias, - linear_weight, - linear_bias, - residual, - eps, - prenorm, - residual_in_fp32, - is_rms_norm, - ) diff --git a/build/torch212-cxx11-cu132-aarch64-linux/ops/triton/linear.py b/build/torch212-cxx11-cu132-aarch64-linux/ops/triton/linear.py deleted file mode 100644 index a8966dbc345ab0e593df0124451ee7be3dae131a..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu132-aarch64-linux/ops/triton/linear.py +++ /dev/null @@ -1,594 +0,0 @@ -# Adapted from https://github.com/ELS-RD/kernl/blob/main/src/kernl/implementations/linear_layer.py -# and https://github.com/openai/triton/blob/master/python/triton/ops/matmul.py -from typing import Optional - -import torch -import triton -import triton.language as tl -from triton.ops.matmul_perf_model import early_config_prune, estimate_matmul_time - -from flash_attn.ops.triton.k_activations import ( - gelu, - gelu_approx, - gelu_approx_grad, - gelu_grad, - squared_relu, - squared_relu_grad, -) - -# CREDITS: Initially inspired by the Triton tutorial on matrix multiplications - - -def init_to_zero(name): - return lambda nargs: nargs[name].zero_() - - -def get_configs_io_bound(): - configs = [] - for num_stages in [2, 3, 4, 5, 6]: - for block_m in [16, 32]: - for block_k in [32, 64]: - for block_n in [32, 64, 128, 256]: - num_warps = 2 if block_n <= 64 else 4 - configs.append( - triton.Config( - { - "BLOCK_M": block_m, - "BLOCK_N": block_n, - "BLOCK_K": block_k, - "SPLIT_K": 1, - }, - num_stages=num_stages, - num_warps=num_warps, - ) - ) - # split_k not used - # for split_k in [2, 4, 8, 16]: - # configs.append(triton.Config( - # {'BLOCK_M': block_m, 'BLOCK_N': block_n, 'BLOCK_K': block_k, 'SPLIT_K': split_k}, - # num_stages=num_stages, num_warps=num_warps, pre_hook=init_to_zero('C'))) - return configs - - -@triton.autotune( - configs=[ - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - # good for int8 - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=4, - num_warps=4, - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - ] - + get_configs_io_bound(), - key=["CACHE_KEY_M", "CACHE_KEY_N", "CACHE_KEY_K"], - prune_configs_by={ - "early_config_prune": early_config_prune, - "perf_model": estimate_matmul_time, - "top_k": 10, - }, -) -@triton.heuristics( - { - "EVEN_K": lambda args: args["K"] % (args["BLOCK_K"] * args["SPLIT_K"]) == 0, - } -) -@triton.jit -def kernel_fwd( - C, # Pointers to matrices - ACT_INPUT, - A, - B, - bias, - # Matrix dimensions - M, - N, - K, - CACHE_KEY_M, - CACHE_KEY_N, - CACHE_KEY_K, - # The stride variables represent how much to increase the ptr by when moving by 1 - # element in a particular dimension. E.g. stride_am is how much to increase a_ptr - # by to get the element one row down (A has M rows) - stride_cm, - # stride_cn, # Assume that stride_cn == 1 - stride_am, - stride_ak, - stride_bn, - stride_bk, - # Meta-parameters - BLOCK_M: tl.constexpr, - GROUP_M: tl.constexpr, - BLOCK_N: tl.constexpr, - BLOCK_K: tl.constexpr, - # split k not used, not performant with activation, kept because early_config_prune is expecting it - SPLIT_K: tl.constexpr, - EVEN_K: tl.constexpr, - A_ROWMAJOR: tl.constexpr, - B_COLMAJOR: tl.constexpr, - BIAS: tl.constexpr, - SAVE_ACT_INPUT: tl.constexpr, - ACTIVATION: tl.constexpr, -): - - """ - Kernel for computing Out = activation(A x W + C) - - Input has shape (M, K) - - Weight has shape (K, N) - - Bias has shape (N,) - - Output has shape (M, N) - - ActInputs (optional) has shape (M, N) - 'ActInputs' optionally saves the A x W + C intermediate for backward computations - This kernel will consolidate over K - """ - - pid = tl.program_id(axis=0) - - grid_m = (M + BLOCK_M - 1) // BLOCK_M - grid_n = (N + BLOCK_N - 1) // BLOCK_N - # re-order program ID for better L2 performance - width = GROUP_M * grid_n - group_id = pid // width - group_size = min(grid_m - group_id * GROUP_M, GROUP_M) - pid_m = group_id * GROUP_M + (pid % group_size) - pid_n = (pid % width) // (group_size) - - # now compute the block that each program will go through - # rm (resp. rn) denotes a range of indices - # for rows (resp. col) of C - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - # trick to avoid masking on M and N axis - ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M) - rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N) - rk = tl.arange(0, BLOCK_K) - - if A_ROWMAJOR: - A = A + (ram[:, None] * stride_am + rk[None, :]) - else: - A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak) - if B_COLMAJOR: - B = B + (rk[:, None] + rbn[None, :] * stride_bn) - else: - B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn) - - acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) - - for k in range(K, 0, -BLOCK_K): - if EVEN_K: - a = tl.load(A) - b = tl.load(B) - else: - a = tl.load(A, mask=rk[None, :] < k, other=0.0) - b = tl.load(B, mask=rk[:, None] < k, other=0.0) - acc += tl.dot(a, b) - - if A_ROWMAJOR: - A += BLOCK_K - else: - A += BLOCK_K * stride_ak - if B_COLMAJOR: - B += BLOCK_K - else: - B += BLOCK_K * stride_bk - - # Putting bias after the matmul (instead of before) is faster, idk why - if BIAS: - bias = tl.load(bias + rn, mask=rn < N, other=0.0).to(tl.float32) - acc += bias[None, :] - - # optional: save the activation inputs - if SAVE_ACT_INPUT: - # act_in_ptrs = ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :] * stride_cn - act_in_ptrs = ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :] - tl.store(act_in_ptrs, acc) - - # optional: fused activation (while the data is in shared memory) - if ACTIVATION == "gelu": - acc = gelu(acc) - elif ACTIVATION == "gelu_approx": - acc = gelu_approx(acc) - elif ACTIVATION == "squared_relu": - acc = squared_relu(acc) - # rematerialize rm and rn to save registers - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - - # write back result - # C = C + rm[:, None] * stride_cm + rn[None, :] * stride_cn - C = C + rm[:, None] * stride_cm + rn[None, :] - mask = (rm < M)[:, None] & (rn < N)[None, :] - tl.store(C, acc) - - -def triton_linear_act( - x: torch.Tensor, - weight: torch.Tensor, - bias: Optional[torch.Tensor] = None, - activation: str = "id", - save_act_input: bool = False, -) -> torch.Tensor: - """ - Compute e = activation(x @ weight.T + bias). - This wrapper kicks the `kernel_fwd` Triton kernel - :param x: input tensor - :param weight: weight matrix - :param bias: an optional bias tensor - :param activation: Activation name. Needs to be a Triton kernel. - :param act_input: an optional tensor to save the activation inputs (for backward) - :return: result tensor - """ - # if torch.is_autocast_enabled(): - # dtype = torch.get_autocast_gpu_dtype() - # x, weight, bias = [a.to(dtype=dtype) for a in [x, weight, bias]] - - assert activation in ["id", "gelu", "gelu_approx", "squared_relu"] - - batch_shape, n = x.shape[:-1], x.shape[-1] - batch_dim = batch_shape.numel() - x_reshaped = x.reshape(batch_dim, n) - - if x_reshaped.stride(0) > 1 and x_reshaped.stride(1) > 1: - x_reshaped = x_reshaped.contiguous() - if weight.stride(0) > 1 and weight.stride(1) > 1: - weight = weight.contiguous() - bias = bias.contiguous() if bias is not None else None - - assert ( - x.dtype == weight.dtype - ), f"Input and weight must have the same dtype, got {x.dtype} and {weight.dtype}" - if bias is not None: - assert ( - x.dtype == bias.dtype - ), f"Input and bias must have the same dtype, got {x.dtype} and {bias.dtype}" - assert ( - x_reshaped.shape[1] == weight.shape[1] - ), f"Incompatible dimensions: {x_reshaped.shape} - {weight.shape}" - - assert ( - bias is None or bias.shape[0] == weight.shape[0] - ), "Incompatible dimensions in between weight and bias" - - M, K = x_reshaped.shape - N, K = weight.shape - - output = torch.empty((M, N), device=x.device, dtype=x.dtype) - act_input = torch.empty_like(output) if save_act_input else None - - # 1D launch kernel where each block gets its own program. - grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),) # noqa - - kernel_fwd[grid]( - output, - act_input, - x_reshaped, - weight, # data ptrs - bias if bias is not None else x, # auto skip bias if not present - M, # shapes - N, - K, - M // 32, # key for triton cache (limit number of compilations) - N // 32, - K // 32, - stride_cm=output.stride(0), # strides - # stride_cn=output.stride(1), - stride_am=x_reshaped.stride(0), - stride_ak=x_reshaped.stride(1), - stride_bk=weight.stride(1), - stride_bn=weight.stride(0), - BIAS=bias is not None, # optional fused bias - SAVE_ACT_INPUT=save_act_input, # optional save activation inputs - ACTIVATION=activation, # optional fused activation - A_ROWMAJOR=x_reshaped.stride(1) == 1, - B_COLMAJOR=weight.stride(1) == 1, - GROUP_M=8, # speed optimization: group the programs - ) - - if not save_act_input: - return output.reshape(*batch_shape, output.shape[-1]) - else: - return ( - output.reshape(*batch_shape, output.shape[-1]), - act_input.reshape(*batch_shape, act_input.shape[-1]), - ) - - -@triton.autotune( - configs=[ - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - # good for int8 - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=4, - num_warps=4, - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - ] - + get_configs_io_bound(), - key=["CACHE_KEY_M", "CACHE_KEY_N", "CACHE_KEY_K"], - prune_configs_by={ - "early_config_prune": early_config_prune, - "perf_model": estimate_matmul_time, - "top_k": 10, - }, -) -@triton.heuristics( - { - "EVEN_K": lambda args: args["K"] % (args["BLOCK_K"] * args["SPLIT_K"]) == 0, - } -) -@triton.jit -def kernel_bwd( - C, # Pointers to matrices - ACT_INPUT, - A, - B, - # Matrix dimensions - M, - N, - K, - CACHE_KEY_M, - CACHE_KEY_N, - CACHE_KEY_K, - # The stride variables represent how much to increase the ptr by when moving by 1 - # element in a particular dimension. E.g. stride_am is how much to increase a_ptr - # by to get the element one row down (A has M rows) - stride_cm, - # stride_cn, # Assume that stride_cn == 1 - stride_am, - stride_ak, - stride_bk, - stride_bn, - # Meta-parameters - BLOCK_M: tl.constexpr, - GROUP_M: tl.constexpr, - BLOCK_N: tl.constexpr, - BLOCK_K: tl.constexpr, - # split k not used, not performant with activation, kept because early_config_prune is expecting it - SPLIT_K: tl.constexpr, - EVEN_K: tl.constexpr, - ACTIVATION: tl.constexpr, -): - - """ - Kernel for computing Out = activation(A x W + C) - - Input has shape (M, K) - - Weight has shape (K, N) - - Output has shape (M, N) - - ActInputs (optional) has shape (M, N) - 'ActInputs' optionally saves the A x W + C intermediate for backward computations - This kernel will consolidate over K - """ - - pid = tl.program_id(axis=0) - - grid_m = (M + BLOCK_M - 1) // BLOCK_M - grid_n = (N + BLOCK_N - 1) // BLOCK_N - # re-order program ID for better L2 performance - width = GROUP_M * grid_n - group_id = pid // width - group_size = min(grid_m - group_id * GROUP_M, GROUP_M) - pid_m = group_id * GROUP_M + (pid % group_size) - pid_n = (pid % width) // (group_size) - - # now compute the block that each program will go through - # rm (resp. rn) denotes a range of indices - # for rows (resp. col) of C - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - # trick to avoid masking on M and N axis - ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M) - rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N) - rk = tl.arange(0, BLOCK_K) - - A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak) - B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn) - - acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) - - for k in range(K, 0, -BLOCK_K): - if EVEN_K: - a = tl.load(A) - b = tl.load(B) - else: - a = tl.load(A, mask=rk[None, :] < k, other=0.0) - b = tl.load(B, mask=rk[:, None] < k, other=0.0) - acc += tl.dot(a, b) - - A += BLOCK_K * stride_ak - B += BLOCK_K * stride_bk - - # optional: fused activation (while the data is in shared memory) - if ACTIVATION != "id": - act_in_ptrs = ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :] - act_input = tl.load(act_in_ptrs).to(acc.dtype) - if ACTIVATION == "gelu": - acc *= gelu_grad(act_input) - elif ACTIVATION == "gelu_approx": - acc *= gelu_approx_grad(act_input) - elif ACTIVATION == "squared_relu": - acc *= squared_relu_grad(act_input) - - # rematerialize rm and rn to save registers - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - - # write back result - C = C + rm[:, None] * stride_cm + rn[None, :] - mask = (rm < M)[:, None] & (rn < N)[None, :] - tl.store(C, acc, mask=mask) - - -def triton_dgrad_act( - grad_output: torch.Tensor, - weight: torch.Tensor, - activation: str = "id", - act_input: Optional[torch.Tensor] = None, -) -> torch.Tensor: - """ - Compute e = activation(grad_output @ weight + bias). - This wrapper kicks the `kernel_fwd` Triton kernel - :param grad_output: input tensor - :param weight: weight matrix - :param activation: Activation name. Needs to be a Triton kernel. - :param act_input: an optional tensor to save the activation inputs (for backward) - :return: result tensor - """ - assert activation in ["id", "gelu", "gelu_approx", "squared_relu"] - - batch_shape, n = grad_output.shape[:-1], grad_output.shape[-1] - batch_dim = batch_shape.numel() - grad_output_reshaped = grad_output.reshape(batch_dim, n) - - if grad_output_reshaped.stride(0) > 1 and grad_output_reshaped.stride(1) > 1: - grad_output_reshaped = grad_output_reshaped.contiguous() - if weight.stride(0) > 1 and weight.stride(1) > 1: - weight = weight.contiguous() - - assert ( - grad_output.dtype == weight.dtype - ), f"grad_output and weight must have the same dtype, got {grad_output.dtype} and {weight.dtype}" - assert ( - grad_output_reshaped.shape[1] == weight.shape[0] - ), f"Incompatible dimensions: {grad_output_reshaped.shape} - {weight.shape}" - if activation != "id": - assert act_input is not None, f"act_input is required for activation {activation}" - - # M, N, K in bwd are different from M, N, K in fwd - M, K = grad_output_reshaped.shape - K, N = weight.shape - - grad_input = torch.empty((M, N), device=grad_output.device, dtype=grad_output.dtype) - - # 1D launch kernel where each block gets its own program. - grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),) # noqa - - kernel_bwd[grid]( - grad_input, - act_input, - grad_output_reshaped, - weight, # data ptrs - M, # shapes - N, - K, - M // 32, # key for triton cache (limit number of compilations) - N // 32, - K // 32, - stride_cm=grad_input.stride(0), # strides - # stride_cn=grad_input.stride(1), - stride_am=grad_output_reshaped.stride(0), - stride_ak=grad_output_reshaped.stride(1), - stride_bk=weight.stride(0), - stride_bn=weight.stride(1), - ACTIVATION=activation, # optional fused activation - GROUP_M=8, # speed optimization: group the programs - ) - - return grad_input.reshape(*batch_shape, grad_input.shape[-1]) diff --git a/build/torch212-cxx11-cu132-aarch64-linux/ops/triton/mlp.py b/build/torch212-cxx11-cu132-aarch64-linux/ops/triton/mlp.py deleted file mode 100644 index 059f4f8a5e174c1f4824e43d313fca18eaa799b8..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu132-aarch64-linux/ops/triton/mlp.py +++ /dev/null @@ -1,149 +0,0 @@ -# The triton fused matmul + sqrelu is faster for fp16 but slower for bf16, compared -# to naive implementation. -import fused_dense_lib as fused_dense_cuda -import torch -import torch.nn as nn -import torch.nn.functional as F - -from flash_attn.utils.torch import custom_fwd, custom_bwd -from flash_attn.ops.activations import sqrelu_bwd, sqrelu_fwd -from flash_attn.ops.triton.linear import triton_dgrad_act, triton_linear_act - - -class FusedDenseSqreluDenseFunc(torch.autograd.Function): - @staticmethod - @custom_fwd - def forward(ctx, x, weight1, bias1, weight2, bias2, checkpoint_lvl=0): - """checkpoint_lvl: - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute act_input and gelu_out in the bwd - """ - if torch.is_autocast_enabled(): - dtype = torch.get_autocast_gpu_dtype() - x, weight1, bias1, weight2, bias2 = [ - a.to(dtype=dtype) for a in [x, weight1, bias1, weight2, bias2] - ] - is_bf16 = x.dtype == torch.bfloat16 - assert checkpoint_lvl in [0, 1, 2] - x = x.contiguous() - weight1 = weight1.contiguous() - bias1 = bias1.contiguous() - weight2 = weight2.contiguous() - bias2 = bias2.contiguous() - batch_shape, n = x.shape[:-1], x.shape[-1] - batch_dim = batch_shape.numel() - if is_bf16: - act_input = fused_dense_cuda.linear_bias_forward( - x.reshape(batch_dim, n), weight1, bias1 - ) - output1 = sqrelu_fwd(act_input) - else: - save_act_input = checkpoint_lvl != 2 - result = triton_linear_act( - x.reshape(batch_dim, n), - weight1, - bias1, - activation="squared_relu", - save_act_input=save_act_input, - ) - if save_act_input: - output1, act_input = result - else: - output1 = result - output2 = fused_dense_cuda.linear_bias_forward(output1, weight2, bias2) - ctx.checkpoint_lvl = checkpoint_lvl - if checkpoint_lvl == 0: - ctx.save_for_backward(x, weight1, bias1, weight2, act_input, output1) - elif checkpoint_lvl == 1: - ctx.save_for_backward(x, weight1, bias1, weight2, act_input) - elif checkpoint_lvl == 2: - ctx.save_for_backward(x, weight1, bias1, weight2) - return output2.reshape(*batch_shape, output2.shape[-1]) - - @staticmethod - @custom_bwd - def backward(ctx, grad_output): - grad_output = grad_output.contiguous() - checkpoint_lvl = ctx.checkpoint_lvl - x, weight1, bias1, weight2, *rest = ctx.saved_tensors - batch_shape, n = x.shape[:-1], x.shape[-1] - batch_dim = batch_shape.numel() - is_bf16 = x.dtype == torch.bfloat16 - if checkpoint_lvl == 0: - act_input, output1 = rest - elif checkpoint_lvl == 1: - (act_input,) = rest - output1 = sqrelu_fwd(act_input) - elif checkpoint_lvl == 2: - if is_bf16: - act_input = fused_dense_cuda.linear_bias_forward( - x.reshape(batch_dim, n), weight1, bias1 - ) - output1 = sqrelu_fwd(act_input) - else: - output1, act_input = triton_linear_act( - x.reshape(batch_dim, n), - weight1, - bias1, - activation="squared_relu", - save_act_input=True, - ) - - if is_bf16: - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad(output1, grad_output) - grad_output1 = grad_output @ weight2 - grad_act_input = sqrelu_bwd(grad_output1, act_input) - grad_input, grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_backward( - x.reshape(batch_dim, n), weight1, grad_act_input - ) - else: - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad(output1, grad_output) - grad_act_input = triton_dgrad_act( - grad_output, weight2, activation="squared_relu", act_input=act_input - ) - grad_input, grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_backward( - x.reshape(batch_dim, n), weight1, grad_act_input - ) - return grad_input.reshape_as(x), grad_weight1, grad_bias1, grad_weight2, grad_bias2, None - - -fused_dense_sqrelu_dense_function = FusedDenseSqreluDenseFunc.apply - - -class FusedDenseSqreluDense(nn.Module): - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - bias1=True, - bias2=True, - checkpoint_lvl=0, - device=None, - dtype=None, - ): - """ - checkpoint_lvl (increasing lvl means slower but more memory saving): - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute gelu_in and gelu_out in the bwd - """ - assert checkpoint_lvl in [0, 1, 2] - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features * 4 - assert bias1 == True, "DenseSqreluDense module without bias is currently not supported" - assert bias2 == True, "DenseSqreluDense module without bias is currently not supported" - self.checkpoint_lvl = checkpoint_lvl - self.fc1 = nn.Linear(in_features, hidden_features, bias=bias1, **factory_kwargs) - self.fc2 = nn.Linear(hidden_features, out_features, bias=bias2, **factory_kwargs) - - def forward(self, x): - assert x.is_cuda - return fused_dense_sqrelu_dense_function( - x, self.fc1.weight, self.fc1.bias, self.fc2.weight, self.fc2.bias, self.checkpoint_lvl - ) diff --git a/build/torch212-cxx11-cu132-aarch64-linux/ops/triton/rotary.py b/build/torch212-cxx11-cu132-aarch64-linux/ops/triton/rotary.py deleted file mode 100644 index 6ad3ab412a084270120167ad32b67ad1720f2c37..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu132-aarch64-linux/ops/triton/rotary.py +++ /dev/null @@ -1,186 +0,0 @@ -# Copyright (c) 2025, Tri Dao. -# As of 2025-04-23, we require triton >= 3.0 - -from typing import Optional, Union - -import torch - -import triton -import triton.language as tl - - -@triton.jit -def rotary_kernel( - OUT, # Pointers to matrices - X, - COS, - SIN, - CU_SEQLENS, - SEQLEN_OFFSETS, # this could be int or a pointer - # Matrix dimensions - seqlen, - nheads, - seqlen_ro, - # strides - stride_out_batch, - stride_out_seqlen, - stride_out_nheads, - stride_out_headdim, - stride_x_batch, - stride_x_seqlen, - stride_x_nheads, - stride_x_headdim, - # Meta-parameters - # We want ROTARY_DIM to be constexpr, otherwise the triton compiler doesn't know that - # the mask is constant every 8 elements, and it will generate LDG.16 instead of LDG.128 - ROTARY_DIM: tl.constexpr, - IS_SEQLEN_OFFSETS_TENSOR: tl.constexpr, - IS_VARLEN: tl.constexpr, - INTERLEAVED: tl.constexpr, - CONJUGATE: tl.constexpr, - BLOCK_H: tl.constexpr, - BLOCK_M: tl.constexpr, -): - BLOCK_K: tl.constexpr = triton.next_power_of_2(ROTARY_DIM) - ROTARY_DIM_HALF = ROTARY_DIM // 2 - pid_head = tl.program_id(axis=0) - pid_m = tl.program_id(axis=1) - pid_batch = tl.program_id(axis=2) - - if not IS_VARLEN: - X = X + pid_batch * stride_x_batch - OUT = OUT + pid_batch * stride_out_batch - else: - start_idx = tl.load(CU_SEQLENS + pid_batch) - seqlen = tl.load(CU_SEQLENS + pid_batch + 1) - start_idx - X = X + start_idx * stride_x_seqlen - OUT = OUT + start_idx * stride_out_seqlen - - if pid_m * BLOCK_M >= seqlen: - return - - rh = pid_head * BLOCK_H + tl.arange(0, BLOCK_H) - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - if not IS_SEQLEN_OFFSETS_TENSOR: - rm_cs = rm + SEQLEN_OFFSETS - else: - rm_cs = rm + tl.load(SEQLEN_OFFSETS + pid_batch) - - rk_half = tl.arange(0, BLOCK_K // 2) - COS = COS + (rm_cs[:, None] * ROTARY_DIM_HALF + rk_half[None, :]) - SIN = SIN + (rm_cs[:, None] * ROTARY_DIM_HALF + rk_half[None, :]) - mask_cs = (rm_cs[:, None] < seqlen_ro) & (rk_half[None, :] < ROTARY_DIM_HALF) - cos = tl.load(COS, mask=mask_cs, other=1.0).to(tl.float32) - sin = tl.load(SIN, mask=mask_cs, other=0.0).to(tl.float32) - if CONJUGATE: - sin = -sin - - if not INTERLEAVED: - # Load the 1st and 2nd halves of X, do calculation, then store to 1st and 2nd halves of OUT - X = X + (rh[:, None, None] * stride_x_nheads + rm[None, :, None] * stride_x_seqlen + rk_half[None, None, :] * stride_x_headdim) - OUT = OUT + (rh[:, None, None] * stride_out_nheads + rm[None, :, None] * stride_out_seqlen + rk_half[None, None, :] * stride_out_headdim) - mask = (rh[:, None, None] < nheads) & (rm[None, :, None] < seqlen) & (rk_half[None, None, :] < ROTARY_DIM_HALF) - x0 = tl.load(X, mask=mask, other=0.0).to(tl.float32) - x1 = tl.load(X + ROTARY_DIM_HALF * stride_x_headdim, mask=mask, other=0.0,).to(tl.float32) - o0 = x0 * cos - x1 * sin - o1 = x0 * sin + x1 * cos - tl.store(OUT, o0, mask=mask) - tl.store(OUT + ROTARY_DIM_HALF * stride_out_headdim, o1, mask=mask) - else: - rk = tl.arange(0, BLOCK_K) - X = X + (rh[:, None, None] * stride_x_nheads + rm[None, :, None] * stride_x_seqlen + rk[None, None, :] * stride_x_headdim) - OUT = OUT + (rh[:, None, None] * stride_out_nheads + rm[None, :, None] * stride_out_seqlen + rk[None, None, :] * stride_out_headdim) - mask = (rh[:, None, None] < nheads) & (rm[None, :, None] < seqlen) & (rk[None, None, :] < ROTARY_DIM) - x = tl.load(X, mask=mask, other=0.0).to(tl.float32) - x0, x1 = tl.split(tl.reshape(x, [BLOCK_H, BLOCK_M, BLOCK_K // 2, 2])) - o0 = x0 * cos - x1 * sin - o1 = x0 * sin + x1 * cos - o = tl.reshape(tl.join(o0, o1), [BLOCK_H, BLOCK_M, BLOCK_K]) - tl.store(OUT, o, mask=mask) - - -def apply_rotary( - x: torch.Tensor, - cos: torch.Tensor, - sin: torch.Tensor, - seqlen_offsets: Union[int, torch.Tensor] = 0, - cu_seqlens: Optional[torch.Tensor] = None, - max_seqlen: Optional[int] = None, - interleaved=False, - inplace=False, - conjugate=False, -) -> torch.Tensor: - """ - Arguments: - x: (batch, seqlen, nheads, headdim) if cu_seqlens is None - else (total_seqlen, nheads, headdim). - cos: (seqlen_ro, rotary_dim / 2) - sin: (seqlen_ro, rotary_dim / 2) - seqlen_offsets: integer or integer tensor of size (batch,) - cu_seqlens: (batch + 1,) or None - max_seqlen: int - Returns: - y: (batch, seqlen, nheads, headdim) - """ - is_varlen = cu_seqlens is not None - if not is_varlen: - batch, seqlen, nheads, headdim = x.shape - else: - assert max_seqlen is not None, "If cu_seqlens is passed in, then max_seqlen must be passed" - total_seqlen, nheads, headdim = x.shape - batch_p_1 = cu_seqlens.shape[0] - batch = batch_p_1 - 1 - seqlen = max_seqlen - seqlen_ro, rotary_dim = cos.shape - assert sin.shape == cos.shape - rotary_dim *= 2 - assert rotary_dim <= headdim, "rotary_dim must be <= headdim" - assert headdim <= 256, "Only support headdim <= 256" - assert seqlen_ro >= seqlen, "seqlen_ro must be >= seqlen" - - cos, sin = cos.contiguous(), sin.contiguous() - if isinstance(seqlen_offsets, torch.Tensor): - assert seqlen_offsets.shape == (batch,) - assert seqlen_offsets.dtype in [torch.int32, torch.int64] - seqlen_offsets = seqlen_offsets.contiguous() - else: - assert seqlen_offsets + seqlen <= seqlen_ro - - output = torch.empty_like(x) if not inplace else x - if rotary_dim < headdim and not inplace: - output[..., rotary_dim:].copy_(x[..., rotary_dim:]) - - grid = lambda META: (triton.cdiv(nheads, META["BLOCK_H"]), triton.cdiv(seqlen, META["BLOCK_M"]), batch) # noqa - BLOCK_M = 8 if rotary_dim <= 128 else 4 - - # Need this, otherwise Triton tries to launch from cuda:0 and we get - # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?) - device_ctx = torch.cuda.device(x.device.index) if x.device.type == 'cuda' else torch.xpu.device(x.device.index) - with device_ctx: - torch.library.wrap_triton(rotary_kernel)[grid]( - output, # data ptrs - x, - cos, - sin, - cu_seqlens, - seqlen_offsets, - seqlen, # shapes - nheads, - seqlen_ro, - output.stride(0) if not is_varlen else 0, # batch_strides if not varlen else 0 - output.stride(-3), # seqlen_stride or total_seqlen_stride - output.stride(-2), # nheads_stride - output.stride(-1), # headdim_stride - x.stride(0) if not is_varlen else 0, # batch_strides if not varlen else 0 - x.stride(-3), # seqlen stride or total_seqlen_stride - x.stride(-2), # nheads stride - x.stride(-1), # headdim stride - rotary_dim, - isinstance(seqlen_offsets, torch.Tensor), - is_varlen, - interleaved, - conjugate, - BLOCK_M=BLOCK_M, - BLOCK_H=2, - ) - return output diff --git a/build/torch212-cxx11-cu132-x86_64-linux/__init__.py b/build/torch212-cxx11-cu132-x86_64-linux/__init__.py deleted file mode 100644 index ecc2f9d896b6c93f90b0a1499856dc0612177422..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu132-x86_64-linux/__init__.py +++ /dev/null @@ -1,393 +0,0 @@ -from typing import Optional, List -import torch -from ._ops import ops as flash_attn_ops -from .flash_attn_interface import ( - flash_attn_func, - flash_attn_kvpacked_func, - flash_attn_qkvpacked_func, - flash_attn_varlen_func, - flash_attn_varlen_kvpacked_func, - flash_attn_varlen_qkvpacked_func, - flash_attn_with_kvcache, -) - - -def fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - return_softmax: bool = False, - gen: Optional[torch.Generator] = None, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention. - - Args: - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Optional output tensor, same shape as q - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - return_softmax: Whether to return softmax weights - gen: Optional random number generator - - Returns: - List of tensors: [output, softmax_lse, (softmax if return_softmax)] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.fwd( - q, - k, - v, - out, - alibi_slopes, - p_dropout, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - gen, - ) - - -def varlen_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - out: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - block_table: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - max_seqlen_q: int = 0, - max_seqlen_k: int = 0, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - zero_tensors: bool = False, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - return_softmax: bool = False, - gen: Optional[torch.Generator] = None, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention with variable sequence lengths. - - Args: - q: Query tensor of shape [total_q, num_heads, head_size] - k: Key tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - v: Value tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1] - cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1] - out: Optional output tensor of shape [total_q, num_heads, head_size] - seqused_k: Optional tensor specifying how many keys to use per batch element [batch_size] - leftpad_k: Optional left padding for keys of shape [batch_size] - block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq] - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - max_seqlen_q: Maximum sequence length for queries - max_seqlen_k: Maximum sequence length for keys - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - zero_tensors: Whether to zero tensors before computation - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - return_softmax: Whether to return softmax weights - gen: Optional random number generator - - Returns: - List of tensors: [output, softmax_lse, (softmax if return_softmax)] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.varlen_fwd( - q, - k, - v, - out, - cu_seqlens_q, - cu_seqlens_k, - seqused_k, - leftpad_k, - block_table, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - p_dropout, - softmax_scale, - zero_tensors, - is_causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - gen, - ) - - -def bwd( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - p_dropout: float = 0.0, - 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, - gen: Optional[torch.Generator] = None, - rng_state: Optional[torch.Tensor] = None, -) -> List[torch.Tensor]: - """ - Backward pass for multi-head attention. - - Args: - dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size] - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size] - softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q] - dq: Optional gradient tensor for queries, same shape as q - dk: Optional gradient tensor for keys, same shape as k - dv: Optional gradient tensor for values, same shape as v - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - deterministic: Whether to use deterministic algorithms - gen: Optional random number generator - rng_state: Optional RNG state from forward pass - - Returns: - List of tensors: [dq, dk, dv] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - alibi_slopes, - p_dropout, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - deterministic, - gen, - rng_state, - ) - - -def varlen_bwd( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - max_seqlen_q: int = 0, - max_seqlen_k: int = 0, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - zero_tensors: bool = False, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - deterministic: bool = False, - gen: Optional[torch.Generator] = None, - rng_state: Optional[torch.Tensor] = None, -) -> List[torch.Tensor]: - """ - Backward pass for multi-head attention with variable sequence lengths. - - Args: - dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size] - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size] - softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q] - cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1] - cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1] - dq: Optional gradient tensor for queries, same shape as q - dk: Optional gradient tensor for keys, same shape as k - dv: Optional gradient tensor for values, same shape as v - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - max_seqlen_q: Maximum sequence length for queries - max_seqlen_k: Maximum sequence length for keys - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - zero_tensors: Whether to zero tensors before computation - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - deterministic: Whether to use deterministic algorithms - gen: Optional random number generator - rng_state: Optional RNG state from forward pass - - Returns: - List of tensors: [dq, dk, dv] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.varlen_bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - p_dropout, - softmax_scale, - zero_tensors, - is_causal, - window_size_left, - window_size_right, - softcap, - deterministic, - gen, - rng_state, - ) - - -def fwd_kvcache( - q: torch.Tensor, - kcache: torch.Tensor, - vcache: torch.Tensor, - k: Optional[torch.Tensor] = None, - v: Optional[torch.Tensor] = None, - seqlens_k: Optional[torch.Tensor] = None, - rotary_cos: Optional[torch.Tensor] = None, - rotary_sin: Optional[torch.Tensor] = None, - cache_batch_idx: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - block_table: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - out: 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, - is_rotary_interleaved: bool = False, - num_splits: int = 1, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention with KV cache. - - Args: - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - kcache: Key cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - vcache: Value cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - k: Optional new keys tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size] - v: Optional new values tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size] - seqlens_k: Optional sequence lengths for keys of shape [batch_size] - rotary_cos: Optional rotary cosine tensor of shape [seqlen_ro, rotary_dim/2] - rotary_sin: Optional rotary sine tensor of shape [seqlen_ro, rotary_dim/2] - cache_batch_idx: Optional indices to index into the KV cache - leftpad_k: Optional left padding for keys of shape [batch_size] - block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq] - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - out: Optional output tensor, same shape as q - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - is_rotary_interleaved: Whether rotary embeddings are interleaved - num_splits: Number of splits for computation - - Returns: - List of tensors: [output, softmax_lse] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.fwd_kvcache( - q, - kcache, - vcache, - k, - v, - seqlens_k, - rotary_cos, - rotary_sin, - cache_batch_idx, - leftpad_k, - block_table, - alibi_slopes, - out, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - is_rotary_interleaved, - num_splits, - ) diff --git a/build/torch212-cxx11-cu132-x86_64-linux/_flash_attn2_cuda_f12afc9.abi3.so b/build/torch212-cxx11-cu132-x86_64-linux/_flash_attn2_cuda_f12afc9.abi3.so deleted file mode 100644 index af175d84da6b7c1c31ba99e7c948bb8e6ce14aae..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu132-x86_64-linux/_flash_attn2_cuda_f12afc9.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:1f178869049bf5a8b87b4fb367cd8bc9113d834a537688d55121d2ab53067610 -size 1024545064 diff --git a/build/torch212-cxx11-cu132-x86_64-linux/_ops.py b/build/torch212-cxx11-cu132-x86_64-linux/_ops.py deleted file mode 100644 index af5b731cb04b5e6ce54e862db0501f3683c3163f..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu132-x86_64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _flash_attn2_cuda_f12afc9 -ops = torch.ops._flash_attn2_cuda_f12afc9 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_flash_attn2_cuda_f12afc9::{op_name}" diff --git a/build/torch212-cxx11-cu132-x86_64-linux/bert_padding.py b/build/torch212-cxx11-cu132-x86_64-linux/bert_padding.py deleted file mode 100644 index 3c2d35159a014a9d03aabead9e52e009168696ea..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu132-x86_64-linux/bert_padding.py +++ /dev/null @@ -1,218 +0,0 @@ -# Adapted from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/padding.py - -import torch -import torch.nn.functional as F -from einops import rearrange, repeat - - -class IndexFirstAxis(torch.autograd.Function): - @staticmethod - def forward(ctx, input, indices): - ctx.save_for_backward(indices) - assert input.ndim >= 2 - ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] - second_dim = other_shape.numel() - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - # return input[indices] - return torch.gather( - rearrange(input, "b ... -> b (...)"), 0, repeat(indices, "z -> z d", d=second_dim) - ).reshape(-1, *other_shape) - - @staticmethod - def backward(ctx, grad_output): - (indices,) = ctx.saved_tensors - assert grad_output.ndim >= 2 - other_shape = grad_output.shape[1:] - grad_output = rearrange(grad_output, "b ... -> b (...)") - grad_input = torch.zeros( - [ctx.first_axis_dim, grad_output.shape[1]], - device=grad_output.device, - dtype=grad_output.dtype, - ) - # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. - # grad_input[indices] = grad_output - grad_input.scatter_(0, repeat(indices, "z -> z d", d=grad_output.shape[1]), grad_output) - return grad_input.reshape(ctx.first_axis_dim, *other_shape), None - - -index_first_axis = IndexFirstAxis.apply - - -class IndexPutFirstAxis(torch.autograd.Function): - @staticmethod - def forward(ctx, values, indices, first_axis_dim): - ctx.save_for_backward(indices) - assert indices.ndim == 1 - assert values.ndim >= 2 - output = torch.zeros( - first_axis_dim, *values.shape[1:], device=values.device, dtype=values.dtype - ) - # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. - output[indices] = values - # output.scatter_(0, repeat(indices, 'z -> z d', d=values.shape[1]), values) - return output - - @staticmethod - def backward(ctx, grad_output): - (indices,) = ctx.saved_tensors - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - grad_values = grad_output[indices] - # grad_values = torch.gather(grad_output, 0, repeat(indices, 'z -> z d', d=grad_output.shape[1])) - return grad_values, None, None - - -index_put_first_axis = IndexPutFirstAxis.apply - - -class IndexFirstAxisResidual(torch.autograd.Function): - @staticmethod - def forward(ctx, input, indices): - ctx.save_for_backward(indices) - assert input.ndim >= 2 - ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] - second_dim = other_shape.numel() - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - output = input[indices] - # We don't want to reshape input (b ... -> b (...)) since it could change the channel_last - # memory format to channel_first. In other words, input might not be contiguous. - # If we don't detach, Pytorch complains about output being a view and is being modified inplace - return output, input.detach() - - @staticmethod - def backward(ctx, grad_output, grad_residual): - (indices,) = ctx.saved_tensors - assert grad_output.ndim >= 2 - other_shape = grad_output.shape[1:] - assert grad_residual.shape[1:] == other_shape - grad_input = grad_residual - # grad_input[indices] += grad_output - indices = indices.reshape(indices.shape[0], *((1,) * (grad_output.ndim - 1))) - indices = indices.expand_as(grad_output) - grad_input.scatter_add_(0, indices, grad_output) - return grad_input.reshape(ctx.first_axis_dim, *other_shape), None - - -index_first_axis_residual = IndexFirstAxisResidual.apply - - -def unpad_input(hidden_states, attention_mask, unused_mask=None): - """ - Arguments: - hidden_states: (batch, seqlen, ...) - attention_mask: (batch, seqlen), bool / int, 1 means valid and 0 means not valid. - unused_mask: (batch, seqlen), bool / int, 1 means the element is allocated but unused. - Return: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens selected in attention_mask + unused_mask. - indices: (total_nnz), the indices of masked tokens from the flattened input sequence. - cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. - max_seqlen_in_batch: int - seqused: (batch), returns the number of tokens selected in attention_mask + unused_mask. - """ - all_masks = (attention_mask + unused_mask) if unused_mask is not None else attention_mask - seqlens_in_batch = all_masks.sum(dim=-1, dtype=torch.int32) - used_seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) - indices = torch.nonzero(all_masks.flatten(), as_tuple=False).flatten() - max_seqlen_in_batch = seqlens_in_batch.max().item() - cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) - # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the - # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim - # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to - # index with integer indices. Moreover, torch's index is a bit slower than it needs to be, - # so we write custom forward and backward to make it a bit faster. - return ( - index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices), - indices, - cu_seqlens, - max_seqlen_in_batch, - used_seqlens_in_batch, - ) - - -def unpad_input_for_concatenated_sequences(hidden_states, attention_mask_in_length): - """ - Supports concatenating short samples in one sequence. The attention_mask_in_length is utilized to mask other short samples. It helps efficient training of variant lengths-based samples (e.g., the supervised fine-tuning task in large language model). - The motivation for this function is explained [here](https://github.com/Dao-AILab/flash-attention/issues/432#issuecomment-1668822286). - - For example, if batch = 3 and seqlen = 6, the attention_mask_in_length is: - ``` - [ - [2, 3, 0, 0, 0, 0], - [3, 2, 0, 0, 0, 0], - [6, 0, 0, 0, 0, 0] - ] - ``` - , which refers to the 3D-attention mask: - ``` - [ - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [0, 0, 1, 0, 0, 0], - [0, 0, 1, 1, 0, 0], - [0, 0, 1, 1, 1, 0], - [0, 0, 0, 0, 0, 1] - ], - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [0, 0, 0, 1, 0, 0], - [0, 0, 0, 1, 1, 0], - [0, 0, 0, 0, 0, 1] - ], - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [1, 1, 1, 1, 0, 0], - [1, 1, 1, 1, 1, 0], - [1, 1, 1, 1, 1, 1] - ] - ] - ```. - - Arguments: - hidden_states: (batch, seqlen, ...) - attention_mask_in_length: (batch, seqlen), int, a nonzero number (e.g., 1, 2, 3, etc.) means length of concatenated sequence in b-th batch, and 0 means none. - Return: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. - indices: (total_nnz), the indices of non-masked tokens from the flattened input sequence. - cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. - max_seqlen_in_batch: int - """ - length = attention_mask_in_length.sum(dim=-1) - seqlen = attention_mask_in_length.size(-1) - attention_mask_2d = torch.arange(seqlen, device=length.device, dtype=length.dtype).expand(len(length), seqlen) < length.unsqueeze(1) - real_indices_idx = torch.nonzero(attention_mask_in_length.flatten(), as_tuple=False).flatten() - seqlens_in_batch = attention_mask_in_length.flatten()[real_indices_idx] - indices = torch.nonzero(attention_mask_2d.flatten(), as_tuple=False).flatten() - max_seqlen_in_batch = seqlens_in_batch.max().item() - cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) - # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the - # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim - # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to - # index with integer indices. Moreover, torch's index is a bit slower than it needs to be, - # so we write custom forward and backward to make it a bit faster. - return ( - index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices), - indices, - cu_seqlens, - max_seqlen_in_batch, - ) - - -def pad_input(hidden_states, indices, batch, seqlen): - """ - Arguments: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. - indices: (total_nnz), the indices that represent the non-masked tokens of the original padded input sequence. - batch: int, batch size for the padded sequence. - seqlen: int, maximum sequence length for the padded sequence. - Return: - hidden_states: (batch, seqlen, ...) - """ - dim = hidden_states.shape[-1] - # output = torch.zeros((batch * seqlen), dim, device=hidden_states.device, dtype=hidden_states.dtype) - # output[indices] = hidden_states - output = index_put_first_axis(hidden_states, indices, batch * seqlen) - return rearrange(output, "(b s) ... -> b s ...", b=batch) diff --git a/build/torch212-cxx11-cu132-x86_64-linux/flash_attn2/__init__.py b/build/torch212-cxx11-cu132-x86_64-linux/flash_attn2/__init__.py deleted file mode 100644 index a9b2672c1cd85b74c1b3ded0fc0b2100e1aeac23..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu132-x86_64-linux/flash_attn2/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -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/torch212-cxx11-cu132-x86_64-linux/flash_attn_interface.py b/build/torch212-cxx11-cu132-x86_64-linux/flash_attn_interface.py deleted file mode 100644 index b014e569dc9812a9c840d1043f2cebe767b4d477..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu132-x86_64-linux/flash_attn_interface.py +++ /dev/null @@ -1,1631 +0,0 @@ -# Copyright (c) 2023, Tri Dao. - -from typing import Optional, Sequence, Tuple, Union - -import torch -import torch.nn as nn -import os - -# # isort: off -# # We need to import the CUDA kernels after importing torch -# USE_TRITON_ROCM = os.getenv("FLASH_ATTENTION_TRITON_AMD_ENABLE", "FALSE") == "TRUE" -# if USE_TRITON_ROCM: -# from .flash_attn_triton_amd import interface_fa as flash_attn -# else: -# import flash_attn_2_cuda as flash_attn - - -from . import _ops -from ._ops import ops as flash_attn -from ._ops import add_op_namespace_prefix - -# # isort: on - -def maybe_contiguous(x): - return x.contiguous() if x is not None and x.stride(-1) != 1 else x - - -def _get_device(): - if torch.xpu.is_available(): - return "xpu" - elif torch.cuda.is_available(): - return "cuda" - else: - return "cpu" - - -def _get_block_size_n(device, head_dim, is_dropout, is_causal): - assert head_dim <= 256 - - if device.type == "xpu": - if head_dim <= 96: - return 64 - elif head_dim <= 128: - return 32 - elif head_dim <= 256: - return 64 - else: - return 32 - - # This should match the block sizes in the CUDA kernel - major, minor = torch.cuda.get_device_capability(device) - is_sm8x = major == 8 and minor > 0 # Only include sm86 and sm89, exclude sm80 (A100) - is_sm80 = major == 8 and minor == 0 - is_sm90 = major == 9 and minor == 0 - if head_dim <= 32: - return 128 - if head_dim <= 64: - return 128 if not is_dropout else 64 - elif head_dim <= 96: - return 64 - elif head_dim <= 128: - if is_sm8x: - return 64 if (not is_dropout and is_causal) else 32 - else: - return 64 if not is_dropout else 32 - elif head_dim <= 192: - return 64 - elif head_dim <= 224: - return 64 - elif head_dim <= 256: - return 64 - - -def round_multiple(x, m): - return (x + m - 1) // m * m - - -# torch.compile() support is only enabled for pytorch >= 2.4 -# The reason for this is that we are using the new custom_op and register_fake -# APIs, which support inplace modification of inputs in the function itself -if torch.__version__ >= "2.4.0": - _torch_custom_op_wrapper = torch.library.custom_op - _torch_register_fake_wrapper = torch.library.register_fake -else: - def noop_custom_op_wrapper(name, fn=None, /, *, mutates_args, device_types=None, schema=None): - def wrap(func): - return func - if fn is None: - return wrap - return fn - def noop_register_fake_wrapper(op, fn=None, /, *, lib=None, _stacklevel=1): - def wrap(func): - return func - if fn is None: - return wrap - return fn - _torch_custom_op_wrapper = noop_custom_op_wrapper - _torch_register_fake_wrapper = noop_register_fake_wrapper - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_forward"), mutates_args=(), device_types=_get_device()) -def _flash_attn_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - return_softmax: bool -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - out, softmax_lse, S_dmask, rng_state = flash_attn.fwd( - q, - k, - v, - None, - alibi_slopes, - dropout_p, - softmax_scale, - causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - None, - ) - return out, softmax_lse, S_dmask, rng_state - - -@_torch_register_fake_wrapper(add_op_namespace_prefix("_flash_attn_forward")) -def _flash_attn_forward_fake( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - return_softmax: bool -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - batch_size, seqlen_q, num_heads, head_size = q.shape - seqlen_k = k.shape[1] - out = torch.empty_like(q) - softmax_lse = torch.empty((batch_size, num_heads, seqlen_q), dtype=torch.float32, device=q.device, layout=q.layout) - p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) - if return_softmax: - p = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128), round_multiple(seqlen_k, 128)), dtype=q.dtype, device=q.device, layout=q.layout) - rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) - - return out, softmax_lse, p, rng_state - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_forward = _ops.ops._flash_attn_forward -else: - _wrapped_flash_attn_forward = _flash_attn_forward - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_varlen_forward"), mutates_args=(), device_types=_get_device()) -def _flash_attn_varlen_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - alibi_slopes: Optional[torch.Tensor] = None, - return_softmax: bool = False, - block_table: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - out, softmax_lse, S_dmask, rng_state = flash_attn.varlen_fwd( - q, - k, - v, - None, - cu_seqlens_q, - cu_seqlens_k, - seqused_k, - leftpad_k, - block_table, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - zero_tensors, - causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - None, - ) - # if out.isnan().any() or softmax_lse.isnan().any(): - # breakpoint() - return out, softmax_lse, S_dmask, rng_state - - -@_torch_register_fake_wrapper(add_op_namespace_prefix("_flash_attn_varlen_forward")) -def _flash_attn_varlen_forward_fake( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - alibi_slopes: Optional[torch.Tensor] = None, - return_softmax: bool = False, - block_table: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - paged_kv = block_table is not None - batch_size = cu_seqlens_q.numel() - 1 - total_q, num_heads, _ = q.shape - - out = torch.empty_like(q) - softmax_lse = torch.empty((num_heads, total_q), dtype=torch.float32, device=q.device, layout=q.layout) - p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) - seqlen_q_rounded = round_multiple(max_seqlen_q, 128) - seqlen_k_rounded = round_multiple(max_seqlen_k, 128) - if return_softmax: - p = torch.empty((batch_size, num_heads, seqlen_q_rounded, seqlen_k_rounded), dtype=q.dtype, device=q.device, layout=q.layout) - rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) - return out, softmax_lse, p, rng_state - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_varlen_forward = _ops.ops._flash_attn_varlen_forward -else: - _wrapped_flash_attn_varlen_forward = _flash_attn_varlen_forward - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_backward"), mutates_args=("dq", "dk", "dv"), device_types=_get_device()) -def _flash_attn_backward( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, -) -> 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)] - ( - dq, - dk, - dv, - softmax_d, - ) = flash_attn.bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - alibi_slopes, - dropout_p, - softmax_scale, - causal, - window_size_left, - window_size_right, - softcap, - deterministic, - None, - rng_state, - ) - return softmax_d - - -@_torch_register_fake_wrapper(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, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, -) -> torch.Tensor: - dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] - if dq is None: - dq = torch.empty_like(q) - if dk is None: - dk = torch.empty_like(k) - if dv is None: - dv = torch.empty_like(v) - batch_size, seqlen_q, num_heads, _ = q.shape - softmax_d = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128)), device=q.device, dtype=torch.float32) - - return softmax_d - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_backward = _ops.ops._flash_attn_backward -else: - _wrapped_flash_attn_backward = _flash_attn_backward - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_varlen_backward"), mutates_args=("dq", "dk", "dv"), device_types=_get_device()) -def _flash_attn_varlen_backward( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> 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)] - ( - dq, - dk, - dv, - softmax_d, - ) = flash_attn.varlen_bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - zero_tensors, - causal, - window_size_left, - window_size_right, - softcap, - deterministic, - None, - rng_state, - ) - # if dk.isnan().any() or dk.isnan().any() or dv.isnan().any() or softmax_d.isnan().any(): - # breakpoint() - return softmax_d - - -@_torch_register_fake_wrapper(add_op_namespace_prefix("_flash_attn_varlen_backward")) -def _flash_attn_varlen_backward_fake( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> torch.Tensor: - dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] - batch_size = cu_seqlens_q.numel() - 1 - total_q, num_heads, _ = q.shape - - if dq is None: - dq = torch.empty_like(q) - if dk is None: - dk = torch.empty_like(k) - if dv is None: - dv = torch.empty_like(v) - softmax_d = torch.empty((num_heads, total_q + 128 * batch_size), device=q.device, dtype=torch.float32) - - return softmax_d - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_varlen_backward = _ops.ops._flash_attn_varlen_backward -else: - _wrapped_flash_attn_varlen_backward = _flash_attn_varlen_backward - - -class FlashAttnQKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and qkv.requires_grad - if softmax_scale is None: - softmax_scale = qkv.shape[-1] ** (-0.5) - q, k, v = qkv[:, :, 0].detach(), qkv[:, :, 1].detach(), qkv[:, :, 2].detach() - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) - dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dqkv[:, :, 0], - dqkv[:, :, 1], - dqkv[:, :, 2], - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dqkv = dqkv[..., : dout.shape[-1]] # We could have padded the head dimension - return dqkv, None, None, None, None, None, None, None, None, None - - -class FlashAttnVarlenQKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - cu_seqlens, - max_seqlen, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and qkv.requires_grad - if softmax_scale is None: - softmax_scale = qkv.shape[-1] ** (-0.5) - q, k, v = qkv[:, 0].detach(), qkv[:, 1].detach(), qkv[:, 2].detach() - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens, - cu_seqlens, - max_seqlen, - max_seqlen, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=None, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens, rng_state) - ctx.dropout_p = dropout_p - ctx.max_seqlen = max_seqlen - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens, rng_state = ctx.saved_tensors - qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) - dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dqkv[:, 0], - dqkv[:, 1], - dqkv[:, 2], - cu_seqlens, - cu_seqlens, - ctx.max_seqlen, - ctx.max_seqlen, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - 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 FlashAttnKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - kv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, kv] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - k, v = kv[:, :, 0].detach(), kv[:, :, 1].detach() - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - dq = torch.empty_like(q) - kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) - dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dkv[:, :, 0], - dkv[:, :, 1], - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dkv = dkv[..., : dout.shape[-1]] - return dq, dkv, None, None, None, None, None, None, None, None, None - - -class FlashAttnVarlenKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, kv] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - k, v = kv[:, 0].detach(), kv[:, 1].detach() - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=None, - ) - if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state - ) - ctx.dropout_p = dropout_p - 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.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors - dq = torch.empty_like(q) - kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) - dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dkv[:, 0], - dkv[:, 1], - cu_seqlens_q, - cu_seqlens_k, - ctx.max_seqlen_q, - ctx.max_seqlen_k, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dkv = dkv[..., : dout.shape[-1]] - return dq, dkv, None, None, None, None, None, None, None, None, None, None, None, None, None - - -class FlashAttnFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - k, - v, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, k, v] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dk = dk[..., : dout.shape[-1]] - dv = dv[..., : dout.shape[-1]] - return dq, dk, dv, 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, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - block_table, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, k, v] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=block_table, - ) - if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state - ) - ctx.dropout_p = dropout_p - 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.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors - dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - ctx.max_seqlen_q, - ctx.max_seqlen_k, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dk = dk[..., : dout.shape[-1]] - dv = dv[..., : dout.shape[-1]] - return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None - - -def flash_attn_qkvpacked_func( - qkv, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # <=0.0 means deactivate - alibi_slopes=None, - deterministic=False, - 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, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_kvpacked_func( - q, - kv, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, -): - """dropout_p should be set to 0.0 during evaluation - If 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 K, V. - 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) - kv: (batch_size, seqlen, 2, 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. - softcap: float. Anything > 0 activates softcapping 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). - 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 FlashAttnKVPackedFunc.apply( - q, - kv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_func( - q, - k, - v, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - 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). - 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 FlashAttnFunc.apply( - q, - k, - v, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_qkvpacked_func( - qkv, - cu_seqlens, - max_seqlen, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - 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_varlen_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_varlen_kvpacked_func and flash_attn_varlen_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: (total, 3, nheads, headdim), where total = total number of tokens in the batch. - cu_seqlens: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into qkv. - max_seqlen: int. Maximum sequence length in the batch. - 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenQKVPackedFunc.apply( - qkv, - cu_seqlens, - max_seqlen, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_kvpacked_func( - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, -): - """dropout_p should be set to 0.0 during evaluation - If 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 K, V. - 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: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch. - kv: (total_k, 2, nheads_k, headdim), where total_k = total number of key tokens in the batch. - cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into q. - cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into kv. - max_seqlen_q: int. Maximum query sequence length in the batch. - max_seqlen_k: int. Maximum key sequence length in the batch. - 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 + 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenKVPackedFunc.apply( - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_func( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, - block_table=None, -): - """dropout_p should be set to 0.0 during evaluation - Supports multi-query and grouped-query attention (MQA/GQA) by passing in K, V 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: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch. - k: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. - v: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. - cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into q. - cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into kv. - max_seqlen_q: int. Maximum query sequence length in the batch. - max_seqlen_k: int. Maximum key sequence length in the batch. - 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 + 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenFunc.apply( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - block_table, - False if q.device.type == "cpu" else torch.is_grad_enabled(), - ) - - -def flash_attn_with_kvcache( - q, - k_cache, - v_cache, - k=None, - v=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, - block_table: Optional[torch.Tensor] = None, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - rotary_interleaved=True, - alibi_slopes=None, - num_splits=0, - 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 block_table, - or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_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) if there's no block_table, - or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_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). Similar to k. - 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. - block_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). - 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. - 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" - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-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) - cache_batch_idx = maybe_contiguous(cache_batch_idx) - block_table = maybe_contiguous(block_table) - out, softmax_lse = flash_attn.fwd_kvcache( - q, - k_cache, - v_cache, - k, - v, - cache_seqlens, - rotary_cos, - rotary_sin, - cache_batch_idx, - cache_leftpad, - block_table, - alibi_slopes, - None, - softmax_scale, - causal, - window_size[0], - window_size[1], - softcap, - rotary_interleaved, - num_splits, - ) - return (out, softmax_lse) if return_softmax_lse else out diff --git a/build/torch212-cxx11-cu132-x86_64-linux/layers/__init__.py b/build/torch212-cxx11-cu132-x86_64-linux/layers/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/build/torch212-cxx11-cu132-x86_64-linux/layers/patch_embed.py b/build/torch212-cxx11-cu132-x86_64-linux/layers/patch_embed.py deleted file mode 100644 index 05562f8e8bcdb58e947c6f402a49eacd2d031871..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu132-x86_64-linux/layers/patch_embed.py +++ /dev/null @@ -1,67 +0,0 @@ -# We use the same API as https://github.com/rwightman/pytorch-image-models/blob/v0.6.11/timm/models/layers/patch_embed.py -# But we use nn.Linear instead of Conv2d and it's about 8x faster. - -from functools import partial - -import torch.nn as nn -from einops import rearrange -from torch import _assert -from torch.nn.modules.utils import _pair - -try: - from flash_attn.ops.fused_dense import FusedDense -except ImportError: - FusedDense = None - - -class PatchEmbed(nn.Module): - """2D Image to Patch Embedding""" - - def __init__( - self, - img_size=224, - patch_size=16, - in_chans=3, - embed_dim=768, - norm_layer=None, - flatten=True, - bias=True, - fused_bias_fc=False, - ): - super().__init__() - img_size = _pair(img_size) - patch_size = _pair(patch_size) - self.img_size = img_size - self.patch_size = patch_size - self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1]) - self.num_patches = self.grid_size[0] * self.grid_size[1] - self.flatten = flatten - if fused_bias_fc and FusedDense is None: - raise ImportError("fused_dense is not installed") - - linear_cls = nn.Linear if not fused_bias_fc or not bias else FusedDense - self.proj = linear_cls(in_chans * patch_size[0] * patch_size[1], embed_dim, bias=bias) - self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() - - def forward(self, x): - _, _, H, W = x.shape - _assert( - H == self.img_size[0], - f"Input image height ({H}) doesn't match model ({self.img_size[0]}).", - ) - _assert( - W == self.img_size[1], - f"Input image width ({W}) doesn't match model ({self.img_size[1]}).", - ) - x = self.proj( - rearrange( - x, - "b c (h p1) (w p2) -> b h w (c p1 p2)", - p1=self.patch_size[0], - p2=self.patch_size[1], - ) - ) - if self.flatten: - x = rearrange(x, "b h w c -> b (h w) c") - x = self.norm(x) - return x diff --git a/build/torch212-cxx11-cu132-x86_64-linux/layers/rotary.py b/build/torch212-cxx11-cu132-x86_64-linux/layers/rotary.py deleted file mode 100644 index d1bfc21fc7de1dd287e8f382847b194a48075981..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu132-x86_64-linux/layers/rotary.py +++ /dev/null @@ -1,483 +0,0 @@ -# Copyright (c) 2025, Tri Dao - -import math -from functools import partial -from typing import Optional, Tuple, Union - -import torch -from torch import Tensor - -from einops import rearrange, repeat -# from flash_attn.ops.triton.rotary import apply_rotary -from ..ops.triton.rotary import apply_rotary - - -def rotate_half(x, interleaved=False): - if not interleaved: - x1, x2 = x.chunk(2, dim=-1) - return torch.cat((-x2, x1), dim=-1) - else: - x1, x2 = x[..., ::2], x[..., 1::2] - return rearrange(torch.stack((-x2, x1), dim=-1), "... d two -> ... (d two)", two=2) - - -def apply_rotary_emb_torch(x, cos, sin, interleaved=False): - """ - x: (batch_size, seqlen, nheads, headdim) - cos, sin: (seqlen, rotary_dim / 2) or (batch_size, seqlen, rotary_dim / 2) - """ - ro_dim = cos.shape[-1] * 2 - assert ro_dim <= x.shape[-1] - cos = repeat(cos, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)") - sin = repeat(sin, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)") - return torch.cat( - [x[..., :ro_dim] * cos + rotate_half(x[..., :ro_dim], interleaved) * sin, x[..., ro_dim:]], - dim=-1, - ) - - -class ApplyRotaryEmb(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x, - cos, - sin, - interleaved=False, - inplace=False, - seqlen_offsets: Union[int, Tensor] = 0, - cu_seqlens: Optional[Tensor] = None, - max_seqlen: Optional[int] = None, - ): - out = apply_rotary( - x, - cos, - sin, - seqlen_offsets=seqlen_offsets, - cu_seqlens=cu_seqlens, - max_seqlen=max_seqlen, - interleaved=interleaved, - inplace=inplace, - ) - if isinstance(seqlen_offsets, int): - ctx.save_for_backward(cos, sin, cu_seqlens) # Can't save int with save_for_backward - ctx.seqlen_offsets = seqlen_offsets - else: - ctx.save_for_backward(cos, sin, cu_seqlens, seqlen_offsets) - ctx.seqlen_offsets = None - ctx.interleaved = interleaved - ctx.inplace = inplace - ctx.max_seqlen = max_seqlen - return out if not inplace else x - - @staticmethod - def backward(ctx, do): - seqlen_offsets = ctx.seqlen_offsets - if seqlen_offsets is None: - cos, sin, cu_seqlens, seqlen_offsets = ctx.saved_tensors - else: - cos, sin, cu_seqlens = ctx.saved_tensors - dx = apply_rotary( - do, - cos, - sin, - seqlen_offsets=seqlen_offsets, - cu_seqlens=cu_seqlens, - max_seqlen=ctx.max_seqlen, - interleaved=ctx.interleaved, - inplace=ctx.inplace, - conjugate=True, - ) - return dx, None, None, None, None, None, None, None - - -def apply_rotary_emb( - x, - cos, - sin, - interleaved=False, - inplace=False, - seqlen_offsets: Union[int, Tensor] = 0, - cu_seqlens: Optional[Tensor] = None, - max_seqlen: Optional[int] = None, -): - """ - Arguments: - x: (batch_size, seqlen, nheads, headdim) if cu_seqlens is None - else (total_seqlen, nheads, headdim) - cos, sin: (seqlen_rotary, rotary_dim / 2) - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead - of 1st half and 2nd half (GPT-NeoX style). - inplace: if True, apply rotary embedding in-place. - seqlen_offsets: (batch_size,) or int. Each sequence in x is shifted by this amount. - Most commonly used in inference when we have KV cache. - cu_seqlens: (batch + 1,) or None - max_seqlen: int - Return: - out: (batch_size, seqlen, nheads, headdim) if cu_seqlens is None - else (total_seqlen, nheads, headdim) - rotary_dim must be <= headdim - Apply rotary embedding to the first rotary_dim of x. - """ - return ApplyRotaryEmb.apply( - x, cos, sin, interleaved, inplace, seqlen_offsets, cu_seqlens, max_seqlen - ) - - -# For backward compatibility -apply_rotary_emb_func = apply_rotary_emb - - -def _apply_rotary_emb_qkv( - qkv, - cos, - sin, - cos_k=None, - sin_k=None, - interleaved=False, - inplace=False, - conjugate=False, - seqlen_offsets: Union[int, Tensor] = 0, - num_heads_q: Optional[int] = None, -): - apply_rotary_fn = partial( - apply_rotary, - interleaved=interleaved, - inplace=inplace, - conjugate=conjugate, - seqlen_offsets=seqlen_offsets - ) - if cos_k is None and sin_k is None and qkv.is_contiguous(): - # Call 1 kernel instead of 2 kernels - # We need qkv to be contiguous so that when we reshape to combine (3, nheads) - # dimensions, we get the same tensor - if qkv.dim() == 5: - batch, seqlen, three, nheads, headdim = qkv.shape - assert three == 3 - # qk = rearrange(qkv[:, :, :2], "b s t h d -> b s (t h) d") - qk = qkv[:, :, :2].reshape(batch, seqlen, -1, headdim) - qk = apply_rotary_fn(qk, cos, sin) - else: - assert qkv.dim() == 4 - assert num_heads_q is not None - num_heads_k = (qkv.shape[2] - num_heads_q) // 2 - assert qkv.shape[2] == num_heads_q + 2 * num_heads_k - qk = qkv[:, :, :num_heads_q + num_heads_k] - qk = apply_rotary_fn(qk, cos, sin) - if not inplace: - if qkv.dim() == 5: - qkv = torch.cat([rearrange(qk, "b s (t h) d -> b s t h d", t=2), qkv[:, :, 2:]], dim=2) - else: - qkv = torch.cat([qk, qkv[:, :, num_heads_q + num_heads_k :]], dim=2) - else: - cos_k = cos if cos_k is None else cos_k - sin_k = sin if sin_k is None else sin_k - if qkv.dim() == 5: - batch, seqlen, three, nheads, headdim = qkv.shape - assert three == 3 - q, k = qkv[:, :, 0], qkv[:, :, 1] - else: - assert qkv.dim() == 4 - assert num_heads_q is not None - num_heads_k = (qkv.shape[2] - num_heads_q) // 2 - assert qkv.shape[2] == num_heads_q + 2 * num_heads_k - q, k = qkv[:, :, :num_heads_q], qkv[:, :, num_heads_q : num_heads_q + num_heads_k] - q = apply_rotary_fn(q, cos, sin) - k = apply_rotary_fn(k, cos_k, sin_k) - if not inplace: - if qkv.dim() == 5: - qkv = torch.stack([q, k, qkv[:, :, 2]], dim=2) - else: - qkv = torch.cat([q, k, qkv[:, :, num_heads_q + num_heads_k:]], dim=2) - return qkv - - -class ApplyRotaryEmbQKV_(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - cos, - sin, - cos_k=None, - sin_k=None, - interleaved=False, - seqlen_offsets: Union[int, torch.Tensor] = 0, - num_heads_q: Optional[int] = None, - ): - # apply_rotary_emb_qkv_inplace( - qkv = _apply_rotary_emb_qkv( - qkv, cos, sin, cos_k, sin_k, interleaved=interleaved, inplace=True, - seqlen_offsets=seqlen_offsets, num_heads_q=num_heads_q, - ) - if isinstance(seqlen_offsets, int): - ctx.save_for_backward(cos, sin, cos_k, sin_k) - ctx.seqlen_offsets = seqlen_offsets - else: - ctx.save_for_backward(cos, sin, cos_k, sin_k, seqlen_offsets) - ctx.seqlen_offsets = None - ctx.interleaved = interleaved - ctx.num_heads_q = num_heads_q - return qkv - - @staticmethod - def backward(ctx, dqkv): - seqlen_offsets = ctx.seqlen_offsets - if seqlen_offsets is None: - cos, sin, cos_k, sin_k, seqlen_offsets = ctx.saved_tensors - else: - cos, sin, cos_k, sin_k = ctx.saved_tensors - dqkv = _apply_rotary_emb_qkv( - dqkv, cos, sin, cos_k, sin_k, interleaved=ctx.interleaved, inplace=True, - seqlen_offsets=seqlen_offsets, num_heads_q=ctx.num_heads_q, conjugate=True, - ) - return dqkv, None, None, None, None, None, None, None - - -def apply_rotary_emb_qkv_( - qkv, - cos, - sin, - cos_k=None, - sin_k=None, - interleaved=False, - seqlen_offsets: Union[int, torch.Tensor] = 0, - num_heads_q: Optional[int] = None, -): - """ - Arguments: - qkv: (batch_size, seqlen, 3, nheads, headdim) or (batch_size, seqlen, num_heads_q + 2 * num_heads_k, headdim). - If qkv has shape (batch_size, seqlen, num_heads_q + 2 * num_heads_k, headdim) (e.g. MQA / GQA), - then num_heads_q must be provided. - cos, sin: (seqlen, rotary_dim / 2) - cos_k, sin_k: (seqlen, rotary_dim / 2), optional - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead of - 1st half and 2nd half (GPT-NeoX style). - seqlen_offsets: (batch_size,) or int. Each sequence in Q and K is shifted by this amount. - Most commonly used in inference when we have KV cache. - Return: - qkv: (batch_size, seqlen, 3, nheads, headdim) or (batch_size, seqlen, num_heads_q + 2 * num_heads_k, headdim) - rotary_dim must be <= headdim - Apply rotary embedding *inplace* to the first rotary_dim of Q and K. - """ - return ApplyRotaryEmbQKV_.apply( - qkv, cos, sin, cos_k, sin_k, interleaved, seqlen_offsets, num_heads_q - ) - - -class ApplyRotaryEmbKV_(torch.autograd.Function): - - @staticmethod - def forward(ctx, kv, cos, sin, interleaved=False, seqlen_offsets: Union[int, torch.Tensor] = 0): - batch, seqlen, two, nheads, headdim = kv.shape - assert two == 2 - k = kv[:, :, 0] - apply_rotary( - k, cos, sin, seqlen_offsets=seqlen_offsets, interleaved=interleaved, inplace=True - ) - if isinstance(seqlen_offsets, int): - ctx.save_for_backward(cos, sin) # Can't save int with save_for_backward - ctx.seqlen_offsets = seqlen_offsets - else: - ctx.save_for_backward(cos, sin, seqlen_offsets) - ctx.seqlen_offsets = None - ctx.interleaved = interleaved - return kv - - @staticmethod - def backward(ctx, dkv): - seqlen_offsets = ctx.seqlen_offsets - if seqlen_offsets is None: - cos, sin, seqlen_offsets = ctx.saved_tensors - else: - cos, sin = ctx.saved_tensors - apply_rotary( - dkv[:, :, 0], - cos, - sin, - seqlen_offsets=seqlen_offsets, - interleaved=ctx.interleaved, - inplace=True, - conjugate=True, - ) - return dkv, None, None, None, None - - -apply_rotary_emb_kv_ = ApplyRotaryEmbKV_.apply - - -def apply_rotary_emb_kv_( - kv, - cos, - sin, - interleaved=False, - seqlen_offsets: Union[int, torch.Tensor] = 0, -): - """ - Arguments: - kv: (batch_size, seqlen, 2, nheads, headdim) - cos, sin: (seqlen, rotary_dim / 2) - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead of - 1st half and 2nd half (GPT-NeoX style). - seqlen_offsets: (batch_size,) or int. Each sequence in Q and K is shifted by this amount. - Most commonly used in inference when we have KV cache. - Return: - kv: (batch_size, seqlen, 2, nheads, headdim) - rotary_dim must be <= headdim - Apply rotary embedding *inplace* to the first rotary_dim of K. - """ - return ApplyRotaryEmbKV_.apply(kv, cos, sin, interleaved, seqlen_offsets) - - -class RotaryEmbedding(torch.nn.Module): - """ - The rotary position embeddings from RoFormer_ (Su et. al). - A crucial insight from the method is that the query and keys are - transformed by rotation matrices which depend on the relative positions. - - Other implementations are available in the Rotary Transformer repo_ and in - GPT-NeoX_, GPT-NeoX was an inspiration - - .. _RoFormer: https://arxiv.org/abs/2104.09864 - .. _repo: https://github.com/ZhuiyiTechnology/roformer - .. _GPT-NeoX: https://github.com/EleutherAI/gpt-neox - - If scale_base is not None, this implements XPos (Sun et al., https://arxiv.org/abs/2212.10554). - A recommended value for scale_base is 512: https://github.com/HazyResearch/flash-attention/issues/96 - Reference: https://github.com/sunyt32/torchscale/blob/main/torchscale/component/xpos_relative_position.py - """ - - def __init__( - self, - dim: int, - base=10000.0, - interleaved=False, - scale_base=None, - device=None, - ): - """ - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead - of 1st half and 2nd half (GPT-NeoX style). - """ - super().__init__() - self.dim = dim - self.base = float(base) - # Generate and save the inverse frequency buffer (non trainable) - inv_freq = self._compute_inv_freq(device) - self.register_buffer("inv_freq", inv_freq, persistent=False) - self.interleaved = interleaved - self.scale_base = scale_base - scale = ( - (torch.arange(0, dim, 2, device=device, dtype=torch.float32) + 0.4 * dim) / (1.4 * dim) - if scale_base is not None - else None - ) - self.register_buffer("scale", scale, persistent=False) - - self._seq_len_cached = 0 - self._cos_cached = None - self._sin_cached = None - self._cos_k_cached = None - self._sin_k_cached = None - - def _compute_inv_freq(self, device=None): - return 1.0 / ( - self.base - ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim) - ) - - def _update_cos_sin_cache(self, seqlen, device=None, dtype=None): - # Reset the tables if the sequence length has changed, - # if we're on a new device (possibly due to tracing for instance), - # or if we're switching from inference mode to training - if ( - seqlen > self._seq_len_cached - or self._cos_cached is None - or self._cos_cached.device != device - or self._cos_cached.dtype != dtype - or (self.training and self._cos_cached.is_inference()) - ): - self._seq_len_cached = seqlen - # We want fp32 here, not self.inv_freq.dtype, since the model could be loaded in bf16 - # And the output of arange can be quite large, so bf16 would lose a lot of precision. - t = torch.arange(seqlen, device=device, dtype=torch.float32) - # We want fp32 here as well since inv_freq will be multiplied with t, and the output - # will be large. Having it in bf16 will lose a lot of precision and cause the - # cos & sin output to change significantly. - # We want to recompute self.inv_freq if it was not loaded in fp32 - if self.inv_freq.dtype != torch.float32: - inv_freq = self._compute_inv_freq(device=device) - else: - inv_freq = self.inv_freq - # Don't do einsum, it converts fp32 to bf16 under AMP - # freqs = torch.einsum("i,j->ij", t, self.inv_freq) - freqs = torch.outer(t, inv_freq) - if self.scale is None: - self._cos_cached = torch.cos(freqs).to(dtype) - self._sin_cached = torch.sin(freqs).to(dtype) - else: - power = ( - torch.arange(seqlen, dtype=self.scale.dtype, device=self.scale.device) - - seqlen // 2 - ) / self.scale_base - scale = self.scale.to(device=power.device) ** rearrange(power, "s -> s 1") - # We want the multiplication by scale to happen in fp32 - self._cos_cached = (torch.cos(freqs) * scale).to(dtype) - self._sin_cached = (torch.sin(freqs) * scale).to(dtype) - self._cos_k_cached = (torch.cos(freqs) / scale).to(dtype) - self._sin_k_cached = (torch.sin(freqs) / scale).to(dtype) - - def forward( - self, - qkv: torch.Tensor, - kv: Optional[torch.Tensor] = None, - seqlen_offset: Union[int, torch.Tensor] = 0, - max_seqlen: Optional[int] = None, - num_heads_q: Optional[int] = None, - ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: - """ - qkv: (batch, seqlen, 3, nheads, headdim) or (batch, seqlen, num_heads_q + 2 * num_heads_k, headdim) - if kv is none, else it's just q of shape (batch, seqlen, nheads, headdim). - If qkv has shape (batch, seqlen, num_heads_q + 2 * num_heads_k, headdim) (e.g. MQA / GQA), - then num_heads_q must be provided. - kv: (batch, seqlen, 2, nheads, headdim) - seqlen_offset: (batch_size,) or int. Each sequence in x is shifted by this amount. - Most commonly used in inference when we have KV cache. - If it's a tensor of shape (batch_size,), then to update the cos / sin cache, one - should pass in max_seqlen, which will update the cos / sin cache up to that length. - Apply rotary embedding *inplace* to qkv and / or kv. - """ - seqlen = qkv.shape[1] - if max_seqlen is not None: - self._update_cos_sin_cache(max_seqlen, device=qkv.device, dtype=qkv.dtype) - elif isinstance(seqlen_offset, int): - self._update_cos_sin_cache(seqlen + seqlen_offset, device=qkv.device, dtype=qkv.dtype) - if kv is None: - return apply_rotary_emb_qkv_( - qkv, - self._cos_cached, - self._sin_cached, - self._cos_k_cached if self.scale is not None else None, - self._sin_k_cached if self.scale is not None else None, - interleaved=self.interleaved, - seqlen_offsets=seqlen_offset, - num_heads_q=num_heads_q, - ) - else: - q = qkv - q = apply_rotary_emb_func( - q, - self._cos_cached, - self._sin_cached, - interleaved=self.interleaved, - inplace=True, - seqlen_offsets=seqlen_offset, - ) - kv = apply_rotary_emb_kv_( - kv, - self._cos_cached if self.scale is None else self._cos_k_cached, - self._sin_cached if self.scale is None else self._sin_k_cached, - interleaved=self.interleaved, - seqlen_offsets=seqlen_offset, - ) - return q, kv diff --git a/build/torch212-cxx11-cu132-x86_64-linux/metadata.json b/build/torch212-cxx11-cu132-x86_64-linux/metadata.json deleted file mode 100644 index b2ac342ddc645b87ee96ab43a9306713ac4858fe..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu132-x86_64-linux/metadata.json +++ /dev/null @@ -1,42 +0,0 @@ -{ - "name": "flash-attn2", - "id": "_flash_attn2_cuda_f12afc9", - "version": 1, - "license": "BSD-3-Clause", - "python-depends": [], - "backend": { - "type": "cuda", - "archs": [ - "10.0", - "12.0", - "8.0", - "9.0" - ] - }, - "digest": { - "algorithm": "sha256", - "files": { - "__init__.py": "Cd9MEdd/WUYnBmfiQpC4SyPqAxrTilTjvsvswNju4zc=", - "_flash_attn2_cuda_f12afc9.abi3.so": "HxeIaQSb9ai4e0+zZ82LyRE9g0pTdojVUSHSq1MGdhA=", - "_ops.py": "GQbmTTeSOU+7Q2nI68Hz+1RvEogyfEFccx1ubUYw+2s=", - "bert_padding.py": "gF1EmsdJ+HpQ86MRQ4VxDw+Sb/RVISdQALdNnoByHlw=", - "flash_attn2/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY=", - "flash_attn_interface.py": "W7a9bHR2ly6QL3aKqsBqxWA7M7ZWYcZPmtsPJu6rAfY=", - "layers/__init__.py": "47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU=", - "layers/patch_embed.py": "H58CgME/qSOPTZLOG08wFgrQS1j34pvNwMPrkTj3Ek4=", - "layers/rotary.py": "rFIdNrZQFNawwLD6d7t5LE2m+M93wwDZ8EuExw+gfD8=", - "ops/__init__.py": "47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU=", - "ops/activations.py": "t5lzNg1In8LP6bKeTnyeMizwqjv27JGbJ6ylPdGvZYg=", - "ops/fused_dense.py": "BqT8VTSCGdLJAjh6WBMkCRFYrdVodPeoElWBvH6AIzA=", - "ops/layer_norm.py": "zr7NXIm+2mtEynTp1CS0fbFGI2Mqdp41dY4AfDWF6EQ=", - "ops/rms_norm.py": "XEnihcj0a4aSz4LO55m5iKGVn4HKTeKN8TIyHjuDgxI=", - "ops/triton/__init__.py": "AbpHGcgLb+kRsJGnwFEktk7uzpZOCcBY74+YBdrKVGs=", - "ops/triton/cross_entropy.py": "y56nmuOtpSNxM7sDZcUSgXKIsIEayfuvrU8bAVog/nw=", - "ops/triton/k_activations.py": "+Z3vIyO4JkqBMipKsPvhzmxljtBdIhJCsl/M+/ESqBo=", - "ops/triton/layer_norm.py": "BBQdJiqopbIS12viUqzx3SvOQCNbDSAYLJXhpyMGyC8=", - "ops/triton/linear.py": "OtRvKz8xdpl+7v3q/ZTaS9fdBt9XrzMyapgRr50uBbM=", - "ops/triton/mlp.py": "uWA+D7HyuRUsOT+4nhi3vvCDNZnNEXwYcHMQ02UREfQ=", - "ops/triton/rotary.py": "0YSrlbep0lJmzLpTiANs0slTFeGcQQqR+el6uzbSLvI=" - } - } -} \ No newline at end of file diff --git a/build/torch212-cxx11-cu132-x86_64-linux/metadata.json.sigstore b/build/torch212-cxx11-cu132-x86_64-linux/metadata.json.sigstore deleted file mode 100644 index 6ae14a110194598f792dc252fc0db351f9500cfb..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu132-x86_64-linux/metadata.json.sigstore +++ /dev/null @@ -1 +0,0 @@ -{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json","verificationMaterial":{"certificate":{"rawBytes":"MIIHTTCCBtKgAwIBAgIUROBXutQ2wYJSAa3Mw8ZeaniqMX8wCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwNjI0MDEwMTI1WhcNMjYwNjI0MDExMTI1WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEAg9LK/qQ2EDJoNBopkmjXQnEfn64FsU8FQ84rRMC1IcgUOt95mqb/zg1hWt7ifUgs3Ha9e8FP2P9vfCG9xuMo6OCBfEwggXtMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUOPTM9Ut+haLrKxJCoPEuHYB4evswHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoZjEyYWZjOTJjNDBiYmQ3NmQ5OTBkOTExZWJhYmYyOGU0OTBhN2JkZTATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoZjEyYWZjOTJjNDBiYmQ3NmQ5OTBkOTExZWJhYmYyOGU0OTBhN2JkZTAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoZjEyYWZjOTJjNDBiYmQ3NmQ5OTBkOTExZWJhYmYyOGU0OTBhN2JkZTAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKGYxMmFmYzkyYzQwYmJkNzZkOTkwZDkxMWViYWJmMjhlNDkwYTdiZGUwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMjgwMzYyNjg5MzUvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBiwYKKwYBBAHWeQIEAgR9BHsAeQB3AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABnvclbc0AAAQDAEgwRgIhAOZ/nlWEkze1Tr4LrOy9WNk+ImwrPiezpHqZWsq/V8kVAiEAqHUogHJojYJCm4br3jI0p6aBE8GTe8gIt8JFwz2S2HQwCgYIKoZIzj0EAwMDaQAwZgIxANFkeJygCR/XweTIZAs/znXR/CTCrmzfgzZcNpPgoO+ZyNzc2SBbpmYlMElaZV9IuQIxALDTH+3IkhacvRQYd6Rs9GOlxps3NlSnoAM3FdkOxHHH+mfae3IOxMoJRo2ULHwgJg=="},"tlogEntries":[{"logIndex":"1933772868","logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="},"kindVersion":{"kind":"hashedrekord","version":"0.0.1"},"integratedTime":"1782262886","inclusionPromise":{"signedEntryTimestamp":"MEUCIEFfplc9gh9WZxnFdNOEzIRzE3/F9xFVROUIpZQ0vwDHAiEAh0S2VcDh+wWX2cC51jnA6UWmnV5JH8m5IGamKuENJb4="},"inclusionProof":{"logIndex":"1811868606","rootHash":"lYonvSjCLuSYo20cbX1mAMT+505vGa6JsRYBsraMvyA=","treeSize":"1811868624","hashes":["yBn2+cFzviif5riB3+SSRANuxQpSLr2U8qqAFYivN64=","TNwbCVtf0RU2HpZgq5Mimz1DpHWSsF83KDNPKYwbTZg=","Gm80NUOPadnjlP23y+qiwahi4wLY3dQ8nu0GFmcywEg=","QBbMRDe9zzr6I9zVPttwd/ct+jddu787Lj4dVgG9678=","gM0iyO5n28cJjb+zB9cRGva2RHpATiyDyZCsiXJLmVA=","fRDNf+kRYjyzOVDoWA7vWqSSBlVL7SVyt7s0K1YwhE4=","CFF8ni2lObKtSCMAFt3g1Dimj2VKZvFg9FIWME3EMkI=","cqpTs9pwOsqeKF/CsB3hZE8E5CNbqjk7SFSHV6F6cKg=","06lZjiqsJYtErb+gRpEGWMQP2LekWHgm3EQTQgIj37Q=","eEf0n4aXyn0P4Z/HNLjdj0hqsSjjouR8dwugBBfshiA=","fqtytk2if0R+M3kWRpiPUrBEhf2z8y5Li6MeOuWmiks=","mzp9m+2At7GZdgfx0m7RwtO9YzltCZ6UhZBy4p+lUdM=","ipWWLDkqay5JAYe6P4Nrqsck/1C0C5uwqT9QbFsiY1E=","COKvK4TO+dkpqUDGbTbSKewajXzAonx1B1YaRb/ymLw=","dQlR+DXGDKXm5V1fupxaMSOwNZ9F1MybXyyxB96yj80=","JKJdxSZwKPO6QFzLTxHCrUyI5sDU/CwY57zDlA9ni6Q=","OgfnQCh7xr8xUTErIw//YAKaKMDwRJWiThvYExQvLJM=","S8Oe9ZicoJqbVuL1I0dIdTgVpLc82Cm2CMOcV83sKVs=","E1O2ilWlpFRLWGdbmJUMDZ2qbGu/eDZ12m6YjQ9/JPA=","nyv2TDlfkz0XfRwjY98zsbUbOkBZ+Tgz7PML0c2uN14=","OMfDh4k9YD0E+t4sjqJsh+vg8D1ueo1cC07exyxmr6s=","+vHKm4wNwGJ46ceyZ8mePD7Vn1SediStNHCStRoTSWQ=","5HwWjuGNDisetbiWM0cWIXM65TIM3EJlqk13ntCQlcU=","NDIRlqxO0IaGr94q86MIU7VCAN3yEEgOHS4loya5lu4=","daxmZaajRpZV+JxHiOYZhJBiSKN5ucqjh2WnGbHhirw=","DOCeoSMovIvLExkhIvisow9AuNXgeWs4ECkyR6EcqYU="],"checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n1811868624\nlYonvSjCLuSYo20cbX1mAMT+505vGa6JsRYBsraMvyA=\n\n— rekor.sigstore.dev wNI9ajBEAiB6G+ID1NuQWdDM+mks51CjbbysLKaqICBbcU9KvpVj/AIgOSpSGwAP/rMyGTqTJOhzSosLN9SUgBcc6afVLfsfRvQ=\n"}},"canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiIzZTFkZjE1YWFiNGRkOGNjNzFmNGI3N2U1MmU4Yzg3MDE2NTA3ODZiMWQ4NDlkMTBjNmZlNDk3ZDU0YWQzNjZkIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FWUNJUUNhTC91M004REwxa1VnVE1qYVh4TGJNOUM3T3hPcmRDcHp6STAzdWVaVUVBSWhBTWxrMno2eGkzT3FSam5adjhPb2duR2hlVGM4U0w0ajZIWnllZkZablg4RCIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFVWRU5EUW5STFowRjNTVUpCWjBsVlVrOUNXSFYwVVRKM1dVcFRRV0V6VFhjNFdtVmhibWx4VFZnNGQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDVxU1RCTlJFVjNUVlJKTVZkb1kwNU5hbGwzVG1wSk1FMUVSWGhOVkVreFYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZCWnpsTVN5OXhVVEpGUkVwdlRrSnZjR3R0YWxoUmJrVm1ialkwUm5OVk9FWlJPRFFLY2xKTlF6RkpZMmRWVDNRNU5XMXhZaTk2WnpGb1YzUTNhV1pWWjNNelNHRTVaVGhHVURKUU9YWm1RMGM1ZUhWTmJ6WlBRMEptUlhkbloxaDBUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlZQVUZSTkNqbFZkQ3RvWVV4eVMzaEtRMjlRUlhWSVdVSTBaWFp6ZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOWFha1Y1V1ZkYWFrOVVTbXBPUkVKcFdXMVJNMDV0VVRWUFZFSnJDazlVUlhoYVYwcG9XVzFaZVU5SFZUQlBWRUpvVGpKS2ExcFVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMXBxUlhsWlYxcHFUMVJLYWs1RVFtbFpiVkV6VG0xUk5VOVVRbXRQVkVWNFdsZEthRmx0V1hrS1QwZFZNRTlVUW1oT01rcHJXbFJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMXBxUlhsWlYxcHFUMVJLYWs1RVFtbFpiVkV6VG0xUk5VOVVRbXNLVDFSRmVGcFhTbWhaYlZsNVQwZFZNRTlVUW1oT01rcHJXbFJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwZFplRTF0Um0wS1dYcHJlVmw2VVhkWmJVcHJUbnBhYTA5VWEzZGFSR3Q0VFZkV2FWbFhTbTFOYW1oc1RrUnJkMWxVWkdsYVIxVjNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOYW1kM1RYcFplVTVxWnpWTmVsVjJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwZDFsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTT1VKSWMwRUtaVkZDTTBGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtNTJZMnhpWXpCQlFVRlJSQXBCUldkM1VtZEphRUZQV2k5dWJGZEZhM3BsTVZSeU5FeHlUM2s1VjA1ckswbHRkM0pRYVdWNmNFaHhXbGR6Y1M5V09HdFdRV2xGUVhGSVZXOW5TRXB2Q21wWlNrTnROR0p5TTJwSk1IQTJZVUpGT0VkVVpUaG5TWFE0U2taM2VqSlRNa2hSZDBObldVbExiMXBKZW1vd1JVRjNUVVJoVVVGM1dtZEplRUZPUm1zS1pVcDVaME5TTDFoM1pWUkpXa0Z6TDNwdVdGSXZRMVJEY20xNlptZDZXbU5PY0ZCbmIwOHJXbmxPZW1NeVUwSmljRzFaYkUxRmJHRmFWamxKZFZGSmVBcEJURVJVU0NzelNXdG9ZV04yVWxGWlpEWlNjemxIVDJ4NGNITXpUbXhUYm05QlRUTkdaR3RQZUVoSVNDdHRabUZsTTBsUGVFMXZTbEp2TWxWTVNIZG5Da3BuUFQwS0xTMHRMUzFGVGtRZ1EwVlNWRWxHU1VOQlZFVXRMUzB0TFFvPSJ9fX19"}],"timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyjADAgEAMIICwQYJKoZIhvcNAQcCoIICsjCCAq4CAQMxDTALBglghkgBZQMEAgEwgbcGCyqGSIb3DQEJEAEEoIGnBIGkMIGhAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgNrRDoD59SGdC+OCHcmsZhIWoCsLwRoUL6671FeiDKLICFGlHYsFL3E9poKAo2c2MQWd8G1SaGA8yMDI2MDYyNDAxMDEyNVowAwIBAaAypDAwLjEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MRUwEwYDVQQDEwxzaWdzdG9yZS10c2GgADGCAdwwggHYAgEBMFEwOTEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MSAwHgYDVQQDExdzaWdzdG9yZS10c2Etc2VsZnNpZ25lZAIUOhNULwyQYe68wUMvy4qOiyojiwwwCwYJYIZIAWUDBAIBoIH8MBoGCSqGSIb3DQEJAzENBgsqhkiG9w0BCRABBDAcBgkqhkiG9w0BCQUxDxcNMjYwNjI0MDEwMTI1WjAvBgkqhkiG9w0BCQQxIgQgFxd60DNIdX8SmVhnpOMRsixXWI0yxlKHIZPYv9KyjdgwgY4GCyqGSIb3DQEJEAIvMX8wfTB7MHkEIIX5J7wHq2LKw7RDVsEO/IGyxog/2nq55thw2dE6zQW3MFUwPaQ7MDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAoGCCqGSM49BAMCBGgwZgIxALkhBs0TBiY/JmY80e2ENjcxW5dT+IofsCwDXEGMK+TCM/99p84k2Yqzf9sJqUDnfQIxANplC4/7wztMJPU5LJpp+dgxuXGvQ1OQEr/BIIhwuTuk/7ZJxT8hdvq5mYrJW5/bsg=="}]}},"messageSignature":{"messageDigest":{"algorithm":"SHA2_256","digest":"Ph3xWqtN2Mxx9Ld+UujIcBZQeGsdhJ0Qxv5JfVStNm0="},"signature":"MEYCIQCaL/u3M8DL1kUgTMjaXxLbM9C7OxOrdCpzzI03ueZUEAIhAMlk2z6xi3OqRjnZv8OognGheTc8SL4j6HZyefFZnX8D"}} \ No newline at end of file diff --git a/build/torch212-cxx11-cu132-x86_64-linux/ops/__init__.py b/build/torch212-cxx11-cu132-x86_64-linux/ops/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/build/torch212-cxx11-cu132-x86_64-linux/ops/activations.py b/build/torch212-cxx11-cu132-x86_64-linux/ops/activations.py deleted file mode 100644 index 7c09649fc41e12d5a360c5672825d8380bc7ec80..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu132-x86_64-linux/ops/activations.py +++ /dev/null @@ -1,135 +0,0 @@ -# Copied from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/model/layers/activations.py -import math - -import torch -import torch.nn as nn -import torch.nn.functional as F - -# 1/sqrt(2*pi)-> 0.3989423 -# 1/sqrt(2) -> 0.70710678 -# sqrt(2/pi) -> 0.79788456 - -# this function is tanh approximation of gelu -# actual gelu is: -# x * 0.5 * (1.0 + torch.erf(x * 0.70710678)) -@torch.jit.script -def bias_gelu(y, bias): - x = bias + y - return (x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))).to(dtype=y.dtype) - - -# gradient of tanh approximation of gelu -# gradient of actual gelu is: -# 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x) -@torch.jit.script -def bias_gelu_back(g, y, bias): - """Assume that y has shape (B, D) and bias has shape (D)""" - x = bias + y - tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)) - # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243 - ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( - 1 + tanh_out - ) - grad_y = ff * g - return grad_y.to(dtype=y.dtype), grad_y.sum(dim=(0), dtype=bias.dtype) - - -class GeLUFunction(torch.autograd.Function): - @staticmethod - # bias is an optional argument - def forward(ctx, input, bias): - ctx.save_for_backward(input, bias) - return bias_gelu(input, bias) - - @staticmethod - def backward(ctx, grad_output): - input, bias = ctx.saved_tensors - tmp = bias_gelu_back(grad_output, input, bias) - return tmp, tmp - - -bias_gelu_impl = GeLUFunction.apply - -# this function is tanh approximation of gelu -# actual gelu is: -# x * 0.5 * (1.0 + torch.erf(x * 0.70710678)) -@torch.jit.script -def gelu_fwd(x): - return (x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))).to(dtype=x.dtype) - - -# gradient of tanh approximation of gelu -# gradient of actual gelu is: -# 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x) -@torch.jit.script -def gelu_bwd(g, x): - tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)) - # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243 - ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( - 1 + tanh_out - ) - return (ff * g).to(dtype=x.dtype) - - -class FastGeLUFunction(torch.autograd.Function): - @staticmethod - # bias is an optional argument - def forward(ctx, input): - ctx.save_for_backward(input) - return gelu_fwd(input) - - @staticmethod - def backward(ctx, grad_output): - (input,) = ctx.saved_tensors - tmp = gelu_bwd(grad_output, input) - return tmp - - -fast_gelu_impl = FastGeLUFunction.apply - - -@torch.jit.script -def relu_bwd(g, x): - return torch.where(x >= 0, g, 0.0).to(dtype=x.dtype) - - -@torch.jit.script -def sqrelu_fwd(x): - r = F.relu(x) - return (r * r).to(dtype=x.dtype) - - -@torch.jit.script -def sqrelu_bwd(g, x): - return (2.0 * g * F.relu(x)).to(dtype=x.dtype) - - -swiglu_fwd_codestring = """ -template T swiglu_fwd(T x, T y) { - return float(x) * float(y) / (1.0f + ::exp(-float(x))); -} -""" -swiglu_bwd_codestring = """ -template void swiglu_bwd(T x, T y, T g, T& dx, T& dy) { - float x_sigmoid = 1.0f / (1.0f + ::exp(-float(x))); - dx = x_sigmoid * (1 + float(x) * (1.0f - x_sigmoid)) * float(g) * float(y); - dy = float(x) * x_sigmoid * float(g); -} -""" -swiglu_fwd = torch.cuda.jiterator._create_jit_fn(swiglu_fwd_codestring) -swiglu_bwd = torch.cuda.jiterator._create_multi_output_jit_fn(swiglu_bwd_codestring, num_outputs=2) - - -class SwiGLUFunction(torch.autograd.Function): - - @staticmethod - def forward(ctx, x, y): - ctx.save_for_backward(x, y) - return swiglu_fwd(x, y) - - @staticmethod - def backward(ctx, dout): - x, y = ctx.saved_tensors - return swiglu_bwd(x, y, dout) - -swiglu = SwiGLUFunction.apply diff --git a/build/torch212-cxx11-cu132-x86_64-linux/ops/fused_dense.py b/build/torch212-cxx11-cu132-x86_64-linux/ops/fused_dense.py deleted file mode 100644 index 6b4033d134e4093fe278f7b3f8c7d3128ce9f36d..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu132-x86_64-linux/ops/fused_dense.py +++ /dev/null @@ -1,688 +0,0 @@ -# Copyright (c) 2023, Tri Dao. -# Inspired by https://github.com/NVIDIA/apex/blob/master/apex/fused_dense/fused_dense.py -# We make it work with pytorch amp and with bfloat16. -# The TensorParallel linear modules are inspired by https://github.com/NVIDIA/apex/blob/master/apex/transformer/tensor_parallel/layers.py -from functools import partial -from typing import Optional - -# import fused_dense_cuda # from apex -import fused_dense_lib as fused_dense_cuda -import torch -import torch.nn as nn -import torch.nn.functional as F -from torch import Tensor -from torch.distributed import ProcessGroup - -from flash_attn.utils.torch import custom_fwd, custom_bwd -from flash_attn.ops.activations import gelu_bwd, relu_bwd, sqrelu_bwd, sqrelu_fwd -from flash_attn.utils.distributed import ( - all_gather_raw, - all_reduce, - all_reduce_raw, - reduce_scatter, - reduce_scatter_raw, -) - - -class FusedDenseFunc(torch.autograd.Function): - @staticmethod - @custom_fwd - def forward( - ctx, x, weight, bias, return_residual=False, process_group=None, sequence_parallel=True - ): - """ - If process_group is not None and sequence_parallel=True, we're doing Tensor Parallel - with sequence parallelism: we do an all_gather_raw of x before doing the matmul. - """ - ctx.compute_weight_gradient = weight.requires_grad - ctx.return_residual = return_residual - ctx.process_group = process_group - ctx.sequence_parallel = sequence_parallel - - if torch.is_autocast_enabled(): - x = x.to(dtype=torch.get_autocast_gpu_dtype()) - x = x.contiguous() - if process_group is not None and sequence_parallel: - # We want to kick off the all_gather early, before weight dtype conversion - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - else: - total_x = x - - if torch.is_autocast_enabled(): - weight = weight.to(dtype=torch.get_autocast_gpu_dtype()) - bias = bias.to(dtype=torch.get_autocast_gpu_dtype()) if bias is not None else None - weight = weight.contiguous() - if process_group is not None and sequence_parallel: - handle_x.wait() - batch_shape, n = total_x.shape[:-1], total_x.shape[-1] - batch_dim = batch_shape.numel() - # https://github.com/pytorch/pytorch/blob/5b51849b48a7dbccd297286cc0110def4706f9e7/aten/src/ATen/native/cuda/Blas.cpp#L174 - if min(batch_dim, n, *weight.shape) > 65535 * 32: - raise RuntimeError("fused_dense only supports matrix dims <= 2M") - output = F.linear(total_x, weight, bias) - if ctx.compute_weight_gradient: - ctx.save_for_backward(x, weight) - else: - ctx.save_for_backward(weight) - return output if not return_residual else (output, x) - - @staticmethod - @custom_bwd - def backward(ctx, grad_output, *args): - grad_output = grad_output.contiguous() - if ctx.return_residual: - (grad_input,) = args - grad_input = grad_input.contiguous() - process_group = ctx.process_group - sequence_parallel = ctx.sequence_parallel - if ctx.compute_weight_gradient: - x, weight = ctx.saved_tensors - if process_group is not None and sequence_parallel: - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - else: - total_x = x - else: - (weight,) = ctx.saved_tensors - total_x = None - batch_shape = grad_output.shape[:-1] - batch_dim = batch_shape.numel() - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - if ctx.needs_input_grad[0]: - if not ctx.return_residual: - grad_input = F.linear(grad_output, weight.t()) - else: - grad_input = torch.addmm( - grad_input.reshape(batch_dim, grad_input.shape[-1]), grad_output, weight - ) - grad_input = grad_input.reshape(*batch_shape, grad_input.shape[-1]) - if process_group is not None: - reduce_fn = reduce_scatter_raw if sequence_parallel else all_reduce_raw - grad_input, handle_grad_input = reduce_fn(grad_input, process_group, async_op=True) - else: - grad_input = None - if ctx.needs_input_grad[1]: - assert ctx.compute_weight_gradient - if process_group is not None and sequence_parallel: - handle_x.wait() - grad_weight, grad_bias = fused_dense_cuda.linear_bias_wgrad( - total_x.reshape(batch_dim, total_x.shape[-1]), grad_output, ctx.needs_input_grad[2] - ) - else: - grad_weight = None - grad_bias = grad_output if ctx.needs_input_grad[2] else None - if process_group is not None and ctx.needs_input_grad[0]: - handle_grad_input.wait() - return grad_input, grad_weight, grad_bias, None, None, None - - -def fused_dense_func( - x: Tensor, - weight: Tensor, - bias: Optional[Tensor] = None, - return_residual: bool = False, - process_group: Optional[ProcessGroup] = None, - sequence_parallel: bool = True, -): - dtype_eligible = x.dtype in [torch.float16, torch.bfloat16] or ( - x.dtype == torch.float32 and torch.is_autocast_enabled() - ) - if x.is_cuda and weight.is_cuda and (bias is None or bias.is_cuda) and dtype_eligible: - return FusedDenseFunc.apply( - x, weight, bias, return_residual, process_group, sequence_parallel - ) - else: - assert process_group is None - out = F.linear(x, weight, bias) - return out if not return_residual else (out, x) - - -class FusedDense(nn.Linear): - def __init__( - self, - in_features: int, - out_features: int, - bias: bool = True, - return_residual: bool = False, - device=None, - dtype=None, - ) -> None: - super().__init__(in_features, out_features, bias=bias, device=device, dtype=dtype) - self.return_residual = return_residual - - def forward(self, x, process_group=None): - """ - If process_group is not None, we're doing Tensor Parallel with sequence parallelism: - we do an all_gather of x before doing the matmul. - """ - return fused_dense_func( - x, - self.weight, - self.bias, - return_residual=self.return_residual, - process_group=process_group, - ) - - -class ColumnParallelLinear(nn.Linear): - def __init__( - self, - in_features: int, - out_features: int, - process_group: ProcessGroup, - bias: bool = True, - sequence_parallel=True, - multiple_of=1, - device=None, - dtype=None, - ) -> None: - world_size = torch.distributed.get_world_size(process_group) - if out_features % multiple_of: - raise ValueError(f"out_features ({out_features}) must be a multiple of {multiple_of}") - multiple = out_features // multiple_of - # We want to split @multiple across world_size, but it could be an uneven split - div = multiple // world_size - mod = multiple % world_size - # The first @mod ranks get @div + 1 copies, the rest get @div copies - local_multiple = div + int(torch.distributed.get_rank(process_group) < mod) - super().__init__( - in_features, local_multiple * multiple_of, bias=bias, device=device, dtype=dtype - ) - self.process_group = process_group - self.sequence_parallel = sequence_parallel - - def forward(self, x): - # If self.sequence_parallel is True, we're doing Tensor Parallel with sequence parallelism: - # we do an all_gather of x before doing the matmul. - # If not, then the input is already gathered. - return fused_dense_func( - x, - self.weight, - self.bias, - process_group=self.process_group, - sequence_parallel=self.sequence_parallel, - ) - - -class RowParallelLinear(nn.Linear): - def __init__( - self, - in_features: int, - out_features: int, - process_group: ProcessGroup, - bias: bool = True, - sequence_parallel=True, - multiple_of=1, - device=None, - dtype=None, - ) -> None: - world_size = torch.distributed.get_world_size(process_group) - rank = torch.distributed.get_rank(process_group) - if in_features % multiple_of: - raise ValueError(f"in_features ({in_features}) must be a multiple of {multiple_of}") - multiple = in_features // multiple_of - # We want to split @multiple across world_size, but it could be an uneven split - div = multiple // world_size - mod = multiple % world_size - # The first @mod ranks get @div + 1 copies, the rest get @div copies - local_multiple = div + int(torch.distributed.get_rank(process_group) < mod) - # Only rank 0 will have bias - super().__init__( - local_multiple * multiple_of, - out_features, - bias=bias and rank == 0, - device=device, - dtype=dtype, - ) - self.process_group = process_group - self.sequence_parallel = sequence_parallel - - def forward(self, x): - """ - We're doing Tensor Parallel with sequence parallelism: we do the matmul and then - a reduce_scatter of the result. - """ - out = fused_dense_func(x, self.weight, self.bias) - reduce_fn = reduce_scatter if self.sequence_parallel else all_reduce - return reduce_fn(out, self.process_group) - - -class FusedMLPFunc(torch.autograd.Function): - @staticmethod - @custom_fwd - def forward( - ctx, - x, - weight1, - bias1, - weight2, - bias2, - activation="gelu_approx", - save_pre_act=True, - return_residual=False, - checkpoint_lvl=0, - heuristic=0, - process_group=None, - sequence_parallel=True, - ): - """ - If process_group is not None and sequence_parallel=True, we're doing Tensor Parallel - with sequence parallelism: we do an all_gather of x before doing the matmul. - If sequence_parallel=False, then the input is already gathered. - - checkpoint_lvl: - 0: no recomputation in the bwd - 1: recompute gelu_out / relu_out in the bwd - 2: recompute pre_act and gelu_out / relu_out in the bwd - """ - assert -1 <= heuristic <= 4 - assert activation in ["gelu_approx", "relu", "sqrelu"] - if activation == "sqrelu": - assert heuristic == -1 - if not save_pre_act: - checkpoint_lvl = 2 - assert checkpoint_lvl in [0, 1, 2] - ctx.return_residual = return_residual - ctx.process_group = process_group - ctx.sequence_parallel = sequence_parallel - ctx.checkpoint_lvl = checkpoint_lvl - ctx.activation = activation - ctx.heuristic = heuristic - - if torch.is_autocast_enabled(): - x = x.to(dtype=torch.get_autocast_gpu_dtype()) - x = x.contiguous() - if process_group is not None and sequence_parallel: - # We want to kick off the all_gather early, before weight dtype conversion - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - else: - total_x = x - - if torch.is_autocast_enabled(): - dtype = torch.get_autocast_gpu_dtype() - weight1, weight2 = [a.to(dtype=dtype) for a in [weight1, weight2]] - bias1 = bias1.to(dtype=dtype) if bias1 is not None else None - bias2 = bias2.to(dtype=dtype) if bias2 is not None else None - weight1 = weight1.contiguous() - bias1 = bias1.contiguous() if bias1 is not None else None - weight2 = weight2.contiguous() - bias2 = bias2.contiguous() if bias2 is not None else None - if process_group is not None and sequence_parallel: - handle_x.wait() - batch_shape, n = total_x.shape[:-1], total_x.shape[-1] - batch_dim = batch_shape.numel() - # https://github.com/pytorch/pytorch/blob/5b51849b48a7dbccd297286cc0110def4706f9e7/aten/src/ATen/native/cuda/Blas.cpp#L174 - if min(batch_dim, n, *weight1.shape, *weight2.shape) > 65535 * 32: - raise RuntimeError("fused_dense only supports matrix dims <= 2M") - if heuristic == -1: - pre_act = F.linear(total_x, weight1, bias1) - activation_fn = ( - partial(F.gelu, approximate="tanh") - if activation == "gelu_approx" - else (sqrelu_fwd if activation == "sqrelu" else F.relu) - ) - with torch.jit.fuser("fuser2"): - output1 = activation_fn(pre_act) - # This is before adding bias1 - # pre_act = F.linear(total_x.reshape(batch_dim, n), weight1) - # with torch.jit.fuser('fuser2'): - # output1 = bias_gelu(pre_act, bias1) - else: - is_gelu = activation == "gelu_approx" - output1, *rest = fused_dense_cuda.linear_act_forward( - total_x.reshape(batch_dim, n), weight1, bias1, is_gelu, save_pre_act, heuristic - ) - if save_pre_act: - pre_act = rest[0] - output2 = F.linear(output1, weight2, bias2) - if checkpoint_lvl == 0 or (checkpoint_lvl == 1 and activation == "relu"): - # For RELU the pre_act is very small (just a bit-mask) so we just save it - ctx.save_for_backward(x, weight1, weight2, pre_act, output1) - elif checkpoint_lvl == 1: - ctx.save_for_backward(x, weight1, weight2, pre_act) - elif checkpoint_lvl == 2: - ctx.save_for_backward(x, weight1, weight2, bias1) - output2 = output2.reshape(*batch_shape, output2.shape[-1]) - return output2 if not return_residual else (output2, x) - - @staticmethod - @custom_bwd - def backward(ctx, grad_output, *args): - grad_output = grad_output.contiguous() - checkpoint_lvl = ctx.checkpoint_lvl - activation = ctx.activation - activation_fn = ( - partial(F.gelu, approximate="tanh") - if activation == "gelu_approx" - else (sqrelu_fwd if activation == "sqrelu" else F.relu) - ) - if ctx.return_residual: - (grad_input,) = args - grad_input = grad_input.contiguous() - process_group = ctx.process_group - sequence_parallel = ctx.sequence_parallel - x, weight1, weight2, *rest = ctx.saved_tensors - if process_group is None or not sequence_parallel: - total_x = x - batch_shape = grad_output.shape[:-1] - batch_dim = batch_shape.numel() - if checkpoint_lvl in [0, 1]: - if process_group is not None and sequence_parallel: - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - if checkpoint_lvl == 0 or (checkpoint_lvl == 1 and activation == "relu"): - pre_act, output1 = rest - elif checkpoint_lvl == 1: - (pre_act,) = rest - with torch.jit.fuser("fuser2"): - output1 = activation_fn(pre_act) - elif checkpoint_lvl == 2: - (bias1,) = rest - if process_group is not None and sequence_parallel: - total_x, _ = all_gather_raw(x, process_group) - if ctx.heuristic == -1: - pre_act = F.linear(total_x, weight1, bias1) - with torch.jit.fuser("fuser2"): - output1 = activation_fn(pre_act) - else: - output1, pre_act = fused_dense_cuda.linear_act_forward( - total_x.reshape(batch_dim, total_x.shape[-1]), - weight1, - bias1, - activation == "gelu_approx", - True, - ctx.heuristic, - ) - - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - output1 = output1.reshape(batch_dim, output1.shape[-1]) - pre_act = pre_act.reshape(batch_dim, pre_act.shape[-1]) - if ctx.needs_input_grad[3]: - grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad( - output1, grad_output, ctx.needs_input_grad[4] - ) - else: - grad_weight2 = None - grad_bias2 = grad_output if ctx.needs_input_grad[4] else None - if ctx.heuristic == -1: - # grad_pre_act = matmul_dgelu(grad_output, weight2, pre_act) - grad_output1 = F.linear(grad_output, weight2.t()) - activation_grad_fn = ( - gelu_bwd - if activation == "gelu_approx" - else (sqrelu_bwd if activation == "sqrelu" else relu_bwd) - ) - with torch.jit.fuser("fuser2"): - grad_pre_act = activation_grad_fn(grad_output1, pre_act) - else: - # The cublasLt epilogue has to compute both gelu/relu grad and bias grad, we can't - # just compute gelu/relu grad - grad_pre_act, grad_bias1 = fused_dense_cuda.bias_act_linear_dgrad_bgrad( - weight2, grad_output, pre_act, activation == "gelu_approx", ctx.heuristic - ) - if not ctx.needs_input_grad[2]: - grad_bias1 = None - if ctx.needs_input_grad[0]: - if not ctx.return_residual: - grad_input = F.linear(grad_pre_act, weight1.t()) - else: - grad_input = torch.addmm( - grad_input.reshape(batch_dim, grad_input.shape[-1]), grad_pre_act, weight1 - ) - grad_input = grad_input.reshape(*batch_shape, grad_input.shape[-1]) - if process_group is not None: - reduce_fn = reduce_scatter_raw if sequence_parallel else all_reduce_raw - grad_input, handle_grad_input = reduce_fn(grad_input, process_group, async_op=True) - else: - grad_input = None - if ctx.heuristic == -1: - if ctx.needs_input_grad[1]: - if process_group is not None and sequence_parallel and checkpoint_lvl != 2: - handle_x.wait() - grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_wgrad( - total_x.reshape(batch_dim, total_x.shape[-1]), - grad_pre_act, - ctx.needs_input_grad[2], - ) - else: - grad_weight1 = None - grad_bias1 = grad_pre_act if ctx.needs_input_grad[2] else None - else: - if ctx.needs_input_grad[1]: - if process_group is not None and sequence_parallel and checkpoint_lvl != 2: - handle_x.wait() - grad_weight1 = F.linear( - grad_pre_act.t(), total_x.reshape(batch_dim, total_x.shape[-1]).t() - ) - else: - grad_weight1 = None - if process_group is not None and ctx.needs_input_grad[0]: - handle_grad_input.wait() - return ( - grad_input, - grad_weight1, - grad_bias1, - grad_weight2, - grad_bias2, - None, - None, - None, - None, - None, - None, - None, - ) - - -def fused_mlp_func( - x: Tensor, - weight1: Tensor, - weight2: Tensor, - bias1: Optional[Tensor] = None, - bias2: Optional[Tensor] = None, - activation: str = "gelu_approx", - save_pre_act: bool = True, - return_residual: bool = False, - checkpoint_lvl: int = 0, - heuristic: int = 0, - process_group: Optional[ProcessGroup] = None, - sequence_parallel: bool = True, -): - assert activation in ["gelu_approx", "relu", "sqrelu"] - dtype_eligible = x.dtype in [torch.float16, torch.bfloat16] or ( - x.dtype == torch.float32 and torch.is_autocast_enabled() - ) - # If we save pre-activation, dimension must be divisible by 128 (relu) or 8 (gelu) - dim_eligible = not save_pre_act or (x.shape[-1] % (128 if activation == "relu" else 8) == 0) - if ( - x.is_cuda - and weight1.is_cuda - and weight2.is_cuda - and (bias1 is None or bias1.is_cuda) - and (bias2 is None or bias2.is_cuda) - and dtype_eligible - and dim_eligible - ): - return FusedMLPFunc.apply( - x, - weight1, - bias1, - weight2, - bias2, - activation, - save_pre_act, - return_residual, - checkpoint_lvl, - heuristic, - process_group, - sequence_parallel, - ) - else: - assert process_group is None - pre_act = F.linear(x, weight1, bias1) - activation_fn = ( - partial(F.gelu, approximate="tanh") - if activation == "gelu_approx" - else partial(F.relu, inplace=True) - ) - output1 = activation_fn(pre_act) - output2 = F.linear(output1, weight2, bias2) - return output2 if not return_residual else (output2, x) - - -class FusedMLP(nn.Module): - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - bias1=True, - bias2=True, - activation="gelu_approx", - return_residual=False, - checkpoint_lvl=0, - heuristic="auto", - device=None, - dtype=None, - ): - """ - If process_group is not None, we're doing Tensor Parallel with sequence parallelism: - we do an all_gather of x before doing the matmul, gelu, then matmul. - Finally we do a reduce_scatter of the output. - - checkpoint_lvl (increasing lvl means slower but more memory saving): - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute pre_act and gelu_out in the bwd - heuristic: - -1: don't fuse gemm + gelu (separate kernel) - 0..4: use this heuristic for the algo section in the fused gemm + gelu - 'auto': heuristic will be picked automatically: - For CUDA >= 11.8, we set heuristic=0 for both fp16 and bf16 for best perf. - For CUDA <= 11.7, we set heuristic=1 for fp16 and heuristic=-1 for bf16. - For H100, we set heuristic=-1 for both fp16 and bf16 as the fused cuBlasLt implementation - is slower than the unfused version. - return_residual: whether to return the input x along with the output. This is for - performance reason: for post-norm architecture, returning the input allows us - to fuse the backward of nn.Linear with the residual connection. - """ - assert checkpoint_lvl in [0, 1, 2] - assert activation in ["gelu_approx", "relu", "sqrelu"] - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features * 4 - self.activation = activation - self.return_residual = return_residual - self.checkpoint_lvl = checkpoint_lvl - self.heuristic = heuristic if activation != "sqrelu" else -1 - self.fc1 = nn.Linear(in_features, hidden_features, bias=bias1, **factory_kwargs) - self.fc2 = nn.Linear(hidden_features, out_features, bias=bias2, **factory_kwargs) - - def forward(self, x, process_group=None): - dtype = x.dtype if not torch.is_autocast_enabled() else torch.get_autocast_gpu_dtype() - if self.heuristic == "auto": - if self.activation == "gelu_approx": - if torch.cuda.get_device_capability("cuda") == (9, 0): - heuristic = -1 - else: - cuda_ver = tuple(map(int, torch.version.cuda.split("."))) - heuristic = 0 if cuda_ver >= (11, 8) else (1 if dtype == torch.float16 else -1) - else: - heuristic = 0 - else: - heuristic = self.heuristic - out = fused_mlp_func( - x, - self.fc1.weight, - self.fc2.weight, - self.fc1.bias, - self.fc2.bias, - activation=self.activation, - save_pre_act=self.training, - return_residual=self.return_residual, - checkpoint_lvl=self.checkpoint_lvl, - heuristic=heuristic, - process_group=process_group, - ) - if self.return_residual: - out, x = out - if process_group is not None: - out = reduce_scatter(out, process_group) - return out if not self.return_residual else (out, x) - - -class ParallelFusedMLP(nn.Module): - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - activation="gelu_approx", - process_group: ProcessGroup = None, - bias1=True, - bias2=True, - sequence_parallel=True, - checkpoint_lvl=0, - heuristic="auto", - device=None, - dtype=None, - ): - """ - process_group is required. We're doing Tensor Parallel with sequence parallelism: - we do an all_gather of x before doing the matmul, gelu, then matmul. - Finally we do a reduce_scatter of the output. - - checkpoint_lvl (increasing lvl means slower but more memory saving): - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute pre_act and gelu_out in the bwd - heuristic: - -1: don't fuse gemm + gelu (separate kernel) - 0..4: use this heuristic for the algo section in the fused gemm + gelu - 'auto': heuristic will be picked automatically: - For CUDA >= 11.8, we set heuristic=0 for both fp16 and bf16 for best perf. - For CUDA <= 11.7, we set heuristic=1 for fp16 and heuristic=-1 for bf16. - """ - assert checkpoint_lvl in [0, 1, 2] - assert activation in ["gelu_approx", "relu", "sqrelu"] - assert process_group is not None - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features * 4 - self.activation = activation - self.process_group = process_group - self.sequence_parallel = sequence_parallel - self.checkpoint_lvl = checkpoint_lvl - self.heuristic = heuristic if activation != "sqrelu" else -1 - self.fc1 = ColumnParallelLinear( - in_features, hidden_features, process_group, bias=bias1, **factory_kwargs - ) - self.fc2 = RowParallelLinear( - hidden_features, out_features, process_group, bias=bias2, **factory_kwargs - ) - - def forward(self, x): - dtype = x.dtype if not torch.is_autocast_enabled() else torch.get_autocast_gpu_dtype() - if self.heuristic == "auto": - if self.activation == "gelu_approx": - cuda_ver = tuple(map(int, torch.version.cuda.split("."))) - heuristic = 0 if cuda_ver >= (11, 8) else (1 if dtype == torch.float16 else -1) - else: - heuristic = 0 - else: - heuristic = self.heuristic - out = fused_mlp_func( - x, - self.fc1.weight, - self.fc2.weight, - self.fc1.bias, - self.fc2.bias, - activation=self.activation, - save_pre_act=self.training, - checkpoint_lvl=self.checkpoint_lvl, - heuristic=heuristic, - process_group=self.process_group, - sequence_parallel=self.sequence_parallel, - ) - reduce_fn = reduce_scatter if self.sequence_parallel else all_reduce - return reduce_fn(out, self.process_group) diff --git a/build/torch212-cxx11-cu132-x86_64-linux/ops/layer_norm.py b/build/torch212-cxx11-cu132-x86_64-linux/ops/layer_norm.py deleted file mode 100644 index 4b6cd798fd02844ef9cd3897f8ab95e490e638bf..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu132-x86_64-linux/ops/layer_norm.py +++ /dev/null @@ -1,800 +0,0 @@ -# Copyright (c) 2022, Tri Dao. -# Adapted from https://github.com/NVIDIA/apex/blob/master/apex/contrib/layer_norm/layer_norm.py - -import dropout_layer_norm -import torch -from torch.nn import init - - -def maybe_align(x, alignment_in_bytes=16): - """Assume that x already has last dim divisible by alignment_in_bytes""" - # TD [2023-07-04] I'm not 100% sure that clone will align the memory - # https://discuss.pytorch.org/t/how-to-ensure-that-tensor-data-ptr-is-aligned-to-16-bytes/183440 - return x if x.data_ptr() % alignment_in_bytes == 0 else x.clone() - - -def _dropout_add_layer_norm_forward( - x0, - residual, - gamma, - beta, - rowscale, - colscale, - dropout_p, - epsilon, - residual_in_fp32=False, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes""" - hidden_size = gamma.numel() - x0mat = x0.view((-1, hidden_size)) - residualmat = residual.view((-1, hidden_size)) if residual is not None else None - rowscale = rowscale.view(-1) if rowscale is not None else None - zmat, xmat, dmask, mu, rsigma = dropout_layer_norm.dropout_add_ln_fwd( - x0mat, - residualmat, - gamma, - beta, - rowscale, - colscale, - None, - None, - dropout_p, - epsilon, - 1.0, - 0, - None, - residual_in_fp32, - is_rms_norm, - ) - # dmask is None if dropout_p == 0.0 - # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype - return zmat, xmat if xmat is not None else x0mat, dmask, mu, rsigma - - -def _dropout_add_layer_norm_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - rowscale, - colscale, - dropout_p, - has_residual, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes - dx == None means that it was a post-norm architecture - (x = drop(x0) + residual was not returned in the fwd). - x0 must not be None if we have colscale. - """ - hidden_size = gamma.numel() - xmat = x.view((-1, hidden_size)) - dzmat = dz.view(xmat.shape) - dxmat = dx.view(xmat.shape) if dx is not None else None - x0mat = x0.view((-1, hidden_size)) if x0 is not None else None - rowscale = rowscale.view(-1) if rowscale is not None else None - if colscale is not None: - assert x0 is not None, "x0 is required to compute the gradient of colscale" - dx0mat, dresidualmat, dgamma, dbeta, _, _, *rest = dropout_layer_norm.dropout_add_ln_bwd( - dzmat, - dxmat, - xmat, - x0mat, - dmask, - mu, - rsigma, - gamma, - rowscale, - colscale, - None, - None, - dropout_p, - 1.0, - 0, - has_residual, - is_rms_norm, - ) - # dresidualmat is None if not has_residual - if colscale is None: - return dx0mat, dresidualmat, dgamma, dbeta - else: - dcolscale = rest[0] - return dx0mat, dresidualmat, dgamma, dbeta, dcolscale - - -def _dropout_add_layer_norm_subset_forward( - x0, - residual, - gamma, - beta, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32=False, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes""" - hidden_size = gamma.numel() - x0mat = x0.view((-1, hidden_size)) - residualmat = residual.view((-1, hidden_size)) if residual is not None else None - x0_subset = x0_subset.view(-1) if x0_subset is not None else None - out_subset = out_subset.view(-1) if out_subset is not None else None - zmat, xmat, dmask, mu, rsigma = dropout_layer_norm.dropout_add_ln_fwd( - x0mat, - residualmat, - gamma, - beta, - None, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - None, - residual_in_fp32, - is_rms_norm, - ) - # dmask is None if dropout_p == 0.0 - # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype - return zmat, xmat if xmat is not None else x0mat, dmask, mu, rsigma - - -def _dropout_add_layer_norm_subset_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - colscale, - x0_subset, - out_subset, - dropout_p, - rowscale_const, - x0_numrows, - has_residual, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes - dx == None means that it was a post-norm architecture - (x = drop(x0) + residual was not returned in the fwd). - x0 must not be None if we have colscale. - """ - hidden_size = gamma.numel() - xmat = x.view((-1, hidden_size)) - dzmat = dz.view(-1, hidden_size) - dxmat = dx.view(xmat.shape) if dx is not None else None - x0mat = x0.view((-1, hidden_size)) if x0 is not None else None - x0_subset = x0_subset.view(-1) if x0_subset is not None else None - out_subset = out_subset.view(-1) if out_subset is not None else None - if colscale is not None: - assert x0 is not None, "x0 is required to compute the gradient of colscale" - dx0mat, dresidualmat, dgamma, dbeta, _, _, *rest = dropout_layer_norm.dropout_add_ln_bwd( - dzmat, - dxmat, - xmat, - x0mat, - dmask, - mu, - rsigma, - gamma, - None, - colscale, - x0_subset, - out_subset, - dropout_p, - rowscale_const, - x0_numrows, - has_residual, - is_rms_norm, - ) - # dresidualmat is None if not has_residual - if colscale is None: - return dx0mat, dresidualmat, dgamma, dbeta - else: - dcolscale = rest[0] - return dx0mat, dresidualmat, dgamma, dbeta, dcolscale - - -def _dropout_add_layer_norm_parallel_residual_forward( - x0, - x1, - residual, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - residual_in_fp32=False, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes""" - hidden_size = gamma0.numel() - x0mat = x0.view((-1, hidden_size)) - x1mat = x1.view((-1, hidden_size)) if x1 is not None else None - residualmat = residual.view((-1, hidden_size)) if residual is not None else None - ( - z0mat, - z1mat, - xmat, - dmask0, - dmask1, - mu, - rsigma, - ) = dropout_layer_norm.dropout_add_ln_parallel_residual_fwd( - x0mat, - x1mat, - residualmat, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - None, - residual_in_fp32, - is_rms_norm, - ) - # dmask0 and dmask1 are None if dropout_p == 0.0 - # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype - return z0mat, z1mat, xmat if xmat is not None else x0mat, dmask0, dmask1, mu, rsigma - - -def _dropout_add_layer_norm_parallel_residual_backward( - dz0, - dz1, - dx, - x, - dmask0, - dmask1, - mu, - rsigma, - gamma0, - gamma1, - dropout_p, - has_x1, - has_residual, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes - dx == None means that it was a post-norm architecture - (x = drop(x0) + residual was not returned in the fwd). - """ - hidden_size = gamma0.numel() - xmat = x.view((-1, hidden_size)) - dz0mat = dz0.view(xmat.shape) - dz1mat = dz1.view(xmat.shape) if dz1 is not None else None - dxmat = dx.view(xmat.shape) if dx is not None else None - ( - dx0mat, - dx1mat, - dresidualmat, - dgamma0, - dbeta0, - dgamma1, - dbeta1, - *rest, - ) = dropout_layer_norm.dropout_add_ln_parallel_residual_bwd( - dz0mat, - dz1mat, - dxmat, - xmat, - dmask0, - dmask1, - mu, - rsigma, - gamma0, - gamma1, - dropout_p, - has_x1, - has_residual, - is_rms_norm, - ) - # dresidualmat is None if not has_residual - return dx0mat, dx1mat, dresidualmat, dgamma0, dbeta0, dgamma1, dbeta1 - - -class DropoutAddLayerNormFn(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x0, - residual, - gamma, - beta, - rowscale, - colscale, - dropout_p, - epsilon, - residual_in_fp32=False, - prenorm=False, - is_rms_norm=False, - return_dmask=False, - ): - x0 = maybe_align(x0.contiguous(), 16) - residual = maybe_align(residual.contiguous(), 16) if residual is not None else None - gamma = maybe_align(gamma.contiguous(), 16) - beta = maybe_align(beta.contiguous(), 16) if beta is not None else None - rowscale = maybe_align(rowscale.contiguous(), 16) if rowscale is not None else None - colscale = maybe_align(colscale.contiguous(), 16) if colscale is not None else None - zmat, xmat, dmask, mu, rsigma = _dropout_add_layer_norm_forward( - x0, - residual, - gamma, - beta, - rowscale, - colscale, - dropout_p, - epsilon, - residual_in_fp32, - is_rms_norm, - ) - # Only need to save x0 if we need to compute gradient wrt colscale - x0_saved = x0 if colscale is not None else None - ctx.save_for_backward( - xmat.view(x0.shape), x0_saved, dmask, gamma, mu, rsigma, rowscale, colscale - ) - ctx.prenorm = prenorm - ctx.dropout_p = dropout_p - ctx.has_residual = residual is not None - ctx.is_rms_norm = is_rms_norm - ctx.has_beta = beta is not None - if not return_dmask: - return ( - zmat.view(x0.shape) if not prenorm else (zmat.view(x0.shape), xmat.view(x0.shape)) - ) - else: - dmask = ( - dmask.view(x0.shape) - if dropout_p > 0.0 - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - ctx.mark_non_differentiable(dmask) - return ( - (zmat.view(x0.shape), dmask) - if not prenorm - else (zmat.view(x0.shape), xmat.view(x0.shape), dmask) - ) - - @staticmethod - def backward(ctx, dz, *args): - # assert dz.is_contiguous() - dz = maybe_align(dz.contiguous(), 16) # this happens! - dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None - x, x0, dmask, gamma, mu, rsigma, rowscale, colscale = ctx.saved_tensors - # x0 is None if colscale is None - dropout_p = ctx.dropout_p - has_residual = ctx.has_residual - dx0mat, dresidualmat, dgamma, dbeta, *rest = _dropout_add_layer_norm_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - rowscale, - colscale, - dropout_p, - has_residual, - ctx.is_rms_norm, - ) - dx0 = dx0mat.view(x.shape) - dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None - dcolscale = rest[0] if colscale is not None else None - return ( - dx0, - dresidual, - dgamma, - dbeta if ctx.has_beta else None, - None, - dcolscale, - None, - None, - None, - None, - None, - None, - ) - - -class DropoutAddLayerNormSubsetFn(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x0, - residual, - gamma, - beta, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32=False, - prenorm=False, - is_rms_norm=False, - return_dmask=False, - ): - x0 = maybe_align(x0.contiguous(), 16) - residual = maybe_align(residual.contiguous(), 16) if residual is not None else None - gamma = maybe_align(gamma.contiguous(), 16) - beta = maybe_align(beta.contiguous(), 16) if beta is not None else None - colscale = maybe_align(colscale.contiguous(), 16) if colscale is not None else None - zmat, xmat, dmask, mu, rsigma = _dropout_add_layer_norm_subset_forward( - x0, - residual, - gamma, - beta, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32, - is_rms_norm, - ) - # Only need to save x0 if we need to compute gradient wrt colscale - x0_saved = x0 if colscale is not None else None - x_shape = (-1, *x0.shape[1:]) - ctx.save_for_backward( - xmat.view(x_shape), x0_saved, dmask, gamma, mu, rsigma, colscale, x0_subset, out_subset - ) - ctx.prenorm = prenorm - ctx.dropout_p = dropout_p - ctx.rowscale_const = rowscale_const - ctx.x0_numrows = x0.shape[:-1].numel() - ctx.has_residual = residual is not None - ctx.is_rms_norm = is_rms_norm - ctx.has_beta = beta is not None - z_shape = (-1, *x0.shape[1:]) - if not return_dmask: - return zmat.view(z_shape) if not prenorm else (zmat.view(z_shape), xmat.view(x0.shape)) - else: - z = zmat.view(z_shape) - dmask = ( - dmask.view(x0.shape) - if dropout_p > 0.0 - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - ctx.mark_non_differentiable(dmask) - return (z, dmask) if not prenorm else (z, xmat.view(x_shape), dmask) - - @staticmethod - def backward(ctx, dz, *args): - # assert dz.is_contiguous() - dz = maybe_align(dz.contiguous(), 16) # this happens! - dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None - x, x0, dmask, gamma, mu, rsigma, colscale, x0_subset, out_subset = ctx.saved_tensors - # x0 is None if colscale is None - dropout_p = ctx.dropout_p - has_residual = ctx.has_residual - dx0mat, dresidualmat, dgamma, dbeta, *rest = _dropout_add_layer_norm_subset_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - colscale, - x0_subset, - out_subset, - dropout_p, - ctx.rowscale_const, - ctx.x0_numrows, - has_residual, - ctx.is_rms_norm, - ) - dx0 = dx0mat.view(-1, *x.shape[1:]) - dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None - dcolscale = rest[0] if colscale is not None else None - return ( - dx0, - dresidual, - dgamma, - dbeta if ctx.has_beta else None, - dcolscale, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) - - -class DropoutAddLayerNormParallelResidualFn(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x0, - x1, - residual, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - residual_in_fp32=False, - prenorm=False, - is_rms_norm=False, - return_dmask=False, - ): - x0 = maybe_align(x0.contiguous(), 16) - x1 = maybe_align(x1.contiguous(), 16) if x1 is not None else None - residual = maybe_align(residual.contiguous(), 16) if residual is not None else None - gamma0 = maybe_align(gamma0.contiguous(), 16) - beta0 = maybe_align(beta0.contiguous(), 16) if beta0 is not None else None - gamma1 = maybe_align(gamma1.contiguous(), 16) if gamma1 is not None else None - beta1 = maybe_align(beta1.contiguous(), 16) if beta1 is not None else None - ( - z0mat, - z1mat, - xmat, - dmask0, - dmask1, - mu, - rsigma, - ) = _dropout_add_layer_norm_parallel_residual_forward( - x0, - x1, - residual, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - residual_in_fp32, - is_rms_norm, - ) - ctx.save_for_backward(xmat.view(x0.shape), dmask0, dmask1, gamma0, gamma1, mu, rsigma) - ctx.prenorm = prenorm - ctx.dropout_p = dropout_p - ctx.has_x1 = x1 is not None - ctx.has_residual = residual is not None - ctx.is_rms_norm = is_rms_norm - ctx.has_beta = beta0 is not None - z = (z0mat.view(x0.shape), z1mat.view(x0.shape) if z1mat is not None else None) - if not return_dmask: - return z if not prenorm else (*z, xmat.view(x0.shape)) - else: - dmask0 = ( - dmask0.view(x0.shape) - if dropout_p > 0.0 - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - dmask1 = ( - dmask1.view(x0.shape) - if dropout_p > 0.0 and x1 is not None - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - ctx.mark_non_differentiable(dmask0) - ctx.mark_non_differentiable(dmask1) - return ( - (*z, dmask0, dmask1) if not prenorm else (*z, xmat.view(x0.shape), dmask0, dmask1) - ) - - @staticmethod - def backward(ctx, dz0, dz1, *args): - dz0 = maybe_align(dz0.contiguous(), 16) # this happens! - dz1 = maybe_align(dz1.contiguous(), 16) if dz1 is not None else None - dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None - x, dmask0, dmask1, gamma0, gamma1, mu, rsigma = ctx.saved_tensors - dropout_p = ctx.dropout_p - has_x1 = ctx.has_x1 - has_residual = ctx.has_residual - ( - dx0mat, - dx1mat, - dresidualmat, - dgamma0, - dbeta0, - dgamma1, - dbeta1, - ) = _dropout_add_layer_norm_parallel_residual_backward( - dz0, - dz1, - dx, - x, - dmask0, - dmask1, - mu, - rsigma, - gamma0, - gamma1, - dropout_p, - has_x1, - has_residual, - ctx.is_rms_norm, - ) - dx0 = dx0mat.view(x.shape) - dx1 = dx1mat.view(x.shape) if dx1mat is not None else None - dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None - return ( - dx0, - dx1, - dresidual, - dgamma0, - dbeta0 if ctx.has_beta else None, - dgamma1, - dbeta1 if ctx.has_beta else None, - None, - None, - None, - None, - None, - None, - ) - - -def layer_norm(x, weight, bias, epsilon): - return DropoutAddLayerNormFn.apply(x, None, weight, bias, None, None, 0.0, epsilon, False) - - -def dropout_add_layer_norm( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - rowscale=None, - layerscale=None, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormFn.apply( - x0, - residual, - weight, - bias, - rowscale, - layerscale, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - False, - return_dropout_mask, - ) - - -def dropout_add_layer_norm_subset( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - layerscale=None, - x0_subset=None, - out_subset=None, - rowscale_const=1.0, - out_numrows=0, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormSubsetFn.apply( - x0, - residual, - weight, - bias, - layerscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32, - prenorm, - False, - return_dropout_mask, - ) - - -def dropout_add_layer_norm_parallel_residual( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormParallelResidualFn.apply( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - False, - return_dropout_mask, - ) - - -class DropoutAddLayerNorm(torch.nn.Module): - def __init__( - self, - hidden_size, - prenorm=False, - p=0.0, - eps=1e-5, - residual_in_fp32=False, - device=None, - dtype=None, - ): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.prenorm = prenorm - self.p = p - self.eps = eps - self.residual_in_fp32 = residual_in_fp32 - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.bias = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.reset_parameters() - - def reset_parameters(self): - init.ones_(self.weight) - init.zeros_(self.bias) - - def forward(self, x0, residual=None): - return dropout_add_layer_norm( - x0, - residual, - self.weight, - self.bias, - self.p if self.training else 0.0, - self.eps, - prenorm=self.prenorm, - residual_in_fp32=self.residual_in_fp32, - ) diff --git a/build/torch212-cxx11-cu132-x86_64-linux/ops/rms_norm.py b/build/torch212-cxx11-cu132-x86_64-linux/ops/rms_norm.py deleted file mode 100644 index 068348d61290e3839dd082b540d898578ba1e8e2..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu132-x86_64-linux/ops/rms_norm.py +++ /dev/null @@ -1,174 +0,0 @@ -# Copyright (c) 2022, Tri Dao. -# Adapted from https://github.com/NVIDIA/apex/blob/master/apex/contrib/layer_norm/layer_norm.py - -import torch -from torch.nn import init - -from flash_attn.ops.layer_norm import ( - DropoutAddLayerNormFn, - DropoutAddLayerNormParallelResidualFn, - DropoutAddLayerNormSubsetFn, -) - - -def rms_norm(x, weight, epsilon): - return DropoutAddLayerNormFn.apply( - x, None, weight, None, None, None, 0.0, epsilon, False, False, True - ) - - -def dropout_add_rms_norm( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - rowscale=None, - layerscale=None, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormFn.apply( - x0, - residual, - weight, - bias, - rowscale, - layerscale, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - True, - return_dropout_mask, - ) - - -def dropout_add_rms_norm_subset( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - layerscale=None, - x0_subset=None, - out_subset=None, - rowscale_const=1.0, - out_numrows=0, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormSubsetFn.apply( - x0, - residual, - weight, - bias, - layerscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32, - prenorm, - True, - return_dropout_mask, - ) - - -def dropout_add_rms_norm_parallel_residual( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormParallelResidualFn.apply( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - True, - return_dropout_mask, - ) - - -class RMSNorm(torch.nn.Module): - def __init__(self, hidden_size, eps=1e-5, device=None, dtype=None): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.eps = eps - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.register_parameter("bias", None) - self.reset_parameters() - - def reset_parameters(self): - init.ones_(self.weight) - - def forward(self, x): - return rms_norm(x, self.weight, self.eps) - - -class DropoutAddRMSNorm(torch.nn.Module): - def __init__( - self, - hidden_size, - prenorm=False, - p=0.0, - eps=1e-5, - residual_in_fp32=False, - device=None, - dtype=None, - ): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.prenorm = prenorm - self.p = p - self.eps = eps - self.residual_in_fp32 = residual_in_fp32 - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.register_parameter("bias", None) - self.reset_parameters() - - def reset_parameters(self): - init.ones_(self.weight) - - def forward(self, x0, residual=None): - return dropout_add_rms_norm( - x0, - residual, - self.weight, - None, - self.p if self.training else 0.0, - self.eps, - prenorm=self.prenorm, - residual_in_fp32=self.residual_in_fp32, - ) diff --git a/build/torch212-cxx11-cu132-x86_64-linux/ops/triton/__init__.py b/build/torch212-cxx11-cu132-x86_64-linux/ops/triton/__init__.py deleted file mode 100644 index 8b137891791fe96927ad78e64b0aad7bded08bdc..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu132-x86_64-linux/ops/triton/__init__.py +++ /dev/null @@ -1 +0,0 @@ - diff --git a/build/torch212-cxx11-cu132-x86_64-linux/ops/triton/cross_entropy.py b/build/torch212-cxx11-cu132-x86_64-linux/ops/triton/cross_entropy.py deleted file mode 100644 index 1b5a415b73f236f3e05fb14b9141959559e18526..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu132-x86_64-linux/ops/triton/cross_entropy.py +++ /dev/null @@ -1,330 +0,0 @@ -# Copyright (c) 2023, Tri Dao. - -from typing import Tuple, Optional, Union - -import torch -import torch.nn.functional as F - -import triton -import triton.language as tl - -# `all_gather_into_tensor` and `reduce_scatter_tensor` are new placeholders for -# `_all_gather_base` and `_reduce_scatter_base`. They require the most recent -# version of PyTorch. The following 2 lines are for backward compatibility with -# older PyTorch. -if "all_gather_into_tensor" not in dir(torch.distributed): - torch.distributed.all_gather_into_tensor = torch.distributed._all_gather_base - - -@triton.heuristics( - { - "HAS_SMOOTHING": lambda args: args["smoothing"] > 0.0, - } -) -@triton.jit -def cross_entropy_fwd_kernel( - loss_ptr, # data ptrs - lse_ptr, - z_loss_ptr, - logits_ptr, - labels_ptr, - smoothing, - logit_scale, - lse_square_scale, - ignore_index, - total_classes, - class_start_idx, # Useful for tensor parallel when each rank only has a subset of classes - n_cols, # shapes - logits_row_stride, # strides - BLOCK_SIZE: tl.constexpr, - HAS_SMOOTHING: tl.constexpr, - # if SPLIT (e.g. tensor parallel), don't include the LSE in the loss since it's not the final LSE - SPLIT: tl.constexpr, - PRECOMPUTED_LSE: tl.constexpr, # If LSE is already computed (also no smoothing and logit_scale == 1.0) -): - row_idx = tl.program_id(0) - logits_ptr = logits_ptr + row_idx * logits_row_stride.to(tl.int64) - sum_logits = 0.0 # For smoothing - if not PRECOMPUTED_LSE: - # Statistics for online softmax - m_i = -float("inf") - l_i = 0.0 - for col_offset in range(0, n_cols, BLOCK_SIZE): - cols = col_offset + tl.arange(0, BLOCK_SIZE) - logits = tl.load(logits_ptr + cols, mask=cols < n_cols, other=-float("inf")).to( - tl.float32 - ) * logit_scale - if HAS_SMOOTHING: - sum_logits += tl.sum(tl.where(cols < n_cols, logits, 0.0)) - m_i_new = tl.maximum(m_i, tl.max(logits)) - l_i = tl.exp(m_i - m_i_new) * l_i + tl.sum(tl.exp(logits - m_i_new)) - m_i = m_i_new - lse = tl.log(l_i) + m_i - tl.store(lse_ptr + row_idx, lse) - else: - lse = tl.load(lse_ptr + row_idx) - label_idx = tl.load(labels_ptr + row_idx) - if label_idx == ignore_index: - loss = 0.0 - z_loss = 0.0 - else: - label_idx -= class_start_idx - if label_idx >= 0 and label_idx < n_cols: - logits_label = tl.load(logits_ptr + label_idx) * logit_scale - if HAS_SMOOTHING: - loss = ( - (lse if not SPLIT else 0.0) - - smoothing * sum_logits / total_classes - - (1 - smoothing) * logits_label - ) - else: - loss = (lse if not SPLIT else 0.0) - logits_label - else: - # If label is out of bounds, we set the CE loss to 0.0. But we still want the smoothing loss - if HAS_SMOOTHING: - loss = smoothing * ((lse if not SPLIT else 0.0) - sum_logits / total_classes) - else: - loss = 0.0 - if not SPLIT: - z_loss = lse_square_scale * lse * lse - loss += z_loss - else: - z_loss = 0.0 - tl.store(loss_ptr + row_idx, loss) - if not SPLIT: - tl.store(z_loss_ptr + row_idx, z_loss) - - -@triton.heuristics( - { - "HAS_SMOOTHING": lambda args: args["smoothing"] > 0.0, - } -) -@triton.jit -def cross_entropy_bwd_kernel( - dlogits_ptr, # data ptrs - dloss_ptr, - logits_ptr, - lse_ptr, - labels_ptr, - smoothing, - logit_scale, - lse_square_scale, - ignore_index, - total_classes, - class_start_idx, # Useful for tensor parallel when each rank only has a subset of classes - n_cols, # shapes - logits_row_stride, # strides - dlogits_row_stride, - dloss_row_stride, - BLOCK_SIZE: tl.constexpr, - HAS_SMOOTHING: tl.constexpr, -): - row_idx = tl.program_id(0) - col_block_idx = tl.program_id(1) - logits_ptr = logits_ptr + row_idx * logits_row_stride.to(tl.int64) - dlogits_ptr = dlogits_ptr + row_idx * dlogits_row_stride.to(tl.int64) - col_offsets = col_block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - label_idx = tl.load(labels_ptr + row_idx) - if label_idx != ignore_index: - dloss = tl.load(dloss_ptr + row_idx * dloss_row_stride) - else: - dloss = 0.0 - logits = tl.load(logits_ptr + col_offsets, mask=col_offsets < n_cols, other=-float("inf")).to( - tl.float32 - ) * logit_scale - lse = tl.load(lse_ptr + row_idx) - probs = tl.exp(logits - lse) - probs += 2.0 * lse_square_scale * lse * probs - label_idx -= class_start_idx - if HAS_SMOOTHING: - smooth_positive = 1.0 - smoothing - smooth_negative = smoothing / total_classes - probs = tl.where(col_offsets == label_idx, probs - smooth_positive, probs) - smooth_negative - else: - probs = tl.where(col_offsets == label_idx, probs - 1.0, probs) - tl.store(dlogits_ptr + col_offsets, (dloss * logit_scale) * probs, mask=col_offsets < n_cols) - - -class CrossEntropyLoss(torch.autograd.Function): - - @staticmethod - def forward( - ctx, - logits, - labels, - precomputed_lse=None, - smoothing=0.0, - logit_scale=1.0, - lse_square_scale=0.0, - ignore_index=-100, - inplace_backward=False, - process_group=None, - ): - # For some reason Triton generates wrong code when labels has dtype long and its address - # is not aligned to 16 bytes. The ld.global.b64 seems to load the wrong label index. - if labels.dtype == torch.long and labels.data_ptr() % 16 != 0: - labels = F.pad(labels, (0, 1))[..., :-1] - assert labels.data_ptr() % 16 == 0 - assert logit_scale > 0.0 - n_rows, n_cols = logits.shape - assert labels.shape == (n_rows,) - world_size = 1 if process_group is None else torch.distributed.get_world_size(process_group) - total_classes = world_size * n_cols - rank = 0 if process_group is None else torch.distributed.get_rank(process_group) - class_start_idx = rank * n_cols - use_precomputed_lse = precomputed_lse is not None and logit_scale == 1.0 and smoothing == 0.0 - - if logits.stride(-1) != 1: - logits = logits.contiguous() - MAX_BLOCK_SIZE = 16 * 1024 - BLOCK_SIZE = min(triton.next_power_of_2(n_cols), MAX_BLOCK_SIZE) - num_warps = ( - 4 - if BLOCK_SIZE < 2048 - else (8 if BLOCK_SIZE < 8192 else (16 if BLOCK_SIZE < 128 * 1024 else 32)) - ) - losses = torch.empty(n_rows, dtype=torch.float, device=logits.device) - if use_precomputed_lse: - assert precomputed_lse.shape == (n_rows,) - lse = precomputed_lse.contiguous() - else: - lse = torch.empty(n_rows, dtype=torch.float, device=logits.device) - z_losses = torch.empty(n_rows, dtype=torch.float, device=logits.device) - # Need this, otherwise Triton tries to launch from cuda:0 and we get - # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?) - with torch.cuda.device(logits.device.index): - cross_entropy_fwd_kernel[(n_rows,)]( - losses, # data ptrs - lse, - z_losses, - logits, - labels, - smoothing, - logit_scale, - lse_square_scale, - ignore_index, - total_classes, - class_start_idx, - n_cols, # shapes - logits.stride(0), # strides - BLOCK_SIZE=BLOCK_SIZE, # constants - SPLIT=world_size > 1, - PRECOMPUTED_LSE=use_precomputed_lse, - num_warps=num_warps, - ) - - if world_size > 1: - # If there's no smoothing, if labels are in the vocab of this partition, losses contains - # - predicted logit, and 0 otherwise. - # If there's smoothing=0.1, for labels in the vocab of this partition, losses contains - # -0.9 * predicted logit - 0.1 * sum logit / total_classes. - # For labels not in the vocab of this partition, losses contains - # -0.1 * sum logit / total_classes. - if world_size > 1: - lse_allgather = torch.empty(world_size, n_rows, dtype=lse.dtype, device=lse.device) - torch.distributed.all_gather_into_tensor(lse_allgather, lse, group=process_group) - handle_losses = torch.distributed.all_reduce( - losses, op=torch.distributed.ReduceOp.SUM, group=process_group, async_op=True - ) - lse = torch.logsumexp(lse_allgather, dim=0) - handle_losses.wait() - # After the allreduce, if there's no smoothing, the total losses are - predicted_logit, - # we just have to add the (global) lse. - # If there's smoothing=0.1, the total losses are - # -0.9 * predicted_logit - 0.1 * sum logit / total_classes. - # Again, we just have to add the (global) lse. - losses += lse - if lse_square_scale != 0.0: - z_losses = lse_square_scale * lse.square() - z_losses.masked_fill_(labels == ignore_index, 0.0) - losses += z_losses - else: - z_losses = torch.zeros_like(losses) - losses.masked_fill_(labels == ignore_index, 0.0) - - ctx.save_for_backward(logits, lse, labels) - ctx.mark_non_differentiable(z_losses) - ctx.smoothing = smoothing - ctx.logit_scale = logit_scale - ctx.lse_square_scale = lse_square_scale - ctx.ignore_index = ignore_index - ctx.total_classes = total_classes - ctx.class_start_idx = class_start_idx - ctx.inplace_backward = inplace_backward - return losses, z_losses - - @staticmethod - def backward(ctx, grad_losses, grad_z_losses): - del grad_z_losses # z_losses are only for logging. - - logits, lse, labels = ctx.saved_tensors - dlogits = logits if ctx.inplace_backward else torch.empty_like(logits) - n_rows, n_cols = logits.shape - BLOCK_SIZE = min(triton.next_power_of_2(n_cols), 4 * 1024) - num_warps = 4 if BLOCK_SIZE < 2048 else (8 if BLOCK_SIZE < 8192 else 16) - grid = lambda META: (n_rows, triton.cdiv(n_cols, META["BLOCK_SIZE"])) # noqa - # Need this, otherwise Triton tries to launch from cuda:0 and we get - # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?) - with torch.cuda.device(logits.device.index): - cross_entropy_bwd_kernel[grid]( - dlogits, # data ptrs - grad_losses, - logits, - lse, - labels, - ctx.smoothing, - ctx.logit_scale, - ctx.lse_square_scale, - ctx.ignore_index, - ctx.total_classes, - ctx.class_start_idx, - n_cols, # shapes - logits.stride(0), # strides - dlogits.stride(0), - grad_losses.stride(0), - BLOCK_SIZE=BLOCK_SIZE, # constants - num_warps=num_warps, - ) - return dlogits, None, None, None, None, None, None, None, None, None - - -def cross_entropy_loss( - logits: torch.Tensor, - labels: torch.Tensor, - precomputed_lse: Optional[torch.Tensor] = None, - label_smoothing: float = 0.0, - logit_scale: float = 1.0, - lse_square_scale: float = 0.0, - ignore_index=-100, - inplace_backward: bool = False, - process_group=None, -) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Arguments: - logits: (batch, vocab_size) - labels: (batch,) - label_smoothing: float - logit_scale: float. Multiply logits by this scale before calculating the loss. - lse_square_scale: float. If > 0, we add lse_square_scale * lse(logits) ^ 2 to the loss. - This is also referred to as "z-loss". - ignore_index: int. If labels == ignore_index, the loss is set to 0.0. - inplace_backward: bool. If True, we do the backward pass in-place by modifying the logits. - This saves memory. - process_group: if not None, we're doing Tensor Parallel: each process is responsible for - one part of the vocab. The loss will be aggregated across processes. - Returns: - losses: (batch,), float - z_losses: (batch,), float - """ - return CrossEntropyLoss.apply( - logits, - labels, - precomputed_lse, - label_smoothing, - logit_scale, - lse_square_scale, - ignore_index, - inplace_backward, - process_group, - ) diff --git a/build/torch212-cxx11-cu132-x86_64-linux/ops/triton/k_activations.py b/build/torch212-cxx11-cu132-x86_64-linux/ops/triton/k_activations.py deleted file mode 100644 index efb83c358eb4a85d069ee340a3c83f418f9a805b..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu132-x86_64-linux/ops/triton/k_activations.py +++ /dev/null @@ -1,162 +0,0 @@ -# Adapted from https://github.com/facebookresearch/xformers/blob/main/xformers/triton/k_activations.py -# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved. -# -# This source code is licensed under the BSD license found in the -# LICENSE file in the root directory of this source tree. - -import math -from enum import Enum -from typing import Optional - -import triton -import triton.language as tl - -_sqrt2pi = math.sqrt(2.0 / math.pi) -_sqrt1_2 = math.sqrt(1.0 / 2) -_gaussian_pdf_normalization = 1.0 / math.sqrt(2 * math.pi) - - -class Activation(str, Enum): - SquaredReLU = "squared_relu" - GeLU = "gelu" - GeLUApprox = "gelu_approx" - LeakyReLU = "leaky_relu" - ReLU = "relu" - - -def get_triton_activation_kernel(activation: Optional[Activation]): - return ( - { - Activation.ReLU: relu, - Activation.LeakyReLU: leaky_relu, - Activation.GeLU: gelu, - Activation.GeLUApprox: gelu_approx, - Activation.SquaredReLU: squared_relu, - }[activation] - if activation - else None - ) - - -def get_triton_activation_bwd_kernel(activation: Optional[Activation]): - return ( - { - Activation.ReLU: relu_grad, - Activation.LeakyReLU: leaky_relu_grad, - Activation.GeLU: gelu_grad, - Activation.GeLUApprox: gelu_approx_grad, - Activation.SquaredReLU: squared_relu_grad, - }[activation] - if activation - else None - ) - - -@triton.jit -def tanh(x): - # Tanh is just a scaled sigmoid - return 2 * tl.sigmoid(2 * x) - 1 - - -@triton.jit -def cosh(x): - exp_x = tl.exp(x) - return (exp_x + 1.0 / exp_x) * 0.5 - - -# a Triton implementation of the most used activations -# See for instance http://arxiv.org/abs/1606.08415 for an overview - -# ReLU -@triton.jit -def relu(x): - """ - ReLU_ activation function - - .. _ReLU: https://pytorch.org/docs/stable/generated/torch.nn.ReLU.html - """ - zero = 0.0 - return tl.where(x >= 0, x, zero.to(x.dtype)) - - -@triton.jit -def relu_grad(x): - # ReLU is different from other activations - # in that it does not require the input to retrospectively compute its gradient - # here the input is the downstream gradient, and we return the upstream gradient directly - zero = 0.0 - one = 1.0 - return tl.where(x >= 0, one.to(x.dtype), zero.to(x.dtype)) - - -@triton.jit -def squared_relu(x): - """ - Squared ReLU activation, as proposed in the Primer_ paper. - - .. _Primer: https://arxiv.org/abs/2109.08668 - """ - x_ = relu(x) - return (x_ * x_).to(x.dtype) - - -@triton.jit -def squared_relu_grad(x): - return tl.where(x >= 0, 2.0 * x, 0.0) - - -# Leaky ReLU -@triton.jit -def leaky_relu(x): - """ - LeakyReLU_ activation - - .. _LeakyReLU: https://pytorch.org/docs/stable/generated/torch.nn.LeakyReLU.html - """ - scale = 0.01 + 0.0 - scale = scale.to(x.dtype) - return tl.where(x >= 0, x, scale * x) - - -@triton.jit -def leaky_relu_grad(x): - min_grad = 0.01 - max_grad = 1 - - min_grad = min_grad.to(x.dtype) - max_grad = max_grad.to(x.dtype) - - return tl.where(x >= 0, max_grad, min_grad) - - -@triton.jit -def gelu(x): - """Gaussian Error Linear Unit (GELU)""" - return x * 0.5 * (1.0 + tl.libdevice.erf(x * _sqrt1_2)) - - -@triton.jit -def gelu_grad(x): - cdf = 0.5 * (1.0 + tl.libdevice.erf(x * _sqrt1_2)) - pdf = tl.exp(-0.5 * x * x) * _gaussian_pdf_normalization - return cdf + x * pdf - - -@triton.jit -def gelu_approx(x): - """ - GeLU_ activation - Gaussian error linear unit, with tanh approximation - - .. _GeLU: https://arxiv.org/pdf/1606.08415.pdf - """ - return 0.5 * x * (1.0 + tanh(_sqrt2pi * x * (1.0 + 0.044715 * x * x))) - - -@triton.jit -def gelu_approx_grad(x): - # CREDITS: Fast implementation proposed in - # https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/fused_bias_gelu.py#L30 - tanh_out = tanh(0.79788456 * x * (1 + 0.044715 * x * x)) - return 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( - 1 + tanh_out - ) diff --git a/build/torch212-cxx11-cu132-x86_64-linux/ops/triton/layer_norm.py b/build/torch212-cxx11-cu132-x86_64-linux/ops/triton/layer_norm.py deleted file mode 100644 index 5e6f0b3697b612e6a19031bbc55ac76ad233dc85..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu132-x86_64-linux/ops/triton/layer_norm.py +++ /dev/null @@ -1,1254 +0,0 @@ -# Copyright (c) 2024, Tri Dao. -# Implement dropout + residual + layer_norm / rms_norm. - -# Based on the Triton LayerNorm tutorial: https://triton-lang.org/main/getting-started/tutorials/05-layer-norm.html -# For the backward pass, we keep weight_grad and bias_grad in registers and accumulate. -# This is faster for dimensions up to 8k, but after that it's much slower due to register spilling. -# The models we train have hidden dim up to 8k anyway (e.g. Llama 70B), so this is fine. - -import math -from typing import Optional, List - -import torch -import torch.nn.functional as F -from torch import Tensor - -import triton -import triton.language as tl - -from flash_attn.utils.torch import custom_fwd, custom_bwd -from flash_attn.utils.library import triton_op - -from ..._ops import add_op_namespace_prefix - - -def maybe_contiguous_lastdim(x): - return x.contiguous() if x is not None and x.stride(-1) != 1 else x - - -def maybe_contiguous(x): - return x.contiguous() if x is not None else None - - -def triton_autotune_configs(): - # Return configs with a valid warp count for the current device - configs = [] - # Maximum threads per block is architecture-dependent in theory, but in reality all are 1024 - max_threads_per_block = 1024 - # Default to warp size 32 if not defined by device - warp_size = getattr(torch.cuda.get_device_properties(torch.cuda.current_device()), "warp_size", 32) - # Autotune for warp counts which are powers of 2 and do not exceed thread per block limit - return [triton.Config({}, num_warps=warp_count) for warp_count in [1, 2, 4, 8, 16, 32] - if warp_count * warp_size <= max_threads_per_block] - # return [triton.Config({}, num_warps=8)] - - -def layer_norm_ref( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - zero_centered_weight=False, - dropout_mask=None, - dropout_mask1=None, - upcast=False, -): - dtype = x.dtype - if upcast: - x = x.float() - weight = weight.float() - bias = bias.float() if bias is not None else None - residual = residual.float() if residual is not None else residual - x1 = x1.float() if x1 is not None else None - weight1 = weight1.float() if weight1 is not None else None - bias1 = bias1.float() if bias1 is not None else None - if zero_centered_weight: - weight = weight + 1.0 - if weight1 is not None: - weight1 = weight1 + 1.0 - if x1 is not None: - assert rowscale is None, "rowscale is not supported with parallel LayerNorm" - if rowscale is not None: - x = x * rowscale[..., None] - if dropout_p > 0.0: - if dropout_mask is not None: - x = x.masked_fill(~dropout_mask, 0.0) / (1.0 - dropout_p) - else: - x = F.dropout(x, p=dropout_p) - if x1 is not None: - if dropout_mask1 is not None: - x1 = x1.masked_fill(~dropout_mask1, 0.0) / (1.0 - dropout_p) - else: - x1 = F.dropout(x1, p=dropout_p) - if x1 is not None: - x = x + x1 - if residual is not None: - x = (x + residual).to(x.dtype) - out = F.layer_norm(x.to(weight.dtype), x.shape[-1:], weight=weight, bias=bias, eps=eps).to( - dtype - ) - if weight1 is None: - return out if not prenorm else (out, x) - else: - out1 = F.layer_norm( - x.to(weight1.dtype), x.shape[-1:], weight=weight1, bias=bias1, eps=eps - ).to(dtype) - return (out, out1) if not prenorm else (out, out1, x) - - -def rms_norm_ref( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - zero_centered_weight=False, - dropout_mask=None, - dropout_mask1=None, - upcast=False, -): - dtype = x.dtype - if upcast: - x = x.float() - weight = weight.float() - bias = bias.float() if bias is not None else None - residual = residual.float() if residual is not None else residual - x1 = x1.float() if x1 is not None else None - weight1 = weight1.float() if weight1 is not None else None - bias1 = bias1.float() if bias1 is not None else None - if zero_centered_weight: - weight = weight + 1.0 - if weight1 is not None: - weight1 = weight1 + 1.0 - if x1 is not None: - assert rowscale is None, "rowscale is not supported with parallel LayerNorm" - if rowscale is not None: - x = x * rowscale[..., None] - if dropout_p > 0.0: - if dropout_mask is not None: - x = x.masked_fill(~dropout_mask, 0.0) / (1.0 - dropout_p) - else: - x = F.dropout(x, p=dropout_p) - if x1 is not None: - if dropout_mask1 is not None: - x1 = x1.masked_fill(~dropout_mask1, 0.0) / (1.0 - dropout_p) - else: - x1 = F.dropout(x1, p=dropout_p) - if x1 is not None: - x = x + x1 - if residual is not None: - x = (x + residual).to(x.dtype) - rstd = 1 / torch.sqrt((x.square()).mean(dim=-1, keepdim=True) + eps) - out = ((x * rstd * weight) + bias if bias is not None else (x * rstd * weight)).to(dtype) - if weight1 is None: - return out if not prenorm else (out, x) - else: - out1 = ((x * rstd * weight1) + bias1 if bias1 is not None else (x * rstd * weight1)).to( - dtype - ) - return (out, out1) if not prenorm else (out, out1, x) - - -@triton.autotune( - configs=triton_autotune_configs(), - key=["N", "HAS_RESIDUAL", "STORE_RESIDUAL_OUT", "IS_RMS_NORM", "HAS_BIAS", "HAS_X1", "HAS_W1", "HAS_B1"], -) -# torch compile doesn't like triton.heuristics, so we set these manually when calling the kernel -# @triton.heuristics({"HAS_BIAS": lambda args: args["B"] is not None}) -# @triton.heuristics({"HAS_RESIDUAL": lambda args: args["RESIDUAL"] is not None}) -# @triton.heuristics({"HAS_X1": lambda args: args["X1"] is not None}) -# @triton.heuristics({"HAS_W1": lambda args: args["W1"] is not None}) -# @triton.heuristics({"HAS_B1": lambda args: args["B1"] is not None}) -@triton.jit -def _layer_norm_fwd_1pass_kernel( - X, # pointer to the input - Y, # pointer to the output - W, # pointer to the weights - B, # pointer to the biases - RESIDUAL, # pointer to the residual - X1, - W1, - B1, - Y1, - RESIDUAL_OUT, # pointer to the residual - ROWSCALE, - SEEDS, # Dropout seeds for each row - DROPOUT_MASK, - DROPOUT_MASK1, - Mean, # pointer to the mean - Rstd, # pointer to the 1/std - stride_x_row, # how much to increase the pointer when moving by 1 row - stride_y_row, - stride_res_row, - stride_res_out_row, - stride_x1_row, - stride_y1_row, - M, # number of rows in X - N, # number of columns in X - eps, # epsilon to avoid division by zero - dropout_p, # Dropout probability - zero_centered_weight, # If true, add 1.0 to the weight - IS_RMS_NORM: tl.constexpr, - BLOCK_N: tl.constexpr, - HAS_RESIDUAL: tl.constexpr, - STORE_RESIDUAL_OUT: tl.constexpr, - HAS_BIAS: tl.constexpr, - HAS_DROPOUT: tl.constexpr, - STORE_DROPOUT_MASK: tl.constexpr, - HAS_ROWSCALE: tl.constexpr, - HAS_X1: tl.constexpr, - HAS_W1: tl.constexpr, - HAS_B1: tl.constexpr, -): - # Map the program id to the row of X and Y it should compute. - row = tl.program_id(0) - X += row * stride_x_row - Y += row * stride_y_row - if HAS_RESIDUAL: - RESIDUAL += row * stride_res_row - if STORE_RESIDUAL_OUT: - RESIDUAL_OUT += row * stride_res_out_row - if HAS_X1: - X1 += row * stride_x1_row - if HAS_W1: - Y1 += row * stride_y1_row - # Compute mean and variance - cols = tl.arange(0, BLOCK_N) - x = tl.load(X + cols, mask=cols < N, other=0.0).to(tl.float32) - if HAS_ROWSCALE: - rowscale = tl.load(ROWSCALE + row).to(tl.float32) - x *= rowscale - if HAS_DROPOUT: - # Compute dropout mask - # 7 rounds is good enough, and reduces register pressure - keep_mask = tl.rand(tl.load(SEEDS + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - x = tl.where(keep_mask, x / (1.0 - dropout_p), 0.0) - if STORE_DROPOUT_MASK: - tl.store(DROPOUT_MASK + row * N + cols, keep_mask, mask=cols < N) - if HAS_X1: - x1 = tl.load(X1 + cols, mask=cols < N, other=0.0).to(tl.float32) - if HAS_ROWSCALE: - rowscale = tl.load(ROWSCALE + M + row).to(tl.float32) - x1 *= rowscale - if HAS_DROPOUT: - # Compute dropout mask - # 7 rounds is good enough, and reduces register pressure - keep_mask = ( - tl.rand(tl.load(SEEDS + M + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - ) - x1 = tl.where(keep_mask, x1 / (1.0 - dropout_p), 0.0) - if STORE_DROPOUT_MASK: - tl.store(DROPOUT_MASK1 + row * N + cols, keep_mask, mask=cols < N) - x += x1 - if HAS_RESIDUAL: - residual = tl.load(RESIDUAL + cols, mask=cols < N, other=0.0).to(tl.float32) - x += residual - if STORE_RESIDUAL_OUT: - tl.store(RESIDUAL_OUT + cols, x, mask=cols < N) - if not IS_RMS_NORM: - mean = tl.sum(x, axis=0) / N - tl.store(Mean + row, mean) - xbar = tl.where(cols < N, x - mean, 0.0) - var = tl.sum(xbar * xbar, axis=0) / N - else: - xbar = tl.where(cols < N, x, 0.0) - var = tl.sum(xbar * xbar, axis=0) / N - rstd = 1 / tl.sqrt(var + eps) - tl.store(Rstd + row, rstd) - # Normalize and apply linear transformation - mask = cols < N - w = tl.load(W + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w += 1.0 - if HAS_BIAS: - b = tl.load(B + cols, mask=mask).to(tl.float32) - x_hat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd - y = x_hat * w + b if HAS_BIAS else x_hat * w - # Write output - tl.store(Y + cols, y, mask=mask) - if HAS_W1: - w1 = tl.load(W1 + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w1 += 1.0 - if HAS_B1: - b1 = tl.load(B1 + cols, mask=mask).to(tl.float32) - y1 = x_hat * w1 + b1 if HAS_B1 else x_hat * w1 - tl.store(Y1 + cols, y1, mask=mask) - - -def _layer_norm_fwd( - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - residual: Optional[Tensor] = None, - x1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - out_dtype: Optional[torch.dtype] = None, - residual_dtype: Optional[torch.dtype] = None, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - return_dropout_mask: bool = False, - out: Optional[Tensor] = None, - residual_out: Optional[Tensor] = None -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - # Need to wrap to handle the case where residual_out is a alias of x, which makes torch.library - # and torch.compile unhappy. Also allocate memory for out and residual_out if they are None - # so that _layer_norm_fwd_impl doesn't have to return them. - if out is None: - out = torch.empty_like(x, dtype=x.dtype if out_dtype is None else out_dtype) - if residual is not None: - residual_dtype = residual.dtype - if residual_out is None and ( - residual is not None - or (residual_dtype is not None and residual_dtype != x.dtype) - or dropout_p > 0.0 - or rowscale is not None - or x1 is not None - ): - residual_out = torch.empty_like( - x, dtype=residual_dtype if residual_dtype is not None else x.dtype - ) - else: - residual_out = None - y1, mean, rstd, seeds, dropout_mask, dropout_mask1 = _layer_norm_fwd_impl( - x, - weight, - bias, - eps, - out, - residual=residual, - x1=x1, - weight1=weight1, - bias1=bias1, - dropout_p=dropout_p, - rowscale=rowscale, - zero_centered_weight=zero_centered_weight, - is_rms_norm=is_rms_norm, - return_dropout_mask=return_dropout_mask, - residual_out=residual_out, - ) - # residual_out is None if residual is None and residual_dtype == input_dtype and dropout_p == 0.0 - if residual_out is None: - residual_out = x - return out, y1, mean, rstd, residual_out, seeds, dropout_mask, dropout_mask1 - - -# [2025-04-28] torch.library.triton_op ignores the schema argument, but here we need the schema -# since we're returning a tuple of tensors -@triton_op(add_op_namespace_prefix("layer_norm_fwd_impl"), mutates_args={"out", "residual_out"}, - schema="(Tensor x, Tensor weight, Tensor bias, float eps, Tensor(a!) out, Tensor? residual, Tensor? x1, Tensor? weight1, Tensor? bias1, float dropout_p, Tensor? rowscale, bool zero_centered_weight, bool is_rms_norm, bool return_dropout_mask, Tensor(a!)? residual_out) -> (Tensor y1, Tensor mean, Tensor rstd, Tensor seeds, Tensor dropout_mask, Tensor dropout_mask1)") -def _layer_norm_fwd_impl( - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - out: Tensor, - residual: Optional[Tensor] = None, - x1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - return_dropout_mask: bool = False, - residual_out: Optional[Tensor] = None -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - M, N = x.shape - assert x.stride(-1) == 1 - if residual is not None: - assert residual.stride(-1) == 1 - assert residual.shape == (M, N) - assert weight.shape == (N,) - assert weight.stride(-1) == 1 - if bias is not None: - assert bias.stride(-1) == 1 - assert bias.shape == (N,) - if x1 is not None: - assert x1.shape == x.shape - assert rowscale is None - assert x1.stride(-1) == 1 - if weight1 is not None: - assert weight1.shape == (N,) - assert weight1.stride(-1) == 1 - if bias1 is not None: - assert bias1.shape == (N,) - assert bias1.stride(-1) == 1 - if rowscale is not None: - assert rowscale.is_contiguous() - assert rowscale.shape == (M,) - assert out.shape == x.shape - assert out.stride(-1) == 1 - if residual_out is not None: - assert residual_out.shape == x.shape - assert residual_out.stride(-1) == 1 - if weight1 is not None: - y1 = torch.empty_like(out) - assert y1.stride(-1) == 1 - else: - y1 = None - mean = torch.empty((M,), dtype=torch.float32, device=x.device) if not is_rms_norm else None - rstd = torch.empty((M,), dtype=torch.float32, device=x.device) - if dropout_p > 0.0: - seeds = torch.randint( - 2**32, (M if x1 is None else 2 * M,), device=x.device, dtype=torch.int64 - ) - else: - seeds = None - if return_dropout_mask and dropout_p > 0.0: - dropout_mask = torch.empty(M, N, device=x.device, dtype=torch.bool) - if x1 is not None: - dropout_mask1 = torch.empty(M, N, device=x.device, dtype=torch.bool) - else: - dropout_mask1 = None - else: - dropout_mask, dropout_mask1 = None, None - # Less than 64KB per feature: enqueue fused kernel - MAX_FUSED_SIZE = 65536 // x.element_size() - BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) - if N > BLOCK_N: - raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") - with torch.cuda.device(x.device.index): - torch.library.wrap_triton(_layer_norm_fwd_1pass_kernel)[(M,)]( - x, - out, - weight, - bias, - residual, - x1, - weight1, - bias1, - y1, - residual_out, - rowscale, - seeds, - dropout_mask, - dropout_mask1, - mean, - rstd, - x.stride(0), - out.stride(0), - residual.stride(0) if residual is not None else 0, - residual_out.stride(0) if residual_out is not None else 0, - x1.stride(0) if x1 is not None else 0, - y1.stride(0) if y1 is not None else 0, - M, - N, - eps, - dropout_p, - # Passing bool make torch inductor very unhappy since it then tries to compare to int_max - int(zero_centered_weight), - is_rms_norm, - BLOCK_N, - residual is not None, - residual_out is not None, - bias is not None, - dropout_p > 0.0, - dropout_mask is not None, - rowscale is not None, - HAS_X1=x1 is not None, - HAS_W1=weight1 is not None, - HAS_B1=bias1 is not None, - ) - return y1, mean, rstd, seeds, dropout_mask, dropout_mask1 - - -@triton.autotune( - configs=triton_autotune_configs(), - key=["N", "HAS_DRESIDUAL", "STORE_DRESIDUAL", "IS_RMS_NORM", "HAS_BIAS", "HAS_DROPOUT"], -) -# torch compile doesn't like triton.heuristics, so we set these manually when calling the kernel -# @triton.heuristics({"HAS_BIAS": lambda args: args["B"] is not None}) -# @triton.heuristics({"HAS_DRESIDUAL": lambda args: args["DRESIDUAL"] is not None}) -# @triton.heuristics({"STORE_DRESIDUAL": lambda args: args["DRESIDUAL_IN"] is not None}) -# @triton.heuristics({"HAS_ROWSCALE": lambda args: args["ROWSCALE"] is not None}) -# @triton.heuristics({"HAS_DY1": lambda args: args["DY1"] is not None}) -# @triton.heuristics({"HAS_DX1": lambda args: args["DX1"] is not None}) -# @triton.heuristics({"HAS_B1": lambda args: args["DB1"] is not None}) -# @triton.heuristics({"RECOMPUTE_OUTPUT": lambda args: args["Y"] is not None}) -@triton.jit -def _layer_norm_bwd_kernel( - X, # pointer to the input - W, # pointer to the weights - B, # pointer to the biases - Y, # pointer to the output to be recomputed - DY, # pointer to the output gradient - DX, # pointer to the input gradient - DW, # pointer to the partial sum of weights gradient - DB, # pointer to the partial sum of biases gradient - DRESIDUAL, - W1, - DY1, - DX1, - DW1, - DB1, - DRESIDUAL_IN, - ROWSCALE, - SEEDS, - Mean, # pointer to the mean - Rstd, # pointer to the 1/std - stride_x_row, # how much to increase the pointer when moving by 1 row - stride_y_row, - stride_dy_row, - stride_dx_row, - stride_dres_row, - stride_dy1_row, - stride_dx1_row, - stride_dres_in_row, - M, # number of rows in X - N, # number of columns in X - eps, # epsilon to avoid division by zero - dropout_p, - zero_centered_weight, - rows_per_program, - IS_RMS_NORM: tl.constexpr, - BLOCK_N: tl.constexpr, - HAS_DRESIDUAL: tl.constexpr, - STORE_DRESIDUAL: tl.constexpr, - HAS_BIAS: tl.constexpr, - HAS_DROPOUT: tl.constexpr, - HAS_ROWSCALE: tl.constexpr, - HAS_DY1: tl.constexpr, - HAS_DX1: tl.constexpr, - HAS_B1: tl.constexpr, - RECOMPUTE_OUTPUT: tl.constexpr, -): - # Map the program id to the elements of X, DX, and DY it should compute. - row_block_id = tl.program_id(0) - row_start = row_block_id * rows_per_program - # Do not early exit if row_start >= M, because we need to write DW and DB - cols = tl.arange(0, BLOCK_N) - mask = cols < N - X += row_start * stride_x_row - if HAS_DRESIDUAL: - DRESIDUAL += row_start * stride_dres_row - if STORE_DRESIDUAL: - DRESIDUAL_IN += row_start * stride_dres_in_row - DY += row_start * stride_dy_row - DX += row_start * stride_dx_row - if HAS_DY1: - DY1 += row_start * stride_dy1_row - if HAS_DX1: - DX1 += row_start * stride_dx1_row - if RECOMPUTE_OUTPUT: - Y += row_start * stride_y_row - w = tl.load(W + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w += 1.0 - if RECOMPUTE_OUTPUT and HAS_BIAS: - b = tl.load(B + cols, mask=mask, other=0.0).to(tl.float32) - if HAS_DY1: - w1 = tl.load(W1 + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w1 += 1.0 - dw = tl.zeros((BLOCK_N,), dtype=tl.float32) - if HAS_BIAS: - db = tl.zeros((BLOCK_N,), dtype=tl.float32) - if HAS_DY1: - dw1 = tl.zeros((BLOCK_N,), dtype=tl.float32) - if HAS_B1: - db1 = tl.zeros((BLOCK_N,), dtype=tl.float32) - row_end = min((row_block_id + 1) * rows_per_program, M) - for row in range(row_start, row_end): - # Load data to SRAM - x = tl.load(X + cols, mask=mask, other=0).to(tl.float32) - dy = tl.load(DY + cols, mask=mask, other=0).to(tl.float32) - if HAS_DY1: - dy1 = tl.load(DY1 + cols, mask=mask, other=0).to(tl.float32) - if not IS_RMS_NORM: - mean = tl.load(Mean + row) - rstd = tl.load(Rstd + row) - # Compute dx - xhat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd - xhat = tl.where(mask, xhat, 0.0) - if RECOMPUTE_OUTPUT: - y = xhat * w + b if HAS_BIAS else xhat * w - tl.store(Y + cols, y, mask=mask) - wdy = w * dy - dw += dy * xhat - if HAS_BIAS: - db += dy - if HAS_DY1: - wdy += w1 * dy1 - dw1 += dy1 * xhat - if HAS_B1: - db1 += dy1 - if not IS_RMS_NORM: - c1 = tl.sum(xhat * wdy, axis=0) / N - c2 = tl.sum(wdy, axis=0) / N - dx = (wdy - (xhat * c1 + c2)) * rstd - else: - c1 = tl.sum(xhat * wdy, axis=0) / N - dx = (wdy - xhat * c1) * rstd - if HAS_DRESIDUAL: - dres = tl.load(DRESIDUAL + cols, mask=mask, other=0).to(tl.float32) - dx += dres - # Write dx - if STORE_DRESIDUAL: - tl.store(DRESIDUAL_IN + cols, dx, mask=mask) - if HAS_DX1: - if HAS_DROPOUT: - keep_mask = ( - tl.rand(tl.load(SEEDS + M + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - ) - dx1 = tl.where(keep_mask, dx / (1.0 - dropout_p), 0.0) - else: - dx1 = dx - tl.store(DX1 + cols, dx1, mask=mask) - if HAS_DROPOUT: - keep_mask = tl.rand(tl.load(SEEDS + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - dx = tl.where(keep_mask, dx / (1.0 - dropout_p), 0.0) - if HAS_ROWSCALE: - rowscale = tl.load(ROWSCALE + row).to(tl.float32) - dx *= rowscale - tl.store(DX + cols, dx, mask=mask) - - X += stride_x_row - if HAS_DRESIDUAL: - DRESIDUAL += stride_dres_row - if STORE_DRESIDUAL: - DRESIDUAL_IN += stride_dres_in_row - if RECOMPUTE_OUTPUT: - Y += stride_y_row - DY += stride_dy_row - DX += stride_dx_row - if HAS_DY1: - DY1 += stride_dy1_row - if HAS_DX1: - DX1 += stride_dx1_row - tl.store(DW + row_block_id * N + cols, dw, mask=mask) - if HAS_BIAS: - tl.store(DB + row_block_id * N + cols, db, mask=mask) - if HAS_DY1: - tl.store(DW1 + row_block_id * N + cols, dw1, mask=mask) - if HAS_B1: - tl.store(DB1 + row_block_id * N + cols, db1, mask=mask) - - -def _layer_norm_bwd( - dy: Tensor, - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - mean: Tensor, - rstd: Tensor, - dresidual: Optional[Tensor] = None, - dy1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - seeds: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - has_residual: bool = False, - has_x1: bool = False, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - x_dtype: Optional[torch.dtype] = None, - recompute_output: bool = False, -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - # Need to wrap to handle the case where dresidual_in or dx1 are aliases of x, - # which makes torch.library unhappy - dx, dw, db, dresidual_in, dx1, dw1, db1, y = _layer_norm_bwd_impl( - dy, - x, - weight, - bias, - eps, - mean, - rstd, - dresidual, - dy1, - weight1, - bias1, - seeds, - dropout_p, - rowscale, - has_residual, - has_x1, - zero_centered_weight, - is_rms_norm, - x_dtype=x_dtype, - recompute_output=recompute_output, - ) - # Don't need to compute dresidual_in separately in this case - if has_residual and dx.dtype == x.dtype and dropout_p == 0.0 and rowscale is None: - dresidual_in = dx - if has_x1 and dropout_p == 0.0: - dx1 = dx - return dx, dw, db, dresidual_in, dx1, dw1, db1, y - - - -@triton_op(add_op_namespace_prefix("layer_norm_bwd_impl"), mutates_args={}, - schema="(Tensor dy, Tensor x, Tensor weight, Tensor bias, float eps, Tensor mean, Tensor rstd, Tensor? dresidual, Tensor? dy1, Tensor? weight1, Tensor? bias1, Tensor? seeds, float dropout_p, Tensor? rowscale, bool has_residual, bool has_x1, bool zero_centered_weight, bool is_rms_norm, ScalarType? x_dtype, bool recompute_output) -> (Tensor dx, Tensor dw, Tensor db, Tensor dresidual_in, Tensor dx1, Tensor dw1, Tensor db1, Tensor y)", - allow_decomposition=False, # Don't let torch.compile trace inside - ) -def _layer_norm_bwd_impl( - dy: Tensor, - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - mean: Tensor, - rstd: Tensor, - dresidual: Optional[Tensor] = None, - dy1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - seeds: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - has_residual: bool = False, - has_x1: bool = False, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - x_dtype: Optional[torch.dtype] = None, - recompute_output: bool = False, -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - M, N = x.shape - assert x.stride(-1) == 1 - dy = maybe_contiguous_lastdim(dy) - assert dy.stride(-1) == 1 - assert dy.shape == (M, N) - if dresidual is not None: - dresidual = maybe_contiguous_lastdim(dresidual) - assert dresidual.stride(-1) == 1 - assert dresidual.shape == (M, N) - assert weight.shape == (N,) - assert weight.stride(-1) == 1 - if bias is not None: - assert bias.stride(-1) == 1 - assert bias.shape == (N,) - if dy1 is not None: - dy1 = maybe_contiguous_lastdim(dy1) - assert weight1 is not None - assert dy1.shape == dy.shape - assert dy1.stride(-1) == 1 - if weight1 is not None: - assert weight1.shape == (N,) - assert weight1.stride(-1) == 1 - if bias1 is not None: - assert bias1.shape == (N,) - assert bias1.stride(-1) == 1 - if seeds is not None: - assert seeds.is_contiguous() - assert seeds.shape == (M if not has_x1 else M * 2,) - if rowscale is not None: - assert rowscale.is_contiguous() - assert rowscale.shape == (M,) - # allocate output - dx = ( - torch.empty_like(x) - if x_dtype is None - else torch.empty(M, N, dtype=x_dtype, device=x.device) - ) - dresidual_in = ( - torch.empty_like(x) - if has_residual - and (dx.dtype != x.dtype or dropout_p > 0.0 or rowscale is not None or has_x1) - else None - ) - dx1 = torch.empty_like(dx) if (has_x1 and dropout_p > 0.0) else None - y = torch.empty(M, N, dtype=dy.dtype, device=dy.device) if recompute_output else None - if recompute_output: - assert weight1 is None, "recompute_output is not supported with parallel LayerNorm" - - # Less than 64KB per feature: enqueue fused kernel - MAX_FUSED_SIZE = 65536 // x.element_size() - BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) - if N > BLOCK_N: - raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") - # Increasing the multiple (e.g. 8) will allow more thread blocks to be launched and hide the - # latency of the gmem reads/writes, but will increase the time of summing up dw / db. - sm_count = torch.cuda.get_device_properties(x.device).multi_processor_count * 8 - _dw = torch.empty((sm_count, N), dtype=torch.float32, device=weight.device) - _db = ( - torch.empty((sm_count, N), dtype=torch.float32, device=bias.device) - if bias is not None - else None - ) - _dw1 = torch.empty_like(_dw) if weight1 is not None else None - _db1 = torch.empty_like(_db) if bias1 is not None else None - rows_per_program = math.ceil(M / sm_count) - grid = (sm_count,) - with torch.cuda.device(x.device.index): - torch.library.wrap_triton(_layer_norm_bwd_kernel)[grid]( - x, - weight, - bias, - y, - dy, - dx, - _dw, - _db, - dresidual, - weight1, - dy1, - dx1, - _dw1, - _db1, - dresidual_in, - rowscale, - seeds, - mean, - rstd, - x.stride(0), - 0 if not recompute_output else y.stride(0), - dy.stride(0), - dx.stride(0), - dresidual.stride(0) if dresidual is not None else 0, - dy1.stride(0) if dy1 is not None else 0, - dx1.stride(0) if dx1 is not None else 0, - dresidual_in.stride(0) if dresidual_in is not None else 0, - M, - N, - eps, - dropout_p, - # Passing bool make torch inductor very unhappy since it then tries to compare to int_max - int(zero_centered_weight), - rows_per_program, - is_rms_norm, - BLOCK_N, - dresidual is not None, - dresidual_in is not None, - bias is not None, - dropout_p > 0.0, - HAS_ROWSCALE=rowscale is not None, - HAS_DY1=dy1 is not None, - HAS_DX1=dx1 is not None, - HAS_B1=bias1 is not None, - RECOMPUTE_OUTPUT=y is not None, - ) - dw = _dw.sum(0).to(weight.dtype) - db = _db.sum(0).to(bias.dtype) if bias is not None else None - dw1 = _dw1.sum(0).to(weight1.dtype) if weight1 is not None else None - db1 = _db1.sum(0).to(bias1.dtype) if bias1 is not None else None - # dresidual_in and dx1 could be None, the wrapper will handle assigning them from dx - return dx, dw, db, dresidual_in, dx1, dw1, db1, y - - -class LayerNormFn(torch.autograd.Function): - - @staticmethod - def forward( - ctx, - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - residual_in_fp32=False, - zero_centered_weight=False, - is_rms_norm=False, - return_dropout_mask=False, - out_dtype=None, - out=None, - residual_out=None - ): - x_shape_og = x.shape - # reshape input data into 2D tensor - x = maybe_contiguous_lastdim(x.reshape(-1, x.shape[-1])) - if residual is not None: - assert residual.shape == x_shape_og - residual = maybe_contiguous_lastdim(residual.reshape(-1, residual.shape[-1])) - if x1 is not None: - assert x1.shape == x_shape_og - assert rowscale is None, "rowscale is not supported with parallel LayerNorm" - x1 = maybe_contiguous_lastdim(x1.reshape(-1, x1.shape[-1])) - weight = weight.contiguous() - bias = maybe_contiguous(bias) - weight1 = maybe_contiguous(weight1) - bias1 = maybe_contiguous(bias1) - if rowscale is not None: - rowscale = rowscale.reshape(-1).contiguous() - residual_dtype = ( - residual.dtype - if residual is not None - else (torch.float32 if residual_in_fp32 else None) - ) - if out is not None: - out = out.reshape(-1, out.shape[-1]) - if residual_out is not None: - residual_out = residual_out.reshape(-1, residual_out.shape[-1]) - y, y1, mean, rstd, residual_out, seeds, dropout_mask, dropout_mask1 = _layer_norm_fwd( - x, - weight, - bias, - eps, - residual, - x1, - weight1, - bias1, - dropout_p=dropout_p, - rowscale=rowscale, - out_dtype=out_dtype, - residual_dtype=residual_dtype, - zero_centered_weight=zero_centered_weight, - is_rms_norm=is_rms_norm, - return_dropout_mask=return_dropout_mask, - out=out, - residual_out=residual_out, - ) - ctx.save_for_backward( - residual_out, weight, bias, weight1, bias1, rowscale, seeds, mean, rstd - ) - ctx.x_shape_og = x_shape_og - ctx.eps = eps - ctx.dropout_p = dropout_p - ctx.is_rms_norm = is_rms_norm - ctx.has_residual = residual is not None - ctx.has_x1 = x1 is not None - ctx.prenorm = prenorm - ctx.x_dtype = x.dtype - ctx.zero_centered_weight = zero_centered_weight - y = y.reshape(x_shape_og) - y1 = y1.reshape(x_shape_og) if y1 is not None else None - residual_out = residual_out.reshape(x_shape_og) if residual_out is not None else None - dropout_mask = dropout_mask.reshape(x_shape_og) if dropout_mask is not None else None - dropout_mask1 = dropout_mask1.reshape(x_shape_og) if dropout_mask1 is not None else None - if not return_dropout_mask: - if weight1 is None: - return y if not prenorm else (y, residual_out) - else: - return (y, y1) if not prenorm else (y, y1, residual_out) - else: - if weight1 is None: - return ( - (y, dropout_mask, dropout_mask1) - if not prenorm - else (y, residual_out, dropout_mask, dropout_mask1) - ) - else: - return ( - (y, y1, dropout_mask, dropout_mask1) - if not prenorm - else (y, y1, residual_out, dropout_mask, dropout_mask1) - ) - - @staticmethod - def backward(ctx, dy, *args): - x, weight, bias, weight1, bias1, rowscale, seeds, mean, rstd = ctx.saved_tensors - dy = dy.reshape(-1, dy.shape[-1]) - if weight1 is not None: - dy1, args = args[0], args[1:] - dy1 = dy1.reshape(-1, dy1.shape[-1]) - assert dy1.shape == x.shape - else: - dy1 = None - if ctx.prenorm: - dresidual = args[0] - dresidual = dresidual.reshape(-1, dresidual.shape[-1]) - assert dresidual.shape == x.shape - else: - dresidual = None - dx, dw, db, dresidual_in, dx1, dw1, db1, _ = _layer_norm_bwd( - dy, - x, - weight, - bias, - ctx.eps, - mean, - rstd, - dresidual, - dy1, - weight1, - bias1, - seeds, - ctx.dropout_p, - rowscale, - ctx.has_residual, - ctx.has_x1, - ctx.zero_centered_weight, - ctx.is_rms_norm, - x_dtype=ctx.x_dtype, - recompute_output=False, - ) - return ( - dx.reshape(ctx.x_shape_og), - dw, - db, - dresidual_in.reshape(ctx.x_shape_og) if ctx.has_residual else None, - dx1.reshape(ctx.x_shape_og) if dx1 is not None else None, - dw1, - db1, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) - - -def layer_norm_fn( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - residual_in_fp32=False, - zero_centered_weight=False, - is_rms_norm=False, - return_dropout_mask=False, - out_dtype=None, - out=None, - residual_out=None -): - return LayerNormFn.apply( - x, - weight, - bias, - residual, - x1, - weight1, - bias1, - eps, - dropout_p, - rowscale, - prenorm, - residual_in_fp32, - zero_centered_weight, - is_rms_norm, - return_dropout_mask, - out_dtype, - out, - residual_out - ) - - -def rms_norm_fn( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - residual_in_fp32=False, - zero_centered_weight=False, - return_dropout_mask=False, - out_dtype=None, - out=None, - residual_out=None -): - return LayerNormFn.apply( - x, - weight, - bias, - residual, - x1, - weight1, - bias1, - eps, - dropout_p, - rowscale, - prenorm, - residual_in_fp32, - zero_centered_weight, - True, - return_dropout_mask, - out_dtype, - out, - residual_out - ) - - -class RMSNorm(torch.nn.Module): - - def __init__(self, hidden_size, eps=1e-5, dropout_p=0.0, zero_centered_weight=False, - device=None, dtype=None): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.eps = eps - if dropout_p > 0.0: - self.drop = torch.nn.Dropout(dropout_p) - else: - self.drop = None - self.zero_centered_weight = zero_centered_weight - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.register_parameter("bias", None) - self.reset_parameters() - - def reset_parameters(self): - if not self.zero_centered_weight: - torch.nn.init.ones_(self.weight) - else: - torch.nn.init.zeros_(self.weight) - - def forward(self, x, residual=None, prenorm=False, residual_in_fp32=False): - return rms_norm_fn( - x, - self.weight, - self.bias, - residual=residual, - eps=self.eps, - dropout_p=self.drop.p if self.drop is not None and self.training else 0.0, - prenorm=prenorm, - residual_in_fp32=residual_in_fp32, - zero_centered_weight=self.zero_centered_weight, - ) - - -class LayerNormLinearFn(torch.autograd.Function): - - @staticmethod - @custom_fwd - def forward( - ctx, - x, - norm_weight, - norm_bias, - linear_weight, - linear_bias, - residual=None, - eps=1e-6, - prenorm=False, - residual_in_fp32=False, - is_rms_norm=False, - ): - x_shape_og = x.shape - # reshape input data into 2D tensor - x = maybe_contiguous_lastdim(x.reshape(-1, x.shape[-1])) - if residual is not None: - assert residual.shape == x_shape_og - residual = maybe_contiguous_lastdim(residual.reshape(-1, residual.shape[-1])) - norm_weight = norm_weight.contiguous() - norm_bias = maybe_contiguous(norm_bias) - residual_dtype = ( - residual.dtype - if residual is not None - else (torch.float32 if residual_in_fp32 else None) - ) - y, _, mean, rstd, residual_out, *rest = _layer_norm_fwd( - x, - norm_weight, - norm_bias, - eps, - residual, - out_dtype=None if not torch.is_autocast_enabled() else torch.get_autocast_dtype("cuda"), - residual_dtype=residual_dtype, - is_rms_norm=is_rms_norm, - ) - y = y.reshape(x_shape_og) - dtype = torch.get_autocast_dtype("cuda") if torch.is_autocast_enabled() else y.dtype - linear_weight = linear_weight.to(dtype) - linear_bias = linear_bias.to(dtype) if linear_bias is not None else None - out = F.linear(y.to(linear_weight.dtype), linear_weight, linear_bias) - # We don't store y, will be recomputed in the backward pass to save memory - ctx.save_for_backward(residual_out, norm_weight, norm_bias, linear_weight, mean, rstd) - ctx.x_shape_og = x_shape_og - ctx.eps = eps - ctx.is_rms_norm = is_rms_norm - ctx.has_residual = residual is not None - ctx.prenorm = prenorm - ctx.x_dtype = x.dtype - ctx.linear_bias_is_none = linear_bias is None - return out if not prenorm else (out, residual_out.reshape(x_shape_og)) - - @staticmethod - @custom_bwd - def backward(ctx, dout, *args): - x, norm_weight, norm_bias, linear_weight, mean, rstd = ctx.saved_tensors - dout = dout.reshape(-1, dout.shape[-1]) - dy = F.linear(dout, linear_weight.t()) - dlinear_bias = None if ctx.linear_bias_is_none else dout.sum(0) - dy = maybe_contiguous_lastdim(dy) - assert dy.shape == x.shape - if ctx.prenorm: - dresidual = args[0] - dresidual = maybe_contiguous_lastdim(dresidual.reshape(-1, dresidual.shape[-1])) - assert dresidual.shape == x.shape - else: - dresidual = None - dx, dnorm_weight, dnorm_bias, dresidual_in, _, _, _, y = _layer_norm_bwd( - dy, - x, - norm_weight, - norm_bias, - ctx.eps, - mean, - rstd, - dresidual=dresidual, - has_residual=ctx.has_residual, - is_rms_norm=ctx.is_rms_norm, - x_dtype=ctx.x_dtype, - recompute_output=True, - ) - dlinear_weight = torch.einsum("bo,bi->oi", dout, y) - return ( - dx.reshape(ctx.x_shape_og), - dnorm_weight, - dnorm_bias, - dlinear_weight, - dlinear_bias, - dresidual_in.reshape(ctx.x_shape_og) if ctx.has_residual else None, - None, - None, - None, - None, - ) - - -def layer_norm_linear_fn( - x, - norm_weight, - norm_bias, - linear_weight, - linear_bias, - residual=None, - eps=1e-6, - prenorm=False, - residual_in_fp32=False, - is_rms_norm=False, -): - return LayerNormLinearFn.apply( - x, - norm_weight, - norm_bias, - linear_weight, - linear_bias, - residual, - eps, - prenorm, - residual_in_fp32, - is_rms_norm, - ) diff --git a/build/torch212-cxx11-cu132-x86_64-linux/ops/triton/linear.py b/build/torch212-cxx11-cu132-x86_64-linux/ops/triton/linear.py deleted file mode 100644 index a8966dbc345ab0e593df0124451ee7be3dae131a..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu132-x86_64-linux/ops/triton/linear.py +++ /dev/null @@ -1,594 +0,0 @@ -# Adapted from https://github.com/ELS-RD/kernl/blob/main/src/kernl/implementations/linear_layer.py -# and https://github.com/openai/triton/blob/master/python/triton/ops/matmul.py -from typing import Optional - -import torch -import triton -import triton.language as tl -from triton.ops.matmul_perf_model import early_config_prune, estimate_matmul_time - -from flash_attn.ops.triton.k_activations import ( - gelu, - gelu_approx, - gelu_approx_grad, - gelu_grad, - squared_relu, - squared_relu_grad, -) - -# CREDITS: Initially inspired by the Triton tutorial on matrix multiplications - - -def init_to_zero(name): - return lambda nargs: nargs[name].zero_() - - -def get_configs_io_bound(): - configs = [] - for num_stages in [2, 3, 4, 5, 6]: - for block_m in [16, 32]: - for block_k in [32, 64]: - for block_n in [32, 64, 128, 256]: - num_warps = 2 if block_n <= 64 else 4 - configs.append( - triton.Config( - { - "BLOCK_M": block_m, - "BLOCK_N": block_n, - "BLOCK_K": block_k, - "SPLIT_K": 1, - }, - num_stages=num_stages, - num_warps=num_warps, - ) - ) - # split_k not used - # for split_k in [2, 4, 8, 16]: - # configs.append(triton.Config( - # {'BLOCK_M': block_m, 'BLOCK_N': block_n, 'BLOCK_K': block_k, 'SPLIT_K': split_k}, - # num_stages=num_stages, num_warps=num_warps, pre_hook=init_to_zero('C'))) - return configs - - -@triton.autotune( - configs=[ - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - # good for int8 - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=4, - num_warps=4, - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - ] - + get_configs_io_bound(), - key=["CACHE_KEY_M", "CACHE_KEY_N", "CACHE_KEY_K"], - prune_configs_by={ - "early_config_prune": early_config_prune, - "perf_model": estimate_matmul_time, - "top_k": 10, - }, -) -@triton.heuristics( - { - "EVEN_K": lambda args: args["K"] % (args["BLOCK_K"] * args["SPLIT_K"]) == 0, - } -) -@triton.jit -def kernel_fwd( - C, # Pointers to matrices - ACT_INPUT, - A, - B, - bias, - # Matrix dimensions - M, - N, - K, - CACHE_KEY_M, - CACHE_KEY_N, - CACHE_KEY_K, - # The stride variables represent how much to increase the ptr by when moving by 1 - # element in a particular dimension. E.g. stride_am is how much to increase a_ptr - # by to get the element one row down (A has M rows) - stride_cm, - # stride_cn, # Assume that stride_cn == 1 - stride_am, - stride_ak, - stride_bn, - stride_bk, - # Meta-parameters - BLOCK_M: tl.constexpr, - GROUP_M: tl.constexpr, - BLOCK_N: tl.constexpr, - BLOCK_K: tl.constexpr, - # split k not used, not performant with activation, kept because early_config_prune is expecting it - SPLIT_K: tl.constexpr, - EVEN_K: tl.constexpr, - A_ROWMAJOR: tl.constexpr, - B_COLMAJOR: tl.constexpr, - BIAS: tl.constexpr, - SAVE_ACT_INPUT: tl.constexpr, - ACTIVATION: tl.constexpr, -): - - """ - Kernel for computing Out = activation(A x W + C) - - Input has shape (M, K) - - Weight has shape (K, N) - - Bias has shape (N,) - - Output has shape (M, N) - - ActInputs (optional) has shape (M, N) - 'ActInputs' optionally saves the A x W + C intermediate for backward computations - This kernel will consolidate over K - """ - - pid = tl.program_id(axis=0) - - grid_m = (M + BLOCK_M - 1) // BLOCK_M - grid_n = (N + BLOCK_N - 1) // BLOCK_N - # re-order program ID for better L2 performance - width = GROUP_M * grid_n - group_id = pid // width - group_size = min(grid_m - group_id * GROUP_M, GROUP_M) - pid_m = group_id * GROUP_M + (pid % group_size) - pid_n = (pid % width) // (group_size) - - # now compute the block that each program will go through - # rm (resp. rn) denotes a range of indices - # for rows (resp. col) of C - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - # trick to avoid masking on M and N axis - ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M) - rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N) - rk = tl.arange(0, BLOCK_K) - - if A_ROWMAJOR: - A = A + (ram[:, None] * stride_am + rk[None, :]) - else: - A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak) - if B_COLMAJOR: - B = B + (rk[:, None] + rbn[None, :] * stride_bn) - else: - B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn) - - acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) - - for k in range(K, 0, -BLOCK_K): - if EVEN_K: - a = tl.load(A) - b = tl.load(B) - else: - a = tl.load(A, mask=rk[None, :] < k, other=0.0) - b = tl.load(B, mask=rk[:, None] < k, other=0.0) - acc += tl.dot(a, b) - - if A_ROWMAJOR: - A += BLOCK_K - else: - A += BLOCK_K * stride_ak - if B_COLMAJOR: - B += BLOCK_K - else: - B += BLOCK_K * stride_bk - - # Putting bias after the matmul (instead of before) is faster, idk why - if BIAS: - bias = tl.load(bias + rn, mask=rn < N, other=0.0).to(tl.float32) - acc += bias[None, :] - - # optional: save the activation inputs - if SAVE_ACT_INPUT: - # act_in_ptrs = ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :] * stride_cn - act_in_ptrs = ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :] - tl.store(act_in_ptrs, acc) - - # optional: fused activation (while the data is in shared memory) - if ACTIVATION == "gelu": - acc = gelu(acc) - elif ACTIVATION == "gelu_approx": - acc = gelu_approx(acc) - elif ACTIVATION == "squared_relu": - acc = squared_relu(acc) - # rematerialize rm and rn to save registers - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - - # write back result - # C = C + rm[:, None] * stride_cm + rn[None, :] * stride_cn - C = C + rm[:, None] * stride_cm + rn[None, :] - mask = (rm < M)[:, None] & (rn < N)[None, :] - tl.store(C, acc) - - -def triton_linear_act( - x: torch.Tensor, - weight: torch.Tensor, - bias: Optional[torch.Tensor] = None, - activation: str = "id", - save_act_input: bool = False, -) -> torch.Tensor: - """ - Compute e = activation(x @ weight.T + bias). - This wrapper kicks the `kernel_fwd` Triton kernel - :param x: input tensor - :param weight: weight matrix - :param bias: an optional bias tensor - :param activation: Activation name. Needs to be a Triton kernel. - :param act_input: an optional tensor to save the activation inputs (for backward) - :return: result tensor - """ - # if torch.is_autocast_enabled(): - # dtype = torch.get_autocast_gpu_dtype() - # x, weight, bias = [a.to(dtype=dtype) for a in [x, weight, bias]] - - assert activation in ["id", "gelu", "gelu_approx", "squared_relu"] - - batch_shape, n = x.shape[:-1], x.shape[-1] - batch_dim = batch_shape.numel() - x_reshaped = x.reshape(batch_dim, n) - - if x_reshaped.stride(0) > 1 and x_reshaped.stride(1) > 1: - x_reshaped = x_reshaped.contiguous() - if weight.stride(0) > 1 and weight.stride(1) > 1: - weight = weight.contiguous() - bias = bias.contiguous() if bias is not None else None - - assert ( - x.dtype == weight.dtype - ), f"Input and weight must have the same dtype, got {x.dtype} and {weight.dtype}" - if bias is not None: - assert ( - x.dtype == bias.dtype - ), f"Input and bias must have the same dtype, got {x.dtype} and {bias.dtype}" - assert ( - x_reshaped.shape[1] == weight.shape[1] - ), f"Incompatible dimensions: {x_reshaped.shape} - {weight.shape}" - - assert ( - bias is None or bias.shape[0] == weight.shape[0] - ), "Incompatible dimensions in between weight and bias" - - M, K = x_reshaped.shape - N, K = weight.shape - - output = torch.empty((M, N), device=x.device, dtype=x.dtype) - act_input = torch.empty_like(output) if save_act_input else None - - # 1D launch kernel where each block gets its own program. - grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),) # noqa - - kernel_fwd[grid]( - output, - act_input, - x_reshaped, - weight, # data ptrs - bias if bias is not None else x, # auto skip bias if not present - M, # shapes - N, - K, - M // 32, # key for triton cache (limit number of compilations) - N // 32, - K // 32, - stride_cm=output.stride(0), # strides - # stride_cn=output.stride(1), - stride_am=x_reshaped.stride(0), - stride_ak=x_reshaped.stride(1), - stride_bk=weight.stride(1), - stride_bn=weight.stride(0), - BIAS=bias is not None, # optional fused bias - SAVE_ACT_INPUT=save_act_input, # optional save activation inputs - ACTIVATION=activation, # optional fused activation - A_ROWMAJOR=x_reshaped.stride(1) == 1, - B_COLMAJOR=weight.stride(1) == 1, - GROUP_M=8, # speed optimization: group the programs - ) - - if not save_act_input: - return output.reshape(*batch_shape, output.shape[-1]) - else: - return ( - output.reshape(*batch_shape, output.shape[-1]), - act_input.reshape(*batch_shape, act_input.shape[-1]), - ) - - -@triton.autotune( - configs=[ - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - # good for int8 - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=4, - num_warps=4, - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - ] - + get_configs_io_bound(), - key=["CACHE_KEY_M", "CACHE_KEY_N", "CACHE_KEY_K"], - prune_configs_by={ - "early_config_prune": early_config_prune, - "perf_model": estimate_matmul_time, - "top_k": 10, - }, -) -@triton.heuristics( - { - "EVEN_K": lambda args: args["K"] % (args["BLOCK_K"] * args["SPLIT_K"]) == 0, - } -) -@triton.jit -def kernel_bwd( - C, # Pointers to matrices - ACT_INPUT, - A, - B, - # Matrix dimensions - M, - N, - K, - CACHE_KEY_M, - CACHE_KEY_N, - CACHE_KEY_K, - # The stride variables represent how much to increase the ptr by when moving by 1 - # element in a particular dimension. E.g. stride_am is how much to increase a_ptr - # by to get the element one row down (A has M rows) - stride_cm, - # stride_cn, # Assume that stride_cn == 1 - stride_am, - stride_ak, - stride_bk, - stride_bn, - # Meta-parameters - BLOCK_M: tl.constexpr, - GROUP_M: tl.constexpr, - BLOCK_N: tl.constexpr, - BLOCK_K: tl.constexpr, - # split k not used, not performant with activation, kept because early_config_prune is expecting it - SPLIT_K: tl.constexpr, - EVEN_K: tl.constexpr, - ACTIVATION: tl.constexpr, -): - - """ - Kernel for computing Out = activation(A x W + C) - - Input has shape (M, K) - - Weight has shape (K, N) - - Output has shape (M, N) - - ActInputs (optional) has shape (M, N) - 'ActInputs' optionally saves the A x W + C intermediate for backward computations - This kernel will consolidate over K - """ - - pid = tl.program_id(axis=0) - - grid_m = (M + BLOCK_M - 1) // BLOCK_M - grid_n = (N + BLOCK_N - 1) // BLOCK_N - # re-order program ID for better L2 performance - width = GROUP_M * grid_n - group_id = pid // width - group_size = min(grid_m - group_id * GROUP_M, GROUP_M) - pid_m = group_id * GROUP_M + (pid % group_size) - pid_n = (pid % width) // (group_size) - - # now compute the block that each program will go through - # rm (resp. rn) denotes a range of indices - # for rows (resp. col) of C - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - # trick to avoid masking on M and N axis - ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M) - rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N) - rk = tl.arange(0, BLOCK_K) - - A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak) - B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn) - - acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) - - for k in range(K, 0, -BLOCK_K): - if EVEN_K: - a = tl.load(A) - b = tl.load(B) - else: - a = tl.load(A, mask=rk[None, :] < k, other=0.0) - b = tl.load(B, mask=rk[:, None] < k, other=0.0) - acc += tl.dot(a, b) - - A += BLOCK_K * stride_ak - B += BLOCK_K * stride_bk - - # optional: fused activation (while the data is in shared memory) - if ACTIVATION != "id": - act_in_ptrs = ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :] - act_input = tl.load(act_in_ptrs).to(acc.dtype) - if ACTIVATION == "gelu": - acc *= gelu_grad(act_input) - elif ACTIVATION == "gelu_approx": - acc *= gelu_approx_grad(act_input) - elif ACTIVATION == "squared_relu": - acc *= squared_relu_grad(act_input) - - # rematerialize rm and rn to save registers - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - - # write back result - C = C + rm[:, None] * stride_cm + rn[None, :] - mask = (rm < M)[:, None] & (rn < N)[None, :] - tl.store(C, acc, mask=mask) - - -def triton_dgrad_act( - grad_output: torch.Tensor, - weight: torch.Tensor, - activation: str = "id", - act_input: Optional[torch.Tensor] = None, -) -> torch.Tensor: - """ - Compute e = activation(grad_output @ weight + bias). - This wrapper kicks the `kernel_fwd` Triton kernel - :param grad_output: input tensor - :param weight: weight matrix - :param activation: Activation name. Needs to be a Triton kernel. - :param act_input: an optional tensor to save the activation inputs (for backward) - :return: result tensor - """ - assert activation in ["id", "gelu", "gelu_approx", "squared_relu"] - - batch_shape, n = grad_output.shape[:-1], grad_output.shape[-1] - batch_dim = batch_shape.numel() - grad_output_reshaped = grad_output.reshape(batch_dim, n) - - if grad_output_reshaped.stride(0) > 1 and grad_output_reshaped.stride(1) > 1: - grad_output_reshaped = grad_output_reshaped.contiguous() - if weight.stride(0) > 1 and weight.stride(1) > 1: - weight = weight.contiguous() - - assert ( - grad_output.dtype == weight.dtype - ), f"grad_output and weight must have the same dtype, got {grad_output.dtype} and {weight.dtype}" - assert ( - grad_output_reshaped.shape[1] == weight.shape[0] - ), f"Incompatible dimensions: {grad_output_reshaped.shape} - {weight.shape}" - if activation != "id": - assert act_input is not None, f"act_input is required for activation {activation}" - - # M, N, K in bwd are different from M, N, K in fwd - M, K = grad_output_reshaped.shape - K, N = weight.shape - - grad_input = torch.empty((M, N), device=grad_output.device, dtype=grad_output.dtype) - - # 1D launch kernel where each block gets its own program. - grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),) # noqa - - kernel_bwd[grid]( - grad_input, - act_input, - grad_output_reshaped, - weight, # data ptrs - M, # shapes - N, - K, - M // 32, # key for triton cache (limit number of compilations) - N // 32, - K // 32, - stride_cm=grad_input.stride(0), # strides - # stride_cn=grad_input.stride(1), - stride_am=grad_output_reshaped.stride(0), - stride_ak=grad_output_reshaped.stride(1), - stride_bk=weight.stride(0), - stride_bn=weight.stride(1), - ACTIVATION=activation, # optional fused activation - GROUP_M=8, # speed optimization: group the programs - ) - - return grad_input.reshape(*batch_shape, grad_input.shape[-1]) diff --git a/build/torch212-cxx11-cu132-x86_64-linux/ops/triton/mlp.py b/build/torch212-cxx11-cu132-x86_64-linux/ops/triton/mlp.py deleted file mode 100644 index 059f4f8a5e174c1f4824e43d313fca18eaa799b8..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu132-x86_64-linux/ops/triton/mlp.py +++ /dev/null @@ -1,149 +0,0 @@ -# The triton fused matmul + sqrelu is faster for fp16 but slower for bf16, compared -# to naive implementation. -import fused_dense_lib as fused_dense_cuda -import torch -import torch.nn as nn -import torch.nn.functional as F - -from flash_attn.utils.torch import custom_fwd, custom_bwd -from flash_attn.ops.activations import sqrelu_bwd, sqrelu_fwd -from flash_attn.ops.triton.linear import triton_dgrad_act, triton_linear_act - - -class FusedDenseSqreluDenseFunc(torch.autograd.Function): - @staticmethod - @custom_fwd - def forward(ctx, x, weight1, bias1, weight2, bias2, checkpoint_lvl=0): - """checkpoint_lvl: - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute act_input and gelu_out in the bwd - """ - if torch.is_autocast_enabled(): - dtype = torch.get_autocast_gpu_dtype() - x, weight1, bias1, weight2, bias2 = [ - a.to(dtype=dtype) for a in [x, weight1, bias1, weight2, bias2] - ] - is_bf16 = x.dtype == torch.bfloat16 - assert checkpoint_lvl in [0, 1, 2] - x = x.contiguous() - weight1 = weight1.contiguous() - bias1 = bias1.contiguous() - weight2 = weight2.contiguous() - bias2 = bias2.contiguous() - batch_shape, n = x.shape[:-1], x.shape[-1] - batch_dim = batch_shape.numel() - if is_bf16: - act_input = fused_dense_cuda.linear_bias_forward( - x.reshape(batch_dim, n), weight1, bias1 - ) - output1 = sqrelu_fwd(act_input) - else: - save_act_input = checkpoint_lvl != 2 - result = triton_linear_act( - x.reshape(batch_dim, n), - weight1, - bias1, - activation="squared_relu", - save_act_input=save_act_input, - ) - if save_act_input: - output1, act_input = result - else: - output1 = result - output2 = fused_dense_cuda.linear_bias_forward(output1, weight2, bias2) - ctx.checkpoint_lvl = checkpoint_lvl - if checkpoint_lvl == 0: - ctx.save_for_backward(x, weight1, bias1, weight2, act_input, output1) - elif checkpoint_lvl == 1: - ctx.save_for_backward(x, weight1, bias1, weight2, act_input) - elif checkpoint_lvl == 2: - ctx.save_for_backward(x, weight1, bias1, weight2) - return output2.reshape(*batch_shape, output2.shape[-1]) - - @staticmethod - @custom_bwd - def backward(ctx, grad_output): - grad_output = grad_output.contiguous() - checkpoint_lvl = ctx.checkpoint_lvl - x, weight1, bias1, weight2, *rest = ctx.saved_tensors - batch_shape, n = x.shape[:-1], x.shape[-1] - batch_dim = batch_shape.numel() - is_bf16 = x.dtype == torch.bfloat16 - if checkpoint_lvl == 0: - act_input, output1 = rest - elif checkpoint_lvl == 1: - (act_input,) = rest - output1 = sqrelu_fwd(act_input) - elif checkpoint_lvl == 2: - if is_bf16: - act_input = fused_dense_cuda.linear_bias_forward( - x.reshape(batch_dim, n), weight1, bias1 - ) - output1 = sqrelu_fwd(act_input) - else: - output1, act_input = triton_linear_act( - x.reshape(batch_dim, n), - weight1, - bias1, - activation="squared_relu", - save_act_input=True, - ) - - if is_bf16: - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad(output1, grad_output) - grad_output1 = grad_output @ weight2 - grad_act_input = sqrelu_bwd(grad_output1, act_input) - grad_input, grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_backward( - x.reshape(batch_dim, n), weight1, grad_act_input - ) - else: - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad(output1, grad_output) - grad_act_input = triton_dgrad_act( - grad_output, weight2, activation="squared_relu", act_input=act_input - ) - grad_input, grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_backward( - x.reshape(batch_dim, n), weight1, grad_act_input - ) - return grad_input.reshape_as(x), grad_weight1, grad_bias1, grad_weight2, grad_bias2, None - - -fused_dense_sqrelu_dense_function = FusedDenseSqreluDenseFunc.apply - - -class FusedDenseSqreluDense(nn.Module): - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - bias1=True, - bias2=True, - checkpoint_lvl=0, - device=None, - dtype=None, - ): - """ - checkpoint_lvl (increasing lvl means slower but more memory saving): - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute gelu_in and gelu_out in the bwd - """ - assert checkpoint_lvl in [0, 1, 2] - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features * 4 - assert bias1 == True, "DenseSqreluDense module without bias is currently not supported" - assert bias2 == True, "DenseSqreluDense module without bias is currently not supported" - self.checkpoint_lvl = checkpoint_lvl - self.fc1 = nn.Linear(in_features, hidden_features, bias=bias1, **factory_kwargs) - self.fc2 = nn.Linear(hidden_features, out_features, bias=bias2, **factory_kwargs) - - def forward(self, x): - assert x.is_cuda - return fused_dense_sqrelu_dense_function( - x, self.fc1.weight, self.fc1.bias, self.fc2.weight, self.fc2.bias, self.checkpoint_lvl - ) diff --git a/build/torch212-cxx11-cu132-x86_64-linux/ops/triton/rotary.py b/build/torch212-cxx11-cu132-x86_64-linux/ops/triton/rotary.py deleted file mode 100644 index 6ad3ab412a084270120167ad32b67ad1720f2c37..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu132-x86_64-linux/ops/triton/rotary.py +++ /dev/null @@ -1,186 +0,0 @@ -# Copyright (c) 2025, Tri Dao. -# As of 2025-04-23, we require triton >= 3.0 - -from typing import Optional, Union - -import torch - -import triton -import triton.language as tl - - -@triton.jit -def rotary_kernel( - OUT, # Pointers to matrices - X, - COS, - SIN, - CU_SEQLENS, - SEQLEN_OFFSETS, # this could be int or a pointer - # Matrix dimensions - seqlen, - nheads, - seqlen_ro, - # strides - stride_out_batch, - stride_out_seqlen, - stride_out_nheads, - stride_out_headdim, - stride_x_batch, - stride_x_seqlen, - stride_x_nheads, - stride_x_headdim, - # Meta-parameters - # We want ROTARY_DIM to be constexpr, otherwise the triton compiler doesn't know that - # the mask is constant every 8 elements, and it will generate LDG.16 instead of LDG.128 - ROTARY_DIM: tl.constexpr, - IS_SEQLEN_OFFSETS_TENSOR: tl.constexpr, - IS_VARLEN: tl.constexpr, - INTERLEAVED: tl.constexpr, - CONJUGATE: tl.constexpr, - BLOCK_H: tl.constexpr, - BLOCK_M: tl.constexpr, -): - BLOCK_K: tl.constexpr = triton.next_power_of_2(ROTARY_DIM) - ROTARY_DIM_HALF = ROTARY_DIM // 2 - pid_head = tl.program_id(axis=0) - pid_m = tl.program_id(axis=1) - pid_batch = tl.program_id(axis=2) - - if not IS_VARLEN: - X = X + pid_batch * stride_x_batch - OUT = OUT + pid_batch * stride_out_batch - else: - start_idx = tl.load(CU_SEQLENS + pid_batch) - seqlen = tl.load(CU_SEQLENS + pid_batch + 1) - start_idx - X = X + start_idx * stride_x_seqlen - OUT = OUT + start_idx * stride_out_seqlen - - if pid_m * BLOCK_M >= seqlen: - return - - rh = pid_head * BLOCK_H + tl.arange(0, BLOCK_H) - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - if not IS_SEQLEN_OFFSETS_TENSOR: - rm_cs = rm + SEQLEN_OFFSETS - else: - rm_cs = rm + tl.load(SEQLEN_OFFSETS + pid_batch) - - rk_half = tl.arange(0, BLOCK_K // 2) - COS = COS + (rm_cs[:, None] * ROTARY_DIM_HALF + rk_half[None, :]) - SIN = SIN + (rm_cs[:, None] * ROTARY_DIM_HALF + rk_half[None, :]) - mask_cs = (rm_cs[:, None] < seqlen_ro) & (rk_half[None, :] < ROTARY_DIM_HALF) - cos = tl.load(COS, mask=mask_cs, other=1.0).to(tl.float32) - sin = tl.load(SIN, mask=mask_cs, other=0.0).to(tl.float32) - if CONJUGATE: - sin = -sin - - if not INTERLEAVED: - # Load the 1st and 2nd halves of X, do calculation, then store to 1st and 2nd halves of OUT - X = X + (rh[:, None, None] * stride_x_nheads + rm[None, :, None] * stride_x_seqlen + rk_half[None, None, :] * stride_x_headdim) - OUT = OUT + (rh[:, None, None] * stride_out_nheads + rm[None, :, None] * stride_out_seqlen + rk_half[None, None, :] * stride_out_headdim) - mask = (rh[:, None, None] < nheads) & (rm[None, :, None] < seqlen) & (rk_half[None, None, :] < ROTARY_DIM_HALF) - x0 = tl.load(X, mask=mask, other=0.0).to(tl.float32) - x1 = tl.load(X + ROTARY_DIM_HALF * stride_x_headdim, mask=mask, other=0.0,).to(tl.float32) - o0 = x0 * cos - x1 * sin - o1 = x0 * sin + x1 * cos - tl.store(OUT, o0, mask=mask) - tl.store(OUT + ROTARY_DIM_HALF * stride_out_headdim, o1, mask=mask) - else: - rk = tl.arange(0, BLOCK_K) - X = X + (rh[:, None, None] * stride_x_nheads + rm[None, :, None] * stride_x_seqlen + rk[None, None, :] * stride_x_headdim) - OUT = OUT + (rh[:, None, None] * stride_out_nheads + rm[None, :, None] * stride_out_seqlen + rk[None, None, :] * stride_out_headdim) - mask = (rh[:, None, None] < nheads) & (rm[None, :, None] < seqlen) & (rk[None, None, :] < ROTARY_DIM) - x = tl.load(X, mask=mask, other=0.0).to(tl.float32) - x0, x1 = tl.split(tl.reshape(x, [BLOCK_H, BLOCK_M, BLOCK_K // 2, 2])) - o0 = x0 * cos - x1 * sin - o1 = x0 * sin + x1 * cos - o = tl.reshape(tl.join(o0, o1), [BLOCK_H, BLOCK_M, BLOCK_K]) - tl.store(OUT, o, mask=mask) - - -def apply_rotary( - x: torch.Tensor, - cos: torch.Tensor, - sin: torch.Tensor, - seqlen_offsets: Union[int, torch.Tensor] = 0, - cu_seqlens: Optional[torch.Tensor] = None, - max_seqlen: Optional[int] = None, - interleaved=False, - inplace=False, - conjugate=False, -) -> torch.Tensor: - """ - Arguments: - x: (batch, seqlen, nheads, headdim) if cu_seqlens is None - else (total_seqlen, nheads, headdim). - cos: (seqlen_ro, rotary_dim / 2) - sin: (seqlen_ro, rotary_dim / 2) - seqlen_offsets: integer or integer tensor of size (batch,) - cu_seqlens: (batch + 1,) or None - max_seqlen: int - Returns: - y: (batch, seqlen, nheads, headdim) - """ - is_varlen = cu_seqlens is not None - if not is_varlen: - batch, seqlen, nheads, headdim = x.shape - else: - assert max_seqlen is not None, "If cu_seqlens is passed in, then max_seqlen must be passed" - total_seqlen, nheads, headdim = x.shape - batch_p_1 = cu_seqlens.shape[0] - batch = batch_p_1 - 1 - seqlen = max_seqlen - seqlen_ro, rotary_dim = cos.shape - assert sin.shape == cos.shape - rotary_dim *= 2 - assert rotary_dim <= headdim, "rotary_dim must be <= headdim" - assert headdim <= 256, "Only support headdim <= 256" - assert seqlen_ro >= seqlen, "seqlen_ro must be >= seqlen" - - cos, sin = cos.contiguous(), sin.contiguous() - if isinstance(seqlen_offsets, torch.Tensor): - assert seqlen_offsets.shape == (batch,) - assert seqlen_offsets.dtype in [torch.int32, torch.int64] - seqlen_offsets = seqlen_offsets.contiguous() - else: - assert seqlen_offsets + seqlen <= seqlen_ro - - output = torch.empty_like(x) if not inplace else x - if rotary_dim < headdim and not inplace: - output[..., rotary_dim:].copy_(x[..., rotary_dim:]) - - grid = lambda META: (triton.cdiv(nheads, META["BLOCK_H"]), triton.cdiv(seqlen, META["BLOCK_M"]), batch) # noqa - BLOCK_M = 8 if rotary_dim <= 128 else 4 - - # Need this, otherwise Triton tries to launch from cuda:0 and we get - # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?) - device_ctx = torch.cuda.device(x.device.index) if x.device.type == 'cuda' else torch.xpu.device(x.device.index) - with device_ctx: - torch.library.wrap_triton(rotary_kernel)[grid]( - output, # data ptrs - x, - cos, - sin, - cu_seqlens, - seqlen_offsets, - seqlen, # shapes - nheads, - seqlen_ro, - output.stride(0) if not is_varlen else 0, # batch_strides if not varlen else 0 - output.stride(-3), # seqlen_stride or total_seqlen_stride - output.stride(-2), # nheads_stride - output.stride(-1), # headdim_stride - x.stride(0) if not is_varlen else 0, # batch_strides if not varlen else 0 - x.stride(-3), # seqlen stride or total_seqlen_stride - x.stride(-2), # nheads stride - x.stride(-1), # headdim stride - rotary_dim, - isinstance(seqlen_offsets, torch.Tensor), - is_varlen, - interleaved, - conjugate, - BLOCK_M=BLOCK_M, - BLOCK_H=2, - ) - return output diff --git a/build/torch212-cxx11-xpu20253-x86_64-linux/__init__.py b/build/torch212-cxx11-xpu20253-x86_64-linux/__init__.py deleted file mode 100644 index ecc2f9d896b6c93f90b0a1499856dc0612177422..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-xpu20253-x86_64-linux/__init__.py +++ /dev/null @@ -1,393 +0,0 @@ -from typing import Optional, List -import torch -from ._ops import ops as flash_attn_ops -from .flash_attn_interface import ( - flash_attn_func, - flash_attn_kvpacked_func, - flash_attn_qkvpacked_func, - flash_attn_varlen_func, - flash_attn_varlen_kvpacked_func, - flash_attn_varlen_qkvpacked_func, - flash_attn_with_kvcache, -) - - -def fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - return_softmax: bool = False, - gen: Optional[torch.Generator] = None, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention. - - Args: - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Optional output tensor, same shape as q - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - return_softmax: Whether to return softmax weights - gen: Optional random number generator - - Returns: - List of tensors: [output, softmax_lse, (softmax if return_softmax)] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.fwd( - q, - k, - v, - out, - alibi_slopes, - p_dropout, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - gen, - ) - - -def varlen_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - out: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - block_table: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - max_seqlen_q: int = 0, - max_seqlen_k: int = 0, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - zero_tensors: bool = False, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - return_softmax: bool = False, - gen: Optional[torch.Generator] = None, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention with variable sequence lengths. - - Args: - q: Query tensor of shape [total_q, num_heads, head_size] - k: Key tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - v: Value tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1] - cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1] - out: Optional output tensor of shape [total_q, num_heads, head_size] - seqused_k: Optional tensor specifying how many keys to use per batch element [batch_size] - leftpad_k: Optional left padding for keys of shape [batch_size] - block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq] - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - max_seqlen_q: Maximum sequence length for queries - max_seqlen_k: Maximum sequence length for keys - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - zero_tensors: Whether to zero tensors before computation - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - return_softmax: Whether to return softmax weights - gen: Optional random number generator - - Returns: - List of tensors: [output, softmax_lse, (softmax if return_softmax)] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.varlen_fwd( - q, - k, - v, - out, - cu_seqlens_q, - cu_seqlens_k, - seqused_k, - leftpad_k, - block_table, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - p_dropout, - softmax_scale, - zero_tensors, - is_causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - gen, - ) - - -def bwd( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - p_dropout: float = 0.0, - 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, - gen: Optional[torch.Generator] = None, - rng_state: Optional[torch.Tensor] = None, -) -> List[torch.Tensor]: - """ - Backward pass for multi-head attention. - - Args: - dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size] - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size] - softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q] - dq: Optional gradient tensor for queries, same shape as q - dk: Optional gradient tensor for keys, same shape as k - dv: Optional gradient tensor for values, same shape as v - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - deterministic: Whether to use deterministic algorithms - gen: Optional random number generator - rng_state: Optional RNG state from forward pass - - Returns: - List of tensors: [dq, dk, dv] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - alibi_slopes, - p_dropout, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - deterministic, - gen, - rng_state, - ) - - -def varlen_bwd( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - max_seqlen_q: int = 0, - max_seqlen_k: int = 0, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - zero_tensors: bool = False, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - deterministic: bool = False, - gen: Optional[torch.Generator] = None, - rng_state: Optional[torch.Tensor] = None, -) -> List[torch.Tensor]: - """ - Backward pass for multi-head attention with variable sequence lengths. - - Args: - dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size] - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size] - softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q] - cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1] - cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1] - dq: Optional gradient tensor for queries, same shape as q - dk: Optional gradient tensor for keys, same shape as k - dv: Optional gradient tensor for values, same shape as v - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - max_seqlen_q: Maximum sequence length for queries - max_seqlen_k: Maximum sequence length for keys - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - zero_tensors: Whether to zero tensors before computation - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - deterministic: Whether to use deterministic algorithms - gen: Optional random number generator - rng_state: Optional RNG state from forward pass - - Returns: - List of tensors: [dq, dk, dv] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.varlen_bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - p_dropout, - softmax_scale, - zero_tensors, - is_causal, - window_size_left, - window_size_right, - softcap, - deterministic, - gen, - rng_state, - ) - - -def fwd_kvcache( - q: torch.Tensor, - kcache: torch.Tensor, - vcache: torch.Tensor, - k: Optional[torch.Tensor] = None, - v: Optional[torch.Tensor] = None, - seqlens_k: Optional[torch.Tensor] = None, - rotary_cos: Optional[torch.Tensor] = None, - rotary_sin: Optional[torch.Tensor] = None, - cache_batch_idx: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - block_table: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - out: 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, - is_rotary_interleaved: bool = False, - num_splits: int = 1, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention with KV cache. - - Args: - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - kcache: Key cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - vcache: Value cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - k: Optional new keys tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size] - v: Optional new values tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size] - seqlens_k: Optional sequence lengths for keys of shape [batch_size] - rotary_cos: Optional rotary cosine tensor of shape [seqlen_ro, rotary_dim/2] - rotary_sin: Optional rotary sine tensor of shape [seqlen_ro, rotary_dim/2] - cache_batch_idx: Optional indices to index into the KV cache - leftpad_k: Optional left padding for keys of shape [batch_size] - block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq] - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - out: Optional output tensor, same shape as q - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - is_rotary_interleaved: Whether rotary embeddings are interleaved - num_splits: Number of splits for computation - - Returns: - List of tensors: [output, softmax_lse] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.fwd_kvcache( - q, - kcache, - vcache, - k, - v, - seqlens_k, - rotary_cos, - rotary_sin, - cache_batch_idx, - leftpad_k, - block_table, - alibi_slopes, - out, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - is_rotary_interleaved, - num_splits, - ) diff --git a/build/torch212-cxx11-xpu20253-x86_64-linux/_flash_attn2_xpu_f12afc9.abi3.so b/build/torch212-cxx11-xpu20253-x86_64-linux/_flash_attn2_xpu_f12afc9.abi3.so deleted file mode 100644 index cfc097fc83fcc31b0b09ccc7440578fc159bb26e..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-xpu20253-x86_64-linux/_flash_attn2_xpu_f12afc9.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:ac4a81b9380257110b219cff74fa5396f84db6e3bb2adc26ebf2a990a3480242 -size 82767408 diff --git a/build/torch212-cxx11-xpu20253-x86_64-linux/_ops.py b/build/torch212-cxx11-xpu20253-x86_64-linux/_ops.py deleted file mode 100644 index ad32616257d5cef8c756337dded818dda050fc0c..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-xpu20253-x86_64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _flash_attn2_xpu_f12afc9 -ops = torch.ops._flash_attn2_xpu_f12afc9 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_flash_attn2_xpu_f12afc9::{op_name}" diff --git a/build/torch212-cxx11-xpu20253-x86_64-linux/bert_padding.py b/build/torch212-cxx11-xpu20253-x86_64-linux/bert_padding.py deleted file mode 100644 index 3c2d35159a014a9d03aabead9e52e009168696ea..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-xpu20253-x86_64-linux/bert_padding.py +++ /dev/null @@ -1,218 +0,0 @@ -# Adapted from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/padding.py - -import torch -import torch.nn.functional as F -from einops import rearrange, repeat - - -class IndexFirstAxis(torch.autograd.Function): - @staticmethod - def forward(ctx, input, indices): - ctx.save_for_backward(indices) - assert input.ndim >= 2 - ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] - second_dim = other_shape.numel() - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - # return input[indices] - return torch.gather( - rearrange(input, "b ... -> b (...)"), 0, repeat(indices, "z -> z d", d=second_dim) - ).reshape(-1, *other_shape) - - @staticmethod - def backward(ctx, grad_output): - (indices,) = ctx.saved_tensors - assert grad_output.ndim >= 2 - other_shape = grad_output.shape[1:] - grad_output = rearrange(grad_output, "b ... -> b (...)") - grad_input = torch.zeros( - [ctx.first_axis_dim, grad_output.shape[1]], - device=grad_output.device, - dtype=grad_output.dtype, - ) - # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. - # grad_input[indices] = grad_output - grad_input.scatter_(0, repeat(indices, "z -> z d", d=grad_output.shape[1]), grad_output) - return grad_input.reshape(ctx.first_axis_dim, *other_shape), None - - -index_first_axis = IndexFirstAxis.apply - - -class IndexPutFirstAxis(torch.autograd.Function): - @staticmethod - def forward(ctx, values, indices, first_axis_dim): - ctx.save_for_backward(indices) - assert indices.ndim == 1 - assert values.ndim >= 2 - output = torch.zeros( - first_axis_dim, *values.shape[1:], device=values.device, dtype=values.dtype - ) - # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. - output[indices] = values - # output.scatter_(0, repeat(indices, 'z -> z d', d=values.shape[1]), values) - return output - - @staticmethod - def backward(ctx, grad_output): - (indices,) = ctx.saved_tensors - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - grad_values = grad_output[indices] - # grad_values = torch.gather(grad_output, 0, repeat(indices, 'z -> z d', d=grad_output.shape[1])) - return grad_values, None, None - - -index_put_first_axis = IndexPutFirstAxis.apply - - -class IndexFirstAxisResidual(torch.autograd.Function): - @staticmethod - def forward(ctx, input, indices): - ctx.save_for_backward(indices) - assert input.ndim >= 2 - ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] - second_dim = other_shape.numel() - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - output = input[indices] - # We don't want to reshape input (b ... -> b (...)) since it could change the channel_last - # memory format to channel_first. In other words, input might not be contiguous. - # If we don't detach, Pytorch complains about output being a view and is being modified inplace - return output, input.detach() - - @staticmethod - def backward(ctx, grad_output, grad_residual): - (indices,) = ctx.saved_tensors - assert grad_output.ndim >= 2 - other_shape = grad_output.shape[1:] - assert grad_residual.shape[1:] == other_shape - grad_input = grad_residual - # grad_input[indices] += grad_output - indices = indices.reshape(indices.shape[0], *((1,) * (grad_output.ndim - 1))) - indices = indices.expand_as(grad_output) - grad_input.scatter_add_(0, indices, grad_output) - return grad_input.reshape(ctx.first_axis_dim, *other_shape), None - - -index_first_axis_residual = IndexFirstAxisResidual.apply - - -def unpad_input(hidden_states, attention_mask, unused_mask=None): - """ - Arguments: - hidden_states: (batch, seqlen, ...) - attention_mask: (batch, seqlen), bool / int, 1 means valid and 0 means not valid. - unused_mask: (batch, seqlen), bool / int, 1 means the element is allocated but unused. - Return: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens selected in attention_mask + unused_mask. - indices: (total_nnz), the indices of masked tokens from the flattened input sequence. - cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. - max_seqlen_in_batch: int - seqused: (batch), returns the number of tokens selected in attention_mask + unused_mask. - """ - all_masks = (attention_mask + unused_mask) if unused_mask is not None else attention_mask - seqlens_in_batch = all_masks.sum(dim=-1, dtype=torch.int32) - used_seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) - indices = torch.nonzero(all_masks.flatten(), as_tuple=False).flatten() - max_seqlen_in_batch = seqlens_in_batch.max().item() - cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) - # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the - # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim - # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to - # index with integer indices. Moreover, torch's index is a bit slower than it needs to be, - # so we write custom forward and backward to make it a bit faster. - return ( - index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices), - indices, - cu_seqlens, - max_seqlen_in_batch, - used_seqlens_in_batch, - ) - - -def unpad_input_for_concatenated_sequences(hidden_states, attention_mask_in_length): - """ - Supports concatenating short samples in one sequence. The attention_mask_in_length is utilized to mask other short samples. It helps efficient training of variant lengths-based samples (e.g., the supervised fine-tuning task in large language model). - The motivation for this function is explained [here](https://github.com/Dao-AILab/flash-attention/issues/432#issuecomment-1668822286). - - For example, if batch = 3 and seqlen = 6, the attention_mask_in_length is: - ``` - [ - [2, 3, 0, 0, 0, 0], - [3, 2, 0, 0, 0, 0], - [6, 0, 0, 0, 0, 0] - ] - ``` - , which refers to the 3D-attention mask: - ``` - [ - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [0, 0, 1, 0, 0, 0], - [0, 0, 1, 1, 0, 0], - [0, 0, 1, 1, 1, 0], - [0, 0, 0, 0, 0, 1] - ], - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [0, 0, 0, 1, 0, 0], - [0, 0, 0, 1, 1, 0], - [0, 0, 0, 0, 0, 1] - ], - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [1, 1, 1, 1, 0, 0], - [1, 1, 1, 1, 1, 0], - [1, 1, 1, 1, 1, 1] - ] - ] - ```. - - Arguments: - hidden_states: (batch, seqlen, ...) - attention_mask_in_length: (batch, seqlen), int, a nonzero number (e.g., 1, 2, 3, etc.) means length of concatenated sequence in b-th batch, and 0 means none. - Return: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. - indices: (total_nnz), the indices of non-masked tokens from the flattened input sequence. - cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. - max_seqlen_in_batch: int - """ - length = attention_mask_in_length.sum(dim=-1) - seqlen = attention_mask_in_length.size(-1) - attention_mask_2d = torch.arange(seqlen, device=length.device, dtype=length.dtype).expand(len(length), seqlen) < length.unsqueeze(1) - real_indices_idx = torch.nonzero(attention_mask_in_length.flatten(), as_tuple=False).flatten() - seqlens_in_batch = attention_mask_in_length.flatten()[real_indices_idx] - indices = torch.nonzero(attention_mask_2d.flatten(), as_tuple=False).flatten() - max_seqlen_in_batch = seqlens_in_batch.max().item() - cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) - # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the - # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim - # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to - # index with integer indices. Moreover, torch's index is a bit slower than it needs to be, - # so we write custom forward and backward to make it a bit faster. - return ( - index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices), - indices, - cu_seqlens, - max_seqlen_in_batch, - ) - - -def pad_input(hidden_states, indices, batch, seqlen): - """ - Arguments: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. - indices: (total_nnz), the indices that represent the non-masked tokens of the original padded input sequence. - batch: int, batch size for the padded sequence. - seqlen: int, maximum sequence length for the padded sequence. - Return: - hidden_states: (batch, seqlen, ...) - """ - dim = hidden_states.shape[-1] - # output = torch.zeros((batch * seqlen), dim, device=hidden_states.device, dtype=hidden_states.dtype) - # output[indices] = hidden_states - output = index_put_first_axis(hidden_states, indices, batch * seqlen) - return rearrange(output, "(b s) ... -> b s ...", b=batch) diff --git a/build/torch212-cxx11-xpu20253-x86_64-linux/flash_attn2/__init__.py b/build/torch212-cxx11-xpu20253-x86_64-linux/flash_attn2/__init__.py deleted file mode 100644 index a9b2672c1cd85b74c1b3ded0fc0b2100e1aeac23..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-xpu20253-x86_64-linux/flash_attn2/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -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/torch212-cxx11-xpu20253-x86_64-linux/flash_attn_interface.py b/build/torch212-cxx11-xpu20253-x86_64-linux/flash_attn_interface.py deleted file mode 100644 index b014e569dc9812a9c840d1043f2cebe767b4d477..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-xpu20253-x86_64-linux/flash_attn_interface.py +++ /dev/null @@ -1,1631 +0,0 @@ -# Copyright (c) 2023, Tri Dao. - -from typing import Optional, Sequence, Tuple, Union - -import torch -import torch.nn as nn -import os - -# # isort: off -# # We need to import the CUDA kernels after importing torch -# USE_TRITON_ROCM = os.getenv("FLASH_ATTENTION_TRITON_AMD_ENABLE", "FALSE") == "TRUE" -# if USE_TRITON_ROCM: -# from .flash_attn_triton_amd import interface_fa as flash_attn -# else: -# import flash_attn_2_cuda as flash_attn - - -from . import _ops -from ._ops import ops as flash_attn -from ._ops import add_op_namespace_prefix - -# # isort: on - -def maybe_contiguous(x): - return x.contiguous() if x is not None and x.stride(-1) != 1 else x - - -def _get_device(): - if torch.xpu.is_available(): - return "xpu" - elif torch.cuda.is_available(): - return "cuda" - else: - return "cpu" - - -def _get_block_size_n(device, head_dim, is_dropout, is_causal): - assert head_dim <= 256 - - if device.type == "xpu": - if head_dim <= 96: - return 64 - elif head_dim <= 128: - return 32 - elif head_dim <= 256: - return 64 - else: - return 32 - - # This should match the block sizes in the CUDA kernel - major, minor = torch.cuda.get_device_capability(device) - is_sm8x = major == 8 and minor > 0 # Only include sm86 and sm89, exclude sm80 (A100) - is_sm80 = major == 8 and minor == 0 - is_sm90 = major == 9 and minor == 0 - if head_dim <= 32: - return 128 - if head_dim <= 64: - return 128 if not is_dropout else 64 - elif head_dim <= 96: - return 64 - elif head_dim <= 128: - if is_sm8x: - return 64 if (not is_dropout and is_causal) else 32 - else: - return 64 if not is_dropout else 32 - elif head_dim <= 192: - return 64 - elif head_dim <= 224: - return 64 - elif head_dim <= 256: - return 64 - - -def round_multiple(x, m): - return (x + m - 1) // m * m - - -# torch.compile() support is only enabled for pytorch >= 2.4 -# The reason for this is that we are using the new custom_op and register_fake -# APIs, which support inplace modification of inputs in the function itself -if torch.__version__ >= "2.4.0": - _torch_custom_op_wrapper = torch.library.custom_op - _torch_register_fake_wrapper = torch.library.register_fake -else: - def noop_custom_op_wrapper(name, fn=None, /, *, mutates_args, device_types=None, schema=None): - def wrap(func): - return func - if fn is None: - return wrap - return fn - def noop_register_fake_wrapper(op, fn=None, /, *, lib=None, _stacklevel=1): - def wrap(func): - return func - if fn is None: - return wrap - return fn - _torch_custom_op_wrapper = noop_custom_op_wrapper - _torch_register_fake_wrapper = noop_register_fake_wrapper - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_forward"), mutates_args=(), device_types=_get_device()) -def _flash_attn_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - return_softmax: bool -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - out, softmax_lse, S_dmask, rng_state = flash_attn.fwd( - q, - k, - v, - None, - alibi_slopes, - dropout_p, - softmax_scale, - causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - None, - ) - return out, softmax_lse, S_dmask, rng_state - - -@_torch_register_fake_wrapper(add_op_namespace_prefix("_flash_attn_forward")) -def _flash_attn_forward_fake( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - return_softmax: bool -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - batch_size, seqlen_q, num_heads, head_size = q.shape - seqlen_k = k.shape[1] - out = torch.empty_like(q) - softmax_lse = torch.empty((batch_size, num_heads, seqlen_q), dtype=torch.float32, device=q.device, layout=q.layout) - p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) - if return_softmax: - p = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128), round_multiple(seqlen_k, 128)), dtype=q.dtype, device=q.device, layout=q.layout) - rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) - - return out, softmax_lse, p, rng_state - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_forward = _ops.ops._flash_attn_forward -else: - _wrapped_flash_attn_forward = _flash_attn_forward - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_varlen_forward"), mutates_args=(), device_types=_get_device()) -def _flash_attn_varlen_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - alibi_slopes: Optional[torch.Tensor] = None, - return_softmax: bool = False, - block_table: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - out, softmax_lse, S_dmask, rng_state = flash_attn.varlen_fwd( - q, - k, - v, - None, - cu_seqlens_q, - cu_seqlens_k, - seqused_k, - leftpad_k, - block_table, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - zero_tensors, - causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - None, - ) - # if out.isnan().any() or softmax_lse.isnan().any(): - # breakpoint() - return out, softmax_lse, S_dmask, rng_state - - -@_torch_register_fake_wrapper(add_op_namespace_prefix("_flash_attn_varlen_forward")) -def _flash_attn_varlen_forward_fake( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - alibi_slopes: Optional[torch.Tensor] = None, - return_softmax: bool = False, - block_table: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - paged_kv = block_table is not None - batch_size = cu_seqlens_q.numel() - 1 - total_q, num_heads, _ = q.shape - - out = torch.empty_like(q) - softmax_lse = torch.empty((num_heads, total_q), dtype=torch.float32, device=q.device, layout=q.layout) - p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) - seqlen_q_rounded = round_multiple(max_seqlen_q, 128) - seqlen_k_rounded = round_multiple(max_seqlen_k, 128) - if return_softmax: - p = torch.empty((batch_size, num_heads, seqlen_q_rounded, seqlen_k_rounded), dtype=q.dtype, device=q.device, layout=q.layout) - rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) - return out, softmax_lse, p, rng_state - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_varlen_forward = _ops.ops._flash_attn_varlen_forward -else: - _wrapped_flash_attn_varlen_forward = _flash_attn_varlen_forward - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_backward"), mutates_args=("dq", "dk", "dv"), device_types=_get_device()) -def _flash_attn_backward( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, -) -> 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)] - ( - dq, - dk, - dv, - softmax_d, - ) = flash_attn.bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - alibi_slopes, - dropout_p, - softmax_scale, - causal, - window_size_left, - window_size_right, - softcap, - deterministic, - None, - rng_state, - ) - return softmax_d - - -@_torch_register_fake_wrapper(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, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, -) -> torch.Tensor: - dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] - if dq is None: - dq = torch.empty_like(q) - if dk is None: - dk = torch.empty_like(k) - if dv is None: - dv = torch.empty_like(v) - batch_size, seqlen_q, num_heads, _ = q.shape - softmax_d = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128)), device=q.device, dtype=torch.float32) - - return softmax_d - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_backward = _ops.ops._flash_attn_backward -else: - _wrapped_flash_attn_backward = _flash_attn_backward - - -@_torch_custom_op_wrapper(add_op_namespace_prefix("_flash_attn_varlen_backward"), mutates_args=("dq", "dk", "dv"), device_types=_get_device()) -def _flash_attn_varlen_backward( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> 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)] - ( - dq, - dk, - dv, - softmax_d, - ) = flash_attn.varlen_bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - zero_tensors, - causal, - window_size_left, - window_size_right, - softcap, - deterministic, - None, - rng_state, - ) - # if dk.isnan().any() or dk.isnan().any() or dv.isnan().any() or softmax_d.isnan().any(): - # breakpoint() - return softmax_d - - -@_torch_register_fake_wrapper(add_op_namespace_prefix("_flash_attn_varlen_backward")) -def _flash_attn_varlen_backward_fake( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> torch.Tensor: - dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] - batch_size = cu_seqlens_q.numel() - 1 - total_q, num_heads, _ = q.shape - - if dq is None: - dq = torch.empty_like(q) - if dk is None: - dk = torch.empty_like(k) - if dv is None: - dv = torch.empty_like(v) - softmax_d = torch.empty((num_heads, total_q + 128 * batch_size), device=q.device, dtype=torch.float32) - - return softmax_d - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_varlen_backward = _ops.ops._flash_attn_varlen_backward -else: - _wrapped_flash_attn_varlen_backward = _flash_attn_varlen_backward - - -class FlashAttnQKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and qkv.requires_grad - if softmax_scale is None: - softmax_scale = qkv.shape[-1] ** (-0.5) - q, k, v = qkv[:, :, 0].detach(), qkv[:, :, 1].detach(), qkv[:, :, 2].detach() - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) - dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dqkv[:, :, 0], - dqkv[:, :, 1], - dqkv[:, :, 2], - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dqkv = dqkv[..., : dout.shape[-1]] # We could have padded the head dimension - return dqkv, None, None, None, None, None, None, None, None, None - - -class FlashAttnVarlenQKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - cu_seqlens, - max_seqlen, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and qkv.requires_grad - if softmax_scale is None: - softmax_scale = qkv.shape[-1] ** (-0.5) - q, k, v = qkv[:, 0].detach(), qkv[:, 1].detach(), qkv[:, 2].detach() - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens, - cu_seqlens, - max_seqlen, - max_seqlen, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=None, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens, rng_state) - ctx.dropout_p = dropout_p - ctx.max_seqlen = max_seqlen - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens, rng_state = ctx.saved_tensors - qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) - dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dqkv[:, 0], - dqkv[:, 1], - dqkv[:, 2], - cu_seqlens, - cu_seqlens, - ctx.max_seqlen, - ctx.max_seqlen, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - 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 FlashAttnKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - kv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, kv] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - k, v = kv[:, :, 0].detach(), kv[:, :, 1].detach() - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - dq = torch.empty_like(q) - kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) - dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dkv[:, :, 0], - dkv[:, :, 1], - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dkv = dkv[..., : dout.shape[-1]] - return dq, dkv, None, None, None, None, None, None, None, None, None - - -class FlashAttnVarlenKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, kv] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - k, v = kv[:, 0].detach(), kv[:, 1].detach() - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=None, - ) - if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state - ) - ctx.dropout_p = dropout_p - 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.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors - dq = torch.empty_like(q) - kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) - dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dkv[:, 0], - dkv[:, 1], - cu_seqlens_q, - cu_seqlens_k, - ctx.max_seqlen_q, - ctx.max_seqlen_k, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dkv = dkv[..., : dout.shape[-1]] - return dq, dkv, None, None, None, None, None, None, None, None, None, None, None, None, None - - -class FlashAttnFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - k, - v, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, k, v] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dk = dk[..., : dout.shape[-1]] - dv = dv[..., : dout.shape[-1]] - return dq, dk, dv, 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, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - block_table, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, k, v] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=block_table, - ) - if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state - ) - ctx.dropout_p = dropout_p - 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.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors - dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - ctx.max_seqlen_q, - ctx.max_seqlen_k, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dk = dk[..., : dout.shape[-1]] - dv = dv[..., : dout.shape[-1]] - return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None - - -def flash_attn_qkvpacked_func( - qkv, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # <=0.0 means deactivate - alibi_slopes=None, - deterministic=False, - 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, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_kvpacked_func( - q, - kv, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, -): - """dropout_p should be set to 0.0 during evaluation - If 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 K, V. - 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) - kv: (batch_size, seqlen, 2, 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. - softcap: float. Anything > 0 activates softcapping 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). - 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 FlashAttnKVPackedFunc.apply( - q, - kv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_func( - q, - k, - v, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - 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). - 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 FlashAttnFunc.apply( - q, - k, - v, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_qkvpacked_func( - qkv, - cu_seqlens, - max_seqlen, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - 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_varlen_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_varlen_kvpacked_func and flash_attn_varlen_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: (total, 3, nheads, headdim), where total = total number of tokens in the batch. - cu_seqlens: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into qkv. - max_seqlen: int. Maximum sequence length in the batch. - 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenQKVPackedFunc.apply( - qkv, - cu_seqlens, - max_seqlen, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_kvpacked_func( - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, -): - """dropout_p should be set to 0.0 during evaluation - If 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 K, V. - 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: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch. - kv: (total_k, 2, nheads_k, headdim), where total_k = total number of key tokens in the batch. - cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into q. - cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into kv. - max_seqlen_q: int. Maximum query sequence length in the batch. - max_seqlen_k: int. Maximum key sequence length in the batch. - 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 + 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenKVPackedFunc.apply( - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_func( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, - block_table=None, -): - """dropout_p should be set to 0.0 during evaluation - Supports multi-query and grouped-query attention (MQA/GQA) by passing in K, V 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: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch. - k: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. - v: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. - cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into q. - cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into kv. - max_seqlen_q: int. Maximum query sequence length in the batch. - max_seqlen_k: int. Maximum key sequence length in the batch. - 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 + 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenFunc.apply( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - block_table, - False if q.device.type == "cpu" else torch.is_grad_enabled(), - ) - - -def flash_attn_with_kvcache( - q, - k_cache, - v_cache, - k=None, - v=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, - block_table: Optional[torch.Tensor] = None, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - rotary_interleaved=True, - alibi_slopes=None, - num_splits=0, - 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 block_table, - or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_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) if there's no block_table, - or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_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). Similar to k. - 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. - block_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). - 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. - 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" - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-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) - cache_batch_idx = maybe_contiguous(cache_batch_idx) - block_table = maybe_contiguous(block_table) - out, softmax_lse = flash_attn.fwd_kvcache( - q, - k_cache, - v_cache, - k, - v, - cache_seqlens, - rotary_cos, - rotary_sin, - cache_batch_idx, - cache_leftpad, - block_table, - alibi_slopes, - None, - softmax_scale, - causal, - window_size[0], - window_size[1], - softcap, - rotary_interleaved, - num_splits, - ) - return (out, softmax_lse) if return_softmax_lse else out diff --git a/build/torch212-cxx11-xpu20253-x86_64-linux/layers/__init__.py b/build/torch212-cxx11-xpu20253-x86_64-linux/layers/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/build/torch212-cxx11-xpu20253-x86_64-linux/layers/patch_embed.py b/build/torch212-cxx11-xpu20253-x86_64-linux/layers/patch_embed.py deleted file mode 100644 index 05562f8e8bcdb58e947c6f402a49eacd2d031871..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-xpu20253-x86_64-linux/layers/patch_embed.py +++ /dev/null @@ -1,67 +0,0 @@ -# We use the same API as https://github.com/rwightman/pytorch-image-models/blob/v0.6.11/timm/models/layers/patch_embed.py -# But we use nn.Linear instead of Conv2d and it's about 8x faster. - -from functools import partial - -import torch.nn as nn -from einops import rearrange -from torch import _assert -from torch.nn.modules.utils import _pair - -try: - from flash_attn.ops.fused_dense import FusedDense -except ImportError: - FusedDense = None - - -class PatchEmbed(nn.Module): - """2D Image to Patch Embedding""" - - def __init__( - self, - img_size=224, - patch_size=16, - in_chans=3, - embed_dim=768, - norm_layer=None, - flatten=True, - bias=True, - fused_bias_fc=False, - ): - super().__init__() - img_size = _pair(img_size) - patch_size = _pair(patch_size) - self.img_size = img_size - self.patch_size = patch_size - self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1]) - self.num_patches = self.grid_size[0] * self.grid_size[1] - self.flatten = flatten - if fused_bias_fc and FusedDense is None: - raise ImportError("fused_dense is not installed") - - linear_cls = nn.Linear if not fused_bias_fc or not bias else FusedDense - self.proj = linear_cls(in_chans * patch_size[0] * patch_size[1], embed_dim, bias=bias) - self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() - - def forward(self, x): - _, _, H, W = x.shape - _assert( - H == self.img_size[0], - f"Input image height ({H}) doesn't match model ({self.img_size[0]}).", - ) - _assert( - W == self.img_size[1], - f"Input image width ({W}) doesn't match model ({self.img_size[1]}).", - ) - x = self.proj( - rearrange( - x, - "b c (h p1) (w p2) -> b h w (c p1 p2)", - p1=self.patch_size[0], - p2=self.patch_size[1], - ) - ) - if self.flatten: - x = rearrange(x, "b h w c -> b (h w) c") - x = self.norm(x) - return x diff --git a/build/torch212-cxx11-xpu20253-x86_64-linux/layers/rotary.py b/build/torch212-cxx11-xpu20253-x86_64-linux/layers/rotary.py deleted file mode 100644 index d1bfc21fc7de1dd287e8f382847b194a48075981..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-xpu20253-x86_64-linux/layers/rotary.py +++ /dev/null @@ -1,483 +0,0 @@ -# Copyright (c) 2025, Tri Dao - -import math -from functools import partial -from typing import Optional, Tuple, Union - -import torch -from torch import Tensor - -from einops import rearrange, repeat -# from flash_attn.ops.triton.rotary import apply_rotary -from ..ops.triton.rotary import apply_rotary - - -def rotate_half(x, interleaved=False): - if not interleaved: - x1, x2 = x.chunk(2, dim=-1) - return torch.cat((-x2, x1), dim=-1) - else: - x1, x2 = x[..., ::2], x[..., 1::2] - return rearrange(torch.stack((-x2, x1), dim=-1), "... d two -> ... (d two)", two=2) - - -def apply_rotary_emb_torch(x, cos, sin, interleaved=False): - """ - x: (batch_size, seqlen, nheads, headdim) - cos, sin: (seqlen, rotary_dim / 2) or (batch_size, seqlen, rotary_dim / 2) - """ - ro_dim = cos.shape[-1] * 2 - assert ro_dim <= x.shape[-1] - cos = repeat(cos, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)") - sin = repeat(sin, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)") - return torch.cat( - [x[..., :ro_dim] * cos + rotate_half(x[..., :ro_dim], interleaved) * sin, x[..., ro_dim:]], - dim=-1, - ) - - -class ApplyRotaryEmb(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x, - cos, - sin, - interleaved=False, - inplace=False, - seqlen_offsets: Union[int, Tensor] = 0, - cu_seqlens: Optional[Tensor] = None, - max_seqlen: Optional[int] = None, - ): - out = apply_rotary( - x, - cos, - sin, - seqlen_offsets=seqlen_offsets, - cu_seqlens=cu_seqlens, - max_seqlen=max_seqlen, - interleaved=interleaved, - inplace=inplace, - ) - if isinstance(seqlen_offsets, int): - ctx.save_for_backward(cos, sin, cu_seqlens) # Can't save int with save_for_backward - ctx.seqlen_offsets = seqlen_offsets - else: - ctx.save_for_backward(cos, sin, cu_seqlens, seqlen_offsets) - ctx.seqlen_offsets = None - ctx.interleaved = interleaved - ctx.inplace = inplace - ctx.max_seqlen = max_seqlen - return out if not inplace else x - - @staticmethod - def backward(ctx, do): - seqlen_offsets = ctx.seqlen_offsets - if seqlen_offsets is None: - cos, sin, cu_seqlens, seqlen_offsets = ctx.saved_tensors - else: - cos, sin, cu_seqlens = ctx.saved_tensors - dx = apply_rotary( - do, - cos, - sin, - seqlen_offsets=seqlen_offsets, - cu_seqlens=cu_seqlens, - max_seqlen=ctx.max_seqlen, - interleaved=ctx.interleaved, - inplace=ctx.inplace, - conjugate=True, - ) - return dx, None, None, None, None, None, None, None - - -def apply_rotary_emb( - x, - cos, - sin, - interleaved=False, - inplace=False, - seqlen_offsets: Union[int, Tensor] = 0, - cu_seqlens: Optional[Tensor] = None, - max_seqlen: Optional[int] = None, -): - """ - Arguments: - x: (batch_size, seqlen, nheads, headdim) if cu_seqlens is None - else (total_seqlen, nheads, headdim) - cos, sin: (seqlen_rotary, rotary_dim / 2) - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead - of 1st half and 2nd half (GPT-NeoX style). - inplace: if True, apply rotary embedding in-place. - seqlen_offsets: (batch_size,) or int. Each sequence in x is shifted by this amount. - Most commonly used in inference when we have KV cache. - cu_seqlens: (batch + 1,) or None - max_seqlen: int - Return: - out: (batch_size, seqlen, nheads, headdim) if cu_seqlens is None - else (total_seqlen, nheads, headdim) - rotary_dim must be <= headdim - Apply rotary embedding to the first rotary_dim of x. - """ - return ApplyRotaryEmb.apply( - x, cos, sin, interleaved, inplace, seqlen_offsets, cu_seqlens, max_seqlen - ) - - -# For backward compatibility -apply_rotary_emb_func = apply_rotary_emb - - -def _apply_rotary_emb_qkv( - qkv, - cos, - sin, - cos_k=None, - sin_k=None, - interleaved=False, - inplace=False, - conjugate=False, - seqlen_offsets: Union[int, Tensor] = 0, - num_heads_q: Optional[int] = None, -): - apply_rotary_fn = partial( - apply_rotary, - interleaved=interleaved, - inplace=inplace, - conjugate=conjugate, - seqlen_offsets=seqlen_offsets - ) - if cos_k is None and sin_k is None and qkv.is_contiguous(): - # Call 1 kernel instead of 2 kernels - # We need qkv to be contiguous so that when we reshape to combine (3, nheads) - # dimensions, we get the same tensor - if qkv.dim() == 5: - batch, seqlen, three, nheads, headdim = qkv.shape - assert three == 3 - # qk = rearrange(qkv[:, :, :2], "b s t h d -> b s (t h) d") - qk = qkv[:, :, :2].reshape(batch, seqlen, -1, headdim) - qk = apply_rotary_fn(qk, cos, sin) - else: - assert qkv.dim() == 4 - assert num_heads_q is not None - num_heads_k = (qkv.shape[2] - num_heads_q) // 2 - assert qkv.shape[2] == num_heads_q + 2 * num_heads_k - qk = qkv[:, :, :num_heads_q + num_heads_k] - qk = apply_rotary_fn(qk, cos, sin) - if not inplace: - if qkv.dim() == 5: - qkv = torch.cat([rearrange(qk, "b s (t h) d -> b s t h d", t=2), qkv[:, :, 2:]], dim=2) - else: - qkv = torch.cat([qk, qkv[:, :, num_heads_q + num_heads_k :]], dim=2) - else: - cos_k = cos if cos_k is None else cos_k - sin_k = sin if sin_k is None else sin_k - if qkv.dim() == 5: - batch, seqlen, three, nheads, headdim = qkv.shape - assert three == 3 - q, k = qkv[:, :, 0], qkv[:, :, 1] - else: - assert qkv.dim() == 4 - assert num_heads_q is not None - num_heads_k = (qkv.shape[2] - num_heads_q) // 2 - assert qkv.shape[2] == num_heads_q + 2 * num_heads_k - q, k = qkv[:, :, :num_heads_q], qkv[:, :, num_heads_q : num_heads_q + num_heads_k] - q = apply_rotary_fn(q, cos, sin) - k = apply_rotary_fn(k, cos_k, sin_k) - if not inplace: - if qkv.dim() == 5: - qkv = torch.stack([q, k, qkv[:, :, 2]], dim=2) - else: - qkv = torch.cat([q, k, qkv[:, :, num_heads_q + num_heads_k:]], dim=2) - return qkv - - -class ApplyRotaryEmbQKV_(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - cos, - sin, - cos_k=None, - sin_k=None, - interleaved=False, - seqlen_offsets: Union[int, torch.Tensor] = 0, - num_heads_q: Optional[int] = None, - ): - # apply_rotary_emb_qkv_inplace( - qkv = _apply_rotary_emb_qkv( - qkv, cos, sin, cos_k, sin_k, interleaved=interleaved, inplace=True, - seqlen_offsets=seqlen_offsets, num_heads_q=num_heads_q, - ) - if isinstance(seqlen_offsets, int): - ctx.save_for_backward(cos, sin, cos_k, sin_k) - ctx.seqlen_offsets = seqlen_offsets - else: - ctx.save_for_backward(cos, sin, cos_k, sin_k, seqlen_offsets) - ctx.seqlen_offsets = None - ctx.interleaved = interleaved - ctx.num_heads_q = num_heads_q - return qkv - - @staticmethod - def backward(ctx, dqkv): - seqlen_offsets = ctx.seqlen_offsets - if seqlen_offsets is None: - cos, sin, cos_k, sin_k, seqlen_offsets = ctx.saved_tensors - else: - cos, sin, cos_k, sin_k = ctx.saved_tensors - dqkv = _apply_rotary_emb_qkv( - dqkv, cos, sin, cos_k, sin_k, interleaved=ctx.interleaved, inplace=True, - seqlen_offsets=seqlen_offsets, num_heads_q=ctx.num_heads_q, conjugate=True, - ) - return dqkv, None, None, None, None, None, None, None - - -def apply_rotary_emb_qkv_( - qkv, - cos, - sin, - cos_k=None, - sin_k=None, - interleaved=False, - seqlen_offsets: Union[int, torch.Tensor] = 0, - num_heads_q: Optional[int] = None, -): - """ - Arguments: - qkv: (batch_size, seqlen, 3, nheads, headdim) or (batch_size, seqlen, num_heads_q + 2 * num_heads_k, headdim). - If qkv has shape (batch_size, seqlen, num_heads_q + 2 * num_heads_k, headdim) (e.g. MQA / GQA), - then num_heads_q must be provided. - cos, sin: (seqlen, rotary_dim / 2) - cos_k, sin_k: (seqlen, rotary_dim / 2), optional - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead of - 1st half and 2nd half (GPT-NeoX style). - seqlen_offsets: (batch_size,) or int. Each sequence in Q and K is shifted by this amount. - Most commonly used in inference when we have KV cache. - Return: - qkv: (batch_size, seqlen, 3, nheads, headdim) or (batch_size, seqlen, num_heads_q + 2 * num_heads_k, headdim) - rotary_dim must be <= headdim - Apply rotary embedding *inplace* to the first rotary_dim of Q and K. - """ - return ApplyRotaryEmbQKV_.apply( - qkv, cos, sin, cos_k, sin_k, interleaved, seqlen_offsets, num_heads_q - ) - - -class ApplyRotaryEmbKV_(torch.autograd.Function): - - @staticmethod - def forward(ctx, kv, cos, sin, interleaved=False, seqlen_offsets: Union[int, torch.Tensor] = 0): - batch, seqlen, two, nheads, headdim = kv.shape - assert two == 2 - k = kv[:, :, 0] - apply_rotary( - k, cos, sin, seqlen_offsets=seqlen_offsets, interleaved=interleaved, inplace=True - ) - if isinstance(seqlen_offsets, int): - ctx.save_for_backward(cos, sin) # Can't save int with save_for_backward - ctx.seqlen_offsets = seqlen_offsets - else: - ctx.save_for_backward(cos, sin, seqlen_offsets) - ctx.seqlen_offsets = None - ctx.interleaved = interleaved - return kv - - @staticmethod - def backward(ctx, dkv): - seqlen_offsets = ctx.seqlen_offsets - if seqlen_offsets is None: - cos, sin, seqlen_offsets = ctx.saved_tensors - else: - cos, sin = ctx.saved_tensors - apply_rotary( - dkv[:, :, 0], - cos, - sin, - seqlen_offsets=seqlen_offsets, - interleaved=ctx.interleaved, - inplace=True, - conjugate=True, - ) - return dkv, None, None, None, None - - -apply_rotary_emb_kv_ = ApplyRotaryEmbKV_.apply - - -def apply_rotary_emb_kv_( - kv, - cos, - sin, - interleaved=False, - seqlen_offsets: Union[int, torch.Tensor] = 0, -): - """ - Arguments: - kv: (batch_size, seqlen, 2, nheads, headdim) - cos, sin: (seqlen, rotary_dim / 2) - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead of - 1st half and 2nd half (GPT-NeoX style). - seqlen_offsets: (batch_size,) or int. Each sequence in Q and K is shifted by this amount. - Most commonly used in inference when we have KV cache. - Return: - kv: (batch_size, seqlen, 2, nheads, headdim) - rotary_dim must be <= headdim - Apply rotary embedding *inplace* to the first rotary_dim of K. - """ - return ApplyRotaryEmbKV_.apply(kv, cos, sin, interleaved, seqlen_offsets) - - -class RotaryEmbedding(torch.nn.Module): - """ - The rotary position embeddings from RoFormer_ (Su et. al). - A crucial insight from the method is that the query and keys are - transformed by rotation matrices which depend on the relative positions. - - Other implementations are available in the Rotary Transformer repo_ and in - GPT-NeoX_, GPT-NeoX was an inspiration - - .. _RoFormer: https://arxiv.org/abs/2104.09864 - .. _repo: https://github.com/ZhuiyiTechnology/roformer - .. _GPT-NeoX: https://github.com/EleutherAI/gpt-neox - - If scale_base is not None, this implements XPos (Sun et al., https://arxiv.org/abs/2212.10554). - A recommended value for scale_base is 512: https://github.com/HazyResearch/flash-attention/issues/96 - Reference: https://github.com/sunyt32/torchscale/blob/main/torchscale/component/xpos_relative_position.py - """ - - def __init__( - self, - dim: int, - base=10000.0, - interleaved=False, - scale_base=None, - device=None, - ): - """ - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead - of 1st half and 2nd half (GPT-NeoX style). - """ - super().__init__() - self.dim = dim - self.base = float(base) - # Generate and save the inverse frequency buffer (non trainable) - inv_freq = self._compute_inv_freq(device) - self.register_buffer("inv_freq", inv_freq, persistent=False) - self.interleaved = interleaved - self.scale_base = scale_base - scale = ( - (torch.arange(0, dim, 2, device=device, dtype=torch.float32) + 0.4 * dim) / (1.4 * dim) - if scale_base is not None - else None - ) - self.register_buffer("scale", scale, persistent=False) - - self._seq_len_cached = 0 - self._cos_cached = None - self._sin_cached = None - self._cos_k_cached = None - self._sin_k_cached = None - - def _compute_inv_freq(self, device=None): - return 1.0 / ( - self.base - ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim) - ) - - def _update_cos_sin_cache(self, seqlen, device=None, dtype=None): - # Reset the tables if the sequence length has changed, - # if we're on a new device (possibly due to tracing for instance), - # or if we're switching from inference mode to training - if ( - seqlen > self._seq_len_cached - or self._cos_cached is None - or self._cos_cached.device != device - or self._cos_cached.dtype != dtype - or (self.training and self._cos_cached.is_inference()) - ): - self._seq_len_cached = seqlen - # We want fp32 here, not self.inv_freq.dtype, since the model could be loaded in bf16 - # And the output of arange can be quite large, so bf16 would lose a lot of precision. - t = torch.arange(seqlen, device=device, dtype=torch.float32) - # We want fp32 here as well since inv_freq will be multiplied with t, and the output - # will be large. Having it in bf16 will lose a lot of precision and cause the - # cos & sin output to change significantly. - # We want to recompute self.inv_freq if it was not loaded in fp32 - if self.inv_freq.dtype != torch.float32: - inv_freq = self._compute_inv_freq(device=device) - else: - inv_freq = self.inv_freq - # Don't do einsum, it converts fp32 to bf16 under AMP - # freqs = torch.einsum("i,j->ij", t, self.inv_freq) - freqs = torch.outer(t, inv_freq) - if self.scale is None: - self._cos_cached = torch.cos(freqs).to(dtype) - self._sin_cached = torch.sin(freqs).to(dtype) - else: - power = ( - torch.arange(seqlen, dtype=self.scale.dtype, device=self.scale.device) - - seqlen // 2 - ) / self.scale_base - scale = self.scale.to(device=power.device) ** rearrange(power, "s -> s 1") - # We want the multiplication by scale to happen in fp32 - self._cos_cached = (torch.cos(freqs) * scale).to(dtype) - self._sin_cached = (torch.sin(freqs) * scale).to(dtype) - self._cos_k_cached = (torch.cos(freqs) / scale).to(dtype) - self._sin_k_cached = (torch.sin(freqs) / scale).to(dtype) - - def forward( - self, - qkv: torch.Tensor, - kv: Optional[torch.Tensor] = None, - seqlen_offset: Union[int, torch.Tensor] = 0, - max_seqlen: Optional[int] = None, - num_heads_q: Optional[int] = None, - ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: - """ - qkv: (batch, seqlen, 3, nheads, headdim) or (batch, seqlen, num_heads_q + 2 * num_heads_k, headdim) - if kv is none, else it's just q of shape (batch, seqlen, nheads, headdim). - If qkv has shape (batch, seqlen, num_heads_q + 2 * num_heads_k, headdim) (e.g. MQA / GQA), - then num_heads_q must be provided. - kv: (batch, seqlen, 2, nheads, headdim) - seqlen_offset: (batch_size,) or int. Each sequence in x is shifted by this amount. - Most commonly used in inference when we have KV cache. - If it's a tensor of shape (batch_size,), then to update the cos / sin cache, one - should pass in max_seqlen, which will update the cos / sin cache up to that length. - Apply rotary embedding *inplace* to qkv and / or kv. - """ - seqlen = qkv.shape[1] - if max_seqlen is not None: - self._update_cos_sin_cache(max_seqlen, device=qkv.device, dtype=qkv.dtype) - elif isinstance(seqlen_offset, int): - self._update_cos_sin_cache(seqlen + seqlen_offset, device=qkv.device, dtype=qkv.dtype) - if kv is None: - return apply_rotary_emb_qkv_( - qkv, - self._cos_cached, - self._sin_cached, - self._cos_k_cached if self.scale is not None else None, - self._sin_k_cached if self.scale is not None else None, - interleaved=self.interleaved, - seqlen_offsets=seqlen_offset, - num_heads_q=num_heads_q, - ) - else: - q = qkv - q = apply_rotary_emb_func( - q, - self._cos_cached, - self._sin_cached, - interleaved=self.interleaved, - inplace=True, - seqlen_offsets=seqlen_offset, - ) - kv = apply_rotary_emb_kv_( - kv, - self._cos_cached if self.scale is None else self._cos_k_cached, - self._sin_cached if self.scale is None else self._sin_k_cached, - interleaved=self.interleaved, - seqlen_offsets=seqlen_offset, - ) - return q, kv diff --git a/build/torch212-cxx11-xpu20253-x86_64-linux/metadata.json b/build/torch212-cxx11-xpu20253-x86_64-linux/metadata.json deleted file mode 100644 index 4f3ed5d5f25d5a0e23a04d466e1a69d635f00ec2..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-xpu20253-x86_64-linux/metadata.json +++ /dev/null @@ -1,36 +0,0 @@ -{ - "name": "flash-attn2", - "id": "_flash_attn2_xpu_f12afc9", - "version": 1, - "license": "BSD-3-Clause", - "python-depends": [], - "backend": { - "type": "xpu" - }, - "digest": { - "algorithm": "sha256", - "files": { - "__init__.py": "Cd9MEdd/WUYnBmfiQpC4SyPqAxrTilTjvsvswNju4zc=", - "_flash_attn2_xpu_f12afc9.abi3.so": "rEqBuTgCVxELIZz/dPpTlvhNtuO7Ktwm6/KpkKNIAkI=", - "_ops.py": "cmwaGTsubLtetqeo95vdShoO62KcV0DrdeNH34KtCyk=", - "bert_padding.py": "gF1EmsdJ+HpQ86MRQ4VxDw+Sb/RVISdQALdNnoByHlw=", - "flash_attn2/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY=", - "flash_attn_interface.py": "W7a9bHR2ly6QL3aKqsBqxWA7M7ZWYcZPmtsPJu6rAfY=", - "layers/__init__.py": "47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU=", - "layers/patch_embed.py": "H58CgME/qSOPTZLOG08wFgrQS1j34pvNwMPrkTj3Ek4=", - "layers/rotary.py": "rFIdNrZQFNawwLD6d7t5LE2m+M93wwDZ8EuExw+gfD8=", - "ops/__init__.py": "47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU=", - "ops/activations.py": "t5lzNg1In8LP6bKeTnyeMizwqjv27JGbJ6ylPdGvZYg=", - "ops/fused_dense.py": "BqT8VTSCGdLJAjh6WBMkCRFYrdVodPeoElWBvH6AIzA=", - "ops/layer_norm.py": "zr7NXIm+2mtEynTp1CS0fbFGI2Mqdp41dY4AfDWF6EQ=", - "ops/rms_norm.py": "XEnihcj0a4aSz4LO55m5iKGVn4HKTeKN8TIyHjuDgxI=", - "ops/triton/__init__.py": "AbpHGcgLb+kRsJGnwFEktk7uzpZOCcBY74+YBdrKVGs=", - "ops/triton/cross_entropy.py": "y56nmuOtpSNxM7sDZcUSgXKIsIEayfuvrU8bAVog/nw=", - "ops/triton/k_activations.py": "+Z3vIyO4JkqBMipKsPvhzmxljtBdIhJCsl/M+/ESqBo=", - "ops/triton/layer_norm.py": "BBQdJiqopbIS12viUqzx3SvOQCNbDSAYLJXhpyMGyC8=", - "ops/triton/linear.py": "OtRvKz8xdpl+7v3q/ZTaS9fdBt9XrzMyapgRr50uBbM=", - "ops/triton/mlp.py": "uWA+D7HyuRUsOT+4nhi3vvCDNZnNEXwYcHMQ02UREfQ=", - "ops/triton/rotary.py": "0YSrlbep0lJmzLpTiANs0slTFeGcQQqR+el6uzbSLvI=" - } - } -} \ No newline at end of file diff --git a/build/torch212-cxx11-xpu20253-x86_64-linux/metadata.json.sigstore b/build/torch212-cxx11-xpu20253-x86_64-linux/metadata.json.sigstore deleted file mode 100644 index f0c3e4aac897e1ad267b091a84bfbf1dfade009e..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-xpu20253-x86_64-linux/metadata.json.sigstore +++ /dev/null @@ -1 +0,0 @@ -{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json","verificationMaterial":{"certificate":{"rawBytes":"MIIHSjCCBtCgAwIBAgIUGUF/X7u70cK2H/KeZosjret/gbgwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwNjIzMTgyMDI1WhcNMjYwNjIzMTgzMDI1WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAE/MB9Qb0P7ZdxWM09J0Q5gaql7AHZOKstfaAzXJt2I7cmoNnCAVrfmt54ksNSQWaKgNFTNqbzK2RT2/YvNrH5hKOCBe8wggXrMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUQXJ8Bx+pb/vqphcU+QaXwsbhR3QwHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoZjEyYWZjOTJjNDBiYmQ3NmQ5OTBkOTExZWJhYmYyOGU0OTBhN2JkZTATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoZjEyYWZjOTJjNDBiYmQ3NmQ5OTBkOTExZWJhYmYyOGU0OTBhN2JkZTAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoZjEyYWZjOTJjNDBiYmQ3NmQ5OTBkOTExZWJhYmYyOGU0OTBhN2JkZTAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKGYxMmFmYzkyYzQwYmJkNzZkOTkwZDkxMWViYWJmMjhlNDkwYTdiZGUwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMjgwMzYyNjg5MzUvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBiQYKKwYBBAHWeQIEAgR7BHkAdwB1AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABnvW2S1QAAAQDAEYwRAIgNJAwIUYC+u/MDKy6I7yXXQXvo6QfXxIxkhDzhuci/HcCICAQpkw0Hhu3r7fgAS2Iz0RHEuABVMqK42FKpF8x8Z1IMAoGCCqGSM49BAMDA2gAMGUCMQDCSMdd3aLnQzH5VNGi4dIziaKgtCFQSBj9qQWeH6zgV8kG4LN8WZ38SVozORTnCDYCMDVfD/UwMX9uI+ZwL/rzNIkJoC72pyBkSOWAGI+dVFCXGatv5/4hUj56ra51u/Uo+w=="},"tlogEntries":[{"logIndex":"1929966493","logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="},"kindVersion":{"kind":"hashedrekord","version":"0.0.1"},"integratedTime":"1782238825","inclusionPromise":{"signedEntryTimestamp":"MEYCIQDGJtY832xPCCKlnLD4VRF0Bkx8FEVWNWXJX2lQWmAdfgIhAIHvP0u1VsI6nmW5NjNCmktZexRnxm1UNxBpYnlOwl//"},"inclusionProof":{"logIndex":"1808062231","rootHash":"M6bBtwKzXMRC6pslsEYCgZwaaFK1esuiWPxAuv0bX8Q=","treeSize":"1808062258","hashes":["IndZ/q0vKORPU7bw4t94MH9ELdEixGn2xVRKyrMsXz4=","VJ9yaFD79SuCwi4dECn+MHPfY62qZWZjJM1JJ/9INMs=","bsYNroONIHTv3B2u4SjRSdn4PTKUbsaIyxnhwcuWW+8=","5Xh3VCXrEdwAWf2Vke2ugmclHkiLukte3XR2iEF1eqs=","jihrxwF4e/WPcFnsqD1wxKl2Wx92ax8o33KYducFSGY=","A0q0E1Q0SHxoHhzrAl3GNMlJc9b+z+UqSns1/jI2ogA=","VrhYX+LlNagV3pSn3KMaGr/lC81bfD7BQz859yCToYY=","jgKZMZUAOznhZow0dOq9j+6kBbYqJbpAwbHuyR0vtw0=","43kco2WsynkkNcu9FNHzcL7ZH810rbB/a5HS62l4P9Q=","Rf4AP9k6JckOtBvSb8IZhPtOdgQbDwTqzHj/Ob7GoR0=","2azvF0ZQXjI0TTkEYlWo40jJTIJKrIaRQSKDw70qaz8=","MXnb5KZ/vcuCLPtNBIq6MAywaatte0rS/H2lJyqLfIY=","QdxFjkGuHUV4B069H6Stxfn75M+rztEbA+7TnHdszN8=","nyv2TDlfkz0XfRwjY98zsbUbOkBZ+Tgz7PML0c2uN14=","OMfDh4k9YD0E+t4sjqJsh+vg8D1ueo1cC07exyxmr6s=","+vHKm4wNwGJ46ceyZ8mePD7Vn1SediStNHCStRoTSWQ=","5HwWjuGNDisetbiWM0cWIXM65TIM3EJlqk13ntCQlcU=","NDIRlqxO0IaGr94q86MIU7VCAN3yEEgOHS4loya5lu4=","daxmZaajRpZV+JxHiOYZhJBiSKN5ucqjh2WnGbHhirw=","DOCeoSMovIvLExkhIvisow9AuNXgeWs4ECkyR6EcqYU="],"checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n1808062258\nM6bBtwKzXMRC6pslsEYCgZwaaFK1esuiWPxAuv0bX8Q=\n\n— rekor.sigstore.dev wNI9ajBFAiEAyYIW1zM4LZL68fLGHuRetqzhvaCO+4EknerHn+XKI6ACIG/4uUTaRLxUiM6yA4rtIiZQEkpbKBD8xfvUSowADP4D\n"}},"canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiJlYTQwM2U5MDQyZWE2NDU0MGFkZGVmZDNhZjk3ZWU3Yjc5MjZjOWI5OTc5MWFhNmUyYmQ2N2UxMjU5YzdlMGQxIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJUUNZTVhUS094V2JoNXBjTWFjNUkzK0ZJOWlpZHBJcDgvSWRLekMzMXJDUnJRSWdXVWZYeGFwUmVhbXV4OVdkc3dZVmFmSzB0aFlDQmZkc2JibHJ5SzQ0Q2ZjPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRha05EUW5SRFowRjNTVUpCWjBsVlIxVkdMMWczZFRjd1kwc3lTQzlMWlZwdmMycHlaWFF2WjJKbmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDVxU1hwTlZHZDVUVVJKTVZkb1kwNU5hbGwzVG1wSmVrMVVaM3BOUkVreFYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVV2VFVJNVVXSXdVRGRhWkhoWFRUQTVTakJSTldkaGNXdzNRVWhhVDB0emRHWmhRWG9LV0VwME1razNZMjF2VG01RFFWWnlabTEwTlRScmMwNVRVVmRoUzJkT1JsUk9jV0o2U3pKU1ZESXZXWFpPY2tnMWFFdFBRMEpsT0hkbloxaHlUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlZSV0VvNENrSjRLM0JpTDNaeGNHaGpWU3RSWVZoM2MySm9Vak5SZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOWFha1Y1V1ZkYWFrOVVTbXBPUkVKcFdXMVJNMDV0VVRWUFZFSnJDazlVUlhoYVYwcG9XVzFaZVU5SFZUQlBWRUpvVGpKS2ExcFVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMXBxUlhsWlYxcHFUMVJLYWs1RVFtbFpiVkV6VG0xUk5VOVVRbXRQVkVWNFdsZEthRmx0V1hrS1QwZFZNRTlVUW1oT01rcHJXbFJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMXBxUlhsWlYxcHFUMVJLYWs1RVFtbFpiVkV6VG0xUk5VOVVRbXNLVDFSRmVGcFhTbWhaYlZsNVQwZFZNRTlVUW1oT01rcHJXbFJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwZFplRTF0Um0wS1dYcHJlVmw2VVhkWmJVcHJUbnBhYTA5VWEzZGFSR3Q0VFZkV2FWbFhTbTFOYW1oc1RrUnJkMWxVWkdsYVIxVjNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOYW1kM1RYcFplVTVxWnpWTmVsVjJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwVVZsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTTjBKSWEwRUtaSGRDTVVGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtNTJWekpUTVZGQlFVRlJSQXBCUlZsM1VrRkpaMDVLUVhkSlZWbERLM1V2VFVSTGVUWkpOM2xZV0ZGWWRtODJVV1pZZUVsNGEyaEVlbWgxWTJrdlNHTkRTVU5CVVhCcmR6QklhSFV6Q25JM1ptZEJVekpKZWpCU1NFVjFRVUpXVFhGTE5ESkdTM0JHT0hnNFdqRkpUVUZ2UjBORGNVZFRUVFE1UWtGTlJFRXlaMEZOUjFWRFRWRkVRMU5OWkdRS00yRk1ibEY2U0RWV1RrZHBOR1JKZW1saFMyZDBRMFpSVTBKcU9YRlJWMlZJTm5wblZqaHJSelJNVGpoWFdqTTRVMVp2ZWs5U1ZHNURSRmxEVFVSV1pncEVMMVYzVFZnNWRVa3JXbmRNTDNKNlRrbHJTbTlETnpKd2VVSnJVMDlYUVVkSksyUldSa05ZUjJGMGRqVXZOR2hWYWpVMmNtRTFNWFV2Vlc4cmR6MDlDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}],"timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyjADAgEAMIICwQYJKoZIhvcNAQcCoIICsjCCAq4CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgxI8zDvem/zyKTdFSqYy6ib/yQXUt6MGNvDecWKLmlP8CFQCyLByKNcEBVO5flM9JzLSAGTeVXxgPMjAyNjA2MjMxODIwMjVaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHbMIIB1wIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MDYyMzE4MjAyNVowLwYJKoZIhvcNAQkEMSIEIFbD/+/Wnhzjw4X7f0g1/a7iqf8YAorG9LVAShbEV894MIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRnMGUCMQCYEfENOE656RQjbnjoqSi9YoIXpl+HJs1aO+F4G/ej9WCPZOeCSMFgQTlTzwOxK0cCMAQ9fGAgliYgTgFYtY4fXBODpp/KHk6dgqx+sp6sfIHFTmbaEtes0fmrgKDDlad7zA=="}]}},"messageSignature":{"messageDigest":{"algorithm":"SHA2_256","digest":"6kA+kELqZFQK3e/Tr5fue3kmybmXkapuK9Z+ElnH4NE="},"signature":"MEUCIQCYMXTKOxWbh5pcMac5I3+FI9iidpIp8/IdKzC31rCRrQIgWUfXxapReamux9WdswYVafK0thYCBfdsbblryK44Cfc="}} \ No newline at end of file diff --git a/build/torch212-cxx11-xpu20253-x86_64-linux/ops/__init__.py b/build/torch212-cxx11-xpu20253-x86_64-linux/ops/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/build/torch212-cxx11-xpu20253-x86_64-linux/ops/activations.py b/build/torch212-cxx11-xpu20253-x86_64-linux/ops/activations.py deleted file mode 100644 index 7c09649fc41e12d5a360c5672825d8380bc7ec80..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-xpu20253-x86_64-linux/ops/activations.py +++ /dev/null @@ -1,135 +0,0 @@ -# Copied from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/model/layers/activations.py -import math - -import torch -import torch.nn as nn -import torch.nn.functional as F - -# 1/sqrt(2*pi)-> 0.3989423 -# 1/sqrt(2) -> 0.70710678 -# sqrt(2/pi) -> 0.79788456 - -# this function is tanh approximation of gelu -# actual gelu is: -# x * 0.5 * (1.0 + torch.erf(x * 0.70710678)) -@torch.jit.script -def bias_gelu(y, bias): - x = bias + y - return (x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))).to(dtype=y.dtype) - - -# gradient of tanh approximation of gelu -# gradient of actual gelu is: -# 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x) -@torch.jit.script -def bias_gelu_back(g, y, bias): - """Assume that y has shape (B, D) and bias has shape (D)""" - x = bias + y - tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)) - # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243 - ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( - 1 + tanh_out - ) - grad_y = ff * g - return grad_y.to(dtype=y.dtype), grad_y.sum(dim=(0), dtype=bias.dtype) - - -class GeLUFunction(torch.autograd.Function): - @staticmethod - # bias is an optional argument - def forward(ctx, input, bias): - ctx.save_for_backward(input, bias) - return bias_gelu(input, bias) - - @staticmethod - def backward(ctx, grad_output): - input, bias = ctx.saved_tensors - tmp = bias_gelu_back(grad_output, input, bias) - return tmp, tmp - - -bias_gelu_impl = GeLUFunction.apply - -# this function is tanh approximation of gelu -# actual gelu is: -# x * 0.5 * (1.0 + torch.erf(x * 0.70710678)) -@torch.jit.script -def gelu_fwd(x): - return (x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))).to(dtype=x.dtype) - - -# gradient of tanh approximation of gelu -# gradient of actual gelu is: -# 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x) -@torch.jit.script -def gelu_bwd(g, x): - tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)) - # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243 - ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( - 1 + tanh_out - ) - return (ff * g).to(dtype=x.dtype) - - -class FastGeLUFunction(torch.autograd.Function): - @staticmethod - # bias is an optional argument - def forward(ctx, input): - ctx.save_for_backward(input) - return gelu_fwd(input) - - @staticmethod - def backward(ctx, grad_output): - (input,) = ctx.saved_tensors - tmp = gelu_bwd(grad_output, input) - return tmp - - -fast_gelu_impl = FastGeLUFunction.apply - - -@torch.jit.script -def relu_bwd(g, x): - return torch.where(x >= 0, g, 0.0).to(dtype=x.dtype) - - -@torch.jit.script -def sqrelu_fwd(x): - r = F.relu(x) - return (r * r).to(dtype=x.dtype) - - -@torch.jit.script -def sqrelu_bwd(g, x): - return (2.0 * g * F.relu(x)).to(dtype=x.dtype) - - -swiglu_fwd_codestring = """ -template T swiglu_fwd(T x, T y) { - return float(x) * float(y) / (1.0f + ::exp(-float(x))); -} -""" -swiglu_bwd_codestring = """ -template void swiglu_bwd(T x, T y, T g, T& dx, T& dy) { - float x_sigmoid = 1.0f / (1.0f + ::exp(-float(x))); - dx = x_sigmoid * (1 + float(x) * (1.0f - x_sigmoid)) * float(g) * float(y); - dy = float(x) * x_sigmoid * float(g); -} -""" -swiglu_fwd = torch.cuda.jiterator._create_jit_fn(swiglu_fwd_codestring) -swiglu_bwd = torch.cuda.jiterator._create_multi_output_jit_fn(swiglu_bwd_codestring, num_outputs=2) - - -class SwiGLUFunction(torch.autograd.Function): - - @staticmethod - def forward(ctx, x, y): - ctx.save_for_backward(x, y) - return swiglu_fwd(x, y) - - @staticmethod - def backward(ctx, dout): - x, y = ctx.saved_tensors - return swiglu_bwd(x, y, dout) - -swiglu = SwiGLUFunction.apply diff --git a/build/torch212-cxx11-xpu20253-x86_64-linux/ops/fused_dense.py b/build/torch212-cxx11-xpu20253-x86_64-linux/ops/fused_dense.py deleted file mode 100644 index 6b4033d134e4093fe278f7b3f8c7d3128ce9f36d..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-xpu20253-x86_64-linux/ops/fused_dense.py +++ /dev/null @@ -1,688 +0,0 @@ -# Copyright (c) 2023, Tri Dao. -# Inspired by https://github.com/NVIDIA/apex/blob/master/apex/fused_dense/fused_dense.py -# We make it work with pytorch amp and with bfloat16. -# The TensorParallel linear modules are inspired by https://github.com/NVIDIA/apex/blob/master/apex/transformer/tensor_parallel/layers.py -from functools import partial -from typing import Optional - -# import fused_dense_cuda # from apex -import fused_dense_lib as fused_dense_cuda -import torch -import torch.nn as nn -import torch.nn.functional as F -from torch import Tensor -from torch.distributed import ProcessGroup - -from flash_attn.utils.torch import custom_fwd, custom_bwd -from flash_attn.ops.activations import gelu_bwd, relu_bwd, sqrelu_bwd, sqrelu_fwd -from flash_attn.utils.distributed import ( - all_gather_raw, - all_reduce, - all_reduce_raw, - reduce_scatter, - reduce_scatter_raw, -) - - -class FusedDenseFunc(torch.autograd.Function): - @staticmethod - @custom_fwd - def forward( - ctx, x, weight, bias, return_residual=False, process_group=None, sequence_parallel=True - ): - """ - If process_group is not None and sequence_parallel=True, we're doing Tensor Parallel - with sequence parallelism: we do an all_gather_raw of x before doing the matmul. - """ - ctx.compute_weight_gradient = weight.requires_grad - ctx.return_residual = return_residual - ctx.process_group = process_group - ctx.sequence_parallel = sequence_parallel - - if torch.is_autocast_enabled(): - x = x.to(dtype=torch.get_autocast_gpu_dtype()) - x = x.contiguous() - if process_group is not None and sequence_parallel: - # We want to kick off the all_gather early, before weight dtype conversion - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - else: - total_x = x - - if torch.is_autocast_enabled(): - weight = weight.to(dtype=torch.get_autocast_gpu_dtype()) - bias = bias.to(dtype=torch.get_autocast_gpu_dtype()) if bias is not None else None - weight = weight.contiguous() - if process_group is not None and sequence_parallel: - handle_x.wait() - batch_shape, n = total_x.shape[:-1], total_x.shape[-1] - batch_dim = batch_shape.numel() - # https://github.com/pytorch/pytorch/blob/5b51849b48a7dbccd297286cc0110def4706f9e7/aten/src/ATen/native/cuda/Blas.cpp#L174 - if min(batch_dim, n, *weight.shape) > 65535 * 32: - raise RuntimeError("fused_dense only supports matrix dims <= 2M") - output = F.linear(total_x, weight, bias) - if ctx.compute_weight_gradient: - ctx.save_for_backward(x, weight) - else: - ctx.save_for_backward(weight) - return output if not return_residual else (output, x) - - @staticmethod - @custom_bwd - def backward(ctx, grad_output, *args): - grad_output = grad_output.contiguous() - if ctx.return_residual: - (grad_input,) = args - grad_input = grad_input.contiguous() - process_group = ctx.process_group - sequence_parallel = ctx.sequence_parallel - if ctx.compute_weight_gradient: - x, weight = ctx.saved_tensors - if process_group is not None and sequence_parallel: - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - else: - total_x = x - else: - (weight,) = ctx.saved_tensors - total_x = None - batch_shape = grad_output.shape[:-1] - batch_dim = batch_shape.numel() - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - if ctx.needs_input_grad[0]: - if not ctx.return_residual: - grad_input = F.linear(grad_output, weight.t()) - else: - grad_input = torch.addmm( - grad_input.reshape(batch_dim, grad_input.shape[-1]), grad_output, weight - ) - grad_input = grad_input.reshape(*batch_shape, grad_input.shape[-1]) - if process_group is not None: - reduce_fn = reduce_scatter_raw if sequence_parallel else all_reduce_raw - grad_input, handle_grad_input = reduce_fn(grad_input, process_group, async_op=True) - else: - grad_input = None - if ctx.needs_input_grad[1]: - assert ctx.compute_weight_gradient - if process_group is not None and sequence_parallel: - handle_x.wait() - grad_weight, grad_bias = fused_dense_cuda.linear_bias_wgrad( - total_x.reshape(batch_dim, total_x.shape[-1]), grad_output, ctx.needs_input_grad[2] - ) - else: - grad_weight = None - grad_bias = grad_output if ctx.needs_input_grad[2] else None - if process_group is not None and ctx.needs_input_grad[0]: - handle_grad_input.wait() - return grad_input, grad_weight, grad_bias, None, None, None - - -def fused_dense_func( - x: Tensor, - weight: Tensor, - bias: Optional[Tensor] = None, - return_residual: bool = False, - process_group: Optional[ProcessGroup] = None, - sequence_parallel: bool = True, -): - dtype_eligible = x.dtype in [torch.float16, torch.bfloat16] or ( - x.dtype == torch.float32 and torch.is_autocast_enabled() - ) - if x.is_cuda and weight.is_cuda and (bias is None or bias.is_cuda) and dtype_eligible: - return FusedDenseFunc.apply( - x, weight, bias, return_residual, process_group, sequence_parallel - ) - else: - assert process_group is None - out = F.linear(x, weight, bias) - return out if not return_residual else (out, x) - - -class FusedDense(nn.Linear): - def __init__( - self, - in_features: int, - out_features: int, - bias: bool = True, - return_residual: bool = False, - device=None, - dtype=None, - ) -> None: - super().__init__(in_features, out_features, bias=bias, device=device, dtype=dtype) - self.return_residual = return_residual - - def forward(self, x, process_group=None): - """ - If process_group is not None, we're doing Tensor Parallel with sequence parallelism: - we do an all_gather of x before doing the matmul. - """ - return fused_dense_func( - x, - self.weight, - self.bias, - return_residual=self.return_residual, - process_group=process_group, - ) - - -class ColumnParallelLinear(nn.Linear): - def __init__( - self, - in_features: int, - out_features: int, - process_group: ProcessGroup, - bias: bool = True, - sequence_parallel=True, - multiple_of=1, - device=None, - dtype=None, - ) -> None: - world_size = torch.distributed.get_world_size(process_group) - if out_features % multiple_of: - raise ValueError(f"out_features ({out_features}) must be a multiple of {multiple_of}") - multiple = out_features // multiple_of - # We want to split @multiple across world_size, but it could be an uneven split - div = multiple // world_size - mod = multiple % world_size - # The first @mod ranks get @div + 1 copies, the rest get @div copies - local_multiple = div + int(torch.distributed.get_rank(process_group) < mod) - super().__init__( - in_features, local_multiple * multiple_of, bias=bias, device=device, dtype=dtype - ) - self.process_group = process_group - self.sequence_parallel = sequence_parallel - - def forward(self, x): - # If self.sequence_parallel is True, we're doing Tensor Parallel with sequence parallelism: - # we do an all_gather of x before doing the matmul. - # If not, then the input is already gathered. - return fused_dense_func( - x, - self.weight, - self.bias, - process_group=self.process_group, - sequence_parallel=self.sequence_parallel, - ) - - -class RowParallelLinear(nn.Linear): - def __init__( - self, - in_features: int, - out_features: int, - process_group: ProcessGroup, - bias: bool = True, - sequence_parallel=True, - multiple_of=1, - device=None, - dtype=None, - ) -> None: - world_size = torch.distributed.get_world_size(process_group) - rank = torch.distributed.get_rank(process_group) - if in_features % multiple_of: - raise ValueError(f"in_features ({in_features}) must be a multiple of {multiple_of}") - multiple = in_features // multiple_of - # We want to split @multiple across world_size, but it could be an uneven split - div = multiple // world_size - mod = multiple % world_size - # The first @mod ranks get @div + 1 copies, the rest get @div copies - local_multiple = div + int(torch.distributed.get_rank(process_group) < mod) - # Only rank 0 will have bias - super().__init__( - local_multiple * multiple_of, - out_features, - bias=bias and rank == 0, - device=device, - dtype=dtype, - ) - self.process_group = process_group - self.sequence_parallel = sequence_parallel - - def forward(self, x): - """ - We're doing Tensor Parallel with sequence parallelism: we do the matmul and then - a reduce_scatter of the result. - """ - out = fused_dense_func(x, self.weight, self.bias) - reduce_fn = reduce_scatter if self.sequence_parallel else all_reduce - return reduce_fn(out, self.process_group) - - -class FusedMLPFunc(torch.autograd.Function): - @staticmethod - @custom_fwd - def forward( - ctx, - x, - weight1, - bias1, - weight2, - bias2, - activation="gelu_approx", - save_pre_act=True, - return_residual=False, - checkpoint_lvl=0, - heuristic=0, - process_group=None, - sequence_parallel=True, - ): - """ - If process_group is not None and sequence_parallel=True, we're doing Tensor Parallel - with sequence parallelism: we do an all_gather of x before doing the matmul. - If sequence_parallel=False, then the input is already gathered. - - checkpoint_lvl: - 0: no recomputation in the bwd - 1: recompute gelu_out / relu_out in the bwd - 2: recompute pre_act and gelu_out / relu_out in the bwd - """ - assert -1 <= heuristic <= 4 - assert activation in ["gelu_approx", "relu", "sqrelu"] - if activation == "sqrelu": - assert heuristic == -1 - if not save_pre_act: - checkpoint_lvl = 2 - assert checkpoint_lvl in [0, 1, 2] - ctx.return_residual = return_residual - ctx.process_group = process_group - ctx.sequence_parallel = sequence_parallel - ctx.checkpoint_lvl = checkpoint_lvl - ctx.activation = activation - ctx.heuristic = heuristic - - if torch.is_autocast_enabled(): - x = x.to(dtype=torch.get_autocast_gpu_dtype()) - x = x.contiguous() - if process_group is not None and sequence_parallel: - # We want to kick off the all_gather early, before weight dtype conversion - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - else: - total_x = x - - if torch.is_autocast_enabled(): - dtype = torch.get_autocast_gpu_dtype() - weight1, weight2 = [a.to(dtype=dtype) for a in [weight1, weight2]] - bias1 = bias1.to(dtype=dtype) if bias1 is not None else None - bias2 = bias2.to(dtype=dtype) if bias2 is not None else None - weight1 = weight1.contiguous() - bias1 = bias1.contiguous() if bias1 is not None else None - weight2 = weight2.contiguous() - bias2 = bias2.contiguous() if bias2 is not None else None - if process_group is not None and sequence_parallel: - handle_x.wait() - batch_shape, n = total_x.shape[:-1], total_x.shape[-1] - batch_dim = batch_shape.numel() - # https://github.com/pytorch/pytorch/blob/5b51849b48a7dbccd297286cc0110def4706f9e7/aten/src/ATen/native/cuda/Blas.cpp#L174 - if min(batch_dim, n, *weight1.shape, *weight2.shape) > 65535 * 32: - raise RuntimeError("fused_dense only supports matrix dims <= 2M") - if heuristic == -1: - pre_act = F.linear(total_x, weight1, bias1) - activation_fn = ( - partial(F.gelu, approximate="tanh") - if activation == "gelu_approx" - else (sqrelu_fwd if activation == "sqrelu" else F.relu) - ) - with torch.jit.fuser("fuser2"): - output1 = activation_fn(pre_act) - # This is before adding bias1 - # pre_act = F.linear(total_x.reshape(batch_dim, n), weight1) - # with torch.jit.fuser('fuser2'): - # output1 = bias_gelu(pre_act, bias1) - else: - is_gelu = activation == "gelu_approx" - output1, *rest = fused_dense_cuda.linear_act_forward( - total_x.reshape(batch_dim, n), weight1, bias1, is_gelu, save_pre_act, heuristic - ) - if save_pre_act: - pre_act = rest[0] - output2 = F.linear(output1, weight2, bias2) - if checkpoint_lvl == 0 or (checkpoint_lvl == 1 and activation == "relu"): - # For RELU the pre_act is very small (just a bit-mask) so we just save it - ctx.save_for_backward(x, weight1, weight2, pre_act, output1) - elif checkpoint_lvl == 1: - ctx.save_for_backward(x, weight1, weight2, pre_act) - elif checkpoint_lvl == 2: - ctx.save_for_backward(x, weight1, weight2, bias1) - output2 = output2.reshape(*batch_shape, output2.shape[-1]) - return output2 if not return_residual else (output2, x) - - @staticmethod - @custom_bwd - def backward(ctx, grad_output, *args): - grad_output = grad_output.contiguous() - checkpoint_lvl = ctx.checkpoint_lvl - activation = ctx.activation - activation_fn = ( - partial(F.gelu, approximate="tanh") - if activation == "gelu_approx" - else (sqrelu_fwd if activation == "sqrelu" else F.relu) - ) - if ctx.return_residual: - (grad_input,) = args - grad_input = grad_input.contiguous() - process_group = ctx.process_group - sequence_parallel = ctx.sequence_parallel - x, weight1, weight2, *rest = ctx.saved_tensors - if process_group is None or not sequence_parallel: - total_x = x - batch_shape = grad_output.shape[:-1] - batch_dim = batch_shape.numel() - if checkpoint_lvl in [0, 1]: - if process_group is not None and sequence_parallel: - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - if checkpoint_lvl == 0 or (checkpoint_lvl == 1 and activation == "relu"): - pre_act, output1 = rest - elif checkpoint_lvl == 1: - (pre_act,) = rest - with torch.jit.fuser("fuser2"): - output1 = activation_fn(pre_act) - elif checkpoint_lvl == 2: - (bias1,) = rest - if process_group is not None and sequence_parallel: - total_x, _ = all_gather_raw(x, process_group) - if ctx.heuristic == -1: - pre_act = F.linear(total_x, weight1, bias1) - with torch.jit.fuser("fuser2"): - output1 = activation_fn(pre_act) - else: - output1, pre_act = fused_dense_cuda.linear_act_forward( - total_x.reshape(batch_dim, total_x.shape[-1]), - weight1, - bias1, - activation == "gelu_approx", - True, - ctx.heuristic, - ) - - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - output1 = output1.reshape(batch_dim, output1.shape[-1]) - pre_act = pre_act.reshape(batch_dim, pre_act.shape[-1]) - if ctx.needs_input_grad[3]: - grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad( - output1, grad_output, ctx.needs_input_grad[4] - ) - else: - grad_weight2 = None - grad_bias2 = grad_output if ctx.needs_input_grad[4] else None - if ctx.heuristic == -1: - # grad_pre_act = matmul_dgelu(grad_output, weight2, pre_act) - grad_output1 = F.linear(grad_output, weight2.t()) - activation_grad_fn = ( - gelu_bwd - if activation == "gelu_approx" - else (sqrelu_bwd if activation == "sqrelu" else relu_bwd) - ) - with torch.jit.fuser("fuser2"): - grad_pre_act = activation_grad_fn(grad_output1, pre_act) - else: - # The cublasLt epilogue has to compute both gelu/relu grad and bias grad, we can't - # just compute gelu/relu grad - grad_pre_act, grad_bias1 = fused_dense_cuda.bias_act_linear_dgrad_bgrad( - weight2, grad_output, pre_act, activation == "gelu_approx", ctx.heuristic - ) - if not ctx.needs_input_grad[2]: - grad_bias1 = None - if ctx.needs_input_grad[0]: - if not ctx.return_residual: - grad_input = F.linear(grad_pre_act, weight1.t()) - else: - grad_input = torch.addmm( - grad_input.reshape(batch_dim, grad_input.shape[-1]), grad_pre_act, weight1 - ) - grad_input = grad_input.reshape(*batch_shape, grad_input.shape[-1]) - if process_group is not None: - reduce_fn = reduce_scatter_raw if sequence_parallel else all_reduce_raw - grad_input, handle_grad_input = reduce_fn(grad_input, process_group, async_op=True) - else: - grad_input = None - if ctx.heuristic == -1: - if ctx.needs_input_grad[1]: - if process_group is not None and sequence_parallel and checkpoint_lvl != 2: - handle_x.wait() - grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_wgrad( - total_x.reshape(batch_dim, total_x.shape[-1]), - grad_pre_act, - ctx.needs_input_grad[2], - ) - else: - grad_weight1 = None - grad_bias1 = grad_pre_act if ctx.needs_input_grad[2] else None - else: - if ctx.needs_input_grad[1]: - if process_group is not None and sequence_parallel and checkpoint_lvl != 2: - handle_x.wait() - grad_weight1 = F.linear( - grad_pre_act.t(), total_x.reshape(batch_dim, total_x.shape[-1]).t() - ) - else: - grad_weight1 = None - if process_group is not None and ctx.needs_input_grad[0]: - handle_grad_input.wait() - return ( - grad_input, - grad_weight1, - grad_bias1, - grad_weight2, - grad_bias2, - None, - None, - None, - None, - None, - None, - None, - ) - - -def fused_mlp_func( - x: Tensor, - weight1: Tensor, - weight2: Tensor, - bias1: Optional[Tensor] = None, - bias2: Optional[Tensor] = None, - activation: str = "gelu_approx", - save_pre_act: bool = True, - return_residual: bool = False, - checkpoint_lvl: int = 0, - heuristic: int = 0, - process_group: Optional[ProcessGroup] = None, - sequence_parallel: bool = True, -): - assert activation in ["gelu_approx", "relu", "sqrelu"] - dtype_eligible = x.dtype in [torch.float16, torch.bfloat16] or ( - x.dtype == torch.float32 and torch.is_autocast_enabled() - ) - # If we save pre-activation, dimension must be divisible by 128 (relu) or 8 (gelu) - dim_eligible = not save_pre_act or (x.shape[-1] % (128 if activation == "relu" else 8) == 0) - if ( - x.is_cuda - and weight1.is_cuda - and weight2.is_cuda - and (bias1 is None or bias1.is_cuda) - and (bias2 is None or bias2.is_cuda) - and dtype_eligible - and dim_eligible - ): - return FusedMLPFunc.apply( - x, - weight1, - bias1, - weight2, - bias2, - activation, - save_pre_act, - return_residual, - checkpoint_lvl, - heuristic, - process_group, - sequence_parallel, - ) - else: - assert process_group is None - pre_act = F.linear(x, weight1, bias1) - activation_fn = ( - partial(F.gelu, approximate="tanh") - if activation == "gelu_approx" - else partial(F.relu, inplace=True) - ) - output1 = activation_fn(pre_act) - output2 = F.linear(output1, weight2, bias2) - return output2 if not return_residual else (output2, x) - - -class FusedMLP(nn.Module): - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - bias1=True, - bias2=True, - activation="gelu_approx", - return_residual=False, - checkpoint_lvl=0, - heuristic="auto", - device=None, - dtype=None, - ): - """ - If process_group is not None, we're doing Tensor Parallel with sequence parallelism: - we do an all_gather of x before doing the matmul, gelu, then matmul. - Finally we do a reduce_scatter of the output. - - checkpoint_lvl (increasing lvl means slower but more memory saving): - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute pre_act and gelu_out in the bwd - heuristic: - -1: don't fuse gemm + gelu (separate kernel) - 0..4: use this heuristic for the algo section in the fused gemm + gelu - 'auto': heuristic will be picked automatically: - For CUDA >= 11.8, we set heuristic=0 for both fp16 and bf16 for best perf. - For CUDA <= 11.7, we set heuristic=1 for fp16 and heuristic=-1 for bf16. - For H100, we set heuristic=-1 for both fp16 and bf16 as the fused cuBlasLt implementation - is slower than the unfused version. - return_residual: whether to return the input x along with the output. This is for - performance reason: for post-norm architecture, returning the input allows us - to fuse the backward of nn.Linear with the residual connection. - """ - assert checkpoint_lvl in [0, 1, 2] - assert activation in ["gelu_approx", "relu", "sqrelu"] - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features * 4 - self.activation = activation - self.return_residual = return_residual - self.checkpoint_lvl = checkpoint_lvl - self.heuristic = heuristic if activation != "sqrelu" else -1 - self.fc1 = nn.Linear(in_features, hidden_features, bias=bias1, **factory_kwargs) - self.fc2 = nn.Linear(hidden_features, out_features, bias=bias2, **factory_kwargs) - - def forward(self, x, process_group=None): - dtype = x.dtype if not torch.is_autocast_enabled() else torch.get_autocast_gpu_dtype() - if self.heuristic == "auto": - if self.activation == "gelu_approx": - if torch.cuda.get_device_capability("cuda") == (9, 0): - heuristic = -1 - else: - cuda_ver = tuple(map(int, torch.version.cuda.split("."))) - heuristic = 0 if cuda_ver >= (11, 8) else (1 if dtype == torch.float16 else -1) - else: - heuristic = 0 - else: - heuristic = self.heuristic - out = fused_mlp_func( - x, - self.fc1.weight, - self.fc2.weight, - self.fc1.bias, - self.fc2.bias, - activation=self.activation, - save_pre_act=self.training, - return_residual=self.return_residual, - checkpoint_lvl=self.checkpoint_lvl, - heuristic=heuristic, - process_group=process_group, - ) - if self.return_residual: - out, x = out - if process_group is not None: - out = reduce_scatter(out, process_group) - return out if not self.return_residual else (out, x) - - -class ParallelFusedMLP(nn.Module): - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - activation="gelu_approx", - process_group: ProcessGroup = None, - bias1=True, - bias2=True, - sequence_parallel=True, - checkpoint_lvl=0, - heuristic="auto", - device=None, - dtype=None, - ): - """ - process_group is required. We're doing Tensor Parallel with sequence parallelism: - we do an all_gather of x before doing the matmul, gelu, then matmul. - Finally we do a reduce_scatter of the output. - - checkpoint_lvl (increasing lvl means slower but more memory saving): - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute pre_act and gelu_out in the bwd - heuristic: - -1: don't fuse gemm + gelu (separate kernel) - 0..4: use this heuristic for the algo section in the fused gemm + gelu - 'auto': heuristic will be picked automatically: - For CUDA >= 11.8, we set heuristic=0 for both fp16 and bf16 for best perf. - For CUDA <= 11.7, we set heuristic=1 for fp16 and heuristic=-1 for bf16. - """ - assert checkpoint_lvl in [0, 1, 2] - assert activation in ["gelu_approx", "relu", "sqrelu"] - assert process_group is not None - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features * 4 - self.activation = activation - self.process_group = process_group - self.sequence_parallel = sequence_parallel - self.checkpoint_lvl = checkpoint_lvl - self.heuristic = heuristic if activation != "sqrelu" else -1 - self.fc1 = ColumnParallelLinear( - in_features, hidden_features, process_group, bias=bias1, **factory_kwargs - ) - self.fc2 = RowParallelLinear( - hidden_features, out_features, process_group, bias=bias2, **factory_kwargs - ) - - def forward(self, x): - dtype = x.dtype if not torch.is_autocast_enabled() else torch.get_autocast_gpu_dtype() - if self.heuristic == "auto": - if self.activation == "gelu_approx": - cuda_ver = tuple(map(int, torch.version.cuda.split("."))) - heuristic = 0 if cuda_ver >= (11, 8) else (1 if dtype == torch.float16 else -1) - else: - heuristic = 0 - else: - heuristic = self.heuristic - out = fused_mlp_func( - x, - self.fc1.weight, - self.fc2.weight, - self.fc1.bias, - self.fc2.bias, - activation=self.activation, - save_pre_act=self.training, - checkpoint_lvl=self.checkpoint_lvl, - heuristic=heuristic, - process_group=self.process_group, - sequence_parallel=self.sequence_parallel, - ) - reduce_fn = reduce_scatter if self.sequence_parallel else all_reduce - return reduce_fn(out, self.process_group) diff --git a/build/torch212-cxx11-xpu20253-x86_64-linux/ops/layer_norm.py b/build/torch212-cxx11-xpu20253-x86_64-linux/ops/layer_norm.py deleted file mode 100644 index 4b6cd798fd02844ef9cd3897f8ab95e490e638bf..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-xpu20253-x86_64-linux/ops/layer_norm.py +++ /dev/null @@ -1,800 +0,0 @@ -# Copyright (c) 2022, Tri Dao. -# Adapted from https://github.com/NVIDIA/apex/blob/master/apex/contrib/layer_norm/layer_norm.py - -import dropout_layer_norm -import torch -from torch.nn import init - - -def maybe_align(x, alignment_in_bytes=16): - """Assume that x already has last dim divisible by alignment_in_bytes""" - # TD [2023-07-04] I'm not 100% sure that clone will align the memory - # https://discuss.pytorch.org/t/how-to-ensure-that-tensor-data-ptr-is-aligned-to-16-bytes/183440 - return x if x.data_ptr() % alignment_in_bytes == 0 else x.clone() - - -def _dropout_add_layer_norm_forward( - x0, - residual, - gamma, - beta, - rowscale, - colscale, - dropout_p, - epsilon, - residual_in_fp32=False, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes""" - hidden_size = gamma.numel() - x0mat = x0.view((-1, hidden_size)) - residualmat = residual.view((-1, hidden_size)) if residual is not None else None - rowscale = rowscale.view(-1) if rowscale is not None else None - zmat, xmat, dmask, mu, rsigma = dropout_layer_norm.dropout_add_ln_fwd( - x0mat, - residualmat, - gamma, - beta, - rowscale, - colscale, - None, - None, - dropout_p, - epsilon, - 1.0, - 0, - None, - residual_in_fp32, - is_rms_norm, - ) - # dmask is None if dropout_p == 0.0 - # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype - return zmat, xmat if xmat is not None else x0mat, dmask, mu, rsigma - - -def _dropout_add_layer_norm_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - rowscale, - colscale, - dropout_p, - has_residual, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes - dx == None means that it was a post-norm architecture - (x = drop(x0) + residual was not returned in the fwd). - x0 must not be None if we have colscale. - """ - hidden_size = gamma.numel() - xmat = x.view((-1, hidden_size)) - dzmat = dz.view(xmat.shape) - dxmat = dx.view(xmat.shape) if dx is not None else None - x0mat = x0.view((-1, hidden_size)) if x0 is not None else None - rowscale = rowscale.view(-1) if rowscale is not None else None - if colscale is not None: - assert x0 is not None, "x0 is required to compute the gradient of colscale" - dx0mat, dresidualmat, dgamma, dbeta, _, _, *rest = dropout_layer_norm.dropout_add_ln_bwd( - dzmat, - dxmat, - xmat, - x0mat, - dmask, - mu, - rsigma, - gamma, - rowscale, - colscale, - None, - None, - dropout_p, - 1.0, - 0, - has_residual, - is_rms_norm, - ) - # dresidualmat is None if not has_residual - if colscale is None: - return dx0mat, dresidualmat, dgamma, dbeta - else: - dcolscale = rest[0] - return dx0mat, dresidualmat, dgamma, dbeta, dcolscale - - -def _dropout_add_layer_norm_subset_forward( - x0, - residual, - gamma, - beta, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32=False, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes""" - hidden_size = gamma.numel() - x0mat = x0.view((-1, hidden_size)) - residualmat = residual.view((-1, hidden_size)) if residual is not None else None - x0_subset = x0_subset.view(-1) if x0_subset is not None else None - out_subset = out_subset.view(-1) if out_subset is not None else None - zmat, xmat, dmask, mu, rsigma = dropout_layer_norm.dropout_add_ln_fwd( - x0mat, - residualmat, - gamma, - beta, - None, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - None, - residual_in_fp32, - is_rms_norm, - ) - # dmask is None if dropout_p == 0.0 - # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype - return zmat, xmat if xmat is not None else x0mat, dmask, mu, rsigma - - -def _dropout_add_layer_norm_subset_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - colscale, - x0_subset, - out_subset, - dropout_p, - rowscale_const, - x0_numrows, - has_residual, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes - dx == None means that it was a post-norm architecture - (x = drop(x0) + residual was not returned in the fwd). - x0 must not be None if we have colscale. - """ - hidden_size = gamma.numel() - xmat = x.view((-1, hidden_size)) - dzmat = dz.view(-1, hidden_size) - dxmat = dx.view(xmat.shape) if dx is not None else None - x0mat = x0.view((-1, hidden_size)) if x0 is not None else None - x0_subset = x0_subset.view(-1) if x0_subset is not None else None - out_subset = out_subset.view(-1) if out_subset is not None else None - if colscale is not None: - assert x0 is not None, "x0 is required to compute the gradient of colscale" - dx0mat, dresidualmat, dgamma, dbeta, _, _, *rest = dropout_layer_norm.dropout_add_ln_bwd( - dzmat, - dxmat, - xmat, - x0mat, - dmask, - mu, - rsigma, - gamma, - None, - colscale, - x0_subset, - out_subset, - dropout_p, - rowscale_const, - x0_numrows, - has_residual, - is_rms_norm, - ) - # dresidualmat is None if not has_residual - if colscale is None: - return dx0mat, dresidualmat, dgamma, dbeta - else: - dcolscale = rest[0] - return dx0mat, dresidualmat, dgamma, dbeta, dcolscale - - -def _dropout_add_layer_norm_parallel_residual_forward( - x0, - x1, - residual, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - residual_in_fp32=False, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes""" - hidden_size = gamma0.numel() - x0mat = x0.view((-1, hidden_size)) - x1mat = x1.view((-1, hidden_size)) if x1 is not None else None - residualmat = residual.view((-1, hidden_size)) if residual is not None else None - ( - z0mat, - z1mat, - xmat, - dmask0, - dmask1, - mu, - rsigma, - ) = dropout_layer_norm.dropout_add_ln_parallel_residual_fwd( - x0mat, - x1mat, - residualmat, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - None, - residual_in_fp32, - is_rms_norm, - ) - # dmask0 and dmask1 are None if dropout_p == 0.0 - # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype - return z0mat, z1mat, xmat if xmat is not None else x0mat, dmask0, dmask1, mu, rsigma - - -def _dropout_add_layer_norm_parallel_residual_backward( - dz0, - dz1, - dx, - x, - dmask0, - dmask1, - mu, - rsigma, - gamma0, - gamma1, - dropout_p, - has_x1, - has_residual, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes - dx == None means that it was a post-norm architecture - (x = drop(x0) + residual was not returned in the fwd). - """ - hidden_size = gamma0.numel() - xmat = x.view((-1, hidden_size)) - dz0mat = dz0.view(xmat.shape) - dz1mat = dz1.view(xmat.shape) if dz1 is not None else None - dxmat = dx.view(xmat.shape) if dx is not None else None - ( - dx0mat, - dx1mat, - dresidualmat, - dgamma0, - dbeta0, - dgamma1, - dbeta1, - *rest, - ) = dropout_layer_norm.dropout_add_ln_parallel_residual_bwd( - dz0mat, - dz1mat, - dxmat, - xmat, - dmask0, - dmask1, - mu, - rsigma, - gamma0, - gamma1, - dropout_p, - has_x1, - has_residual, - is_rms_norm, - ) - # dresidualmat is None if not has_residual - return dx0mat, dx1mat, dresidualmat, dgamma0, dbeta0, dgamma1, dbeta1 - - -class DropoutAddLayerNormFn(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x0, - residual, - gamma, - beta, - rowscale, - colscale, - dropout_p, - epsilon, - residual_in_fp32=False, - prenorm=False, - is_rms_norm=False, - return_dmask=False, - ): - x0 = maybe_align(x0.contiguous(), 16) - residual = maybe_align(residual.contiguous(), 16) if residual is not None else None - gamma = maybe_align(gamma.contiguous(), 16) - beta = maybe_align(beta.contiguous(), 16) if beta is not None else None - rowscale = maybe_align(rowscale.contiguous(), 16) if rowscale is not None else None - colscale = maybe_align(colscale.contiguous(), 16) if colscale is not None else None - zmat, xmat, dmask, mu, rsigma = _dropout_add_layer_norm_forward( - x0, - residual, - gamma, - beta, - rowscale, - colscale, - dropout_p, - epsilon, - residual_in_fp32, - is_rms_norm, - ) - # Only need to save x0 if we need to compute gradient wrt colscale - x0_saved = x0 if colscale is not None else None - ctx.save_for_backward( - xmat.view(x0.shape), x0_saved, dmask, gamma, mu, rsigma, rowscale, colscale - ) - ctx.prenorm = prenorm - ctx.dropout_p = dropout_p - ctx.has_residual = residual is not None - ctx.is_rms_norm = is_rms_norm - ctx.has_beta = beta is not None - if not return_dmask: - return ( - zmat.view(x0.shape) if not prenorm else (zmat.view(x0.shape), xmat.view(x0.shape)) - ) - else: - dmask = ( - dmask.view(x0.shape) - if dropout_p > 0.0 - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - ctx.mark_non_differentiable(dmask) - return ( - (zmat.view(x0.shape), dmask) - if not prenorm - else (zmat.view(x0.shape), xmat.view(x0.shape), dmask) - ) - - @staticmethod - def backward(ctx, dz, *args): - # assert dz.is_contiguous() - dz = maybe_align(dz.contiguous(), 16) # this happens! - dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None - x, x0, dmask, gamma, mu, rsigma, rowscale, colscale = ctx.saved_tensors - # x0 is None if colscale is None - dropout_p = ctx.dropout_p - has_residual = ctx.has_residual - dx0mat, dresidualmat, dgamma, dbeta, *rest = _dropout_add_layer_norm_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - rowscale, - colscale, - dropout_p, - has_residual, - ctx.is_rms_norm, - ) - dx0 = dx0mat.view(x.shape) - dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None - dcolscale = rest[0] if colscale is not None else None - return ( - dx0, - dresidual, - dgamma, - dbeta if ctx.has_beta else None, - None, - dcolscale, - None, - None, - None, - None, - None, - None, - ) - - -class DropoutAddLayerNormSubsetFn(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x0, - residual, - gamma, - beta, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32=False, - prenorm=False, - is_rms_norm=False, - return_dmask=False, - ): - x0 = maybe_align(x0.contiguous(), 16) - residual = maybe_align(residual.contiguous(), 16) if residual is not None else None - gamma = maybe_align(gamma.contiguous(), 16) - beta = maybe_align(beta.contiguous(), 16) if beta is not None else None - colscale = maybe_align(colscale.contiguous(), 16) if colscale is not None else None - zmat, xmat, dmask, mu, rsigma = _dropout_add_layer_norm_subset_forward( - x0, - residual, - gamma, - beta, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32, - is_rms_norm, - ) - # Only need to save x0 if we need to compute gradient wrt colscale - x0_saved = x0 if colscale is not None else None - x_shape = (-1, *x0.shape[1:]) - ctx.save_for_backward( - xmat.view(x_shape), x0_saved, dmask, gamma, mu, rsigma, colscale, x0_subset, out_subset - ) - ctx.prenorm = prenorm - ctx.dropout_p = dropout_p - ctx.rowscale_const = rowscale_const - ctx.x0_numrows = x0.shape[:-1].numel() - ctx.has_residual = residual is not None - ctx.is_rms_norm = is_rms_norm - ctx.has_beta = beta is not None - z_shape = (-1, *x0.shape[1:]) - if not return_dmask: - return zmat.view(z_shape) if not prenorm else (zmat.view(z_shape), xmat.view(x0.shape)) - else: - z = zmat.view(z_shape) - dmask = ( - dmask.view(x0.shape) - if dropout_p > 0.0 - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - ctx.mark_non_differentiable(dmask) - return (z, dmask) if not prenorm else (z, xmat.view(x_shape), dmask) - - @staticmethod - def backward(ctx, dz, *args): - # assert dz.is_contiguous() - dz = maybe_align(dz.contiguous(), 16) # this happens! - dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None - x, x0, dmask, gamma, mu, rsigma, colscale, x0_subset, out_subset = ctx.saved_tensors - # x0 is None if colscale is None - dropout_p = ctx.dropout_p - has_residual = ctx.has_residual - dx0mat, dresidualmat, dgamma, dbeta, *rest = _dropout_add_layer_norm_subset_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - colscale, - x0_subset, - out_subset, - dropout_p, - ctx.rowscale_const, - ctx.x0_numrows, - has_residual, - ctx.is_rms_norm, - ) - dx0 = dx0mat.view(-1, *x.shape[1:]) - dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None - dcolscale = rest[0] if colscale is not None else None - return ( - dx0, - dresidual, - dgamma, - dbeta if ctx.has_beta else None, - dcolscale, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) - - -class DropoutAddLayerNormParallelResidualFn(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x0, - x1, - residual, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - residual_in_fp32=False, - prenorm=False, - is_rms_norm=False, - return_dmask=False, - ): - x0 = maybe_align(x0.contiguous(), 16) - x1 = maybe_align(x1.contiguous(), 16) if x1 is not None else None - residual = maybe_align(residual.contiguous(), 16) if residual is not None else None - gamma0 = maybe_align(gamma0.contiguous(), 16) - beta0 = maybe_align(beta0.contiguous(), 16) if beta0 is not None else None - gamma1 = maybe_align(gamma1.contiguous(), 16) if gamma1 is not None else None - beta1 = maybe_align(beta1.contiguous(), 16) if beta1 is not None else None - ( - z0mat, - z1mat, - xmat, - dmask0, - dmask1, - mu, - rsigma, - ) = _dropout_add_layer_norm_parallel_residual_forward( - x0, - x1, - residual, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - residual_in_fp32, - is_rms_norm, - ) - ctx.save_for_backward(xmat.view(x0.shape), dmask0, dmask1, gamma0, gamma1, mu, rsigma) - ctx.prenorm = prenorm - ctx.dropout_p = dropout_p - ctx.has_x1 = x1 is not None - ctx.has_residual = residual is not None - ctx.is_rms_norm = is_rms_norm - ctx.has_beta = beta0 is not None - z = (z0mat.view(x0.shape), z1mat.view(x0.shape) if z1mat is not None else None) - if not return_dmask: - return z if not prenorm else (*z, xmat.view(x0.shape)) - else: - dmask0 = ( - dmask0.view(x0.shape) - if dropout_p > 0.0 - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - dmask1 = ( - dmask1.view(x0.shape) - if dropout_p > 0.0 and x1 is not None - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - ctx.mark_non_differentiable(dmask0) - ctx.mark_non_differentiable(dmask1) - return ( - (*z, dmask0, dmask1) if not prenorm else (*z, xmat.view(x0.shape), dmask0, dmask1) - ) - - @staticmethod - def backward(ctx, dz0, dz1, *args): - dz0 = maybe_align(dz0.contiguous(), 16) # this happens! - dz1 = maybe_align(dz1.contiguous(), 16) if dz1 is not None else None - dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None - x, dmask0, dmask1, gamma0, gamma1, mu, rsigma = ctx.saved_tensors - dropout_p = ctx.dropout_p - has_x1 = ctx.has_x1 - has_residual = ctx.has_residual - ( - dx0mat, - dx1mat, - dresidualmat, - dgamma0, - dbeta0, - dgamma1, - dbeta1, - ) = _dropout_add_layer_norm_parallel_residual_backward( - dz0, - dz1, - dx, - x, - dmask0, - dmask1, - mu, - rsigma, - gamma0, - gamma1, - dropout_p, - has_x1, - has_residual, - ctx.is_rms_norm, - ) - dx0 = dx0mat.view(x.shape) - dx1 = dx1mat.view(x.shape) if dx1mat is not None else None - dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None - return ( - dx0, - dx1, - dresidual, - dgamma0, - dbeta0 if ctx.has_beta else None, - dgamma1, - dbeta1 if ctx.has_beta else None, - None, - None, - None, - None, - None, - None, - ) - - -def layer_norm(x, weight, bias, epsilon): - return DropoutAddLayerNormFn.apply(x, None, weight, bias, None, None, 0.0, epsilon, False) - - -def dropout_add_layer_norm( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - rowscale=None, - layerscale=None, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormFn.apply( - x0, - residual, - weight, - bias, - rowscale, - layerscale, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - False, - return_dropout_mask, - ) - - -def dropout_add_layer_norm_subset( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - layerscale=None, - x0_subset=None, - out_subset=None, - rowscale_const=1.0, - out_numrows=0, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormSubsetFn.apply( - x0, - residual, - weight, - bias, - layerscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32, - prenorm, - False, - return_dropout_mask, - ) - - -def dropout_add_layer_norm_parallel_residual( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormParallelResidualFn.apply( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - False, - return_dropout_mask, - ) - - -class DropoutAddLayerNorm(torch.nn.Module): - def __init__( - self, - hidden_size, - prenorm=False, - p=0.0, - eps=1e-5, - residual_in_fp32=False, - device=None, - dtype=None, - ): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.prenorm = prenorm - self.p = p - self.eps = eps - self.residual_in_fp32 = residual_in_fp32 - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.bias = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.reset_parameters() - - def reset_parameters(self): - init.ones_(self.weight) - init.zeros_(self.bias) - - def forward(self, x0, residual=None): - return dropout_add_layer_norm( - x0, - residual, - self.weight, - self.bias, - self.p if self.training else 0.0, - self.eps, - prenorm=self.prenorm, - residual_in_fp32=self.residual_in_fp32, - ) diff --git a/build/torch212-cxx11-xpu20253-x86_64-linux/ops/rms_norm.py b/build/torch212-cxx11-xpu20253-x86_64-linux/ops/rms_norm.py deleted file mode 100644 index 068348d61290e3839dd082b540d898578ba1e8e2..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-xpu20253-x86_64-linux/ops/rms_norm.py +++ /dev/null @@ -1,174 +0,0 @@ -# Copyright (c) 2022, Tri Dao. -# Adapted from https://github.com/NVIDIA/apex/blob/master/apex/contrib/layer_norm/layer_norm.py - -import torch -from torch.nn import init - -from flash_attn.ops.layer_norm import ( - DropoutAddLayerNormFn, - DropoutAddLayerNormParallelResidualFn, - DropoutAddLayerNormSubsetFn, -) - - -def rms_norm(x, weight, epsilon): - return DropoutAddLayerNormFn.apply( - x, None, weight, None, None, None, 0.0, epsilon, False, False, True - ) - - -def dropout_add_rms_norm( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - rowscale=None, - layerscale=None, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormFn.apply( - x0, - residual, - weight, - bias, - rowscale, - layerscale, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - True, - return_dropout_mask, - ) - - -def dropout_add_rms_norm_subset( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - layerscale=None, - x0_subset=None, - out_subset=None, - rowscale_const=1.0, - out_numrows=0, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormSubsetFn.apply( - x0, - residual, - weight, - bias, - layerscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32, - prenorm, - True, - return_dropout_mask, - ) - - -def dropout_add_rms_norm_parallel_residual( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormParallelResidualFn.apply( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - True, - return_dropout_mask, - ) - - -class RMSNorm(torch.nn.Module): - def __init__(self, hidden_size, eps=1e-5, device=None, dtype=None): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.eps = eps - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.register_parameter("bias", None) - self.reset_parameters() - - def reset_parameters(self): - init.ones_(self.weight) - - def forward(self, x): - return rms_norm(x, self.weight, self.eps) - - -class DropoutAddRMSNorm(torch.nn.Module): - def __init__( - self, - hidden_size, - prenorm=False, - p=0.0, - eps=1e-5, - residual_in_fp32=False, - device=None, - dtype=None, - ): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.prenorm = prenorm - self.p = p - self.eps = eps - self.residual_in_fp32 = residual_in_fp32 - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.register_parameter("bias", None) - self.reset_parameters() - - def reset_parameters(self): - init.ones_(self.weight) - - def forward(self, x0, residual=None): - return dropout_add_rms_norm( - x0, - residual, - self.weight, - None, - self.p if self.training else 0.0, - self.eps, - prenorm=self.prenorm, - residual_in_fp32=self.residual_in_fp32, - ) diff --git a/build/torch212-cxx11-xpu20253-x86_64-linux/ops/triton/__init__.py b/build/torch212-cxx11-xpu20253-x86_64-linux/ops/triton/__init__.py deleted file mode 100644 index 8b137891791fe96927ad78e64b0aad7bded08bdc..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-xpu20253-x86_64-linux/ops/triton/__init__.py +++ /dev/null @@ -1 +0,0 @@ - diff --git a/build/torch212-cxx11-xpu20253-x86_64-linux/ops/triton/cross_entropy.py b/build/torch212-cxx11-xpu20253-x86_64-linux/ops/triton/cross_entropy.py deleted file mode 100644 index 1b5a415b73f236f3e05fb14b9141959559e18526..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-xpu20253-x86_64-linux/ops/triton/cross_entropy.py +++ /dev/null @@ -1,330 +0,0 @@ -# Copyright (c) 2023, Tri Dao. - -from typing import Tuple, Optional, Union - -import torch -import torch.nn.functional as F - -import triton -import triton.language as tl - -# `all_gather_into_tensor` and `reduce_scatter_tensor` are new placeholders for -# `_all_gather_base` and `_reduce_scatter_base`. They require the most recent -# version of PyTorch. The following 2 lines are for backward compatibility with -# older PyTorch. -if "all_gather_into_tensor" not in dir(torch.distributed): - torch.distributed.all_gather_into_tensor = torch.distributed._all_gather_base - - -@triton.heuristics( - { - "HAS_SMOOTHING": lambda args: args["smoothing"] > 0.0, - } -) -@triton.jit -def cross_entropy_fwd_kernel( - loss_ptr, # data ptrs - lse_ptr, - z_loss_ptr, - logits_ptr, - labels_ptr, - smoothing, - logit_scale, - lse_square_scale, - ignore_index, - total_classes, - class_start_idx, # Useful for tensor parallel when each rank only has a subset of classes - n_cols, # shapes - logits_row_stride, # strides - BLOCK_SIZE: tl.constexpr, - HAS_SMOOTHING: tl.constexpr, - # if SPLIT (e.g. tensor parallel), don't include the LSE in the loss since it's not the final LSE - SPLIT: tl.constexpr, - PRECOMPUTED_LSE: tl.constexpr, # If LSE is already computed (also no smoothing and logit_scale == 1.0) -): - row_idx = tl.program_id(0) - logits_ptr = logits_ptr + row_idx * logits_row_stride.to(tl.int64) - sum_logits = 0.0 # For smoothing - if not PRECOMPUTED_LSE: - # Statistics for online softmax - m_i = -float("inf") - l_i = 0.0 - for col_offset in range(0, n_cols, BLOCK_SIZE): - cols = col_offset + tl.arange(0, BLOCK_SIZE) - logits = tl.load(logits_ptr + cols, mask=cols < n_cols, other=-float("inf")).to( - tl.float32 - ) * logit_scale - if HAS_SMOOTHING: - sum_logits += tl.sum(tl.where(cols < n_cols, logits, 0.0)) - m_i_new = tl.maximum(m_i, tl.max(logits)) - l_i = tl.exp(m_i - m_i_new) * l_i + tl.sum(tl.exp(logits - m_i_new)) - m_i = m_i_new - lse = tl.log(l_i) + m_i - tl.store(lse_ptr + row_idx, lse) - else: - lse = tl.load(lse_ptr + row_idx) - label_idx = tl.load(labels_ptr + row_idx) - if label_idx == ignore_index: - loss = 0.0 - z_loss = 0.0 - else: - label_idx -= class_start_idx - if label_idx >= 0 and label_idx < n_cols: - logits_label = tl.load(logits_ptr + label_idx) * logit_scale - if HAS_SMOOTHING: - loss = ( - (lse if not SPLIT else 0.0) - - smoothing * sum_logits / total_classes - - (1 - smoothing) * logits_label - ) - else: - loss = (lse if not SPLIT else 0.0) - logits_label - else: - # If label is out of bounds, we set the CE loss to 0.0. But we still want the smoothing loss - if HAS_SMOOTHING: - loss = smoothing * ((lse if not SPLIT else 0.0) - sum_logits / total_classes) - else: - loss = 0.0 - if not SPLIT: - z_loss = lse_square_scale * lse * lse - loss += z_loss - else: - z_loss = 0.0 - tl.store(loss_ptr + row_idx, loss) - if not SPLIT: - tl.store(z_loss_ptr + row_idx, z_loss) - - -@triton.heuristics( - { - "HAS_SMOOTHING": lambda args: args["smoothing"] > 0.0, - } -) -@triton.jit -def cross_entropy_bwd_kernel( - dlogits_ptr, # data ptrs - dloss_ptr, - logits_ptr, - lse_ptr, - labels_ptr, - smoothing, - logit_scale, - lse_square_scale, - ignore_index, - total_classes, - class_start_idx, # Useful for tensor parallel when each rank only has a subset of classes - n_cols, # shapes - logits_row_stride, # strides - dlogits_row_stride, - dloss_row_stride, - BLOCK_SIZE: tl.constexpr, - HAS_SMOOTHING: tl.constexpr, -): - row_idx = tl.program_id(0) - col_block_idx = tl.program_id(1) - logits_ptr = logits_ptr + row_idx * logits_row_stride.to(tl.int64) - dlogits_ptr = dlogits_ptr + row_idx * dlogits_row_stride.to(tl.int64) - col_offsets = col_block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - label_idx = tl.load(labels_ptr + row_idx) - if label_idx != ignore_index: - dloss = tl.load(dloss_ptr + row_idx * dloss_row_stride) - else: - dloss = 0.0 - logits = tl.load(logits_ptr + col_offsets, mask=col_offsets < n_cols, other=-float("inf")).to( - tl.float32 - ) * logit_scale - lse = tl.load(lse_ptr + row_idx) - probs = tl.exp(logits - lse) - probs += 2.0 * lse_square_scale * lse * probs - label_idx -= class_start_idx - if HAS_SMOOTHING: - smooth_positive = 1.0 - smoothing - smooth_negative = smoothing / total_classes - probs = tl.where(col_offsets == label_idx, probs - smooth_positive, probs) - smooth_negative - else: - probs = tl.where(col_offsets == label_idx, probs - 1.0, probs) - tl.store(dlogits_ptr + col_offsets, (dloss * logit_scale) * probs, mask=col_offsets < n_cols) - - -class CrossEntropyLoss(torch.autograd.Function): - - @staticmethod - def forward( - ctx, - logits, - labels, - precomputed_lse=None, - smoothing=0.0, - logit_scale=1.0, - lse_square_scale=0.0, - ignore_index=-100, - inplace_backward=False, - process_group=None, - ): - # For some reason Triton generates wrong code when labels has dtype long and its address - # is not aligned to 16 bytes. The ld.global.b64 seems to load the wrong label index. - if labels.dtype == torch.long and labels.data_ptr() % 16 != 0: - labels = F.pad(labels, (0, 1))[..., :-1] - assert labels.data_ptr() % 16 == 0 - assert logit_scale > 0.0 - n_rows, n_cols = logits.shape - assert labels.shape == (n_rows,) - world_size = 1 if process_group is None else torch.distributed.get_world_size(process_group) - total_classes = world_size * n_cols - rank = 0 if process_group is None else torch.distributed.get_rank(process_group) - class_start_idx = rank * n_cols - use_precomputed_lse = precomputed_lse is not None and logit_scale == 1.0 and smoothing == 0.0 - - if logits.stride(-1) != 1: - logits = logits.contiguous() - MAX_BLOCK_SIZE = 16 * 1024 - BLOCK_SIZE = min(triton.next_power_of_2(n_cols), MAX_BLOCK_SIZE) - num_warps = ( - 4 - if BLOCK_SIZE < 2048 - else (8 if BLOCK_SIZE < 8192 else (16 if BLOCK_SIZE < 128 * 1024 else 32)) - ) - losses = torch.empty(n_rows, dtype=torch.float, device=logits.device) - if use_precomputed_lse: - assert precomputed_lse.shape == (n_rows,) - lse = precomputed_lse.contiguous() - else: - lse = torch.empty(n_rows, dtype=torch.float, device=logits.device) - z_losses = torch.empty(n_rows, dtype=torch.float, device=logits.device) - # Need this, otherwise Triton tries to launch from cuda:0 and we get - # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?) - with torch.cuda.device(logits.device.index): - cross_entropy_fwd_kernel[(n_rows,)]( - losses, # data ptrs - lse, - z_losses, - logits, - labels, - smoothing, - logit_scale, - lse_square_scale, - ignore_index, - total_classes, - class_start_idx, - n_cols, # shapes - logits.stride(0), # strides - BLOCK_SIZE=BLOCK_SIZE, # constants - SPLIT=world_size > 1, - PRECOMPUTED_LSE=use_precomputed_lse, - num_warps=num_warps, - ) - - if world_size > 1: - # If there's no smoothing, if labels are in the vocab of this partition, losses contains - # - predicted logit, and 0 otherwise. - # If there's smoothing=0.1, for labels in the vocab of this partition, losses contains - # -0.9 * predicted logit - 0.1 * sum logit / total_classes. - # For labels not in the vocab of this partition, losses contains - # -0.1 * sum logit / total_classes. - if world_size > 1: - lse_allgather = torch.empty(world_size, n_rows, dtype=lse.dtype, device=lse.device) - torch.distributed.all_gather_into_tensor(lse_allgather, lse, group=process_group) - handle_losses = torch.distributed.all_reduce( - losses, op=torch.distributed.ReduceOp.SUM, group=process_group, async_op=True - ) - lse = torch.logsumexp(lse_allgather, dim=0) - handle_losses.wait() - # After the allreduce, if there's no smoothing, the total losses are - predicted_logit, - # we just have to add the (global) lse. - # If there's smoothing=0.1, the total losses are - # -0.9 * predicted_logit - 0.1 * sum logit / total_classes. - # Again, we just have to add the (global) lse. - losses += lse - if lse_square_scale != 0.0: - z_losses = lse_square_scale * lse.square() - z_losses.masked_fill_(labels == ignore_index, 0.0) - losses += z_losses - else: - z_losses = torch.zeros_like(losses) - losses.masked_fill_(labels == ignore_index, 0.0) - - ctx.save_for_backward(logits, lse, labels) - ctx.mark_non_differentiable(z_losses) - ctx.smoothing = smoothing - ctx.logit_scale = logit_scale - ctx.lse_square_scale = lse_square_scale - ctx.ignore_index = ignore_index - ctx.total_classes = total_classes - ctx.class_start_idx = class_start_idx - ctx.inplace_backward = inplace_backward - return losses, z_losses - - @staticmethod - def backward(ctx, grad_losses, grad_z_losses): - del grad_z_losses # z_losses are only for logging. - - logits, lse, labels = ctx.saved_tensors - dlogits = logits if ctx.inplace_backward else torch.empty_like(logits) - n_rows, n_cols = logits.shape - BLOCK_SIZE = min(triton.next_power_of_2(n_cols), 4 * 1024) - num_warps = 4 if BLOCK_SIZE < 2048 else (8 if BLOCK_SIZE < 8192 else 16) - grid = lambda META: (n_rows, triton.cdiv(n_cols, META["BLOCK_SIZE"])) # noqa - # Need this, otherwise Triton tries to launch from cuda:0 and we get - # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?) - with torch.cuda.device(logits.device.index): - cross_entropy_bwd_kernel[grid]( - dlogits, # data ptrs - grad_losses, - logits, - lse, - labels, - ctx.smoothing, - ctx.logit_scale, - ctx.lse_square_scale, - ctx.ignore_index, - ctx.total_classes, - ctx.class_start_idx, - n_cols, # shapes - logits.stride(0), # strides - dlogits.stride(0), - grad_losses.stride(0), - BLOCK_SIZE=BLOCK_SIZE, # constants - num_warps=num_warps, - ) - return dlogits, None, None, None, None, None, None, None, None, None - - -def cross_entropy_loss( - logits: torch.Tensor, - labels: torch.Tensor, - precomputed_lse: Optional[torch.Tensor] = None, - label_smoothing: float = 0.0, - logit_scale: float = 1.0, - lse_square_scale: float = 0.0, - ignore_index=-100, - inplace_backward: bool = False, - process_group=None, -) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Arguments: - logits: (batch, vocab_size) - labels: (batch,) - label_smoothing: float - logit_scale: float. Multiply logits by this scale before calculating the loss. - lse_square_scale: float. If > 0, we add lse_square_scale * lse(logits) ^ 2 to the loss. - This is also referred to as "z-loss". - ignore_index: int. If labels == ignore_index, the loss is set to 0.0. - inplace_backward: bool. If True, we do the backward pass in-place by modifying the logits. - This saves memory. - process_group: if not None, we're doing Tensor Parallel: each process is responsible for - one part of the vocab. The loss will be aggregated across processes. - Returns: - losses: (batch,), float - z_losses: (batch,), float - """ - return CrossEntropyLoss.apply( - logits, - labels, - precomputed_lse, - label_smoothing, - logit_scale, - lse_square_scale, - ignore_index, - inplace_backward, - process_group, - ) diff --git a/build/torch212-cxx11-xpu20253-x86_64-linux/ops/triton/k_activations.py b/build/torch212-cxx11-xpu20253-x86_64-linux/ops/triton/k_activations.py deleted file mode 100644 index efb83c358eb4a85d069ee340a3c83f418f9a805b..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-xpu20253-x86_64-linux/ops/triton/k_activations.py +++ /dev/null @@ -1,162 +0,0 @@ -# Adapted from https://github.com/facebookresearch/xformers/blob/main/xformers/triton/k_activations.py -# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved. -# -# This source code is licensed under the BSD license found in the -# LICENSE file in the root directory of this source tree. - -import math -from enum import Enum -from typing import Optional - -import triton -import triton.language as tl - -_sqrt2pi = math.sqrt(2.0 / math.pi) -_sqrt1_2 = math.sqrt(1.0 / 2) -_gaussian_pdf_normalization = 1.0 / math.sqrt(2 * math.pi) - - -class Activation(str, Enum): - SquaredReLU = "squared_relu" - GeLU = "gelu" - GeLUApprox = "gelu_approx" - LeakyReLU = "leaky_relu" - ReLU = "relu" - - -def get_triton_activation_kernel(activation: Optional[Activation]): - return ( - { - Activation.ReLU: relu, - Activation.LeakyReLU: leaky_relu, - Activation.GeLU: gelu, - Activation.GeLUApprox: gelu_approx, - Activation.SquaredReLU: squared_relu, - }[activation] - if activation - else None - ) - - -def get_triton_activation_bwd_kernel(activation: Optional[Activation]): - return ( - { - Activation.ReLU: relu_grad, - Activation.LeakyReLU: leaky_relu_grad, - Activation.GeLU: gelu_grad, - Activation.GeLUApprox: gelu_approx_grad, - Activation.SquaredReLU: squared_relu_grad, - }[activation] - if activation - else None - ) - - -@triton.jit -def tanh(x): - # Tanh is just a scaled sigmoid - return 2 * tl.sigmoid(2 * x) - 1 - - -@triton.jit -def cosh(x): - exp_x = tl.exp(x) - return (exp_x + 1.0 / exp_x) * 0.5 - - -# a Triton implementation of the most used activations -# See for instance http://arxiv.org/abs/1606.08415 for an overview - -# ReLU -@triton.jit -def relu(x): - """ - ReLU_ activation function - - .. _ReLU: https://pytorch.org/docs/stable/generated/torch.nn.ReLU.html - """ - zero = 0.0 - return tl.where(x >= 0, x, zero.to(x.dtype)) - - -@triton.jit -def relu_grad(x): - # ReLU is different from other activations - # in that it does not require the input to retrospectively compute its gradient - # here the input is the downstream gradient, and we return the upstream gradient directly - zero = 0.0 - one = 1.0 - return tl.where(x >= 0, one.to(x.dtype), zero.to(x.dtype)) - - -@triton.jit -def squared_relu(x): - """ - Squared ReLU activation, as proposed in the Primer_ paper. - - .. _Primer: https://arxiv.org/abs/2109.08668 - """ - x_ = relu(x) - return (x_ * x_).to(x.dtype) - - -@triton.jit -def squared_relu_grad(x): - return tl.where(x >= 0, 2.0 * x, 0.0) - - -# Leaky ReLU -@triton.jit -def leaky_relu(x): - """ - LeakyReLU_ activation - - .. _LeakyReLU: https://pytorch.org/docs/stable/generated/torch.nn.LeakyReLU.html - """ - scale = 0.01 + 0.0 - scale = scale.to(x.dtype) - return tl.where(x >= 0, x, scale * x) - - -@triton.jit -def leaky_relu_grad(x): - min_grad = 0.01 - max_grad = 1 - - min_grad = min_grad.to(x.dtype) - max_grad = max_grad.to(x.dtype) - - return tl.where(x >= 0, max_grad, min_grad) - - -@triton.jit -def gelu(x): - """Gaussian Error Linear Unit (GELU)""" - return x * 0.5 * (1.0 + tl.libdevice.erf(x * _sqrt1_2)) - - -@triton.jit -def gelu_grad(x): - cdf = 0.5 * (1.0 + tl.libdevice.erf(x * _sqrt1_2)) - pdf = tl.exp(-0.5 * x * x) * _gaussian_pdf_normalization - return cdf + x * pdf - - -@triton.jit -def gelu_approx(x): - """ - GeLU_ activation - Gaussian error linear unit, with tanh approximation - - .. _GeLU: https://arxiv.org/pdf/1606.08415.pdf - """ - return 0.5 * x * (1.0 + tanh(_sqrt2pi * x * (1.0 + 0.044715 * x * x))) - - -@triton.jit -def gelu_approx_grad(x): - # CREDITS: Fast implementation proposed in - # https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/fused_bias_gelu.py#L30 - tanh_out = tanh(0.79788456 * x * (1 + 0.044715 * x * x)) - return 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( - 1 + tanh_out - ) diff --git a/build/torch212-cxx11-xpu20253-x86_64-linux/ops/triton/layer_norm.py b/build/torch212-cxx11-xpu20253-x86_64-linux/ops/triton/layer_norm.py deleted file mode 100644 index 5e6f0b3697b612e6a19031bbc55ac76ad233dc85..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-xpu20253-x86_64-linux/ops/triton/layer_norm.py +++ /dev/null @@ -1,1254 +0,0 @@ -# Copyright (c) 2024, Tri Dao. -# Implement dropout + residual + layer_norm / rms_norm. - -# Based on the Triton LayerNorm tutorial: https://triton-lang.org/main/getting-started/tutorials/05-layer-norm.html -# For the backward pass, we keep weight_grad and bias_grad in registers and accumulate. -# This is faster for dimensions up to 8k, but after that it's much slower due to register spilling. -# The models we train have hidden dim up to 8k anyway (e.g. Llama 70B), so this is fine. - -import math -from typing import Optional, List - -import torch -import torch.nn.functional as F -from torch import Tensor - -import triton -import triton.language as tl - -from flash_attn.utils.torch import custom_fwd, custom_bwd -from flash_attn.utils.library import triton_op - -from ..._ops import add_op_namespace_prefix - - -def maybe_contiguous_lastdim(x): - return x.contiguous() if x is not None and x.stride(-1) != 1 else x - - -def maybe_contiguous(x): - return x.contiguous() if x is not None else None - - -def triton_autotune_configs(): - # Return configs with a valid warp count for the current device - configs = [] - # Maximum threads per block is architecture-dependent in theory, but in reality all are 1024 - max_threads_per_block = 1024 - # Default to warp size 32 if not defined by device - warp_size = getattr(torch.cuda.get_device_properties(torch.cuda.current_device()), "warp_size", 32) - # Autotune for warp counts which are powers of 2 and do not exceed thread per block limit - return [triton.Config({}, num_warps=warp_count) for warp_count in [1, 2, 4, 8, 16, 32] - if warp_count * warp_size <= max_threads_per_block] - # return [triton.Config({}, num_warps=8)] - - -def layer_norm_ref( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - zero_centered_weight=False, - dropout_mask=None, - dropout_mask1=None, - upcast=False, -): - dtype = x.dtype - if upcast: - x = x.float() - weight = weight.float() - bias = bias.float() if bias is not None else None - residual = residual.float() if residual is not None else residual - x1 = x1.float() if x1 is not None else None - weight1 = weight1.float() if weight1 is not None else None - bias1 = bias1.float() if bias1 is not None else None - if zero_centered_weight: - weight = weight + 1.0 - if weight1 is not None: - weight1 = weight1 + 1.0 - if x1 is not None: - assert rowscale is None, "rowscale is not supported with parallel LayerNorm" - if rowscale is not None: - x = x * rowscale[..., None] - if dropout_p > 0.0: - if dropout_mask is not None: - x = x.masked_fill(~dropout_mask, 0.0) / (1.0 - dropout_p) - else: - x = F.dropout(x, p=dropout_p) - if x1 is not None: - if dropout_mask1 is not None: - x1 = x1.masked_fill(~dropout_mask1, 0.0) / (1.0 - dropout_p) - else: - x1 = F.dropout(x1, p=dropout_p) - if x1 is not None: - x = x + x1 - if residual is not None: - x = (x + residual).to(x.dtype) - out = F.layer_norm(x.to(weight.dtype), x.shape[-1:], weight=weight, bias=bias, eps=eps).to( - dtype - ) - if weight1 is None: - return out if not prenorm else (out, x) - else: - out1 = F.layer_norm( - x.to(weight1.dtype), x.shape[-1:], weight=weight1, bias=bias1, eps=eps - ).to(dtype) - return (out, out1) if not prenorm else (out, out1, x) - - -def rms_norm_ref( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - zero_centered_weight=False, - dropout_mask=None, - dropout_mask1=None, - upcast=False, -): - dtype = x.dtype - if upcast: - x = x.float() - weight = weight.float() - bias = bias.float() if bias is not None else None - residual = residual.float() if residual is not None else residual - x1 = x1.float() if x1 is not None else None - weight1 = weight1.float() if weight1 is not None else None - bias1 = bias1.float() if bias1 is not None else None - if zero_centered_weight: - weight = weight + 1.0 - if weight1 is not None: - weight1 = weight1 + 1.0 - if x1 is not None: - assert rowscale is None, "rowscale is not supported with parallel LayerNorm" - if rowscale is not None: - x = x * rowscale[..., None] - if dropout_p > 0.0: - if dropout_mask is not None: - x = x.masked_fill(~dropout_mask, 0.0) / (1.0 - dropout_p) - else: - x = F.dropout(x, p=dropout_p) - if x1 is not None: - if dropout_mask1 is not None: - x1 = x1.masked_fill(~dropout_mask1, 0.0) / (1.0 - dropout_p) - else: - x1 = F.dropout(x1, p=dropout_p) - if x1 is not None: - x = x + x1 - if residual is not None: - x = (x + residual).to(x.dtype) - rstd = 1 / torch.sqrt((x.square()).mean(dim=-1, keepdim=True) + eps) - out = ((x * rstd * weight) + bias if bias is not None else (x * rstd * weight)).to(dtype) - if weight1 is None: - return out if not prenorm else (out, x) - else: - out1 = ((x * rstd * weight1) + bias1 if bias1 is not None else (x * rstd * weight1)).to( - dtype - ) - return (out, out1) if not prenorm else (out, out1, x) - - -@triton.autotune( - configs=triton_autotune_configs(), - key=["N", "HAS_RESIDUAL", "STORE_RESIDUAL_OUT", "IS_RMS_NORM", "HAS_BIAS", "HAS_X1", "HAS_W1", "HAS_B1"], -) -# torch compile doesn't like triton.heuristics, so we set these manually when calling the kernel -# @triton.heuristics({"HAS_BIAS": lambda args: args["B"] is not None}) -# @triton.heuristics({"HAS_RESIDUAL": lambda args: args["RESIDUAL"] is not None}) -# @triton.heuristics({"HAS_X1": lambda args: args["X1"] is not None}) -# @triton.heuristics({"HAS_W1": lambda args: args["W1"] is not None}) -# @triton.heuristics({"HAS_B1": lambda args: args["B1"] is not None}) -@triton.jit -def _layer_norm_fwd_1pass_kernel( - X, # pointer to the input - Y, # pointer to the output - W, # pointer to the weights - B, # pointer to the biases - RESIDUAL, # pointer to the residual - X1, - W1, - B1, - Y1, - RESIDUAL_OUT, # pointer to the residual - ROWSCALE, - SEEDS, # Dropout seeds for each row - DROPOUT_MASK, - DROPOUT_MASK1, - Mean, # pointer to the mean - Rstd, # pointer to the 1/std - stride_x_row, # how much to increase the pointer when moving by 1 row - stride_y_row, - stride_res_row, - stride_res_out_row, - stride_x1_row, - stride_y1_row, - M, # number of rows in X - N, # number of columns in X - eps, # epsilon to avoid division by zero - dropout_p, # Dropout probability - zero_centered_weight, # If true, add 1.0 to the weight - IS_RMS_NORM: tl.constexpr, - BLOCK_N: tl.constexpr, - HAS_RESIDUAL: tl.constexpr, - STORE_RESIDUAL_OUT: tl.constexpr, - HAS_BIAS: tl.constexpr, - HAS_DROPOUT: tl.constexpr, - STORE_DROPOUT_MASK: tl.constexpr, - HAS_ROWSCALE: tl.constexpr, - HAS_X1: tl.constexpr, - HAS_W1: tl.constexpr, - HAS_B1: tl.constexpr, -): - # Map the program id to the row of X and Y it should compute. - row = tl.program_id(0) - X += row * stride_x_row - Y += row * stride_y_row - if HAS_RESIDUAL: - RESIDUAL += row * stride_res_row - if STORE_RESIDUAL_OUT: - RESIDUAL_OUT += row * stride_res_out_row - if HAS_X1: - X1 += row * stride_x1_row - if HAS_W1: - Y1 += row * stride_y1_row - # Compute mean and variance - cols = tl.arange(0, BLOCK_N) - x = tl.load(X + cols, mask=cols < N, other=0.0).to(tl.float32) - if HAS_ROWSCALE: - rowscale = tl.load(ROWSCALE + row).to(tl.float32) - x *= rowscale - if HAS_DROPOUT: - # Compute dropout mask - # 7 rounds is good enough, and reduces register pressure - keep_mask = tl.rand(tl.load(SEEDS + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - x = tl.where(keep_mask, x / (1.0 - dropout_p), 0.0) - if STORE_DROPOUT_MASK: - tl.store(DROPOUT_MASK + row * N + cols, keep_mask, mask=cols < N) - if HAS_X1: - x1 = tl.load(X1 + cols, mask=cols < N, other=0.0).to(tl.float32) - if HAS_ROWSCALE: - rowscale = tl.load(ROWSCALE + M + row).to(tl.float32) - x1 *= rowscale - if HAS_DROPOUT: - # Compute dropout mask - # 7 rounds is good enough, and reduces register pressure - keep_mask = ( - tl.rand(tl.load(SEEDS + M + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - ) - x1 = tl.where(keep_mask, x1 / (1.0 - dropout_p), 0.0) - if STORE_DROPOUT_MASK: - tl.store(DROPOUT_MASK1 + row * N + cols, keep_mask, mask=cols < N) - x += x1 - if HAS_RESIDUAL: - residual = tl.load(RESIDUAL + cols, mask=cols < N, other=0.0).to(tl.float32) - x += residual - if STORE_RESIDUAL_OUT: - tl.store(RESIDUAL_OUT + cols, x, mask=cols < N) - if not IS_RMS_NORM: - mean = tl.sum(x, axis=0) / N - tl.store(Mean + row, mean) - xbar = tl.where(cols < N, x - mean, 0.0) - var = tl.sum(xbar * xbar, axis=0) / N - else: - xbar = tl.where(cols < N, x, 0.0) - var = tl.sum(xbar * xbar, axis=0) / N - rstd = 1 / tl.sqrt(var + eps) - tl.store(Rstd + row, rstd) - # Normalize and apply linear transformation - mask = cols < N - w = tl.load(W + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w += 1.0 - if HAS_BIAS: - b = tl.load(B + cols, mask=mask).to(tl.float32) - x_hat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd - y = x_hat * w + b if HAS_BIAS else x_hat * w - # Write output - tl.store(Y + cols, y, mask=mask) - if HAS_W1: - w1 = tl.load(W1 + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w1 += 1.0 - if HAS_B1: - b1 = tl.load(B1 + cols, mask=mask).to(tl.float32) - y1 = x_hat * w1 + b1 if HAS_B1 else x_hat * w1 - tl.store(Y1 + cols, y1, mask=mask) - - -def _layer_norm_fwd( - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - residual: Optional[Tensor] = None, - x1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - out_dtype: Optional[torch.dtype] = None, - residual_dtype: Optional[torch.dtype] = None, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - return_dropout_mask: bool = False, - out: Optional[Tensor] = None, - residual_out: Optional[Tensor] = None -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - # Need to wrap to handle the case where residual_out is a alias of x, which makes torch.library - # and torch.compile unhappy. Also allocate memory for out and residual_out if they are None - # so that _layer_norm_fwd_impl doesn't have to return them. - if out is None: - out = torch.empty_like(x, dtype=x.dtype if out_dtype is None else out_dtype) - if residual is not None: - residual_dtype = residual.dtype - if residual_out is None and ( - residual is not None - or (residual_dtype is not None and residual_dtype != x.dtype) - or dropout_p > 0.0 - or rowscale is not None - or x1 is not None - ): - residual_out = torch.empty_like( - x, dtype=residual_dtype if residual_dtype is not None else x.dtype - ) - else: - residual_out = None - y1, mean, rstd, seeds, dropout_mask, dropout_mask1 = _layer_norm_fwd_impl( - x, - weight, - bias, - eps, - out, - residual=residual, - x1=x1, - weight1=weight1, - bias1=bias1, - dropout_p=dropout_p, - rowscale=rowscale, - zero_centered_weight=zero_centered_weight, - is_rms_norm=is_rms_norm, - return_dropout_mask=return_dropout_mask, - residual_out=residual_out, - ) - # residual_out is None if residual is None and residual_dtype == input_dtype and dropout_p == 0.0 - if residual_out is None: - residual_out = x - return out, y1, mean, rstd, residual_out, seeds, dropout_mask, dropout_mask1 - - -# [2025-04-28] torch.library.triton_op ignores the schema argument, but here we need the schema -# since we're returning a tuple of tensors -@triton_op(add_op_namespace_prefix("layer_norm_fwd_impl"), mutates_args={"out", "residual_out"}, - schema="(Tensor x, Tensor weight, Tensor bias, float eps, Tensor(a!) out, Tensor? residual, Tensor? x1, Tensor? weight1, Tensor? bias1, float dropout_p, Tensor? rowscale, bool zero_centered_weight, bool is_rms_norm, bool return_dropout_mask, Tensor(a!)? residual_out) -> (Tensor y1, Tensor mean, Tensor rstd, Tensor seeds, Tensor dropout_mask, Tensor dropout_mask1)") -def _layer_norm_fwd_impl( - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - out: Tensor, - residual: Optional[Tensor] = None, - x1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - return_dropout_mask: bool = False, - residual_out: Optional[Tensor] = None -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - M, N = x.shape - assert x.stride(-1) == 1 - if residual is not None: - assert residual.stride(-1) == 1 - assert residual.shape == (M, N) - assert weight.shape == (N,) - assert weight.stride(-1) == 1 - if bias is not None: - assert bias.stride(-1) == 1 - assert bias.shape == (N,) - if x1 is not None: - assert x1.shape == x.shape - assert rowscale is None - assert x1.stride(-1) == 1 - if weight1 is not None: - assert weight1.shape == (N,) - assert weight1.stride(-1) == 1 - if bias1 is not None: - assert bias1.shape == (N,) - assert bias1.stride(-1) == 1 - if rowscale is not None: - assert rowscale.is_contiguous() - assert rowscale.shape == (M,) - assert out.shape == x.shape - assert out.stride(-1) == 1 - if residual_out is not None: - assert residual_out.shape == x.shape - assert residual_out.stride(-1) == 1 - if weight1 is not None: - y1 = torch.empty_like(out) - assert y1.stride(-1) == 1 - else: - y1 = None - mean = torch.empty((M,), dtype=torch.float32, device=x.device) if not is_rms_norm else None - rstd = torch.empty((M,), dtype=torch.float32, device=x.device) - if dropout_p > 0.0: - seeds = torch.randint( - 2**32, (M if x1 is None else 2 * M,), device=x.device, dtype=torch.int64 - ) - else: - seeds = None - if return_dropout_mask and dropout_p > 0.0: - dropout_mask = torch.empty(M, N, device=x.device, dtype=torch.bool) - if x1 is not None: - dropout_mask1 = torch.empty(M, N, device=x.device, dtype=torch.bool) - else: - dropout_mask1 = None - else: - dropout_mask, dropout_mask1 = None, None - # Less than 64KB per feature: enqueue fused kernel - MAX_FUSED_SIZE = 65536 // x.element_size() - BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) - if N > BLOCK_N: - raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") - with torch.cuda.device(x.device.index): - torch.library.wrap_triton(_layer_norm_fwd_1pass_kernel)[(M,)]( - x, - out, - weight, - bias, - residual, - x1, - weight1, - bias1, - y1, - residual_out, - rowscale, - seeds, - dropout_mask, - dropout_mask1, - mean, - rstd, - x.stride(0), - out.stride(0), - residual.stride(0) if residual is not None else 0, - residual_out.stride(0) if residual_out is not None else 0, - x1.stride(0) if x1 is not None else 0, - y1.stride(0) if y1 is not None else 0, - M, - N, - eps, - dropout_p, - # Passing bool make torch inductor very unhappy since it then tries to compare to int_max - int(zero_centered_weight), - is_rms_norm, - BLOCK_N, - residual is not None, - residual_out is not None, - bias is not None, - dropout_p > 0.0, - dropout_mask is not None, - rowscale is not None, - HAS_X1=x1 is not None, - HAS_W1=weight1 is not None, - HAS_B1=bias1 is not None, - ) - return y1, mean, rstd, seeds, dropout_mask, dropout_mask1 - - -@triton.autotune( - configs=triton_autotune_configs(), - key=["N", "HAS_DRESIDUAL", "STORE_DRESIDUAL", "IS_RMS_NORM", "HAS_BIAS", "HAS_DROPOUT"], -) -# torch compile doesn't like triton.heuristics, so we set these manually when calling the kernel -# @triton.heuristics({"HAS_BIAS": lambda args: args["B"] is not None}) -# @triton.heuristics({"HAS_DRESIDUAL": lambda args: args["DRESIDUAL"] is not None}) -# @triton.heuristics({"STORE_DRESIDUAL": lambda args: args["DRESIDUAL_IN"] is not None}) -# @triton.heuristics({"HAS_ROWSCALE": lambda args: args["ROWSCALE"] is not None}) -# @triton.heuristics({"HAS_DY1": lambda args: args["DY1"] is not None}) -# @triton.heuristics({"HAS_DX1": lambda args: args["DX1"] is not None}) -# @triton.heuristics({"HAS_B1": lambda args: args["DB1"] is not None}) -# @triton.heuristics({"RECOMPUTE_OUTPUT": lambda args: args["Y"] is not None}) -@triton.jit -def _layer_norm_bwd_kernel( - X, # pointer to the input - W, # pointer to the weights - B, # pointer to the biases - Y, # pointer to the output to be recomputed - DY, # pointer to the output gradient - DX, # pointer to the input gradient - DW, # pointer to the partial sum of weights gradient - DB, # pointer to the partial sum of biases gradient - DRESIDUAL, - W1, - DY1, - DX1, - DW1, - DB1, - DRESIDUAL_IN, - ROWSCALE, - SEEDS, - Mean, # pointer to the mean - Rstd, # pointer to the 1/std - stride_x_row, # how much to increase the pointer when moving by 1 row - stride_y_row, - stride_dy_row, - stride_dx_row, - stride_dres_row, - stride_dy1_row, - stride_dx1_row, - stride_dres_in_row, - M, # number of rows in X - N, # number of columns in X - eps, # epsilon to avoid division by zero - dropout_p, - zero_centered_weight, - rows_per_program, - IS_RMS_NORM: tl.constexpr, - BLOCK_N: tl.constexpr, - HAS_DRESIDUAL: tl.constexpr, - STORE_DRESIDUAL: tl.constexpr, - HAS_BIAS: tl.constexpr, - HAS_DROPOUT: tl.constexpr, - HAS_ROWSCALE: tl.constexpr, - HAS_DY1: tl.constexpr, - HAS_DX1: tl.constexpr, - HAS_B1: tl.constexpr, - RECOMPUTE_OUTPUT: tl.constexpr, -): - # Map the program id to the elements of X, DX, and DY it should compute. - row_block_id = tl.program_id(0) - row_start = row_block_id * rows_per_program - # Do not early exit if row_start >= M, because we need to write DW and DB - cols = tl.arange(0, BLOCK_N) - mask = cols < N - X += row_start * stride_x_row - if HAS_DRESIDUAL: - DRESIDUAL += row_start * stride_dres_row - if STORE_DRESIDUAL: - DRESIDUAL_IN += row_start * stride_dres_in_row - DY += row_start * stride_dy_row - DX += row_start * stride_dx_row - if HAS_DY1: - DY1 += row_start * stride_dy1_row - if HAS_DX1: - DX1 += row_start * stride_dx1_row - if RECOMPUTE_OUTPUT: - Y += row_start * stride_y_row - w = tl.load(W + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w += 1.0 - if RECOMPUTE_OUTPUT and HAS_BIAS: - b = tl.load(B + cols, mask=mask, other=0.0).to(tl.float32) - if HAS_DY1: - w1 = tl.load(W1 + cols, mask=mask).to(tl.float32) - if zero_centered_weight: - w1 += 1.0 - dw = tl.zeros((BLOCK_N,), dtype=tl.float32) - if HAS_BIAS: - db = tl.zeros((BLOCK_N,), dtype=tl.float32) - if HAS_DY1: - dw1 = tl.zeros((BLOCK_N,), dtype=tl.float32) - if HAS_B1: - db1 = tl.zeros((BLOCK_N,), dtype=tl.float32) - row_end = min((row_block_id + 1) * rows_per_program, M) - for row in range(row_start, row_end): - # Load data to SRAM - x = tl.load(X + cols, mask=mask, other=0).to(tl.float32) - dy = tl.load(DY + cols, mask=mask, other=0).to(tl.float32) - if HAS_DY1: - dy1 = tl.load(DY1 + cols, mask=mask, other=0).to(tl.float32) - if not IS_RMS_NORM: - mean = tl.load(Mean + row) - rstd = tl.load(Rstd + row) - # Compute dx - xhat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd - xhat = tl.where(mask, xhat, 0.0) - if RECOMPUTE_OUTPUT: - y = xhat * w + b if HAS_BIAS else xhat * w - tl.store(Y + cols, y, mask=mask) - wdy = w * dy - dw += dy * xhat - if HAS_BIAS: - db += dy - if HAS_DY1: - wdy += w1 * dy1 - dw1 += dy1 * xhat - if HAS_B1: - db1 += dy1 - if not IS_RMS_NORM: - c1 = tl.sum(xhat * wdy, axis=0) / N - c2 = tl.sum(wdy, axis=0) / N - dx = (wdy - (xhat * c1 + c2)) * rstd - else: - c1 = tl.sum(xhat * wdy, axis=0) / N - dx = (wdy - xhat * c1) * rstd - if HAS_DRESIDUAL: - dres = tl.load(DRESIDUAL + cols, mask=mask, other=0).to(tl.float32) - dx += dres - # Write dx - if STORE_DRESIDUAL: - tl.store(DRESIDUAL_IN + cols, dx, mask=mask) - if HAS_DX1: - if HAS_DROPOUT: - keep_mask = ( - tl.rand(tl.load(SEEDS + M + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - ) - dx1 = tl.where(keep_mask, dx / (1.0 - dropout_p), 0.0) - else: - dx1 = dx - tl.store(DX1 + cols, dx1, mask=mask) - if HAS_DROPOUT: - keep_mask = tl.rand(tl.load(SEEDS + row).to(tl.uint32), cols, n_rounds=7) > dropout_p - dx = tl.where(keep_mask, dx / (1.0 - dropout_p), 0.0) - if HAS_ROWSCALE: - rowscale = tl.load(ROWSCALE + row).to(tl.float32) - dx *= rowscale - tl.store(DX + cols, dx, mask=mask) - - X += stride_x_row - if HAS_DRESIDUAL: - DRESIDUAL += stride_dres_row - if STORE_DRESIDUAL: - DRESIDUAL_IN += stride_dres_in_row - if RECOMPUTE_OUTPUT: - Y += stride_y_row - DY += stride_dy_row - DX += stride_dx_row - if HAS_DY1: - DY1 += stride_dy1_row - if HAS_DX1: - DX1 += stride_dx1_row - tl.store(DW + row_block_id * N + cols, dw, mask=mask) - if HAS_BIAS: - tl.store(DB + row_block_id * N + cols, db, mask=mask) - if HAS_DY1: - tl.store(DW1 + row_block_id * N + cols, dw1, mask=mask) - if HAS_B1: - tl.store(DB1 + row_block_id * N + cols, db1, mask=mask) - - -def _layer_norm_bwd( - dy: Tensor, - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - mean: Tensor, - rstd: Tensor, - dresidual: Optional[Tensor] = None, - dy1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - seeds: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - has_residual: bool = False, - has_x1: bool = False, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - x_dtype: Optional[torch.dtype] = None, - recompute_output: bool = False, -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - # Need to wrap to handle the case where dresidual_in or dx1 are aliases of x, - # which makes torch.library unhappy - dx, dw, db, dresidual_in, dx1, dw1, db1, y = _layer_norm_bwd_impl( - dy, - x, - weight, - bias, - eps, - mean, - rstd, - dresidual, - dy1, - weight1, - bias1, - seeds, - dropout_p, - rowscale, - has_residual, - has_x1, - zero_centered_weight, - is_rms_norm, - x_dtype=x_dtype, - recompute_output=recompute_output, - ) - # Don't need to compute dresidual_in separately in this case - if has_residual and dx.dtype == x.dtype and dropout_p == 0.0 and rowscale is None: - dresidual_in = dx - if has_x1 and dropout_p == 0.0: - dx1 = dx - return dx, dw, db, dresidual_in, dx1, dw1, db1, y - - - -@triton_op(add_op_namespace_prefix("layer_norm_bwd_impl"), mutates_args={}, - schema="(Tensor dy, Tensor x, Tensor weight, Tensor bias, float eps, Tensor mean, Tensor rstd, Tensor? dresidual, Tensor? dy1, Tensor? weight1, Tensor? bias1, Tensor? seeds, float dropout_p, Tensor? rowscale, bool has_residual, bool has_x1, bool zero_centered_weight, bool is_rms_norm, ScalarType? x_dtype, bool recompute_output) -> (Tensor dx, Tensor dw, Tensor db, Tensor dresidual_in, Tensor dx1, Tensor dw1, Tensor db1, Tensor y)", - allow_decomposition=False, # Don't let torch.compile trace inside - ) -def _layer_norm_bwd_impl( - dy: Tensor, - x: Tensor, - weight: Tensor, - bias: Tensor, - eps: float, - mean: Tensor, - rstd: Tensor, - dresidual: Optional[Tensor] = None, - dy1: Optional[Tensor] = None, - weight1: Optional[Tensor] = None, - bias1: Optional[Tensor] = None, - seeds: Optional[Tensor] = None, - dropout_p: float = 0.0, - rowscale: Optional[Tensor] = None, - has_residual: bool = False, - has_x1: bool = False, - zero_centered_weight: bool = False, - is_rms_norm: bool = False, - x_dtype: Optional[torch.dtype] = None, - recompute_output: bool = False, -) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): - M, N = x.shape - assert x.stride(-1) == 1 - dy = maybe_contiguous_lastdim(dy) - assert dy.stride(-1) == 1 - assert dy.shape == (M, N) - if dresidual is not None: - dresidual = maybe_contiguous_lastdim(dresidual) - assert dresidual.stride(-1) == 1 - assert dresidual.shape == (M, N) - assert weight.shape == (N,) - assert weight.stride(-1) == 1 - if bias is not None: - assert bias.stride(-1) == 1 - assert bias.shape == (N,) - if dy1 is not None: - dy1 = maybe_contiguous_lastdim(dy1) - assert weight1 is not None - assert dy1.shape == dy.shape - assert dy1.stride(-1) == 1 - if weight1 is not None: - assert weight1.shape == (N,) - assert weight1.stride(-1) == 1 - if bias1 is not None: - assert bias1.shape == (N,) - assert bias1.stride(-1) == 1 - if seeds is not None: - assert seeds.is_contiguous() - assert seeds.shape == (M if not has_x1 else M * 2,) - if rowscale is not None: - assert rowscale.is_contiguous() - assert rowscale.shape == (M,) - # allocate output - dx = ( - torch.empty_like(x) - if x_dtype is None - else torch.empty(M, N, dtype=x_dtype, device=x.device) - ) - dresidual_in = ( - torch.empty_like(x) - if has_residual - and (dx.dtype != x.dtype or dropout_p > 0.0 or rowscale is not None or has_x1) - else None - ) - dx1 = torch.empty_like(dx) if (has_x1 and dropout_p > 0.0) else None - y = torch.empty(M, N, dtype=dy.dtype, device=dy.device) if recompute_output else None - if recompute_output: - assert weight1 is None, "recompute_output is not supported with parallel LayerNorm" - - # Less than 64KB per feature: enqueue fused kernel - MAX_FUSED_SIZE = 65536 // x.element_size() - BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) - if N > BLOCK_N: - raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") - # Increasing the multiple (e.g. 8) will allow more thread blocks to be launched and hide the - # latency of the gmem reads/writes, but will increase the time of summing up dw / db. - sm_count = torch.cuda.get_device_properties(x.device).multi_processor_count * 8 - _dw = torch.empty((sm_count, N), dtype=torch.float32, device=weight.device) - _db = ( - torch.empty((sm_count, N), dtype=torch.float32, device=bias.device) - if bias is not None - else None - ) - _dw1 = torch.empty_like(_dw) if weight1 is not None else None - _db1 = torch.empty_like(_db) if bias1 is not None else None - rows_per_program = math.ceil(M / sm_count) - grid = (sm_count,) - with torch.cuda.device(x.device.index): - torch.library.wrap_triton(_layer_norm_bwd_kernel)[grid]( - x, - weight, - bias, - y, - dy, - dx, - _dw, - _db, - dresidual, - weight1, - dy1, - dx1, - _dw1, - _db1, - dresidual_in, - rowscale, - seeds, - mean, - rstd, - x.stride(0), - 0 if not recompute_output else y.stride(0), - dy.stride(0), - dx.stride(0), - dresidual.stride(0) if dresidual is not None else 0, - dy1.stride(0) if dy1 is not None else 0, - dx1.stride(0) if dx1 is not None else 0, - dresidual_in.stride(0) if dresidual_in is not None else 0, - M, - N, - eps, - dropout_p, - # Passing bool make torch inductor very unhappy since it then tries to compare to int_max - int(zero_centered_weight), - rows_per_program, - is_rms_norm, - BLOCK_N, - dresidual is not None, - dresidual_in is not None, - bias is not None, - dropout_p > 0.0, - HAS_ROWSCALE=rowscale is not None, - HAS_DY1=dy1 is not None, - HAS_DX1=dx1 is not None, - HAS_B1=bias1 is not None, - RECOMPUTE_OUTPUT=y is not None, - ) - dw = _dw.sum(0).to(weight.dtype) - db = _db.sum(0).to(bias.dtype) if bias is not None else None - dw1 = _dw1.sum(0).to(weight1.dtype) if weight1 is not None else None - db1 = _db1.sum(0).to(bias1.dtype) if bias1 is not None else None - # dresidual_in and dx1 could be None, the wrapper will handle assigning them from dx - return dx, dw, db, dresidual_in, dx1, dw1, db1, y - - -class LayerNormFn(torch.autograd.Function): - - @staticmethod - def forward( - ctx, - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - residual_in_fp32=False, - zero_centered_weight=False, - is_rms_norm=False, - return_dropout_mask=False, - out_dtype=None, - out=None, - residual_out=None - ): - x_shape_og = x.shape - # reshape input data into 2D tensor - x = maybe_contiguous_lastdim(x.reshape(-1, x.shape[-1])) - if residual is not None: - assert residual.shape == x_shape_og - residual = maybe_contiguous_lastdim(residual.reshape(-1, residual.shape[-1])) - if x1 is not None: - assert x1.shape == x_shape_og - assert rowscale is None, "rowscale is not supported with parallel LayerNorm" - x1 = maybe_contiguous_lastdim(x1.reshape(-1, x1.shape[-1])) - weight = weight.contiguous() - bias = maybe_contiguous(bias) - weight1 = maybe_contiguous(weight1) - bias1 = maybe_contiguous(bias1) - if rowscale is not None: - rowscale = rowscale.reshape(-1).contiguous() - residual_dtype = ( - residual.dtype - if residual is not None - else (torch.float32 if residual_in_fp32 else None) - ) - if out is not None: - out = out.reshape(-1, out.shape[-1]) - if residual_out is not None: - residual_out = residual_out.reshape(-1, residual_out.shape[-1]) - y, y1, mean, rstd, residual_out, seeds, dropout_mask, dropout_mask1 = _layer_norm_fwd( - x, - weight, - bias, - eps, - residual, - x1, - weight1, - bias1, - dropout_p=dropout_p, - rowscale=rowscale, - out_dtype=out_dtype, - residual_dtype=residual_dtype, - zero_centered_weight=zero_centered_weight, - is_rms_norm=is_rms_norm, - return_dropout_mask=return_dropout_mask, - out=out, - residual_out=residual_out, - ) - ctx.save_for_backward( - residual_out, weight, bias, weight1, bias1, rowscale, seeds, mean, rstd - ) - ctx.x_shape_og = x_shape_og - ctx.eps = eps - ctx.dropout_p = dropout_p - ctx.is_rms_norm = is_rms_norm - ctx.has_residual = residual is not None - ctx.has_x1 = x1 is not None - ctx.prenorm = prenorm - ctx.x_dtype = x.dtype - ctx.zero_centered_weight = zero_centered_weight - y = y.reshape(x_shape_og) - y1 = y1.reshape(x_shape_og) if y1 is not None else None - residual_out = residual_out.reshape(x_shape_og) if residual_out is not None else None - dropout_mask = dropout_mask.reshape(x_shape_og) if dropout_mask is not None else None - dropout_mask1 = dropout_mask1.reshape(x_shape_og) if dropout_mask1 is not None else None - if not return_dropout_mask: - if weight1 is None: - return y if not prenorm else (y, residual_out) - else: - return (y, y1) if not prenorm else (y, y1, residual_out) - else: - if weight1 is None: - return ( - (y, dropout_mask, dropout_mask1) - if not prenorm - else (y, residual_out, dropout_mask, dropout_mask1) - ) - else: - return ( - (y, y1, dropout_mask, dropout_mask1) - if not prenorm - else (y, y1, residual_out, dropout_mask, dropout_mask1) - ) - - @staticmethod - def backward(ctx, dy, *args): - x, weight, bias, weight1, bias1, rowscale, seeds, mean, rstd = ctx.saved_tensors - dy = dy.reshape(-1, dy.shape[-1]) - if weight1 is not None: - dy1, args = args[0], args[1:] - dy1 = dy1.reshape(-1, dy1.shape[-1]) - assert dy1.shape == x.shape - else: - dy1 = None - if ctx.prenorm: - dresidual = args[0] - dresidual = dresidual.reshape(-1, dresidual.shape[-1]) - assert dresidual.shape == x.shape - else: - dresidual = None - dx, dw, db, dresidual_in, dx1, dw1, db1, _ = _layer_norm_bwd( - dy, - x, - weight, - bias, - ctx.eps, - mean, - rstd, - dresidual, - dy1, - weight1, - bias1, - seeds, - ctx.dropout_p, - rowscale, - ctx.has_residual, - ctx.has_x1, - ctx.zero_centered_weight, - ctx.is_rms_norm, - x_dtype=ctx.x_dtype, - recompute_output=False, - ) - return ( - dx.reshape(ctx.x_shape_og), - dw, - db, - dresidual_in.reshape(ctx.x_shape_og) if ctx.has_residual else None, - dx1.reshape(ctx.x_shape_og) if dx1 is not None else None, - dw1, - db1, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) - - -def layer_norm_fn( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - residual_in_fp32=False, - zero_centered_weight=False, - is_rms_norm=False, - return_dropout_mask=False, - out_dtype=None, - out=None, - residual_out=None -): - return LayerNormFn.apply( - x, - weight, - bias, - residual, - x1, - weight1, - bias1, - eps, - dropout_p, - rowscale, - prenorm, - residual_in_fp32, - zero_centered_weight, - is_rms_norm, - return_dropout_mask, - out_dtype, - out, - residual_out - ) - - -def rms_norm_fn( - x, - weight, - bias, - residual=None, - x1=None, - weight1=None, - bias1=None, - eps=1e-6, - dropout_p=0.0, - rowscale=None, - prenorm=False, - residual_in_fp32=False, - zero_centered_weight=False, - return_dropout_mask=False, - out_dtype=None, - out=None, - residual_out=None -): - return LayerNormFn.apply( - x, - weight, - bias, - residual, - x1, - weight1, - bias1, - eps, - dropout_p, - rowscale, - prenorm, - residual_in_fp32, - zero_centered_weight, - True, - return_dropout_mask, - out_dtype, - out, - residual_out - ) - - -class RMSNorm(torch.nn.Module): - - def __init__(self, hidden_size, eps=1e-5, dropout_p=0.0, zero_centered_weight=False, - device=None, dtype=None): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.eps = eps - if dropout_p > 0.0: - self.drop = torch.nn.Dropout(dropout_p) - else: - self.drop = None - self.zero_centered_weight = zero_centered_weight - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.register_parameter("bias", None) - self.reset_parameters() - - def reset_parameters(self): - if not self.zero_centered_weight: - torch.nn.init.ones_(self.weight) - else: - torch.nn.init.zeros_(self.weight) - - def forward(self, x, residual=None, prenorm=False, residual_in_fp32=False): - return rms_norm_fn( - x, - self.weight, - self.bias, - residual=residual, - eps=self.eps, - dropout_p=self.drop.p if self.drop is not None and self.training else 0.0, - prenorm=prenorm, - residual_in_fp32=residual_in_fp32, - zero_centered_weight=self.zero_centered_weight, - ) - - -class LayerNormLinearFn(torch.autograd.Function): - - @staticmethod - @custom_fwd - def forward( - ctx, - x, - norm_weight, - norm_bias, - linear_weight, - linear_bias, - residual=None, - eps=1e-6, - prenorm=False, - residual_in_fp32=False, - is_rms_norm=False, - ): - x_shape_og = x.shape - # reshape input data into 2D tensor - x = maybe_contiguous_lastdim(x.reshape(-1, x.shape[-1])) - if residual is not None: - assert residual.shape == x_shape_og - residual = maybe_contiguous_lastdim(residual.reshape(-1, residual.shape[-1])) - norm_weight = norm_weight.contiguous() - norm_bias = maybe_contiguous(norm_bias) - residual_dtype = ( - residual.dtype - if residual is not None - else (torch.float32 if residual_in_fp32 else None) - ) - y, _, mean, rstd, residual_out, *rest = _layer_norm_fwd( - x, - norm_weight, - norm_bias, - eps, - residual, - out_dtype=None if not torch.is_autocast_enabled() else torch.get_autocast_dtype("cuda"), - residual_dtype=residual_dtype, - is_rms_norm=is_rms_norm, - ) - y = y.reshape(x_shape_og) - dtype = torch.get_autocast_dtype("cuda") if torch.is_autocast_enabled() else y.dtype - linear_weight = linear_weight.to(dtype) - linear_bias = linear_bias.to(dtype) if linear_bias is not None else None - out = F.linear(y.to(linear_weight.dtype), linear_weight, linear_bias) - # We don't store y, will be recomputed in the backward pass to save memory - ctx.save_for_backward(residual_out, norm_weight, norm_bias, linear_weight, mean, rstd) - ctx.x_shape_og = x_shape_og - ctx.eps = eps - ctx.is_rms_norm = is_rms_norm - ctx.has_residual = residual is not None - ctx.prenorm = prenorm - ctx.x_dtype = x.dtype - ctx.linear_bias_is_none = linear_bias is None - return out if not prenorm else (out, residual_out.reshape(x_shape_og)) - - @staticmethod - @custom_bwd - def backward(ctx, dout, *args): - x, norm_weight, norm_bias, linear_weight, mean, rstd = ctx.saved_tensors - dout = dout.reshape(-1, dout.shape[-1]) - dy = F.linear(dout, linear_weight.t()) - dlinear_bias = None if ctx.linear_bias_is_none else dout.sum(0) - dy = maybe_contiguous_lastdim(dy) - assert dy.shape == x.shape - if ctx.prenorm: - dresidual = args[0] - dresidual = maybe_contiguous_lastdim(dresidual.reshape(-1, dresidual.shape[-1])) - assert dresidual.shape == x.shape - else: - dresidual = None - dx, dnorm_weight, dnorm_bias, dresidual_in, _, _, _, y = _layer_norm_bwd( - dy, - x, - norm_weight, - norm_bias, - ctx.eps, - mean, - rstd, - dresidual=dresidual, - has_residual=ctx.has_residual, - is_rms_norm=ctx.is_rms_norm, - x_dtype=ctx.x_dtype, - recompute_output=True, - ) - dlinear_weight = torch.einsum("bo,bi->oi", dout, y) - return ( - dx.reshape(ctx.x_shape_og), - dnorm_weight, - dnorm_bias, - dlinear_weight, - dlinear_bias, - dresidual_in.reshape(ctx.x_shape_og) if ctx.has_residual else None, - None, - None, - None, - None, - ) - - -def layer_norm_linear_fn( - x, - norm_weight, - norm_bias, - linear_weight, - linear_bias, - residual=None, - eps=1e-6, - prenorm=False, - residual_in_fp32=False, - is_rms_norm=False, -): - return LayerNormLinearFn.apply( - x, - norm_weight, - norm_bias, - linear_weight, - linear_bias, - residual, - eps, - prenorm, - residual_in_fp32, - is_rms_norm, - ) diff --git a/build/torch212-cxx11-xpu20253-x86_64-linux/ops/triton/linear.py b/build/torch212-cxx11-xpu20253-x86_64-linux/ops/triton/linear.py deleted file mode 100644 index a8966dbc345ab0e593df0124451ee7be3dae131a..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-xpu20253-x86_64-linux/ops/triton/linear.py +++ /dev/null @@ -1,594 +0,0 @@ -# Adapted from https://github.com/ELS-RD/kernl/blob/main/src/kernl/implementations/linear_layer.py -# and https://github.com/openai/triton/blob/master/python/triton/ops/matmul.py -from typing import Optional - -import torch -import triton -import triton.language as tl -from triton.ops.matmul_perf_model import early_config_prune, estimate_matmul_time - -from flash_attn.ops.triton.k_activations import ( - gelu, - gelu_approx, - gelu_approx_grad, - gelu_grad, - squared_relu, - squared_relu_grad, -) - -# CREDITS: Initially inspired by the Triton tutorial on matrix multiplications - - -def init_to_zero(name): - return lambda nargs: nargs[name].zero_() - - -def get_configs_io_bound(): - configs = [] - for num_stages in [2, 3, 4, 5, 6]: - for block_m in [16, 32]: - for block_k in [32, 64]: - for block_n in [32, 64, 128, 256]: - num_warps = 2 if block_n <= 64 else 4 - configs.append( - triton.Config( - { - "BLOCK_M": block_m, - "BLOCK_N": block_n, - "BLOCK_K": block_k, - "SPLIT_K": 1, - }, - num_stages=num_stages, - num_warps=num_warps, - ) - ) - # split_k not used - # for split_k in [2, 4, 8, 16]: - # configs.append(triton.Config( - # {'BLOCK_M': block_m, 'BLOCK_N': block_n, 'BLOCK_K': block_k, 'SPLIT_K': split_k}, - # num_stages=num_stages, num_warps=num_warps, pre_hook=init_to_zero('C'))) - return configs - - -@triton.autotune( - configs=[ - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - # good for int8 - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=4, - num_warps=4, - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - ] - + get_configs_io_bound(), - key=["CACHE_KEY_M", "CACHE_KEY_N", "CACHE_KEY_K"], - prune_configs_by={ - "early_config_prune": early_config_prune, - "perf_model": estimate_matmul_time, - "top_k": 10, - }, -) -@triton.heuristics( - { - "EVEN_K": lambda args: args["K"] % (args["BLOCK_K"] * args["SPLIT_K"]) == 0, - } -) -@triton.jit -def kernel_fwd( - C, # Pointers to matrices - ACT_INPUT, - A, - B, - bias, - # Matrix dimensions - M, - N, - K, - CACHE_KEY_M, - CACHE_KEY_N, - CACHE_KEY_K, - # The stride variables represent how much to increase the ptr by when moving by 1 - # element in a particular dimension. E.g. stride_am is how much to increase a_ptr - # by to get the element one row down (A has M rows) - stride_cm, - # stride_cn, # Assume that stride_cn == 1 - stride_am, - stride_ak, - stride_bn, - stride_bk, - # Meta-parameters - BLOCK_M: tl.constexpr, - GROUP_M: tl.constexpr, - BLOCK_N: tl.constexpr, - BLOCK_K: tl.constexpr, - # split k not used, not performant with activation, kept because early_config_prune is expecting it - SPLIT_K: tl.constexpr, - EVEN_K: tl.constexpr, - A_ROWMAJOR: tl.constexpr, - B_COLMAJOR: tl.constexpr, - BIAS: tl.constexpr, - SAVE_ACT_INPUT: tl.constexpr, - ACTIVATION: tl.constexpr, -): - - """ - Kernel for computing Out = activation(A x W + C) - - Input has shape (M, K) - - Weight has shape (K, N) - - Bias has shape (N,) - - Output has shape (M, N) - - ActInputs (optional) has shape (M, N) - 'ActInputs' optionally saves the A x W + C intermediate for backward computations - This kernel will consolidate over K - """ - - pid = tl.program_id(axis=0) - - grid_m = (M + BLOCK_M - 1) // BLOCK_M - grid_n = (N + BLOCK_N - 1) // BLOCK_N - # re-order program ID for better L2 performance - width = GROUP_M * grid_n - group_id = pid // width - group_size = min(grid_m - group_id * GROUP_M, GROUP_M) - pid_m = group_id * GROUP_M + (pid % group_size) - pid_n = (pid % width) // (group_size) - - # now compute the block that each program will go through - # rm (resp. rn) denotes a range of indices - # for rows (resp. col) of C - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - # trick to avoid masking on M and N axis - ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M) - rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N) - rk = tl.arange(0, BLOCK_K) - - if A_ROWMAJOR: - A = A + (ram[:, None] * stride_am + rk[None, :]) - else: - A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak) - if B_COLMAJOR: - B = B + (rk[:, None] + rbn[None, :] * stride_bn) - else: - B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn) - - acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) - - for k in range(K, 0, -BLOCK_K): - if EVEN_K: - a = tl.load(A) - b = tl.load(B) - else: - a = tl.load(A, mask=rk[None, :] < k, other=0.0) - b = tl.load(B, mask=rk[:, None] < k, other=0.0) - acc += tl.dot(a, b) - - if A_ROWMAJOR: - A += BLOCK_K - else: - A += BLOCK_K * stride_ak - if B_COLMAJOR: - B += BLOCK_K - else: - B += BLOCK_K * stride_bk - - # Putting bias after the matmul (instead of before) is faster, idk why - if BIAS: - bias = tl.load(bias + rn, mask=rn < N, other=0.0).to(tl.float32) - acc += bias[None, :] - - # optional: save the activation inputs - if SAVE_ACT_INPUT: - # act_in_ptrs = ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :] * stride_cn - act_in_ptrs = ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :] - tl.store(act_in_ptrs, acc) - - # optional: fused activation (while the data is in shared memory) - if ACTIVATION == "gelu": - acc = gelu(acc) - elif ACTIVATION == "gelu_approx": - acc = gelu_approx(acc) - elif ACTIVATION == "squared_relu": - acc = squared_relu(acc) - # rematerialize rm and rn to save registers - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - - # write back result - # C = C + rm[:, None] * stride_cm + rn[None, :] * stride_cn - C = C + rm[:, None] * stride_cm + rn[None, :] - mask = (rm < M)[:, None] & (rn < N)[None, :] - tl.store(C, acc) - - -def triton_linear_act( - x: torch.Tensor, - weight: torch.Tensor, - bias: Optional[torch.Tensor] = None, - activation: str = "id", - save_act_input: bool = False, -) -> torch.Tensor: - """ - Compute e = activation(x @ weight.T + bias). - This wrapper kicks the `kernel_fwd` Triton kernel - :param x: input tensor - :param weight: weight matrix - :param bias: an optional bias tensor - :param activation: Activation name. Needs to be a Triton kernel. - :param act_input: an optional tensor to save the activation inputs (for backward) - :return: result tensor - """ - # if torch.is_autocast_enabled(): - # dtype = torch.get_autocast_gpu_dtype() - # x, weight, bias = [a.to(dtype=dtype) for a in [x, weight, bias]] - - assert activation in ["id", "gelu", "gelu_approx", "squared_relu"] - - batch_shape, n = x.shape[:-1], x.shape[-1] - batch_dim = batch_shape.numel() - x_reshaped = x.reshape(batch_dim, n) - - if x_reshaped.stride(0) > 1 and x_reshaped.stride(1) > 1: - x_reshaped = x_reshaped.contiguous() - if weight.stride(0) > 1 and weight.stride(1) > 1: - weight = weight.contiguous() - bias = bias.contiguous() if bias is not None else None - - assert ( - x.dtype == weight.dtype - ), f"Input and weight must have the same dtype, got {x.dtype} and {weight.dtype}" - if bias is not None: - assert ( - x.dtype == bias.dtype - ), f"Input and bias must have the same dtype, got {x.dtype} and {bias.dtype}" - assert ( - x_reshaped.shape[1] == weight.shape[1] - ), f"Incompatible dimensions: {x_reshaped.shape} - {weight.shape}" - - assert ( - bias is None or bias.shape[0] == weight.shape[0] - ), "Incompatible dimensions in between weight and bias" - - M, K = x_reshaped.shape - N, K = weight.shape - - output = torch.empty((M, N), device=x.device, dtype=x.dtype) - act_input = torch.empty_like(output) if save_act_input else None - - # 1D launch kernel where each block gets its own program. - grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),) # noqa - - kernel_fwd[grid]( - output, - act_input, - x_reshaped, - weight, # data ptrs - bias if bias is not None else x, # auto skip bias if not present - M, # shapes - N, - K, - M // 32, # key for triton cache (limit number of compilations) - N // 32, - K // 32, - stride_cm=output.stride(0), # strides - # stride_cn=output.stride(1), - stride_am=x_reshaped.stride(0), - stride_ak=x_reshaped.stride(1), - stride_bk=weight.stride(1), - stride_bn=weight.stride(0), - BIAS=bias is not None, # optional fused bias - SAVE_ACT_INPUT=save_act_input, # optional save activation inputs - ACTIVATION=activation, # optional fused activation - A_ROWMAJOR=x_reshaped.stride(1) == 1, - B_COLMAJOR=weight.stride(1) == 1, - GROUP_M=8, # speed optimization: group the programs - ) - - if not save_act_input: - return output.reshape(*batch_shape, output.shape[-1]) - else: - return ( - output.reshape(*batch_shape, output.shape[-1]), - act_input.reshape(*batch_shape, act_input.shape[-1]), - ) - - -@triton.autotune( - configs=[ - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - # good for int8 - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=4, - num_warps=4, - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - ] - + get_configs_io_bound(), - key=["CACHE_KEY_M", "CACHE_KEY_N", "CACHE_KEY_K"], - prune_configs_by={ - "early_config_prune": early_config_prune, - "perf_model": estimate_matmul_time, - "top_k": 10, - }, -) -@triton.heuristics( - { - "EVEN_K": lambda args: args["K"] % (args["BLOCK_K"] * args["SPLIT_K"]) == 0, - } -) -@triton.jit -def kernel_bwd( - C, # Pointers to matrices - ACT_INPUT, - A, - B, - # Matrix dimensions - M, - N, - K, - CACHE_KEY_M, - CACHE_KEY_N, - CACHE_KEY_K, - # The stride variables represent how much to increase the ptr by when moving by 1 - # element in a particular dimension. E.g. stride_am is how much to increase a_ptr - # by to get the element one row down (A has M rows) - stride_cm, - # stride_cn, # Assume that stride_cn == 1 - stride_am, - stride_ak, - stride_bk, - stride_bn, - # Meta-parameters - BLOCK_M: tl.constexpr, - GROUP_M: tl.constexpr, - BLOCK_N: tl.constexpr, - BLOCK_K: tl.constexpr, - # split k not used, not performant with activation, kept because early_config_prune is expecting it - SPLIT_K: tl.constexpr, - EVEN_K: tl.constexpr, - ACTIVATION: tl.constexpr, -): - - """ - Kernel for computing Out = activation(A x W + C) - - Input has shape (M, K) - - Weight has shape (K, N) - - Output has shape (M, N) - - ActInputs (optional) has shape (M, N) - 'ActInputs' optionally saves the A x W + C intermediate for backward computations - This kernel will consolidate over K - """ - - pid = tl.program_id(axis=0) - - grid_m = (M + BLOCK_M - 1) // BLOCK_M - grid_n = (N + BLOCK_N - 1) // BLOCK_N - # re-order program ID for better L2 performance - width = GROUP_M * grid_n - group_id = pid // width - group_size = min(grid_m - group_id * GROUP_M, GROUP_M) - pid_m = group_id * GROUP_M + (pid % group_size) - pid_n = (pid % width) // (group_size) - - # now compute the block that each program will go through - # rm (resp. rn) denotes a range of indices - # for rows (resp. col) of C - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - # trick to avoid masking on M and N axis - ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M) - rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N) - rk = tl.arange(0, BLOCK_K) - - A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak) - B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn) - - acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) - - for k in range(K, 0, -BLOCK_K): - if EVEN_K: - a = tl.load(A) - b = tl.load(B) - else: - a = tl.load(A, mask=rk[None, :] < k, other=0.0) - b = tl.load(B, mask=rk[:, None] < k, other=0.0) - acc += tl.dot(a, b) - - A += BLOCK_K * stride_ak - B += BLOCK_K * stride_bk - - # optional: fused activation (while the data is in shared memory) - if ACTIVATION != "id": - act_in_ptrs = ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :] - act_input = tl.load(act_in_ptrs).to(acc.dtype) - if ACTIVATION == "gelu": - acc *= gelu_grad(act_input) - elif ACTIVATION == "gelu_approx": - acc *= gelu_approx_grad(act_input) - elif ACTIVATION == "squared_relu": - acc *= squared_relu_grad(act_input) - - # rematerialize rm and rn to save registers - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - - # write back result - C = C + rm[:, None] * stride_cm + rn[None, :] - mask = (rm < M)[:, None] & (rn < N)[None, :] - tl.store(C, acc, mask=mask) - - -def triton_dgrad_act( - grad_output: torch.Tensor, - weight: torch.Tensor, - activation: str = "id", - act_input: Optional[torch.Tensor] = None, -) -> torch.Tensor: - """ - Compute e = activation(grad_output @ weight + bias). - This wrapper kicks the `kernel_fwd` Triton kernel - :param grad_output: input tensor - :param weight: weight matrix - :param activation: Activation name. Needs to be a Triton kernel. - :param act_input: an optional tensor to save the activation inputs (for backward) - :return: result tensor - """ - assert activation in ["id", "gelu", "gelu_approx", "squared_relu"] - - batch_shape, n = grad_output.shape[:-1], grad_output.shape[-1] - batch_dim = batch_shape.numel() - grad_output_reshaped = grad_output.reshape(batch_dim, n) - - if grad_output_reshaped.stride(0) > 1 and grad_output_reshaped.stride(1) > 1: - grad_output_reshaped = grad_output_reshaped.contiguous() - if weight.stride(0) > 1 and weight.stride(1) > 1: - weight = weight.contiguous() - - assert ( - grad_output.dtype == weight.dtype - ), f"grad_output and weight must have the same dtype, got {grad_output.dtype} and {weight.dtype}" - assert ( - grad_output_reshaped.shape[1] == weight.shape[0] - ), f"Incompatible dimensions: {grad_output_reshaped.shape} - {weight.shape}" - if activation != "id": - assert act_input is not None, f"act_input is required for activation {activation}" - - # M, N, K in bwd are different from M, N, K in fwd - M, K = grad_output_reshaped.shape - K, N = weight.shape - - grad_input = torch.empty((M, N), device=grad_output.device, dtype=grad_output.dtype) - - # 1D launch kernel where each block gets its own program. - grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),) # noqa - - kernel_bwd[grid]( - grad_input, - act_input, - grad_output_reshaped, - weight, # data ptrs - M, # shapes - N, - K, - M // 32, # key for triton cache (limit number of compilations) - N // 32, - K // 32, - stride_cm=grad_input.stride(0), # strides - # stride_cn=grad_input.stride(1), - stride_am=grad_output_reshaped.stride(0), - stride_ak=grad_output_reshaped.stride(1), - stride_bk=weight.stride(0), - stride_bn=weight.stride(1), - ACTIVATION=activation, # optional fused activation - GROUP_M=8, # speed optimization: group the programs - ) - - return grad_input.reshape(*batch_shape, grad_input.shape[-1]) diff --git a/build/torch212-cxx11-xpu20253-x86_64-linux/ops/triton/mlp.py b/build/torch212-cxx11-xpu20253-x86_64-linux/ops/triton/mlp.py deleted file mode 100644 index 059f4f8a5e174c1f4824e43d313fca18eaa799b8..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-xpu20253-x86_64-linux/ops/triton/mlp.py +++ /dev/null @@ -1,149 +0,0 @@ -# The triton fused matmul + sqrelu is faster for fp16 but slower for bf16, compared -# to naive implementation. -import fused_dense_lib as fused_dense_cuda -import torch -import torch.nn as nn -import torch.nn.functional as F - -from flash_attn.utils.torch import custom_fwd, custom_bwd -from flash_attn.ops.activations import sqrelu_bwd, sqrelu_fwd -from flash_attn.ops.triton.linear import triton_dgrad_act, triton_linear_act - - -class FusedDenseSqreluDenseFunc(torch.autograd.Function): - @staticmethod - @custom_fwd - def forward(ctx, x, weight1, bias1, weight2, bias2, checkpoint_lvl=0): - """checkpoint_lvl: - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute act_input and gelu_out in the bwd - """ - if torch.is_autocast_enabled(): - dtype = torch.get_autocast_gpu_dtype() - x, weight1, bias1, weight2, bias2 = [ - a.to(dtype=dtype) for a in [x, weight1, bias1, weight2, bias2] - ] - is_bf16 = x.dtype == torch.bfloat16 - assert checkpoint_lvl in [0, 1, 2] - x = x.contiguous() - weight1 = weight1.contiguous() - bias1 = bias1.contiguous() - weight2 = weight2.contiguous() - bias2 = bias2.contiguous() - batch_shape, n = x.shape[:-1], x.shape[-1] - batch_dim = batch_shape.numel() - if is_bf16: - act_input = fused_dense_cuda.linear_bias_forward( - x.reshape(batch_dim, n), weight1, bias1 - ) - output1 = sqrelu_fwd(act_input) - else: - save_act_input = checkpoint_lvl != 2 - result = triton_linear_act( - x.reshape(batch_dim, n), - weight1, - bias1, - activation="squared_relu", - save_act_input=save_act_input, - ) - if save_act_input: - output1, act_input = result - else: - output1 = result - output2 = fused_dense_cuda.linear_bias_forward(output1, weight2, bias2) - ctx.checkpoint_lvl = checkpoint_lvl - if checkpoint_lvl == 0: - ctx.save_for_backward(x, weight1, bias1, weight2, act_input, output1) - elif checkpoint_lvl == 1: - ctx.save_for_backward(x, weight1, bias1, weight2, act_input) - elif checkpoint_lvl == 2: - ctx.save_for_backward(x, weight1, bias1, weight2) - return output2.reshape(*batch_shape, output2.shape[-1]) - - @staticmethod - @custom_bwd - def backward(ctx, grad_output): - grad_output = grad_output.contiguous() - checkpoint_lvl = ctx.checkpoint_lvl - x, weight1, bias1, weight2, *rest = ctx.saved_tensors - batch_shape, n = x.shape[:-1], x.shape[-1] - batch_dim = batch_shape.numel() - is_bf16 = x.dtype == torch.bfloat16 - if checkpoint_lvl == 0: - act_input, output1 = rest - elif checkpoint_lvl == 1: - (act_input,) = rest - output1 = sqrelu_fwd(act_input) - elif checkpoint_lvl == 2: - if is_bf16: - act_input = fused_dense_cuda.linear_bias_forward( - x.reshape(batch_dim, n), weight1, bias1 - ) - output1 = sqrelu_fwd(act_input) - else: - output1, act_input = triton_linear_act( - x.reshape(batch_dim, n), - weight1, - bias1, - activation="squared_relu", - save_act_input=True, - ) - - if is_bf16: - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad(output1, grad_output) - grad_output1 = grad_output @ weight2 - grad_act_input = sqrelu_bwd(grad_output1, act_input) - grad_input, grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_backward( - x.reshape(batch_dim, n), weight1, grad_act_input - ) - else: - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad(output1, grad_output) - grad_act_input = triton_dgrad_act( - grad_output, weight2, activation="squared_relu", act_input=act_input - ) - grad_input, grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_backward( - x.reshape(batch_dim, n), weight1, grad_act_input - ) - return grad_input.reshape_as(x), grad_weight1, grad_bias1, grad_weight2, grad_bias2, None - - -fused_dense_sqrelu_dense_function = FusedDenseSqreluDenseFunc.apply - - -class FusedDenseSqreluDense(nn.Module): - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - bias1=True, - bias2=True, - checkpoint_lvl=0, - device=None, - dtype=None, - ): - """ - checkpoint_lvl (increasing lvl means slower but more memory saving): - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute gelu_in and gelu_out in the bwd - """ - assert checkpoint_lvl in [0, 1, 2] - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features * 4 - assert bias1 == True, "DenseSqreluDense module without bias is currently not supported" - assert bias2 == True, "DenseSqreluDense module without bias is currently not supported" - self.checkpoint_lvl = checkpoint_lvl - self.fc1 = nn.Linear(in_features, hidden_features, bias=bias1, **factory_kwargs) - self.fc2 = nn.Linear(hidden_features, out_features, bias=bias2, **factory_kwargs) - - def forward(self, x): - assert x.is_cuda - return fused_dense_sqrelu_dense_function( - x, self.fc1.weight, self.fc1.bias, self.fc2.weight, self.fc2.bias, self.checkpoint_lvl - ) diff --git a/build/torch212-cxx11-xpu20253-x86_64-linux/ops/triton/rotary.py b/build/torch212-cxx11-xpu20253-x86_64-linux/ops/triton/rotary.py deleted file mode 100644 index 6ad3ab412a084270120167ad32b67ad1720f2c37..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-xpu20253-x86_64-linux/ops/triton/rotary.py +++ /dev/null @@ -1,186 +0,0 @@ -# Copyright (c) 2025, Tri Dao. -# As of 2025-04-23, we require triton >= 3.0 - -from typing import Optional, Union - -import torch - -import triton -import triton.language as tl - - -@triton.jit -def rotary_kernel( - OUT, # Pointers to matrices - X, - COS, - SIN, - CU_SEQLENS, - SEQLEN_OFFSETS, # this could be int or a pointer - # Matrix dimensions - seqlen, - nheads, - seqlen_ro, - # strides - stride_out_batch, - stride_out_seqlen, - stride_out_nheads, - stride_out_headdim, - stride_x_batch, - stride_x_seqlen, - stride_x_nheads, - stride_x_headdim, - # Meta-parameters - # We want ROTARY_DIM to be constexpr, otherwise the triton compiler doesn't know that - # the mask is constant every 8 elements, and it will generate LDG.16 instead of LDG.128 - ROTARY_DIM: tl.constexpr, - IS_SEQLEN_OFFSETS_TENSOR: tl.constexpr, - IS_VARLEN: tl.constexpr, - INTERLEAVED: tl.constexpr, - CONJUGATE: tl.constexpr, - BLOCK_H: tl.constexpr, - BLOCK_M: tl.constexpr, -): - BLOCK_K: tl.constexpr = triton.next_power_of_2(ROTARY_DIM) - ROTARY_DIM_HALF = ROTARY_DIM // 2 - pid_head = tl.program_id(axis=0) - pid_m = tl.program_id(axis=1) - pid_batch = tl.program_id(axis=2) - - if not IS_VARLEN: - X = X + pid_batch * stride_x_batch - OUT = OUT + pid_batch * stride_out_batch - else: - start_idx = tl.load(CU_SEQLENS + pid_batch) - seqlen = tl.load(CU_SEQLENS + pid_batch + 1) - start_idx - X = X + start_idx * stride_x_seqlen - OUT = OUT + start_idx * stride_out_seqlen - - if pid_m * BLOCK_M >= seqlen: - return - - rh = pid_head * BLOCK_H + tl.arange(0, BLOCK_H) - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - if not IS_SEQLEN_OFFSETS_TENSOR: - rm_cs = rm + SEQLEN_OFFSETS - else: - rm_cs = rm + tl.load(SEQLEN_OFFSETS + pid_batch) - - rk_half = tl.arange(0, BLOCK_K // 2) - COS = COS + (rm_cs[:, None] * ROTARY_DIM_HALF + rk_half[None, :]) - SIN = SIN + (rm_cs[:, None] * ROTARY_DIM_HALF + rk_half[None, :]) - mask_cs = (rm_cs[:, None] < seqlen_ro) & (rk_half[None, :] < ROTARY_DIM_HALF) - cos = tl.load(COS, mask=mask_cs, other=1.0).to(tl.float32) - sin = tl.load(SIN, mask=mask_cs, other=0.0).to(tl.float32) - if CONJUGATE: - sin = -sin - - if not INTERLEAVED: - # Load the 1st and 2nd halves of X, do calculation, then store to 1st and 2nd halves of OUT - X = X + (rh[:, None, None] * stride_x_nheads + rm[None, :, None] * stride_x_seqlen + rk_half[None, None, :] * stride_x_headdim) - OUT = OUT + (rh[:, None, None] * stride_out_nheads + rm[None, :, None] * stride_out_seqlen + rk_half[None, None, :] * stride_out_headdim) - mask = (rh[:, None, None] < nheads) & (rm[None, :, None] < seqlen) & (rk_half[None, None, :] < ROTARY_DIM_HALF) - x0 = tl.load(X, mask=mask, other=0.0).to(tl.float32) - x1 = tl.load(X + ROTARY_DIM_HALF * stride_x_headdim, mask=mask, other=0.0,).to(tl.float32) - o0 = x0 * cos - x1 * sin - o1 = x0 * sin + x1 * cos - tl.store(OUT, o0, mask=mask) - tl.store(OUT + ROTARY_DIM_HALF * stride_out_headdim, o1, mask=mask) - else: - rk = tl.arange(0, BLOCK_K) - X = X + (rh[:, None, None] * stride_x_nheads + rm[None, :, None] * stride_x_seqlen + rk[None, None, :] * stride_x_headdim) - OUT = OUT + (rh[:, None, None] * stride_out_nheads + rm[None, :, None] * stride_out_seqlen + rk[None, None, :] * stride_out_headdim) - mask = (rh[:, None, None] < nheads) & (rm[None, :, None] < seqlen) & (rk[None, None, :] < ROTARY_DIM) - x = tl.load(X, mask=mask, other=0.0).to(tl.float32) - x0, x1 = tl.split(tl.reshape(x, [BLOCK_H, BLOCK_M, BLOCK_K // 2, 2])) - o0 = x0 * cos - x1 * sin - o1 = x0 * sin + x1 * cos - o = tl.reshape(tl.join(o0, o1), [BLOCK_H, BLOCK_M, BLOCK_K]) - tl.store(OUT, o, mask=mask) - - -def apply_rotary( - x: torch.Tensor, - cos: torch.Tensor, - sin: torch.Tensor, - seqlen_offsets: Union[int, torch.Tensor] = 0, - cu_seqlens: Optional[torch.Tensor] = None, - max_seqlen: Optional[int] = None, - interleaved=False, - inplace=False, - conjugate=False, -) -> torch.Tensor: - """ - Arguments: - x: (batch, seqlen, nheads, headdim) if cu_seqlens is None - else (total_seqlen, nheads, headdim). - cos: (seqlen_ro, rotary_dim / 2) - sin: (seqlen_ro, rotary_dim / 2) - seqlen_offsets: integer or integer tensor of size (batch,) - cu_seqlens: (batch + 1,) or None - max_seqlen: int - Returns: - y: (batch, seqlen, nheads, headdim) - """ - is_varlen = cu_seqlens is not None - if not is_varlen: - batch, seqlen, nheads, headdim = x.shape - else: - assert max_seqlen is not None, "If cu_seqlens is passed in, then max_seqlen must be passed" - total_seqlen, nheads, headdim = x.shape - batch_p_1 = cu_seqlens.shape[0] - batch = batch_p_1 - 1 - seqlen = max_seqlen - seqlen_ro, rotary_dim = cos.shape - assert sin.shape == cos.shape - rotary_dim *= 2 - assert rotary_dim <= headdim, "rotary_dim must be <= headdim" - assert headdim <= 256, "Only support headdim <= 256" - assert seqlen_ro >= seqlen, "seqlen_ro must be >= seqlen" - - cos, sin = cos.contiguous(), sin.contiguous() - if isinstance(seqlen_offsets, torch.Tensor): - assert seqlen_offsets.shape == (batch,) - assert seqlen_offsets.dtype in [torch.int32, torch.int64] - seqlen_offsets = seqlen_offsets.contiguous() - else: - assert seqlen_offsets + seqlen <= seqlen_ro - - output = torch.empty_like(x) if not inplace else x - if rotary_dim < headdim and not inplace: - output[..., rotary_dim:].copy_(x[..., rotary_dim:]) - - grid = lambda META: (triton.cdiv(nheads, META["BLOCK_H"]), triton.cdiv(seqlen, META["BLOCK_M"]), batch) # noqa - BLOCK_M = 8 if rotary_dim <= 128 else 4 - - # Need this, otherwise Triton tries to launch from cuda:0 and we get - # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?) - device_ctx = torch.cuda.device(x.device.index) if x.device.type == 'cuda' else torch.xpu.device(x.device.index) - with device_ctx: - torch.library.wrap_triton(rotary_kernel)[grid]( - output, # data ptrs - x, - cos, - sin, - cu_seqlens, - seqlen_offsets, - seqlen, # shapes - nheads, - seqlen_ro, - output.stride(0) if not is_varlen else 0, # batch_strides if not varlen else 0 - output.stride(-3), # seqlen_stride or total_seqlen_stride - output.stride(-2), # nheads_stride - output.stride(-1), # headdim_stride - x.stride(0) if not is_varlen else 0, # batch_strides if not varlen else 0 - x.stride(-3), # seqlen stride or total_seqlen_stride - x.stride(-2), # nheads stride - x.stride(-1), # headdim stride - rotary_dim, - isinstance(seqlen_offsets, torch.Tensor), - is_varlen, - interleaved, - conjugate, - BLOCK_M=BLOCK_M, - BLOCK_H=2, - ) - return output diff --git a/build/torch27-cxx11-cu118-x86_64-linux/flash_attn2/_flash_attn_9e27194.abi3.so b/build/torch27-cxx11-cu118-x86_64-linux/flash_attn2/_flash_attn_9e27194.abi3.so index 2403670ba5a46be2941aa992204e38d996ee5af9..c1d1644ba2cd76bbe0a8671045765eb553f18933 100644 --- a/build/torch27-cxx11-cu118-x86_64-linux/flash_attn2/_flash_attn_9e27194.abi3.so +++ b/build/torch27-cxx11-cu118-x86_64-linux/flash_attn2/_flash_attn_9e27194.abi3.so @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:76aafd307b4f2f2db05f8e450c68d004b5061ce726b6172a5d3411ddf54f338d +oid sha256:760f19632acec698b69e0898d336e1d5aa26d87318d26ab653366cc8d5a8eec7 size 445273544 diff --git a/build/torch27-cxx11-cu126-x86_64-linux/flash_attn2/_flash_attn_9e27194.abi3.so b/build/torch27-cxx11-cu126-x86_64-linux/flash_attn2/_flash_attn_9e27194.abi3.so index 6d9a26841b52a2fdbbae47c6bbd412a2f1ce3a62..822ed49c8580926df9eafb3735f13ce45da10c6c 100644 --- a/build/torch27-cxx11-cu126-x86_64-linux/flash_attn2/_flash_attn_9e27194.abi3.so +++ b/build/torch27-cxx11-cu126-x86_64-linux/flash_attn2/_flash_attn_9e27194.abi3.so @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:0d130f4d2250a8025dacd5094474507c4f721d1f92ebfd51aaa048b23d346b3d +oid sha256:655d9a3bc163e73eaaa59be60819cd9a5b5b68624035e5e810c295f537c6d260 size 448639288 diff --git a/build/torch27-cxx11-cu128-x86_64-linux/flash_attn2/_flash_attn_9e27194.abi3.so b/build/torch27-cxx11-cu128-x86_64-linux/flash_attn2/_flash_attn_9e27194.abi3.so index 4f19f529c124fc09c9711d4c92e9cf171283dab5..841d19e037dacf8a264d6b3068d5837d66fb5e90 100644 --- a/build/torch27-cxx11-cu128-x86_64-linux/flash_attn2/_flash_attn_9e27194.abi3.so +++ b/build/torch27-cxx11-cu128-x86_64-linux/flash_attn2/_flash_attn_9e27194.abi3.so @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:0974763439908e223d9e3627dfff3f636e9d91f830952a35a914af55ee4d0af2 +oid sha256:d75e5b23edcf1f4c17e88a6610b83eba6eac8293ee52ba297ffb713f894a7e78 size 1037635040 diff --git a/build/torch28-cxx11-cu126-x86_64-linux/_flash_attn2_588b404.abi3.so b/build/torch28-cxx11-cu126-x86_64-linux/_flash_attn2_588b404.abi3.so index 22858e0d0e26b016467f345907ec5df43d1ef5f6..c48a9eb21c747827b9564d3d8b293c1845931a5f 100644 --- a/build/torch28-cxx11-cu126-x86_64-linux/_flash_attn2_588b404.abi3.so +++ b/build/torch28-cxx11-cu126-x86_64-linux/_flash_attn2_588b404.abi3.so @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:5b9440c7b7253af45d06c853a7878bf154705e0e8ad8192b7391790a42d3f3c0 +oid sha256:7833489dd07c22325910d36194943fcf8808f4a5a2e7acda97344ac79af3e433 size 448648680 diff --git a/build/torch28-cxx11-cu128-x86_64-linux/_flash_attn2_588b404.abi3.so b/build/torch28-cxx11-cu128-x86_64-linux/_flash_attn2_588b404.abi3.so index b0aab259e84d65028189f7c55ea1d3bdab5570e9..eb0efe07334dd9d89465bbc03fa0cb8b7daac26d 100644 --- a/build/torch28-cxx11-cu128-x86_64-linux/_flash_attn2_588b404.abi3.so +++ b/build/torch28-cxx11-cu128-x86_64-linux/_flash_attn2_588b404.abi3.so @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:966860bc62695941fa02c2966ab7dad4fafc5b13487e231b7781523eef745dab +oid sha256:e6c2cb7d67321c8bba85b79fb132b0c15b38702403733ca864f35ee2b51dba5c size 1037644552 diff --git a/build/torch28-cxx11-cu129-x86_64-linux/_flash_attn2_588b404.abi3.so b/build/torch28-cxx11-cu129-x86_64-linux/_flash_attn2_588b404.abi3.so index 39b88d2a614f7c996c341d5dcfd24352638e94f8..33ed1b4de6c814a734f349894b693e7825ee3f99 100644 --- a/build/torch28-cxx11-cu129-x86_64-linux/_flash_attn2_588b404.abi3.so +++ b/build/torch28-cxx11-cu129-x86_64-linux/_flash_attn2_588b404.abi3.so @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:8ea86598d250413c68ab4cf8bb2eaf99ff138099dcccee119fd3410d1c619287 +oid sha256:e5c631fd1f3538394320256ac62a5faeb85383051342e26a4e61b79b2b2eeb0e size 1042944096 diff --git a/build/torch29-cxx11-cpu-x86_64-linux/_flash_attn2_588b404.abi3.so b/build/torch29-cxx11-cpu-x86_64-linux/_flash_attn2_588b404.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..6d06e695d9ef36b118ee97b0c9a5853a145bfe0b --- /dev/null +++ b/build/torch29-cxx11-cpu-x86_64-linux/_flash_attn2_588b404.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:dfab2b6517d74cf640d3da51f60a787dc42fa98774e10bf52d0a265cc5423f53 +size 239416 diff --git a/build/torch29-cxx11-cpu-x86_64-linux/_flash_attn2_cpu_abda3a0.abi3.so b/build/torch29-cxx11-cpu-x86_64-linux/_flash_attn2_cpu_abda3a0.abi3.so deleted file mode 100644 index 0738988b3dcd6d27094a4b39aec7de5f6dce9325..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cpu-x86_64-linux/_flash_attn2_cpu_abda3a0.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:f48a5de0256294c0baf17e17ff3eafa7c818a6b70c869d6eec74cdff07b994f1 -size 1932200 diff --git a/build/torch29-cxx11-cpu-x86_64-linux/_ops.py b/build/torch29-cxx11-cpu-x86_64-linux/_ops.py index e86c4dfc843a98b9566275b46b630435fcd980b6..03506d8a18d17ca7296a5f1b7a18ec7dc4ad73bb 100644 --- a/build/torch29-cxx11-cpu-x86_64-linux/_ops.py +++ b/build/torch29-cxx11-cpu-x86_64-linux/_ops.py @@ -1,9 +1,9 @@ import torch -from . import _flash_attn2_cpu_abda3a0 -ops = torch.ops._flash_attn2_cpu_abda3a0 +from . import _flash_attn2_588b404 +ops = torch.ops._flash_attn2_588b404 def add_op_namespace_prefix(op_name: str): """ Prefix op by namespace. """ - return f"_flash_attn2_cpu_abda3a0::{op_name}" + return f"_flash_attn2_588b404::{op_name}" \ No newline at end of file diff --git a/build/torch29-cxx11-cpu-x86_64-linux/flash_attn_interface.py b/build/torch29-cxx11-cpu-x86_64-linux/flash_attn_interface.py index 6a0157f25449eb90c6c688bc8ef3f4fcf84d46dc..66e8d3665834b141d703998496b3da6fe1f6fe06 100644 --- a/build/torch29-cxx11-cpu-x86_64-linux/flash_attn_interface.py +++ b/build/torch29-cxx11-cpu-x86_64-linux/flash_attn_interface.py @@ -31,14 +31,12 @@ def _get_device(): else: return "cpu" +_XPU_AVAILABLE = torch.xpu.is_available() if hasattr(torch, "xpu") else False # TODO remove hasattr check when bwd is supported on XPU -def _get_block_size_n(device, head_dim, is_dropout, is_causal): - assert head_dim <= 256 - - if device.type == "xpu": - return 64 +def _get_block_size_n(device, head_dim, is_dropout, is_causal): # This should match the block sizes in the CUDA kernel + assert head_dim <= 256 major, minor = torch.cuda.get_device_capability(device) is_sm8x = major == 8 and minor > 0 # Only include sm86 and sm89, exclude sm80 (A100) is_sm80 = major == 8 and minor == 0 @@ -1066,7 +1064,7 @@ def flash_attn_qkvpacked_func( alibi_slopes, deterministic, return_attn_probs, - torch.is_grad_enabled(), + False if _XPU_AVAILABLE else torch.is_grad_enabled(), ) @@ -1144,7 +1142,7 @@ def flash_attn_kvpacked_func( alibi_slopes, deterministic, return_attn_probs, - torch.is_grad_enabled(), + False if _XPU_AVAILABLE else torch.is_grad_enabled(), ) @@ -1221,7 +1219,7 @@ def flash_attn_func( alibi_slopes, deterministic, return_attn_probs, - torch.is_grad_enabled(), + False if _XPU_AVAILABLE else torch.is_grad_enabled(), ) @@ -1287,7 +1285,7 @@ def flash_attn_varlen_qkvpacked_func( alibi_slopes, deterministic, return_attn_probs, - torch.is_grad_enabled(), + False if _XPU_AVAILABLE else torch.is_grad_enabled(), ) @@ -1379,7 +1377,7 @@ def flash_attn_varlen_kvpacked_func( alibi_slopes, deterministic, return_attn_probs, - torch.is_grad_enabled(), + False if _XPU_AVAILABLE else torch.is_grad_enabled(), ) @@ -1473,7 +1471,7 @@ def flash_attn_varlen_func( deterministic, return_attn_probs, block_table, - False if q.device.type == "cpu" else torch.is_grad_enabled(), + False if _XPU_AVAILABLE or q.device.type == "cpu" else torch.is_grad_enabled(), ) diff --git a/build/torch29-cxx11-cpu-x86_64-linux/metadata.json b/build/torch29-cxx11-cpu-x86_64-linux/metadata.json index de6520c18deaab0372f91d85948970c48240031c..9cf5deed9898dce769f4cc73913d3530b92a0bd8 100644 --- a/build/torch29-cxx11-cpu-x86_64-linux/metadata.json +++ b/build/torch29-cxx11-cpu-x86_64-linux/metadata.json @@ -1,5 +1,4 @@ { "version": 1, - "license": "BSD-3-Clause", "python-depends": [] } \ No newline at end of file diff --git a/build/torch29-cxx11-cu126-aarch64-linux/__init__.py b/build/torch29-cxx11-cu126-aarch64-linux/__init__.py deleted file mode 100644 index ecc2f9d896b6c93f90b0a1499856dc0612177422..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-aarch64-linux/__init__.py +++ /dev/null @@ -1,393 +0,0 @@ -from typing import Optional, List -import torch -from ._ops import ops as flash_attn_ops -from .flash_attn_interface import ( - flash_attn_func, - flash_attn_kvpacked_func, - flash_attn_qkvpacked_func, - flash_attn_varlen_func, - flash_attn_varlen_kvpacked_func, - flash_attn_varlen_qkvpacked_func, - flash_attn_with_kvcache, -) - - -def fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - return_softmax: bool = False, - gen: Optional[torch.Generator] = None, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention. - - Args: - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Optional output tensor, same shape as q - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - return_softmax: Whether to return softmax weights - gen: Optional random number generator - - Returns: - List of tensors: [output, softmax_lse, (softmax if return_softmax)] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.fwd( - q, - k, - v, - out, - alibi_slopes, - p_dropout, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - gen, - ) - - -def varlen_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - out: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - block_table: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - max_seqlen_q: int = 0, - max_seqlen_k: int = 0, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - zero_tensors: bool = False, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - return_softmax: bool = False, - gen: Optional[torch.Generator] = None, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention with variable sequence lengths. - - Args: - q: Query tensor of shape [total_q, num_heads, head_size] - k: Key tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - v: Value tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1] - cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1] - out: Optional output tensor of shape [total_q, num_heads, head_size] - seqused_k: Optional tensor specifying how many keys to use per batch element [batch_size] - leftpad_k: Optional left padding for keys of shape [batch_size] - block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq] - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - max_seqlen_q: Maximum sequence length for queries - max_seqlen_k: Maximum sequence length for keys - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - zero_tensors: Whether to zero tensors before computation - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - return_softmax: Whether to return softmax weights - gen: Optional random number generator - - Returns: - List of tensors: [output, softmax_lse, (softmax if return_softmax)] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.varlen_fwd( - q, - k, - v, - out, - cu_seqlens_q, - cu_seqlens_k, - seqused_k, - leftpad_k, - block_table, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - p_dropout, - softmax_scale, - zero_tensors, - is_causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - gen, - ) - - -def bwd( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - p_dropout: float = 0.0, - 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, - gen: Optional[torch.Generator] = None, - rng_state: Optional[torch.Tensor] = None, -) -> List[torch.Tensor]: - """ - Backward pass for multi-head attention. - - Args: - dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size] - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size] - softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q] - dq: Optional gradient tensor for queries, same shape as q - dk: Optional gradient tensor for keys, same shape as k - dv: Optional gradient tensor for values, same shape as v - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - deterministic: Whether to use deterministic algorithms - gen: Optional random number generator - rng_state: Optional RNG state from forward pass - - Returns: - List of tensors: [dq, dk, dv] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - alibi_slopes, - p_dropout, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - deterministic, - gen, - rng_state, - ) - - -def varlen_bwd( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - max_seqlen_q: int = 0, - max_seqlen_k: int = 0, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - zero_tensors: bool = False, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - deterministic: bool = False, - gen: Optional[torch.Generator] = None, - rng_state: Optional[torch.Tensor] = None, -) -> List[torch.Tensor]: - """ - Backward pass for multi-head attention with variable sequence lengths. - - Args: - dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size] - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size] - softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q] - cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1] - cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1] - dq: Optional gradient tensor for queries, same shape as q - dk: Optional gradient tensor for keys, same shape as k - dv: Optional gradient tensor for values, same shape as v - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - max_seqlen_q: Maximum sequence length for queries - max_seqlen_k: Maximum sequence length for keys - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - zero_tensors: Whether to zero tensors before computation - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - deterministic: Whether to use deterministic algorithms - gen: Optional random number generator - rng_state: Optional RNG state from forward pass - - Returns: - List of tensors: [dq, dk, dv] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.varlen_bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - p_dropout, - softmax_scale, - zero_tensors, - is_causal, - window_size_left, - window_size_right, - softcap, - deterministic, - gen, - rng_state, - ) - - -def fwd_kvcache( - q: torch.Tensor, - kcache: torch.Tensor, - vcache: torch.Tensor, - k: Optional[torch.Tensor] = None, - v: Optional[torch.Tensor] = None, - seqlens_k: Optional[torch.Tensor] = None, - rotary_cos: Optional[torch.Tensor] = None, - rotary_sin: Optional[torch.Tensor] = None, - cache_batch_idx: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - block_table: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - out: 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, - is_rotary_interleaved: bool = False, - num_splits: int = 1, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention with KV cache. - - Args: - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - kcache: Key cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - vcache: Value cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - k: Optional new keys tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size] - v: Optional new values tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size] - seqlens_k: Optional sequence lengths for keys of shape [batch_size] - rotary_cos: Optional rotary cosine tensor of shape [seqlen_ro, rotary_dim/2] - rotary_sin: Optional rotary sine tensor of shape [seqlen_ro, rotary_dim/2] - cache_batch_idx: Optional indices to index into the KV cache - leftpad_k: Optional left padding for keys of shape [batch_size] - block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq] - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - out: Optional output tensor, same shape as q - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - is_rotary_interleaved: Whether rotary embeddings are interleaved - num_splits: Number of splits for computation - - Returns: - List of tensors: [output, softmax_lse] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.fwd_kvcache( - q, - kcache, - vcache, - k, - v, - seqlens_k, - rotary_cos, - rotary_sin, - cache_batch_idx, - leftpad_k, - block_table, - alibi_slopes, - out, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - is_rotary_interleaved, - num_splits, - ) diff --git a/build/torch29-cxx11-cu126-aarch64-linux/_flash_attn2_cuda_abda3a0.abi3.so b/build/torch29-cxx11-cu126-aarch64-linux/_flash_attn2_cuda_abda3a0.abi3.so deleted file mode 100644 index 6f6b2c52cc0c23c3fc413da45f13ede362300956..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-aarch64-linux/_flash_attn2_cuda_abda3a0.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:e3e9388d04fee347a39220864d3274bfa46a0fa54640ad0c3201425bb3a8cb75 -size 448464272 diff --git a/build/torch29-cxx11-cu126-aarch64-linux/_ops.py b/build/torch29-cxx11-cu126-aarch64-linux/_ops.py deleted file mode 100644 index fb32434e09e651a658beb38be8c101e5db89374f..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-aarch64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _flash_attn2_cuda_abda3a0 -ops = torch.ops._flash_attn2_cuda_abda3a0 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_flash_attn2_cuda_abda3a0::{op_name}" diff --git a/build/torch29-cxx11-cu126-aarch64-linux/bert_padding.py b/build/torch29-cxx11-cu126-aarch64-linux/bert_padding.py deleted file mode 100644 index 3c2d35159a014a9d03aabead9e52e009168696ea..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-aarch64-linux/bert_padding.py +++ /dev/null @@ -1,218 +0,0 @@ -# Adapted from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/padding.py - -import torch -import torch.nn.functional as F -from einops import rearrange, repeat - - -class IndexFirstAxis(torch.autograd.Function): - @staticmethod - def forward(ctx, input, indices): - ctx.save_for_backward(indices) - assert input.ndim >= 2 - ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] - second_dim = other_shape.numel() - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - # return input[indices] - return torch.gather( - rearrange(input, "b ... -> b (...)"), 0, repeat(indices, "z -> z d", d=second_dim) - ).reshape(-1, *other_shape) - - @staticmethod - def backward(ctx, grad_output): - (indices,) = ctx.saved_tensors - assert grad_output.ndim >= 2 - other_shape = grad_output.shape[1:] - grad_output = rearrange(grad_output, "b ... -> b (...)") - grad_input = torch.zeros( - [ctx.first_axis_dim, grad_output.shape[1]], - device=grad_output.device, - dtype=grad_output.dtype, - ) - # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. - # grad_input[indices] = grad_output - grad_input.scatter_(0, repeat(indices, "z -> z d", d=grad_output.shape[1]), grad_output) - return grad_input.reshape(ctx.first_axis_dim, *other_shape), None - - -index_first_axis = IndexFirstAxis.apply - - -class IndexPutFirstAxis(torch.autograd.Function): - @staticmethod - def forward(ctx, values, indices, first_axis_dim): - ctx.save_for_backward(indices) - assert indices.ndim == 1 - assert values.ndim >= 2 - output = torch.zeros( - first_axis_dim, *values.shape[1:], device=values.device, dtype=values.dtype - ) - # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. - output[indices] = values - # output.scatter_(0, repeat(indices, 'z -> z d', d=values.shape[1]), values) - return output - - @staticmethod - def backward(ctx, grad_output): - (indices,) = ctx.saved_tensors - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - grad_values = grad_output[indices] - # grad_values = torch.gather(grad_output, 0, repeat(indices, 'z -> z d', d=grad_output.shape[1])) - return grad_values, None, None - - -index_put_first_axis = IndexPutFirstAxis.apply - - -class IndexFirstAxisResidual(torch.autograd.Function): - @staticmethod - def forward(ctx, input, indices): - ctx.save_for_backward(indices) - assert input.ndim >= 2 - ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] - second_dim = other_shape.numel() - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - output = input[indices] - # We don't want to reshape input (b ... -> b (...)) since it could change the channel_last - # memory format to channel_first. In other words, input might not be contiguous. - # If we don't detach, Pytorch complains about output being a view and is being modified inplace - return output, input.detach() - - @staticmethod - def backward(ctx, grad_output, grad_residual): - (indices,) = ctx.saved_tensors - assert grad_output.ndim >= 2 - other_shape = grad_output.shape[1:] - assert grad_residual.shape[1:] == other_shape - grad_input = grad_residual - # grad_input[indices] += grad_output - indices = indices.reshape(indices.shape[0], *((1,) * (grad_output.ndim - 1))) - indices = indices.expand_as(grad_output) - grad_input.scatter_add_(0, indices, grad_output) - return grad_input.reshape(ctx.first_axis_dim, *other_shape), None - - -index_first_axis_residual = IndexFirstAxisResidual.apply - - -def unpad_input(hidden_states, attention_mask, unused_mask=None): - """ - Arguments: - hidden_states: (batch, seqlen, ...) - attention_mask: (batch, seqlen), bool / int, 1 means valid and 0 means not valid. - unused_mask: (batch, seqlen), bool / int, 1 means the element is allocated but unused. - Return: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens selected in attention_mask + unused_mask. - indices: (total_nnz), the indices of masked tokens from the flattened input sequence. - cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. - max_seqlen_in_batch: int - seqused: (batch), returns the number of tokens selected in attention_mask + unused_mask. - """ - all_masks = (attention_mask + unused_mask) if unused_mask is not None else attention_mask - seqlens_in_batch = all_masks.sum(dim=-1, dtype=torch.int32) - used_seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) - indices = torch.nonzero(all_masks.flatten(), as_tuple=False).flatten() - max_seqlen_in_batch = seqlens_in_batch.max().item() - cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) - # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the - # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim - # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to - # index with integer indices. Moreover, torch's index is a bit slower than it needs to be, - # so we write custom forward and backward to make it a bit faster. - return ( - index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices), - indices, - cu_seqlens, - max_seqlen_in_batch, - used_seqlens_in_batch, - ) - - -def unpad_input_for_concatenated_sequences(hidden_states, attention_mask_in_length): - """ - Supports concatenating short samples in one sequence. The attention_mask_in_length is utilized to mask other short samples. It helps efficient training of variant lengths-based samples (e.g., the supervised fine-tuning task in large language model). - The motivation for this function is explained [here](https://github.com/Dao-AILab/flash-attention/issues/432#issuecomment-1668822286). - - For example, if batch = 3 and seqlen = 6, the attention_mask_in_length is: - ``` - [ - [2, 3, 0, 0, 0, 0], - [3, 2, 0, 0, 0, 0], - [6, 0, 0, 0, 0, 0] - ] - ``` - , which refers to the 3D-attention mask: - ``` - [ - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [0, 0, 1, 0, 0, 0], - [0, 0, 1, 1, 0, 0], - [0, 0, 1, 1, 1, 0], - [0, 0, 0, 0, 0, 1] - ], - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [0, 0, 0, 1, 0, 0], - [0, 0, 0, 1, 1, 0], - [0, 0, 0, 0, 0, 1] - ], - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [1, 1, 1, 1, 0, 0], - [1, 1, 1, 1, 1, 0], - [1, 1, 1, 1, 1, 1] - ] - ] - ```. - - Arguments: - hidden_states: (batch, seqlen, ...) - attention_mask_in_length: (batch, seqlen), int, a nonzero number (e.g., 1, 2, 3, etc.) means length of concatenated sequence in b-th batch, and 0 means none. - Return: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. - indices: (total_nnz), the indices of non-masked tokens from the flattened input sequence. - cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. - max_seqlen_in_batch: int - """ - length = attention_mask_in_length.sum(dim=-1) - seqlen = attention_mask_in_length.size(-1) - attention_mask_2d = torch.arange(seqlen, device=length.device, dtype=length.dtype).expand(len(length), seqlen) < length.unsqueeze(1) - real_indices_idx = torch.nonzero(attention_mask_in_length.flatten(), as_tuple=False).flatten() - seqlens_in_batch = attention_mask_in_length.flatten()[real_indices_idx] - indices = torch.nonzero(attention_mask_2d.flatten(), as_tuple=False).flatten() - max_seqlen_in_batch = seqlens_in_batch.max().item() - cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) - # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the - # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim - # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to - # index with integer indices. Moreover, torch's index is a bit slower than it needs to be, - # so we write custom forward and backward to make it a bit faster. - return ( - index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices), - indices, - cu_seqlens, - max_seqlen_in_batch, - ) - - -def pad_input(hidden_states, indices, batch, seqlen): - """ - Arguments: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. - indices: (total_nnz), the indices that represent the non-masked tokens of the original padded input sequence. - batch: int, batch size for the padded sequence. - seqlen: int, maximum sequence length for the padded sequence. - Return: - hidden_states: (batch, seqlen, ...) - """ - dim = hidden_states.shape[-1] - # output = torch.zeros((batch * seqlen), dim, device=hidden_states.device, dtype=hidden_states.dtype) - # output[indices] = hidden_states - output = index_put_first_axis(hidden_states, indices, batch * seqlen) - return rearrange(output, "(b s) ... -> b s ...", b=batch) diff --git a/build/torch29-cxx11-cu126-aarch64-linux/flash_attn2/__init__.py b/build/torch29-cxx11-cu126-aarch64-linux/flash_attn2/__init__.py deleted file mode 100644 index 03dbc1afe1cf156661a2b1b22003cd5f599a0309..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-aarch64-linux/flash_attn2/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -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-cu126-aarch64-linux/flash_attn_interface.py b/build/torch29-cxx11-cu126-aarch64-linux/flash_attn_interface.py deleted file mode 100644 index 6a0157f25449eb90c6c688bc8ef3f4fcf84d46dc..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-aarch64-linux/flash_attn_interface.py +++ /dev/null @@ -1,1622 +0,0 @@ -# Copyright (c) 2023, Tri Dao. - -from typing import Optional, Sequence, Tuple, Union - -import torch -import torch.nn as nn -import os - -# # isort: off -# # We need to import the CUDA kernels after importing torch -# USE_TRITON_ROCM = os.getenv("FLASH_ATTENTION_TRITON_AMD_ENABLE", "FALSE") == "TRUE" -# if USE_TRITON_ROCM: -# from .flash_attn_triton_amd import interface_fa as flash_attn -# else: -# import flash_attn_2_cuda as flash_attn - - -from ._ops import ops as flash_attn - -# # isort: on - -def maybe_contiguous(x): - return x.contiguous() if x is not None and x.stride(-1) != 1 else x - - -def _get_device(): - if torch.xpu.is_available(): - return "xpu" - elif torch.cuda.is_available(): - return "cuda" - else: - return "cpu" - - -def _get_block_size_n(device, head_dim, is_dropout, is_causal): - assert head_dim <= 256 - - if device.type == "xpu": - return 64 - - # This should match the block sizes in the CUDA kernel - major, minor = torch.cuda.get_device_capability(device) - is_sm8x = major == 8 and minor > 0 # Only include sm86 and sm89, exclude sm80 (A100) - is_sm80 = major == 8 and minor == 0 - is_sm90 = major == 9 and minor == 0 - if head_dim <= 32: - return 128 - if head_dim <= 64: - return 128 if not is_dropout else 64 - elif head_dim <= 96: - return 64 - elif head_dim <= 128: - if is_sm8x: - return 64 if (not is_dropout and is_causal) else 32 - else: - return 64 if not is_dropout else 32 - elif head_dim <= 192: - return 64 - elif head_dim <= 224: - return 64 - elif head_dim <= 256: - return 64 - - -def round_multiple(x, m): - return (x + m - 1) // m * m - - -# torch.compile() support is only enabled for pytorch >= 2.4 -# The reason for this is that we are using the new custom_op and register_fake -# APIs, which support inplace modification of inputs in the function itself -if torch.__version__ >= "2.4.0": - _torch_custom_op_wrapper = torch.library.custom_op - _torch_register_fake_wrapper = torch.library.register_fake -else: - def noop_custom_op_wrapper(name, fn=None, /, *, mutates_args, device_types=None, schema=None): - def wrap(func): - return func - if fn is None: - return wrap - return fn - def noop_register_fake_wrapper(op, fn=None, /, *, lib=None, _stacklevel=1): - def wrap(func): - return func - if fn is None: - return wrap - return fn - _torch_custom_op_wrapper = noop_custom_op_wrapper - _torch_register_fake_wrapper = noop_register_fake_wrapper - - -@_torch_custom_op_wrapper("flash_attn::_flash_attn_forward", mutates_args=(), device_types=_get_device()) -def _flash_attn_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - return_softmax: bool -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - out, softmax_lse, S_dmask, rng_state = flash_attn.fwd( - q, - k, - v, - None, - alibi_slopes, - dropout_p, - softmax_scale, - causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - None, - ) - return out, softmax_lse, S_dmask, rng_state - - -@_torch_register_fake_wrapper("flash_attn::_flash_attn_forward") -def _flash_attn_forward_fake( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - return_softmax: bool -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - batch_size, seqlen_q, num_heads, head_size = q.shape - seqlen_k = k.shape[1] - out = torch.empty_like(q) - softmax_lse = torch.empty((batch_size, num_heads, seqlen_q), dtype=torch.float32, device=q.device, layout=q.layout) - p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) - if return_softmax: - p = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128), round_multiple(seqlen_k, 128)), dtype=q.dtype, device=q.device, layout=q.layout) - rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) - - return out, softmax_lse, p, rng_state - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_forward = torch.ops.flash_attn._flash_attn_forward -else: - _wrapped_flash_attn_forward = _flash_attn_forward - - -@_torch_custom_op_wrapper("flash_attn::_flash_attn_varlen_forward", mutates_args=(), device_types=_get_device()) -def _flash_attn_varlen_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - alibi_slopes: Optional[torch.Tensor] = None, - return_softmax: bool = False, - block_table: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - out, softmax_lse, S_dmask, rng_state = flash_attn.varlen_fwd( - q, - k, - v, - None, - cu_seqlens_q, - cu_seqlens_k, - seqused_k, - leftpad_k, - block_table, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - zero_tensors, - causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - None, - ) - # if out.isnan().any() or softmax_lse.isnan().any(): - # breakpoint() - return out, softmax_lse, S_dmask, rng_state - - -@_torch_register_fake_wrapper("flash_attn::_flash_attn_varlen_forward") -def _flash_attn_varlen_forward_fake( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - alibi_slopes: Optional[torch.Tensor] = None, - return_softmax: bool = False, - block_table: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - paged_kv = block_table is not None - batch_size = cu_seqlens_q.numel() - 1 - total_q, num_heads, _ = q.shape - - out = torch.empty_like(q) - softmax_lse = torch.empty((num_heads, total_q), dtype=torch.float32, device=q.device, layout=q.layout) - p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) - seqlen_q_rounded = round_multiple(max_seqlen_q, 128) - seqlen_k_rounded = round_multiple(max_seqlen_k, 128) - if return_softmax: - p = torch.empty((batch_size, num_heads, seqlen_q_rounded, seqlen_k_rounded), dtype=q.dtype, device=q.device, layout=q.layout) - rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) - return out, softmax_lse, p, rng_state - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_varlen_forward = torch.ops.flash_attn._flash_attn_varlen_forward -else: - _wrapped_flash_attn_varlen_forward = _flash_attn_varlen_forward - - -@_torch_custom_op_wrapper("flash_attn::_flash_attn_backward", mutates_args=("dq", "dk", "dv"), device_types=_get_device()) -def _flash_attn_backward( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, -) -> 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)] - ( - dq, - dk, - dv, - softmax_d, - ) = flash_attn.bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - alibi_slopes, - dropout_p, - softmax_scale, - causal, - window_size_left, - window_size_right, - softcap, - deterministic, - None, - rng_state, - ) - return softmax_d - - -@_torch_register_fake_wrapper("flash_attn::_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, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, -) -> torch.Tensor: - dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] - if dq is None: - dq = torch.empty_like(q) - if dk is None: - dk = torch.empty_like(k) - if dv is None: - dv = torch.empty_like(v) - batch_size, seqlen_q, num_heads, _ = q.shape - softmax_d = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128)), device=q.device, dtype=torch.float32) - - return softmax_d - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_backward = torch.ops.flash_attn._flash_attn_backward -else: - _wrapped_flash_attn_backward = _flash_attn_backward - - -@_torch_custom_op_wrapper("flash_attn::_flash_attn_varlen_backward", mutates_args=("dq", "dk", "dv"), device_types=_get_device()) -def _flash_attn_varlen_backward( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> 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)] - ( - dq, - dk, - dv, - softmax_d, - ) = flash_attn.varlen_bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - zero_tensors, - causal, - window_size_left, - window_size_right, - softcap, - deterministic, - None, - rng_state, - ) - # if dk.isnan().any() or dk.isnan().any() or dv.isnan().any() or softmax_d.isnan().any(): - # breakpoint() - return softmax_d - - -@_torch_register_fake_wrapper("flash_attn::_flash_attn_varlen_backward") -def _flash_attn_varlen_backward_fake( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> torch.Tensor: - dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] - batch_size = cu_seqlens_q.numel() - 1 - total_q, num_heads, _ = q.shape - - if dq is None: - dq = torch.empty_like(q) - if dk is None: - dk = torch.empty_like(k) - if dv is None: - dv = torch.empty_like(v) - softmax_d = torch.empty((num_heads, total_q + 128 * batch_size), device=q.device, dtype=torch.float32) - - return softmax_d - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_varlen_backward = torch.ops.flash_attn._flash_attn_varlen_backward -else: - _wrapped_flash_attn_varlen_backward = _flash_attn_varlen_backward - - -class FlashAttnQKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and qkv.requires_grad - if softmax_scale is None: - softmax_scale = qkv.shape[-1] ** (-0.5) - q, k, v = qkv[:, :, 0].detach(), qkv[:, :, 1].detach(), qkv[:, :, 2].detach() - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) - dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dqkv[:, :, 0], - dqkv[:, :, 1], - dqkv[:, :, 2], - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dqkv = dqkv[..., : dout.shape[-1]] # We could have padded the head dimension - return dqkv, None, None, None, None, None, None, None, None, None - - -class FlashAttnVarlenQKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - cu_seqlens, - max_seqlen, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and qkv.requires_grad - if softmax_scale is None: - softmax_scale = qkv.shape[-1] ** (-0.5) - q, k, v = qkv[:, 0].detach(), qkv[:, 1].detach(), qkv[:, 2].detach() - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens, - cu_seqlens, - max_seqlen, - max_seqlen, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=None, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens, rng_state) - ctx.dropout_p = dropout_p - ctx.max_seqlen = max_seqlen - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens, rng_state = ctx.saved_tensors - qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) - dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dqkv[:, 0], - dqkv[:, 1], - dqkv[:, 2], - cu_seqlens, - cu_seqlens, - ctx.max_seqlen, - ctx.max_seqlen, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - 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 FlashAttnKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - kv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, kv] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - k, v = kv[:, :, 0].detach(), kv[:, :, 1].detach() - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - dq = torch.empty_like(q) - kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) - dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dkv[:, :, 0], - dkv[:, :, 1], - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dkv = dkv[..., : dout.shape[-1]] - return dq, dkv, None, None, None, None, None, None, None, None, None - - -class FlashAttnVarlenKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, kv] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - k, v = kv[:, 0].detach(), kv[:, 1].detach() - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=None, - ) - if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state - ) - ctx.dropout_p = dropout_p - 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.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors - dq = torch.empty_like(q) - kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) - dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dkv[:, 0], - dkv[:, 1], - cu_seqlens_q, - cu_seqlens_k, - ctx.max_seqlen_q, - ctx.max_seqlen_k, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dkv = dkv[..., : dout.shape[-1]] - return dq, dkv, None, None, None, None, None, None, None, None, None, None, None, None, None - - -class FlashAttnFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - k, - v, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, k, v] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dk = dk[..., : dout.shape[-1]] - dv = dv[..., : dout.shape[-1]] - return dq, dk, dv, 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, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - block_table, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, k, v] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=block_table, - ) - if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state - ) - ctx.dropout_p = dropout_p - 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.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors - dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - ctx.max_seqlen_q, - ctx.max_seqlen_k, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dk = dk[..., : dout.shape[-1]] - dv = dv[..., : dout.shape[-1]] - return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None - - -def flash_attn_qkvpacked_func( - qkv, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # <=0.0 means deactivate - alibi_slopes=None, - deterministic=False, - 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, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_kvpacked_func( - q, - kv, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, -): - """dropout_p should be set to 0.0 during evaluation - If 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 K, V. - 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) - kv: (batch_size, seqlen, 2, 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. - softcap: float. Anything > 0 activates softcapping 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). - 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 FlashAttnKVPackedFunc.apply( - q, - kv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_func( - q, - k, - v, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - 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). - 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 FlashAttnFunc.apply( - q, - k, - v, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_qkvpacked_func( - qkv, - cu_seqlens, - max_seqlen, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - 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_varlen_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_varlen_kvpacked_func and flash_attn_varlen_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: (total, 3, nheads, headdim), where total = total number of tokens in the batch. - cu_seqlens: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into qkv. - max_seqlen: int. Maximum sequence length in the batch. - 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenQKVPackedFunc.apply( - qkv, - cu_seqlens, - max_seqlen, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_kvpacked_func( - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, -): - """dropout_p should be set to 0.0 during evaluation - If 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 K, V. - 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: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch. - kv: (total_k, 2, nheads_k, headdim), where total_k = total number of key tokens in the batch. - cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into q. - cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into kv. - max_seqlen_q: int. Maximum query sequence length in the batch. - max_seqlen_k: int. Maximum key sequence length in the batch. - 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 + 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenKVPackedFunc.apply( - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_func( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, - block_table=None, -): - """dropout_p should be set to 0.0 during evaluation - Supports multi-query and grouped-query attention (MQA/GQA) by passing in K, V 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: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch. - k: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. - v: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. - cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into q. - cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into kv. - max_seqlen_q: int. Maximum query sequence length in the batch. - max_seqlen_k: int. Maximum key sequence length in the batch. - 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 + 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenFunc.apply( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - block_table, - False if q.device.type == "cpu" else torch.is_grad_enabled(), - ) - - -def flash_attn_with_kvcache( - q, - k_cache, - v_cache, - k=None, - v=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, - block_table: Optional[torch.Tensor] = None, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - rotary_interleaved=True, - alibi_slopes=None, - num_splits=0, - 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 block_table, - or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_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) if there's no block_table, - or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_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). Similar to k. - 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. - block_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). - 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. - 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" - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-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) - cache_batch_idx = maybe_contiguous(cache_batch_idx) - block_table = maybe_contiguous(block_table) - out, softmax_lse = flash_attn.fwd_kvcache( - q, - k_cache, - v_cache, - k, - v, - cache_seqlens, - rotary_cos, - rotary_sin, - cache_batch_idx, - cache_leftpad, - block_table, - alibi_slopes, - None, - softmax_scale, - causal, - window_size[0], - window_size[1], - softcap, - rotary_interleaved, - num_splits, - ) - return (out, softmax_lse) if return_softmax_lse else out diff --git a/build/torch29-cxx11-cu126-aarch64-linux/layers/__init__.py b/build/torch29-cxx11-cu126-aarch64-linux/layers/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/build/torch29-cxx11-cu126-aarch64-linux/layers/patch_embed.py b/build/torch29-cxx11-cu126-aarch64-linux/layers/patch_embed.py deleted file mode 100644 index 05562f8e8bcdb58e947c6f402a49eacd2d031871..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-aarch64-linux/layers/patch_embed.py +++ /dev/null @@ -1,67 +0,0 @@ -# We use the same API as https://github.com/rwightman/pytorch-image-models/blob/v0.6.11/timm/models/layers/patch_embed.py -# But we use nn.Linear instead of Conv2d and it's about 8x faster. - -from functools import partial - -import torch.nn as nn -from einops import rearrange -from torch import _assert -from torch.nn.modules.utils import _pair - -try: - from flash_attn.ops.fused_dense import FusedDense -except ImportError: - FusedDense = None - - -class PatchEmbed(nn.Module): - """2D Image to Patch Embedding""" - - def __init__( - self, - img_size=224, - patch_size=16, - in_chans=3, - embed_dim=768, - norm_layer=None, - flatten=True, - bias=True, - fused_bias_fc=False, - ): - super().__init__() - img_size = _pair(img_size) - patch_size = _pair(patch_size) - self.img_size = img_size - self.patch_size = patch_size - self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1]) - self.num_patches = self.grid_size[0] * self.grid_size[1] - self.flatten = flatten - if fused_bias_fc and FusedDense is None: - raise ImportError("fused_dense is not installed") - - linear_cls = nn.Linear if not fused_bias_fc or not bias else FusedDense - self.proj = linear_cls(in_chans * patch_size[0] * patch_size[1], embed_dim, bias=bias) - self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() - - def forward(self, x): - _, _, H, W = x.shape - _assert( - H == self.img_size[0], - f"Input image height ({H}) doesn't match model ({self.img_size[0]}).", - ) - _assert( - W == self.img_size[1], - f"Input image width ({W}) doesn't match model ({self.img_size[1]}).", - ) - x = self.proj( - rearrange( - x, - "b c (h p1) (w p2) -> b h w (c p1 p2)", - p1=self.patch_size[0], - p2=self.patch_size[1], - ) - ) - if self.flatten: - x = rearrange(x, "b h w c -> b (h w) c") - x = self.norm(x) - return x diff --git a/build/torch29-cxx11-cu126-aarch64-linux/layers/rotary.py b/build/torch29-cxx11-cu126-aarch64-linux/layers/rotary.py deleted file mode 100644 index d1bfc21fc7de1dd287e8f382847b194a48075981..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-aarch64-linux/layers/rotary.py +++ /dev/null @@ -1,483 +0,0 @@ -# Copyright (c) 2025, Tri Dao - -import math -from functools import partial -from typing import Optional, Tuple, Union - -import torch -from torch import Tensor - -from einops import rearrange, repeat -# from flash_attn.ops.triton.rotary import apply_rotary -from ..ops.triton.rotary import apply_rotary - - -def rotate_half(x, interleaved=False): - if not interleaved: - x1, x2 = x.chunk(2, dim=-1) - return torch.cat((-x2, x1), dim=-1) - else: - x1, x2 = x[..., ::2], x[..., 1::2] - return rearrange(torch.stack((-x2, x1), dim=-1), "... d two -> ... (d two)", two=2) - - -def apply_rotary_emb_torch(x, cos, sin, interleaved=False): - """ - x: (batch_size, seqlen, nheads, headdim) - cos, sin: (seqlen, rotary_dim / 2) or (batch_size, seqlen, rotary_dim / 2) - """ - ro_dim = cos.shape[-1] * 2 - assert ro_dim <= x.shape[-1] - cos = repeat(cos, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)") - sin = repeat(sin, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)") - return torch.cat( - [x[..., :ro_dim] * cos + rotate_half(x[..., :ro_dim], interleaved) * sin, x[..., ro_dim:]], - dim=-1, - ) - - -class ApplyRotaryEmb(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x, - cos, - sin, - interleaved=False, - inplace=False, - seqlen_offsets: Union[int, Tensor] = 0, - cu_seqlens: Optional[Tensor] = None, - max_seqlen: Optional[int] = None, - ): - out = apply_rotary( - x, - cos, - sin, - seqlen_offsets=seqlen_offsets, - cu_seqlens=cu_seqlens, - max_seqlen=max_seqlen, - interleaved=interleaved, - inplace=inplace, - ) - if isinstance(seqlen_offsets, int): - ctx.save_for_backward(cos, sin, cu_seqlens) # Can't save int with save_for_backward - ctx.seqlen_offsets = seqlen_offsets - else: - ctx.save_for_backward(cos, sin, cu_seqlens, seqlen_offsets) - ctx.seqlen_offsets = None - ctx.interleaved = interleaved - ctx.inplace = inplace - ctx.max_seqlen = max_seqlen - return out if not inplace else x - - @staticmethod - def backward(ctx, do): - seqlen_offsets = ctx.seqlen_offsets - if seqlen_offsets is None: - cos, sin, cu_seqlens, seqlen_offsets = ctx.saved_tensors - else: - cos, sin, cu_seqlens = ctx.saved_tensors - dx = apply_rotary( - do, - cos, - sin, - seqlen_offsets=seqlen_offsets, - cu_seqlens=cu_seqlens, - max_seqlen=ctx.max_seqlen, - interleaved=ctx.interleaved, - inplace=ctx.inplace, - conjugate=True, - ) - return dx, None, None, None, None, None, None, None - - -def apply_rotary_emb( - x, - cos, - sin, - interleaved=False, - inplace=False, - seqlen_offsets: Union[int, Tensor] = 0, - cu_seqlens: Optional[Tensor] = None, - max_seqlen: Optional[int] = None, -): - """ - Arguments: - x: (batch_size, seqlen, nheads, headdim) if cu_seqlens is None - else (total_seqlen, nheads, headdim) - cos, sin: (seqlen_rotary, rotary_dim / 2) - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead - of 1st half and 2nd half (GPT-NeoX style). - inplace: if True, apply rotary embedding in-place. - seqlen_offsets: (batch_size,) or int. Each sequence in x is shifted by this amount. - Most commonly used in inference when we have KV cache. - cu_seqlens: (batch + 1,) or None - max_seqlen: int - Return: - out: (batch_size, seqlen, nheads, headdim) if cu_seqlens is None - else (total_seqlen, nheads, headdim) - rotary_dim must be <= headdim - Apply rotary embedding to the first rotary_dim of x. - """ - return ApplyRotaryEmb.apply( - x, cos, sin, interleaved, inplace, seqlen_offsets, cu_seqlens, max_seqlen - ) - - -# For backward compatibility -apply_rotary_emb_func = apply_rotary_emb - - -def _apply_rotary_emb_qkv( - qkv, - cos, - sin, - cos_k=None, - sin_k=None, - interleaved=False, - inplace=False, - conjugate=False, - seqlen_offsets: Union[int, Tensor] = 0, - num_heads_q: Optional[int] = None, -): - apply_rotary_fn = partial( - apply_rotary, - interleaved=interleaved, - inplace=inplace, - conjugate=conjugate, - seqlen_offsets=seqlen_offsets - ) - if cos_k is None and sin_k is None and qkv.is_contiguous(): - # Call 1 kernel instead of 2 kernels - # We need qkv to be contiguous so that when we reshape to combine (3, nheads) - # dimensions, we get the same tensor - if qkv.dim() == 5: - batch, seqlen, three, nheads, headdim = qkv.shape - assert three == 3 - # qk = rearrange(qkv[:, :, :2], "b s t h d -> b s (t h) d") - qk = qkv[:, :, :2].reshape(batch, seqlen, -1, headdim) - qk = apply_rotary_fn(qk, cos, sin) - else: - assert qkv.dim() == 4 - assert num_heads_q is not None - num_heads_k = (qkv.shape[2] - num_heads_q) // 2 - assert qkv.shape[2] == num_heads_q + 2 * num_heads_k - qk = qkv[:, :, :num_heads_q + num_heads_k] - qk = apply_rotary_fn(qk, cos, sin) - if not inplace: - if qkv.dim() == 5: - qkv = torch.cat([rearrange(qk, "b s (t h) d -> b s t h d", t=2), qkv[:, :, 2:]], dim=2) - else: - qkv = torch.cat([qk, qkv[:, :, num_heads_q + num_heads_k :]], dim=2) - else: - cos_k = cos if cos_k is None else cos_k - sin_k = sin if sin_k is None else sin_k - if qkv.dim() == 5: - batch, seqlen, three, nheads, headdim = qkv.shape - assert three == 3 - q, k = qkv[:, :, 0], qkv[:, :, 1] - else: - assert qkv.dim() == 4 - assert num_heads_q is not None - num_heads_k = (qkv.shape[2] - num_heads_q) // 2 - assert qkv.shape[2] == num_heads_q + 2 * num_heads_k - q, k = qkv[:, :, :num_heads_q], qkv[:, :, num_heads_q : num_heads_q + num_heads_k] - q = apply_rotary_fn(q, cos, sin) - k = apply_rotary_fn(k, cos_k, sin_k) - if not inplace: - if qkv.dim() == 5: - qkv = torch.stack([q, k, qkv[:, :, 2]], dim=2) - else: - qkv = torch.cat([q, k, qkv[:, :, num_heads_q + num_heads_k:]], dim=2) - return qkv - - -class ApplyRotaryEmbQKV_(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - cos, - sin, - cos_k=None, - sin_k=None, - interleaved=False, - seqlen_offsets: Union[int, torch.Tensor] = 0, - num_heads_q: Optional[int] = None, - ): - # apply_rotary_emb_qkv_inplace( - qkv = _apply_rotary_emb_qkv( - qkv, cos, sin, cos_k, sin_k, interleaved=interleaved, inplace=True, - seqlen_offsets=seqlen_offsets, num_heads_q=num_heads_q, - ) - if isinstance(seqlen_offsets, int): - ctx.save_for_backward(cos, sin, cos_k, sin_k) - ctx.seqlen_offsets = seqlen_offsets - else: - ctx.save_for_backward(cos, sin, cos_k, sin_k, seqlen_offsets) - ctx.seqlen_offsets = None - ctx.interleaved = interleaved - ctx.num_heads_q = num_heads_q - return qkv - - @staticmethod - def backward(ctx, dqkv): - seqlen_offsets = ctx.seqlen_offsets - if seqlen_offsets is None: - cos, sin, cos_k, sin_k, seqlen_offsets = ctx.saved_tensors - else: - cos, sin, cos_k, sin_k = ctx.saved_tensors - dqkv = _apply_rotary_emb_qkv( - dqkv, cos, sin, cos_k, sin_k, interleaved=ctx.interleaved, inplace=True, - seqlen_offsets=seqlen_offsets, num_heads_q=ctx.num_heads_q, conjugate=True, - ) - return dqkv, None, None, None, None, None, None, None - - -def apply_rotary_emb_qkv_( - qkv, - cos, - sin, - cos_k=None, - sin_k=None, - interleaved=False, - seqlen_offsets: Union[int, torch.Tensor] = 0, - num_heads_q: Optional[int] = None, -): - """ - Arguments: - qkv: (batch_size, seqlen, 3, nheads, headdim) or (batch_size, seqlen, num_heads_q + 2 * num_heads_k, headdim). - If qkv has shape (batch_size, seqlen, num_heads_q + 2 * num_heads_k, headdim) (e.g. MQA / GQA), - then num_heads_q must be provided. - cos, sin: (seqlen, rotary_dim / 2) - cos_k, sin_k: (seqlen, rotary_dim / 2), optional - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead of - 1st half and 2nd half (GPT-NeoX style). - seqlen_offsets: (batch_size,) or int. Each sequence in Q and K is shifted by this amount. - Most commonly used in inference when we have KV cache. - Return: - qkv: (batch_size, seqlen, 3, nheads, headdim) or (batch_size, seqlen, num_heads_q + 2 * num_heads_k, headdim) - rotary_dim must be <= headdim - Apply rotary embedding *inplace* to the first rotary_dim of Q and K. - """ - return ApplyRotaryEmbQKV_.apply( - qkv, cos, sin, cos_k, sin_k, interleaved, seqlen_offsets, num_heads_q - ) - - -class ApplyRotaryEmbKV_(torch.autograd.Function): - - @staticmethod - def forward(ctx, kv, cos, sin, interleaved=False, seqlen_offsets: Union[int, torch.Tensor] = 0): - batch, seqlen, two, nheads, headdim = kv.shape - assert two == 2 - k = kv[:, :, 0] - apply_rotary( - k, cos, sin, seqlen_offsets=seqlen_offsets, interleaved=interleaved, inplace=True - ) - if isinstance(seqlen_offsets, int): - ctx.save_for_backward(cos, sin) # Can't save int with save_for_backward - ctx.seqlen_offsets = seqlen_offsets - else: - ctx.save_for_backward(cos, sin, seqlen_offsets) - ctx.seqlen_offsets = None - ctx.interleaved = interleaved - return kv - - @staticmethod - def backward(ctx, dkv): - seqlen_offsets = ctx.seqlen_offsets - if seqlen_offsets is None: - cos, sin, seqlen_offsets = ctx.saved_tensors - else: - cos, sin = ctx.saved_tensors - apply_rotary( - dkv[:, :, 0], - cos, - sin, - seqlen_offsets=seqlen_offsets, - interleaved=ctx.interleaved, - inplace=True, - conjugate=True, - ) - return dkv, None, None, None, None - - -apply_rotary_emb_kv_ = ApplyRotaryEmbKV_.apply - - -def apply_rotary_emb_kv_( - kv, - cos, - sin, - interleaved=False, - seqlen_offsets: Union[int, torch.Tensor] = 0, -): - """ - Arguments: - kv: (batch_size, seqlen, 2, nheads, headdim) - cos, sin: (seqlen, rotary_dim / 2) - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead of - 1st half and 2nd half (GPT-NeoX style). - seqlen_offsets: (batch_size,) or int. Each sequence in Q and K is shifted by this amount. - Most commonly used in inference when we have KV cache. - Return: - kv: (batch_size, seqlen, 2, nheads, headdim) - rotary_dim must be <= headdim - Apply rotary embedding *inplace* to the first rotary_dim of K. - """ - return ApplyRotaryEmbKV_.apply(kv, cos, sin, interleaved, seqlen_offsets) - - -class RotaryEmbedding(torch.nn.Module): - """ - The rotary position embeddings from RoFormer_ (Su et. al). - A crucial insight from the method is that the query and keys are - transformed by rotation matrices which depend on the relative positions. - - Other implementations are available in the Rotary Transformer repo_ and in - GPT-NeoX_, GPT-NeoX was an inspiration - - .. _RoFormer: https://arxiv.org/abs/2104.09864 - .. _repo: https://github.com/ZhuiyiTechnology/roformer - .. _GPT-NeoX: https://github.com/EleutherAI/gpt-neox - - If scale_base is not None, this implements XPos (Sun et al., https://arxiv.org/abs/2212.10554). - A recommended value for scale_base is 512: https://github.com/HazyResearch/flash-attention/issues/96 - Reference: https://github.com/sunyt32/torchscale/blob/main/torchscale/component/xpos_relative_position.py - """ - - def __init__( - self, - dim: int, - base=10000.0, - interleaved=False, - scale_base=None, - device=None, - ): - """ - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead - of 1st half and 2nd half (GPT-NeoX style). - """ - super().__init__() - self.dim = dim - self.base = float(base) - # Generate and save the inverse frequency buffer (non trainable) - inv_freq = self._compute_inv_freq(device) - self.register_buffer("inv_freq", inv_freq, persistent=False) - self.interleaved = interleaved - self.scale_base = scale_base - scale = ( - (torch.arange(0, dim, 2, device=device, dtype=torch.float32) + 0.4 * dim) / (1.4 * dim) - if scale_base is not None - else None - ) - self.register_buffer("scale", scale, persistent=False) - - self._seq_len_cached = 0 - self._cos_cached = None - self._sin_cached = None - self._cos_k_cached = None - self._sin_k_cached = None - - def _compute_inv_freq(self, device=None): - return 1.0 / ( - self.base - ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim) - ) - - def _update_cos_sin_cache(self, seqlen, device=None, dtype=None): - # Reset the tables if the sequence length has changed, - # if we're on a new device (possibly due to tracing for instance), - # or if we're switching from inference mode to training - if ( - seqlen > self._seq_len_cached - or self._cos_cached is None - or self._cos_cached.device != device - or self._cos_cached.dtype != dtype - or (self.training and self._cos_cached.is_inference()) - ): - self._seq_len_cached = seqlen - # We want fp32 here, not self.inv_freq.dtype, since the model could be loaded in bf16 - # And the output of arange can be quite large, so bf16 would lose a lot of precision. - t = torch.arange(seqlen, device=device, dtype=torch.float32) - # We want fp32 here as well since inv_freq will be multiplied with t, and the output - # will be large. Having it in bf16 will lose a lot of precision and cause the - # cos & sin output to change significantly. - # We want to recompute self.inv_freq if it was not loaded in fp32 - if self.inv_freq.dtype != torch.float32: - inv_freq = self._compute_inv_freq(device=device) - else: - inv_freq = self.inv_freq - # Don't do einsum, it converts fp32 to bf16 under AMP - # freqs = torch.einsum("i,j->ij", t, self.inv_freq) - freqs = torch.outer(t, inv_freq) - if self.scale is None: - self._cos_cached = torch.cos(freqs).to(dtype) - self._sin_cached = torch.sin(freqs).to(dtype) - else: - power = ( - torch.arange(seqlen, dtype=self.scale.dtype, device=self.scale.device) - - seqlen // 2 - ) / self.scale_base - scale = self.scale.to(device=power.device) ** rearrange(power, "s -> s 1") - # We want the multiplication by scale to happen in fp32 - self._cos_cached = (torch.cos(freqs) * scale).to(dtype) - self._sin_cached = (torch.sin(freqs) * scale).to(dtype) - self._cos_k_cached = (torch.cos(freqs) / scale).to(dtype) - self._sin_k_cached = (torch.sin(freqs) / scale).to(dtype) - - def forward( - self, - qkv: torch.Tensor, - kv: Optional[torch.Tensor] = None, - seqlen_offset: Union[int, torch.Tensor] = 0, - max_seqlen: Optional[int] = None, - num_heads_q: Optional[int] = None, - ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: - """ - qkv: (batch, seqlen, 3, nheads, headdim) or (batch, seqlen, num_heads_q + 2 * num_heads_k, headdim) - if kv is none, else it's just q of shape (batch, seqlen, nheads, headdim). - If qkv has shape (batch, seqlen, num_heads_q + 2 * num_heads_k, headdim) (e.g. MQA / GQA), - then num_heads_q must be provided. - kv: (batch, seqlen, 2, nheads, headdim) - seqlen_offset: (batch_size,) or int. Each sequence in x is shifted by this amount. - Most commonly used in inference when we have KV cache. - If it's a tensor of shape (batch_size,), then to update the cos / sin cache, one - should pass in max_seqlen, which will update the cos / sin cache up to that length. - Apply rotary embedding *inplace* to qkv and / or kv. - """ - seqlen = qkv.shape[1] - if max_seqlen is not None: - self._update_cos_sin_cache(max_seqlen, device=qkv.device, dtype=qkv.dtype) - elif isinstance(seqlen_offset, int): - self._update_cos_sin_cache(seqlen + seqlen_offset, device=qkv.device, dtype=qkv.dtype) - if kv is None: - return apply_rotary_emb_qkv_( - qkv, - self._cos_cached, - self._sin_cached, - self._cos_k_cached if self.scale is not None else None, - self._sin_k_cached if self.scale is not None else None, - interleaved=self.interleaved, - seqlen_offsets=seqlen_offset, - num_heads_q=num_heads_q, - ) - else: - q = qkv - q = apply_rotary_emb_func( - q, - self._cos_cached, - self._sin_cached, - interleaved=self.interleaved, - inplace=True, - seqlen_offsets=seqlen_offset, - ) - kv = apply_rotary_emb_kv_( - kv, - self._cos_cached if self.scale is None else self._cos_k_cached, - self._sin_cached if self.scale is None else self._sin_k_cached, - interleaved=self.interleaved, - seqlen_offsets=seqlen_offset, - ) - return q, kv diff --git a/build/torch29-cxx11-cu126-aarch64-linux/metadata.json b/build/torch29-cxx11-cu126-aarch64-linux/metadata.json deleted file mode 100644 index 3514c62e2b286f3c5245cd885503179f49edf4f5..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-aarch64-linux/metadata.json +++ /dev/null @@ -1,12 +0,0 @@ -{ - "version": 1, - "license": "BSD-3-Clause", - "python-depends": [], - "backend": { - "type": "cuda", - "archs": [ - "8.0", - "9.0" - ] - } -} diff --git a/build/torch29-cxx11-cu126-aarch64-linux/ops/__init__.py b/build/torch29-cxx11-cu126-aarch64-linux/ops/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/build/torch29-cxx11-cu126-aarch64-linux/ops/activations.py b/build/torch29-cxx11-cu126-aarch64-linux/ops/activations.py deleted file mode 100644 index 7c09649fc41e12d5a360c5672825d8380bc7ec80..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-aarch64-linux/ops/activations.py +++ /dev/null @@ -1,135 +0,0 @@ -# Copied from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/model/layers/activations.py -import math - -import torch -import torch.nn as nn -import torch.nn.functional as F - -# 1/sqrt(2*pi)-> 0.3989423 -# 1/sqrt(2) -> 0.70710678 -# sqrt(2/pi) -> 0.79788456 - -# this function is tanh approximation of gelu -# actual gelu is: -# x * 0.5 * (1.0 + torch.erf(x * 0.70710678)) -@torch.jit.script -def bias_gelu(y, bias): - x = bias + y - return (x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))).to(dtype=y.dtype) - - -# gradient of tanh approximation of gelu -# gradient of actual gelu is: -# 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x) -@torch.jit.script -def bias_gelu_back(g, y, bias): - """Assume that y has shape (B, D) and bias has shape (D)""" - x = bias + y - tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)) - # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243 - ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( - 1 + tanh_out - ) - grad_y = ff * g - return grad_y.to(dtype=y.dtype), grad_y.sum(dim=(0), dtype=bias.dtype) - - -class GeLUFunction(torch.autograd.Function): - @staticmethod - # bias is an optional argument - def forward(ctx, input, bias): - ctx.save_for_backward(input, bias) - return bias_gelu(input, bias) - - @staticmethod - def backward(ctx, grad_output): - input, bias = ctx.saved_tensors - tmp = bias_gelu_back(grad_output, input, bias) - return tmp, tmp - - -bias_gelu_impl = GeLUFunction.apply - -# this function is tanh approximation of gelu -# actual gelu is: -# x * 0.5 * (1.0 + torch.erf(x * 0.70710678)) -@torch.jit.script -def gelu_fwd(x): - return (x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))).to(dtype=x.dtype) - - -# gradient of tanh approximation of gelu -# gradient of actual gelu is: -# 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x) -@torch.jit.script -def gelu_bwd(g, x): - tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)) - # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243 - ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( - 1 + tanh_out - ) - return (ff * g).to(dtype=x.dtype) - - -class FastGeLUFunction(torch.autograd.Function): - @staticmethod - # bias is an optional argument - def forward(ctx, input): - ctx.save_for_backward(input) - return gelu_fwd(input) - - @staticmethod - def backward(ctx, grad_output): - (input,) = ctx.saved_tensors - tmp = gelu_bwd(grad_output, input) - return tmp - - -fast_gelu_impl = FastGeLUFunction.apply - - -@torch.jit.script -def relu_bwd(g, x): - return torch.where(x >= 0, g, 0.0).to(dtype=x.dtype) - - -@torch.jit.script -def sqrelu_fwd(x): - r = F.relu(x) - return (r * r).to(dtype=x.dtype) - - -@torch.jit.script -def sqrelu_bwd(g, x): - return (2.0 * g * F.relu(x)).to(dtype=x.dtype) - - -swiglu_fwd_codestring = """ -template T swiglu_fwd(T x, T y) { - return float(x) * float(y) / (1.0f + ::exp(-float(x))); -} -""" -swiglu_bwd_codestring = """ -template void swiglu_bwd(T x, T y, T g, T& dx, T& dy) { - float x_sigmoid = 1.0f / (1.0f + ::exp(-float(x))); - dx = x_sigmoid * (1 + float(x) * (1.0f - x_sigmoid)) * float(g) * float(y); - dy = float(x) * x_sigmoid * float(g); -} -""" -swiglu_fwd = torch.cuda.jiterator._create_jit_fn(swiglu_fwd_codestring) -swiglu_bwd = torch.cuda.jiterator._create_multi_output_jit_fn(swiglu_bwd_codestring, num_outputs=2) - - -class SwiGLUFunction(torch.autograd.Function): - - @staticmethod - def forward(ctx, x, y): - ctx.save_for_backward(x, y) - return swiglu_fwd(x, y) - - @staticmethod - def backward(ctx, dout): - x, y = ctx.saved_tensors - return swiglu_bwd(x, y, dout) - -swiglu = SwiGLUFunction.apply diff --git a/build/torch29-cxx11-cu126-aarch64-linux/ops/fused_dense.py b/build/torch29-cxx11-cu126-aarch64-linux/ops/fused_dense.py deleted file mode 100644 index 6b4033d134e4093fe278f7b3f8c7d3128ce9f36d..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-aarch64-linux/ops/fused_dense.py +++ /dev/null @@ -1,688 +0,0 @@ -# Copyright (c) 2023, Tri Dao. -# Inspired by https://github.com/NVIDIA/apex/blob/master/apex/fused_dense/fused_dense.py -# We make it work with pytorch amp and with bfloat16. -# The TensorParallel linear modules are inspired by https://github.com/NVIDIA/apex/blob/master/apex/transformer/tensor_parallel/layers.py -from functools import partial -from typing import Optional - -# import fused_dense_cuda # from apex -import fused_dense_lib as fused_dense_cuda -import torch -import torch.nn as nn -import torch.nn.functional as F -from torch import Tensor -from torch.distributed import ProcessGroup - -from flash_attn.utils.torch import custom_fwd, custom_bwd -from flash_attn.ops.activations import gelu_bwd, relu_bwd, sqrelu_bwd, sqrelu_fwd -from flash_attn.utils.distributed import ( - all_gather_raw, - all_reduce, - all_reduce_raw, - reduce_scatter, - reduce_scatter_raw, -) - - -class FusedDenseFunc(torch.autograd.Function): - @staticmethod - @custom_fwd - def forward( - ctx, x, weight, bias, return_residual=False, process_group=None, sequence_parallel=True - ): - """ - If process_group is not None and sequence_parallel=True, we're doing Tensor Parallel - with sequence parallelism: we do an all_gather_raw of x before doing the matmul. - """ - ctx.compute_weight_gradient = weight.requires_grad - ctx.return_residual = return_residual - ctx.process_group = process_group - ctx.sequence_parallel = sequence_parallel - - if torch.is_autocast_enabled(): - x = x.to(dtype=torch.get_autocast_gpu_dtype()) - x = x.contiguous() - if process_group is not None and sequence_parallel: - # We want to kick off the all_gather early, before weight dtype conversion - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - else: - total_x = x - - if torch.is_autocast_enabled(): - weight = weight.to(dtype=torch.get_autocast_gpu_dtype()) - bias = bias.to(dtype=torch.get_autocast_gpu_dtype()) if bias is not None else None - weight = weight.contiguous() - if process_group is not None and sequence_parallel: - handle_x.wait() - batch_shape, n = total_x.shape[:-1], total_x.shape[-1] - batch_dim = batch_shape.numel() - # https://github.com/pytorch/pytorch/blob/5b51849b48a7dbccd297286cc0110def4706f9e7/aten/src/ATen/native/cuda/Blas.cpp#L174 - if min(batch_dim, n, *weight.shape) > 65535 * 32: - raise RuntimeError("fused_dense only supports matrix dims <= 2M") - output = F.linear(total_x, weight, bias) - if ctx.compute_weight_gradient: - ctx.save_for_backward(x, weight) - else: - ctx.save_for_backward(weight) - return output if not return_residual else (output, x) - - @staticmethod - @custom_bwd - def backward(ctx, grad_output, *args): - grad_output = grad_output.contiguous() - if ctx.return_residual: - (grad_input,) = args - grad_input = grad_input.contiguous() - process_group = ctx.process_group - sequence_parallel = ctx.sequence_parallel - if ctx.compute_weight_gradient: - x, weight = ctx.saved_tensors - if process_group is not None and sequence_parallel: - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - else: - total_x = x - else: - (weight,) = ctx.saved_tensors - total_x = None - batch_shape = grad_output.shape[:-1] - batch_dim = batch_shape.numel() - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - if ctx.needs_input_grad[0]: - if not ctx.return_residual: - grad_input = F.linear(grad_output, weight.t()) - else: - grad_input = torch.addmm( - grad_input.reshape(batch_dim, grad_input.shape[-1]), grad_output, weight - ) - grad_input = grad_input.reshape(*batch_shape, grad_input.shape[-1]) - if process_group is not None: - reduce_fn = reduce_scatter_raw if sequence_parallel else all_reduce_raw - grad_input, handle_grad_input = reduce_fn(grad_input, process_group, async_op=True) - else: - grad_input = None - if ctx.needs_input_grad[1]: - assert ctx.compute_weight_gradient - if process_group is not None and sequence_parallel: - handle_x.wait() - grad_weight, grad_bias = fused_dense_cuda.linear_bias_wgrad( - total_x.reshape(batch_dim, total_x.shape[-1]), grad_output, ctx.needs_input_grad[2] - ) - else: - grad_weight = None - grad_bias = grad_output if ctx.needs_input_grad[2] else None - if process_group is not None and ctx.needs_input_grad[0]: - handle_grad_input.wait() - return grad_input, grad_weight, grad_bias, None, None, None - - -def fused_dense_func( - x: Tensor, - weight: Tensor, - bias: Optional[Tensor] = None, - return_residual: bool = False, - process_group: Optional[ProcessGroup] = None, - sequence_parallel: bool = True, -): - dtype_eligible = x.dtype in [torch.float16, torch.bfloat16] or ( - x.dtype == torch.float32 and torch.is_autocast_enabled() - ) - if x.is_cuda and weight.is_cuda and (bias is None or bias.is_cuda) and dtype_eligible: - return FusedDenseFunc.apply( - x, weight, bias, return_residual, process_group, sequence_parallel - ) - else: - assert process_group is None - out = F.linear(x, weight, bias) - return out if not return_residual else (out, x) - - -class FusedDense(nn.Linear): - def __init__( - self, - in_features: int, - out_features: int, - bias: bool = True, - return_residual: bool = False, - device=None, - dtype=None, - ) -> None: - super().__init__(in_features, out_features, bias=bias, device=device, dtype=dtype) - self.return_residual = return_residual - - def forward(self, x, process_group=None): - """ - If process_group is not None, we're doing Tensor Parallel with sequence parallelism: - we do an all_gather of x before doing the matmul. - """ - return fused_dense_func( - x, - self.weight, - self.bias, - return_residual=self.return_residual, - process_group=process_group, - ) - - -class ColumnParallelLinear(nn.Linear): - def __init__( - self, - in_features: int, - out_features: int, - process_group: ProcessGroup, - bias: bool = True, - sequence_parallel=True, - multiple_of=1, - device=None, - dtype=None, - ) -> None: - world_size = torch.distributed.get_world_size(process_group) - if out_features % multiple_of: - raise ValueError(f"out_features ({out_features}) must be a multiple of {multiple_of}") - multiple = out_features // multiple_of - # We want to split @multiple across world_size, but it could be an uneven split - div = multiple // world_size - mod = multiple % world_size - # The first @mod ranks get @div + 1 copies, the rest get @div copies - local_multiple = div + int(torch.distributed.get_rank(process_group) < mod) - super().__init__( - in_features, local_multiple * multiple_of, bias=bias, device=device, dtype=dtype - ) - self.process_group = process_group - self.sequence_parallel = sequence_parallel - - def forward(self, x): - # If self.sequence_parallel is True, we're doing Tensor Parallel with sequence parallelism: - # we do an all_gather of x before doing the matmul. - # If not, then the input is already gathered. - return fused_dense_func( - x, - self.weight, - self.bias, - process_group=self.process_group, - sequence_parallel=self.sequence_parallel, - ) - - -class RowParallelLinear(nn.Linear): - def __init__( - self, - in_features: int, - out_features: int, - process_group: ProcessGroup, - bias: bool = True, - sequence_parallel=True, - multiple_of=1, - device=None, - dtype=None, - ) -> None: - world_size = torch.distributed.get_world_size(process_group) - rank = torch.distributed.get_rank(process_group) - if in_features % multiple_of: - raise ValueError(f"in_features ({in_features}) must be a multiple of {multiple_of}") - multiple = in_features // multiple_of - # We want to split @multiple across world_size, but it could be an uneven split - div = multiple // world_size - mod = multiple % world_size - # The first @mod ranks get @div + 1 copies, the rest get @div copies - local_multiple = div + int(torch.distributed.get_rank(process_group) < mod) - # Only rank 0 will have bias - super().__init__( - local_multiple * multiple_of, - out_features, - bias=bias and rank == 0, - device=device, - dtype=dtype, - ) - self.process_group = process_group - self.sequence_parallel = sequence_parallel - - def forward(self, x): - """ - We're doing Tensor Parallel with sequence parallelism: we do the matmul and then - a reduce_scatter of the result. - """ - out = fused_dense_func(x, self.weight, self.bias) - reduce_fn = reduce_scatter if self.sequence_parallel else all_reduce - return reduce_fn(out, self.process_group) - - -class FusedMLPFunc(torch.autograd.Function): - @staticmethod - @custom_fwd - def forward( - ctx, - x, - weight1, - bias1, - weight2, - bias2, - activation="gelu_approx", - save_pre_act=True, - return_residual=False, - checkpoint_lvl=0, - heuristic=0, - process_group=None, - sequence_parallel=True, - ): - """ - If process_group is not None and sequence_parallel=True, we're doing Tensor Parallel - with sequence parallelism: we do an all_gather of x before doing the matmul. - If sequence_parallel=False, then the input is already gathered. - - checkpoint_lvl: - 0: no recomputation in the bwd - 1: recompute gelu_out / relu_out in the bwd - 2: recompute pre_act and gelu_out / relu_out in the bwd - """ - assert -1 <= heuristic <= 4 - assert activation in ["gelu_approx", "relu", "sqrelu"] - if activation == "sqrelu": - assert heuristic == -1 - if not save_pre_act: - checkpoint_lvl = 2 - assert checkpoint_lvl in [0, 1, 2] - ctx.return_residual = return_residual - ctx.process_group = process_group - ctx.sequence_parallel = sequence_parallel - ctx.checkpoint_lvl = checkpoint_lvl - ctx.activation = activation - ctx.heuristic = heuristic - - if torch.is_autocast_enabled(): - x = x.to(dtype=torch.get_autocast_gpu_dtype()) - x = x.contiguous() - if process_group is not None and sequence_parallel: - # We want to kick off the all_gather early, before weight dtype conversion - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - else: - total_x = x - - if torch.is_autocast_enabled(): - dtype = torch.get_autocast_gpu_dtype() - weight1, weight2 = [a.to(dtype=dtype) for a in [weight1, weight2]] - bias1 = bias1.to(dtype=dtype) if bias1 is not None else None - bias2 = bias2.to(dtype=dtype) if bias2 is not None else None - weight1 = weight1.contiguous() - bias1 = bias1.contiguous() if bias1 is not None else None - weight2 = weight2.contiguous() - bias2 = bias2.contiguous() if bias2 is not None else None - if process_group is not None and sequence_parallel: - handle_x.wait() - batch_shape, n = total_x.shape[:-1], total_x.shape[-1] - batch_dim = batch_shape.numel() - # https://github.com/pytorch/pytorch/blob/5b51849b48a7dbccd297286cc0110def4706f9e7/aten/src/ATen/native/cuda/Blas.cpp#L174 - if min(batch_dim, n, *weight1.shape, *weight2.shape) > 65535 * 32: - raise RuntimeError("fused_dense only supports matrix dims <= 2M") - if heuristic == -1: - pre_act = F.linear(total_x, weight1, bias1) - activation_fn = ( - partial(F.gelu, approximate="tanh") - if activation == "gelu_approx" - else (sqrelu_fwd if activation == "sqrelu" else F.relu) - ) - with torch.jit.fuser("fuser2"): - output1 = activation_fn(pre_act) - # This is before adding bias1 - # pre_act = F.linear(total_x.reshape(batch_dim, n), weight1) - # with torch.jit.fuser('fuser2'): - # output1 = bias_gelu(pre_act, bias1) - else: - is_gelu = activation == "gelu_approx" - output1, *rest = fused_dense_cuda.linear_act_forward( - total_x.reshape(batch_dim, n), weight1, bias1, is_gelu, save_pre_act, heuristic - ) - if save_pre_act: - pre_act = rest[0] - output2 = F.linear(output1, weight2, bias2) - if checkpoint_lvl == 0 or (checkpoint_lvl == 1 and activation == "relu"): - # For RELU the pre_act is very small (just a bit-mask) so we just save it - ctx.save_for_backward(x, weight1, weight2, pre_act, output1) - elif checkpoint_lvl == 1: - ctx.save_for_backward(x, weight1, weight2, pre_act) - elif checkpoint_lvl == 2: - ctx.save_for_backward(x, weight1, weight2, bias1) - output2 = output2.reshape(*batch_shape, output2.shape[-1]) - return output2 if not return_residual else (output2, x) - - @staticmethod - @custom_bwd - def backward(ctx, grad_output, *args): - grad_output = grad_output.contiguous() - checkpoint_lvl = ctx.checkpoint_lvl - activation = ctx.activation - activation_fn = ( - partial(F.gelu, approximate="tanh") - if activation == "gelu_approx" - else (sqrelu_fwd if activation == "sqrelu" else F.relu) - ) - if ctx.return_residual: - (grad_input,) = args - grad_input = grad_input.contiguous() - process_group = ctx.process_group - sequence_parallel = ctx.sequence_parallel - x, weight1, weight2, *rest = ctx.saved_tensors - if process_group is None or not sequence_parallel: - total_x = x - batch_shape = grad_output.shape[:-1] - batch_dim = batch_shape.numel() - if checkpoint_lvl in [0, 1]: - if process_group is not None and sequence_parallel: - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - if checkpoint_lvl == 0 or (checkpoint_lvl == 1 and activation == "relu"): - pre_act, output1 = rest - elif checkpoint_lvl == 1: - (pre_act,) = rest - with torch.jit.fuser("fuser2"): - output1 = activation_fn(pre_act) - elif checkpoint_lvl == 2: - (bias1,) = rest - if process_group is not None and sequence_parallel: - total_x, _ = all_gather_raw(x, process_group) - if ctx.heuristic == -1: - pre_act = F.linear(total_x, weight1, bias1) - with torch.jit.fuser("fuser2"): - output1 = activation_fn(pre_act) - else: - output1, pre_act = fused_dense_cuda.linear_act_forward( - total_x.reshape(batch_dim, total_x.shape[-1]), - weight1, - bias1, - activation == "gelu_approx", - True, - ctx.heuristic, - ) - - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - output1 = output1.reshape(batch_dim, output1.shape[-1]) - pre_act = pre_act.reshape(batch_dim, pre_act.shape[-1]) - if ctx.needs_input_grad[3]: - grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad( - output1, grad_output, ctx.needs_input_grad[4] - ) - else: - grad_weight2 = None - grad_bias2 = grad_output if ctx.needs_input_grad[4] else None - if ctx.heuristic == -1: - # grad_pre_act = matmul_dgelu(grad_output, weight2, pre_act) - grad_output1 = F.linear(grad_output, weight2.t()) - activation_grad_fn = ( - gelu_bwd - if activation == "gelu_approx" - else (sqrelu_bwd if activation == "sqrelu" else relu_bwd) - ) - with torch.jit.fuser("fuser2"): - grad_pre_act = activation_grad_fn(grad_output1, pre_act) - else: - # The cublasLt epilogue has to compute both gelu/relu grad and bias grad, we can't - # just compute gelu/relu grad - grad_pre_act, grad_bias1 = fused_dense_cuda.bias_act_linear_dgrad_bgrad( - weight2, grad_output, pre_act, activation == "gelu_approx", ctx.heuristic - ) - if not ctx.needs_input_grad[2]: - grad_bias1 = None - if ctx.needs_input_grad[0]: - if not ctx.return_residual: - grad_input = F.linear(grad_pre_act, weight1.t()) - else: - grad_input = torch.addmm( - grad_input.reshape(batch_dim, grad_input.shape[-1]), grad_pre_act, weight1 - ) - grad_input = grad_input.reshape(*batch_shape, grad_input.shape[-1]) - if process_group is not None: - reduce_fn = reduce_scatter_raw if sequence_parallel else all_reduce_raw - grad_input, handle_grad_input = reduce_fn(grad_input, process_group, async_op=True) - else: - grad_input = None - if ctx.heuristic == -1: - if ctx.needs_input_grad[1]: - if process_group is not None and sequence_parallel and checkpoint_lvl != 2: - handle_x.wait() - grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_wgrad( - total_x.reshape(batch_dim, total_x.shape[-1]), - grad_pre_act, - ctx.needs_input_grad[2], - ) - else: - grad_weight1 = None - grad_bias1 = grad_pre_act if ctx.needs_input_grad[2] else None - else: - if ctx.needs_input_grad[1]: - if process_group is not None and sequence_parallel and checkpoint_lvl != 2: - handle_x.wait() - grad_weight1 = F.linear( - grad_pre_act.t(), total_x.reshape(batch_dim, total_x.shape[-1]).t() - ) - else: - grad_weight1 = None - if process_group is not None and ctx.needs_input_grad[0]: - handle_grad_input.wait() - return ( - grad_input, - grad_weight1, - grad_bias1, - grad_weight2, - grad_bias2, - None, - None, - None, - None, - None, - None, - None, - ) - - -def fused_mlp_func( - x: Tensor, - weight1: Tensor, - weight2: Tensor, - bias1: Optional[Tensor] = None, - bias2: Optional[Tensor] = None, - activation: str = "gelu_approx", - save_pre_act: bool = True, - return_residual: bool = False, - checkpoint_lvl: int = 0, - heuristic: int = 0, - process_group: Optional[ProcessGroup] = None, - sequence_parallel: bool = True, -): - assert activation in ["gelu_approx", "relu", "sqrelu"] - dtype_eligible = x.dtype in [torch.float16, torch.bfloat16] or ( - x.dtype == torch.float32 and torch.is_autocast_enabled() - ) - # If we save pre-activation, dimension must be divisible by 128 (relu) or 8 (gelu) - dim_eligible = not save_pre_act or (x.shape[-1] % (128 if activation == "relu" else 8) == 0) - if ( - x.is_cuda - and weight1.is_cuda - and weight2.is_cuda - and (bias1 is None or bias1.is_cuda) - and (bias2 is None or bias2.is_cuda) - and dtype_eligible - and dim_eligible - ): - return FusedMLPFunc.apply( - x, - weight1, - bias1, - weight2, - bias2, - activation, - save_pre_act, - return_residual, - checkpoint_lvl, - heuristic, - process_group, - sequence_parallel, - ) - else: - assert process_group is None - pre_act = F.linear(x, weight1, bias1) - activation_fn = ( - partial(F.gelu, approximate="tanh") - if activation == "gelu_approx" - else partial(F.relu, inplace=True) - ) - output1 = activation_fn(pre_act) - output2 = F.linear(output1, weight2, bias2) - return output2 if not return_residual else (output2, x) - - -class FusedMLP(nn.Module): - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - bias1=True, - bias2=True, - activation="gelu_approx", - return_residual=False, - checkpoint_lvl=0, - heuristic="auto", - device=None, - dtype=None, - ): - """ - If process_group is not None, we're doing Tensor Parallel with sequence parallelism: - we do an all_gather of x before doing the matmul, gelu, then matmul. - Finally we do a reduce_scatter of the output. - - checkpoint_lvl (increasing lvl means slower but more memory saving): - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute pre_act and gelu_out in the bwd - heuristic: - -1: don't fuse gemm + gelu (separate kernel) - 0..4: use this heuristic for the algo section in the fused gemm + gelu - 'auto': heuristic will be picked automatically: - For CUDA >= 11.8, we set heuristic=0 for both fp16 and bf16 for best perf. - For CUDA <= 11.7, we set heuristic=1 for fp16 and heuristic=-1 for bf16. - For H100, we set heuristic=-1 for both fp16 and bf16 as the fused cuBlasLt implementation - is slower than the unfused version. - return_residual: whether to return the input x along with the output. This is for - performance reason: for post-norm architecture, returning the input allows us - to fuse the backward of nn.Linear with the residual connection. - """ - assert checkpoint_lvl in [0, 1, 2] - assert activation in ["gelu_approx", "relu", "sqrelu"] - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features * 4 - self.activation = activation - self.return_residual = return_residual - self.checkpoint_lvl = checkpoint_lvl - self.heuristic = heuristic if activation != "sqrelu" else -1 - self.fc1 = nn.Linear(in_features, hidden_features, bias=bias1, **factory_kwargs) - self.fc2 = nn.Linear(hidden_features, out_features, bias=bias2, **factory_kwargs) - - def forward(self, x, process_group=None): - dtype = x.dtype if not torch.is_autocast_enabled() else torch.get_autocast_gpu_dtype() - if self.heuristic == "auto": - if self.activation == "gelu_approx": - if torch.cuda.get_device_capability("cuda") == (9, 0): - heuristic = -1 - else: - cuda_ver = tuple(map(int, torch.version.cuda.split("."))) - heuristic = 0 if cuda_ver >= (11, 8) else (1 if dtype == torch.float16 else -1) - else: - heuristic = 0 - else: - heuristic = self.heuristic - out = fused_mlp_func( - x, - self.fc1.weight, - self.fc2.weight, - self.fc1.bias, - self.fc2.bias, - activation=self.activation, - save_pre_act=self.training, - return_residual=self.return_residual, - checkpoint_lvl=self.checkpoint_lvl, - heuristic=heuristic, - process_group=process_group, - ) - if self.return_residual: - out, x = out - if process_group is not None: - out = reduce_scatter(out, process_group) - return out if not self.return_residual else (out, x) - - -class ParallelFusedMLP(nn.Module): - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - activation="gelu_approx", - process_group: ProcessGroup = None, - bias1=True, - bias2=True, - sequence_parallel=True, - checkpoint_lvl=0, - heuristic="auto", - device=None, - dtype=None, - ): - """ - process_group is required. We're doing Tensor Parallel with sequence parallelism: - we do an all_gather of x before doing the matmul, gelu, then matmul. - Finally we do a reduce_scatter of the output. - - checkpoint_lvl (increasing lvl means slower but more memory saving): - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute pre_act and gelu_out in the bwd - heuristic: - -1: don't fuse gemm + gelu (separate kernel) - 0..4: use this heuristic for the algo section in the fused gemm + gelu - 'auto': heuristic will be picked automatically: - For CUDA >= 11.8, we set heuristic=0 for both fp16 and bf16 for best perf. - For CUDA <= 11.7, we set heuristic=1 for fp16 and heuristic=-1 for bf16. - """ - assert checkpoint_lvl in [0, 1, 2] - assert activation in ["gelu_approx", "relu", "sqrelu"] - assert process_group is not None - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features * 4 - self.activation = activation - self.process_group = process_group - self.sequence_parallel = sequence_parallel - self.checkpoint_lvl = checkpoint_lvl - self.heuristic = heuristic if activation != "sqrelu" else -1 - self.fc1 = ColumnParallelLinear( - in_features, hidden_features, process_group, bias=bias1, **factory_kwargs - ) - self.fc2 = RowParallelLinear( - hidden_features, out_features, process_group, bias=bias2, **factory_kwargs - ) - - def forward(self, x): - dtype = x.dtype if not torch.is_autocast_enabled() else torch.get_autocast_gpu_dtype() - if self.heuristic == "auto": - if self.activation == "gelu_approx": - cuda_ver = tuple(map(int, torch.version.cuda.split("."))) - heuristic = 0 if cuda_ver >= (11, 8) else (1 if dtype == torch.float16 else -1) - else: - heuristic = 0 - else: - heuristic = self.heuristic - out = fused_mlp_func( - x, - self.fc1.weight, - self.fc2.weight, - self.fc1.bias, - self.fc2.bias, - activation=self.activation, - save_pre_act=self.training, - checkpoint_lvl=self.checkpoint_lvl, - heuristic=heuristic, - process_group=self.process_group, - sequence_parallel=self.sequence_parallel, - ) - reduce_fn = reduce_scatter if self.sequence_parallel else all_reduce - return reduce_fn(out, self.process_group) diff --git a/build/torch29-cxx11-cu126-aarch64-linux/ops/layer_norm.py b/build/torch29-cxx11-cu126-aarch64-linux/ops/layer_norm.py deleted file mode 100644 index 4b6cd798fd02844ef9cd3897f8ab95e490e638bf..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-aarch64-linux/ops/layer_norm.py +++ /dev/null @@ -1,800 +0,0 @@ -# Copyright (c) 2022, Tri Dao. -# Adapted from https://github.com/NVIDIA/apex/blob/master/apex/contrib/layer_norm/layer_norm.py - -import dropout_layer_norm -import torch -from torch.nn import init - - -def maybe_align(x, alignment_in_bytes=16): - """Assume that x already has last dim divisible by alignment_in_bytes""" - # TD [2023-07-04] I'm not 100% sure that clone will align the memory - # https://discuss.pytorch.org/t/how-to-ensure-that-tensor-data-ptr-is-aligned-to-16-bytes/183440 - return x if x.data_ptr() % alignment_in_bytes == 0 else x.clone() - - -def _dropout_add_layer_norm_forward( - x0, - residual, - gamma, - beta, - rowscale, - colscale, - dropout_p, - epsilon, - residual_in_fp32=False, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes""" - hidden_size = gamma.numel() - x0mat = x0.view((-1, hidden_size)) - residualmat = residual.view((-1, hidden_size)) if residual is not None else None - rowscale = rowscale.view(-1) if rowscale is not None else None - zmat, xmat, dmask, mu, rsigma = dropout_layer_norm.dropout_add_ln_fwd( - x0mat, - residualmat, - gamma, - beta, - rowscale, - colscale, - None, - None, - dropout_p, - epsilon, - 1.0, - 0, - None, - residual_in_fp32, - is_rms_norm, - ) - # dmask is None if dropout_p == 0.0 - # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype - return zmat, xmat if xmat is not None else x0mat, dmask, mu, rsigma - - -def _dropout_add_layer_norm_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - rowscale, - colscale, - dropout_p, - has_residual, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes - dx == None means that it was a post-norm architecture - (x = drop(x0) + residual was not returned in the fwd). - x0 must not be None if we have colscale. - """ - hidden_size = gamma.numel() - xmat = x.view((-1, hidden_size)) - dzmat = dz.view(xmat.shape) - dxmat = dx.view(xmat.shape) if dx is not None else None - x0mat = x0.view((-1, hidden_size)) if x0 is not None else None - rowscale = rowscale.view(-1) if rowscale is not None else None - if colscale is not None: - assert x0 is not None, "x0 is required to compute the gradient of colscale" - dx0mat, dresidualmat, dgamma, dbeta, _, _, *rest = dropout_layer_norm.dropout_add_ln_bwd( - dzmat, - dxmat, - xmat, - x0mat, - dmask, - mu, - rsigma, - gamma, - rowscale, - colscale, - None, - None, - dropout_p, - 1.0, - 0, - has_residual, - is_rms_norm, - ) - # dresidualmat is None if not has_residual - if colscale is None: - return dx0mat, dresidualmat, dgamma, dbeta - else: - dcolscale = rest[0] - return dx0mat, dresidualmat, dgamma, dbeta, dcolscale - - -def _dropout_add_layer_norm_subset_forward( - x0, - residual, - gamma, - beta, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32=False, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes""" - hidden_size = gamma.numel() - x0mat = x0.view((-1, hidden_size)) - residualmat = residual.view((-1, hidden_size)) if residual is not None else None - x0_subset = x0_subset.view(-1) if x0_subset is not None else None - out_subset = out_subset.view(-1) if out_subset is not None else None - zmat, xmat, dmask, mu, rsigma = dropout_layer_norm.dropout_add_ln_fwd( - x0mat, - residualmat, - gamma, - beta, - None, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - None, - residual_in_fp32, - is_rms_norm, - ) - # dmask is None if dropout_p == 0.0 - # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype - return zmat, xmat if xmat is not None else x0mat, dmask, mu, rsigma - - -def _dropout_add_layer_norm_subset_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - colscale, - x0_subset, - out_subset, - dropout_p, - rowscale_const, - x0_numrows, - has_residual, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes - dx == None means that it was a post-norm architecture - (x = drop(x0) + residual was not returned in the fwd). - x0 must not be None if we have colscale. - """ - hidden_size = gamma.numel() - xmat = x.view((-1, hidden_size)) - dzmat = dz.view(-1, hidden_size) - dxmat = dx.view(xmat.shape) if dx is not None else None - x0mat = x0.view((-1, hidden_size)) if x0 is not None else None - x0_subset = x0_subset.view(-1) if x0_subset is not None else None - out_subset = out_subset.view(-1) if out_subset is not None else None - if colscale is not None: - assert x0 is not None, "x0 is required to compute the gradient of colscale" - dx0mat, dresidualmat, dgamma, dbeta, _, _, *rest = dropout_layer_norm.dropout_add_ln_bwd( - dzmat, - dxmat, - xmat, - x0mat, - dmask, - mu, - rsigma, - gamma, - None, - colscale, - x0_subset, - out_subset, - dropout_p, - rowscale_const, - x0_numrows, - has_residual, - is_rms_norm, - ) - # dresidualmat is None if not has_residual - if colscale is None: - return dx0mat, dresidualmat, dgamma, dbeta - else: - dcolscale = rest[0] - return dx0mat, dresidualmat, dgamma, dbeta, dcolscale - - -def _dropout_add_layer_norm_parallel_residual_forward( - x0, - x1, - residual, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - residual_in_fp32=False, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes""" - hidden_size = gamma0.numel() - x0mat = x0.view((-1, hidden_size)) - x1mat = x1.view((-1, hidden_size)) if x1 is not None else None - residualmat = residual.view((-1, hidden_size)) if residual is not None else None - ( - z0mat, - z1mat, - xmat, - dmask0, - dmask1, - mu, - rsigma, - ) = dropout_layer_norm.dropout_add_ln_parallel_residual_fwd( - x0mat, - x1mat, - residualmat, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - None, - residual_in_fp32, - is_rms_norm, - ) - # dmask0 and dmask1 are None if dropout_p == 0.0 - # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype - return z0mat, z1mat, xmat if xmat is not None else x0mat, dmask0, dmask1, mu, rsigma - - -def _dropout_add_layer_norm_parallel_residual_backward( - dz0, - dz1, - dx, - x, - dmask0, - dmask1, - mu, - rsigma, - gamma0, - gamma1, - dropout_p, - has_x1, - has_residual, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes - dx == None means that it was a post-norm architecture - (x = drop(x0) + residual was not returned in the fwd). - """ - hidden_size = gamma0.numel() - xmat = x.view((-1, hidden_size)) - dz0mat = dz0.view(xmat.shape) - dz1mat = dz1.view(xmat.shape) if dz1 is not None else None - dxmat = dx.view(xmat.shape) if dx is not None else None - ( - dx0mat, - dx1mat, - dresidualmat, - dgamma0, - dbeta0, - dgamma1, - dbeta1, - *rest, - ) = dropout_layer_norm.dropout_add_ln_parallel_residual_bwd( - dz0mat, - dz1mat, - dxmat, - xmat, - dmask0, - dmask1, - mu, - rsigma, - gamma0, - gamma1, - dropout_p, - has_x1, - has_residual, - is_rms_norm, - ) - # dresidualmat is None if not has_residual - return dx0mat, dx1mat, dresidualmat, dgamma0, dbeta0, dgamma1, dbeta1 - - -class DropoutAddLayerNormFn(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x0, - residual, - gamma, - beta, - rowscale, - colscale, - dropout_p, - epsilon, - residual_in_fp32=False, - prenorm=False, - is_rms_norm=False, - return_dmask=False, - ): - x0 = maybe_align(x0.contiguous(), 16) - residual = maybe_align(residual.contiguous(), 16) if residual is not None else None - gamma = maybe_align(gamma.contiguous(), 16) - beta = maybe_align(beta.contiguous(), 16) if beta is not None else None - rowscale = maybe_align(rowscale.contiguous(), 16) if rowscale is not None else None - colscale = maybe_align(colscale.contiguous(), 16) if colscale is not None else None - zmat, xmat, dmask, mu, rsigma = _dropout_add_layer_norm_forward( - x0, - residual, - gamma, - beta, - rowscale, - colscale, - dropout_p, - epsilon, - residual_in_fp32, - is_rms_norm, - ) - # Only need to save x0 if we need to compute gradient wrt colscale - x0_saved = x0 if colscale is not None else None - ctx.save_for_backward( - xmat.view(x0.shape), x0_saved, dmask, gamma, mu, rsigma, rowscale, colscale - ) - ctx.prenorm = prenorm - ctx.dropout_p = dropout_p - ctx.has_residual = residual is not None - ctx.is_rms_norm = is_rms_norm - ctx.has_beta = beta is not None - if not return_dmask: - return ( - zmat.view(x0.shape) if not prenorm else (zmat.view(x0.shape), xmat.view(x0.shape)) - ) - else: - dmask = ( - dmask.view(x0.shape) - if dropout_p > 0.0 - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - ctx.mark_non_differentiable(dmask) - return ( - (zmat.view(x0.shape), dmask) - if not prenorm - else (zmat.view(x0.shape), xmat.view(x0.shape), dmask) - ) - - @staticmethod - def backward(ctx, dz, *args): - # assert dz.is_contiguous() - dz = maybe_align(dz.contiguous(), 16) # this happens! - dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None - x, x0, dmask, gamma, mu, rsigma, rowscale, colscale = ctx.saved_tensors - # x0 is None if colscale is None - dropout_p = ctx.dropout_p - has_residual = ctx.has_residual - dx0mat, dresidualmat, dgamma, dbeta, *rest = _dropout_add_layer_norm_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - rowscale, - colscale, - dropout_p, - has_residual, - ctx.is_rms_norm, - ) - dx0 = dx0mat.view(x.shape) - dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None - dcolscale = rest[0] if colscale is not None else None - return ( - dx0, - dresidual, - dgamma, - dbeta if ctx.has_beta else None, - None, - dcolscale, - None, - None, - None, - None, - None, - None, - ) - - -class DropoutAddLayerNormSubsetFn(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x0, - residual, - gamma, - beta, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32=False, - prenorm=False, - is_rms_norm=False, - return_dmask=False, - ): - x0 = maybe_align(x0.contiguous(), 16) - residual = maybe_align(residual.contiguous(), 16) if residual is not None else None - gamma = maybe_align(gamma.contiguous(), 16) - beta = maybe_align(beta.contiguous(), 16) if beta is not None else None - colscale = maybe_align(colscale.contiguous(), 16) if colscale is not None else None - zmat, xmat, dmask, mu, rsigma = _dropout_add_layer_norm_subset_forward( - x0, - residual, - gamma, - beta, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32, - is_rms_norm, - ) - # Only need to save x0 if we need to compute gradient wrt colscale - x0_saved = x0 if colscale is not None else None - x_shape = (-1, *x0.shape[1:]) - ctx.save_for_backward( - xmat.view(x_shape), x0_saved, dmask, gamma, mu, rsigma, colscale, x0_subset, out_subset - ) - ctx.prenorm = prenorm - ctx.dropout_p = dropout_p - ctx.rowscale_const = rowscale_const - ctx.x0_numrows = x0.shape[:-1].numel() - ctx.has_residual = residual is not None - ctx.is_rms_norm = is_rms_norm - ctx.has_beta = beta is not None - z_shape = (-1, *x0.shape[1:]) - if not return_dmask: - return zmat.view(z_shape) if not prenorm else (zmat.view(z_shape), xmat.view(x0.shape)) - else: - z = zmat.view(z_shape) - dmask = ( - dmask.view(x0.shape) - if dropout_p > 0.0 - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - ctx.mark_non_differentiable(dmask) - return (z, dmask) if not prenorm else (z, xmat.view(x_shape), dmask) - - @staticmethod - def backward(ctx, dz, *args): - # assert dz.is_contiguous() - dz = maybe_align(dz.contiguous(), 16) # this happens! - dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None - x, x0, dmask, gamma, mu, rsigma, colscale, x0_subset, out_subset = ctx.saved_tensors - # x0 is None if colscale is None - dropout_p = ctx.dropout_p - has_residual = ctx.has_residual - dx0mat, dresidualmat, dgamma, dbeta, *rest = _dropout_add_layer_norm_subset_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - colscale, - x0_subset, - out_subset, - dropout_p, - ctx.rowscale_const, - ctx.x0_numrows, - has_residual, - ctx.is_rms_norm, - ) - dx0 = dx0mat.view(-1, *x.shape[1:]) - dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None - dcolscale = rest[0] if colscale is not None else None - return ( - dx0, - dresidual, - dgamma, - dbeta if ctx.has_beta else None, - dcolscale, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) - - -class DropoutAddLayerNormParallelResidualFn(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x0, - x1, - residual, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - residual_in_fp32=False, - prenorm=False, - is_rms_norm=False, - return_dmask=False, - ): - x0 = maybe_align(x0.contiguous(), 16) - x1 = maybe_align(x1.contiguous(), 16) if x1 is not None else None - residual = maybe_align(residual.contiguous(), 16) if residual is not None else None - gamma0 = maybe_align(gamma0.contiguous(), 16) - beta0 = maybe_align(beta0.contiguous(), 16) if beta0 is not None else None - gamma1 = maybe_align(gamma1.contiguous(), 16) if gamma1 is not None else None - beta1 = maybe_align(beta1.contiguous(), 16) if beta1 is not None else None - ( - z0mat, - z1mat, - xmat, - dmask0, - dmask1, - mu, - rsigma, - ) = _dropout_add_layer_norm_parallel_residual_forward( - x0, - x1, - residual, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - residual_in_fp32, - is_rms_norm, - ) - ctx.save_for_backward(xmat.view(x0.shape), dmask0, dmask1, gamma0, gamma1, mu, rsigma) - ctx.prenorm = prenorm - ctx.dropout_p = dropout_p - ctx.has_x1 = x1 is not None - ctx.has_residual = residual is not None - ctx.is_rms_norm = is_rms_norm - ctx.has_beta = beta0 is not None - z = (z0mat.view(x0.shape), z1mat.view(x0.shape) if z1mat is not None else None) - if not return_dmask: - return z if not prenorm else (*z, xmat.view(x0.shape)) - else: - dmask0 = ( - dmask0.view(x0.shape) - if dropout_p > 0.0 - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - dmask1 = ( - dmask1.view(x0.shape) - if dropout_p > 0.0 and x1 is not None - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - ctx.mark_non_differentiable(dmask0) - ctx.mark_non_differentiable(dmask1) - return ( - (*z, dmask0, dmask1) if not prenorm else (*z, xmat.view(x0.shape), dmask0, dmask1) - ) - - @staticmethod - def backward(ctx, dz0, dz1, *args): - dz0 = maybe_align(dz0.contiguous(), 16) # this happens! - dz1 = maybe_align(dz1.contiguous(), 16) if dz1 is not None else None - dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None - x, dmask0, dmask1, gamma0, gamma1, mu, rsigma = ctx.saved_tensors - dropout_p = ctx.dropout_p - has_x1 = ctx.has_x1 - has_residual = ctx.has_residual - ( - dx0mat, - dx1mat, - dresidualmat, - dgamma0, - dbeta0, - dgamma1, - dbeta1, - ) = _dropout_add_layer_norm_parallel_residual_backward( - dz0, - dz1, - dx, - x, - dmask0, - dmask1, - mu, - rsigma, - gamma0, - gamma1, - dropout_p, - has_x1, - has_residual, - ctx.is_rms_norm, - ) - dx0 = dx0mat.view(x.shape) - dx1 = dx1mat.view(x.shape) if dx1mat is not None else None - dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None - return ( - dx0, - dx1, - dresidual, - dgamma0, - dbeta0 if ctx.has_beta else None, - dgamma1, - dbeta1 if ctx.has_beta else None, - None, - None, - None, - None, - None, - None, - ) - - -def layer_norm(x, weight, bias, epsilon): - return DropoutAddLayerNormFn.apply(x, None, weight, bias, None, None, 0.0, epsilon, False) - - -def dropout_add_layer_norm( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - rowscale=None, - layerscale=None, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormFn.apply( - x0, - residual, - weight, - bias, - rowscale, - layerscale, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - False, - return_dropout_mask, - ) - - -def dropout_add_layer_norm_subset( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - layerscale=None, - x0_subset=None, - out_subset=None, - rowscale_const=1.0, - out_numrows=0, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormSubsetFn.apply( - x0, - residual, - weight, - bias, - layerscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32, - prenorm, - False, - return_dropout_mask, - ) - - -def dropout_add_layer_norm_parallel_residual( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormParallelResidualFn.apply( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - False, - return_dropout_mask, - ) - - -class DropoutAddLayerNorm(torch.nn.Module): - def __init__( - self, - hidden_size, - prenorm=False, - p=0.0, - eps=1e-5, - residual_in_fp32=False, - device=None, - dtype=None, - ): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.prenorm = prenorm - self.p = p - self.eps = eps - self.residual_in_fp32 = residual_in_fp32 - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.bias = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.reset_parameters() - - def reset_parameters(self): - init.ones_(self.weight) - init.zeros_(self.bias) - - def forward(self, x0, residual=None): - return dropout_add_layer_norm( - x0, - residual, - self.weight, - self.bias, - self.p if self.training else 0.0, - self.eps, - prenorm=self.prenorm, - residual_in_fp32=self.residual_in_fp32, - ) diff --git a/build/torch29-cxx11-cu126-aarch64-linux/ops/rms_norm.py b/build/torch29-cxx11-cu126-aarch64-linux/ops/rms_norm.py deleted file mode 100644 index 068348d61290e3839dd082b540d898578ba1e8e2..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-aarch64-linux/ops/rms_norm.py +++ /dev/null @@ -1,174 +0,0 @@ -# Copyright (c) 2022, Tri Dao. -# Adapted from https://github.com/NVIDIA/apex/blob/master/apex/contrib/layer_norm/layer_norm.py - -import torch -from torch.nn import init - -from flash_attn.ops.layer_norm import ( - DropoutAddLayerNormFn, - DropoutAddLayerNormParallelResidualFn, - DropoutAddLayerNormSubsetFn, -) - - -def rms_norm(x, weight, epsilon): - return DropoutAddLayerNormFn.apply( - x, None, weight, None, None, None, 0.0, epsilon, False, False, True - ) - - -def dropout_add_rms_norm( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - rowscale=None, - layerscale=None, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormFn.apply( - x0, - residual, - weight, - bias, - rowscale, - layerscale, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - True, - return_dropout_mask, - ) - - -def dropout_add_rms_norm_subset( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - layerscale=None, - x0_subset=None, - out_subset=None, - rowscale_const=1.0, - out_numrows=0, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormSubsetFn.apply( - x0, - residual, - weight, - bias, - layerscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32, - prenorm, - True, - return_dropout_mask, - ) - - -def dropout_add_rms_norm_parallel_residual( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormParallelResidualFn.apply( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - True, - return_dropout_mask, - ) - - -class RMSNorm(torch.nn.Module): - def __init__(self, hidden_size, eps=1e-5, device=None, dtype=None): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.eps = eps - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.register_parameter("bias", None) - self.reset_parameters() - - def reset_parameters(self): - init.ones_(self.weight) - - def forward(self, x): - return rms_norm(x, self.weight, self.eps) - - -class DropoutAddRMSNorm(torch.nn.Module): - def __init__( - self, - hidden_size, - prenorm=False, - p=0.0, - eps=1e-5, - residual_in_fp32=False, - device=None, - dtype=None, - ): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.prenorm = prenorm - self.p = p - self.eps = eps - self.residual_in_fp32 = residual_in_fp32 - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.register_parameter("bias", None) - self.reset_parameters() - - def reset_parameters(self): - init.ones_(self.weight) - - def forward(self, x0, residual=None): - return dropout_add_rms_norm( - x0, - residual, - self.weight, - None, - self.p if self.training else 0.0, - self.eps, - prenorm=self.prenorm, - residual_in_fp32=self.residual_in_fp32, - ) diff --git a/build/torch29-cxx11-cu126-aarch64-linux/ops/triton/__init__.py b/build/torch29-cxx11-cu126-aarch64-linux/ops/triton/__init__.py deleted file mode 100644 index 8b137891791fe96927ad78e64b0aad7bded08bdc..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-aarch64-linux/ops/triton/__init__.py +++ /dev/null @@ -1 +0,0 @@ - diff --git a/build/torch29-cxx11-cu126-aarch64-linux/ops/triton/cross_entropy.py b/build/torch29-cxx11-cu126-aarch64-linux/ops/triton/cross_entropy.py deleted file mode 100644 index 1b5a415b73f236f3e05fb14b9141959559e18526..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-aarch64-linux/ops/triton/cross_entropy.py +++ /dev/null @@ -1,330 +0,0 @@ -# Copyright (c) 2023, Tri Dao. - -from typing import Tuple, Optional, Union - -import torch -import torch.nn.functional as F - -import triton -import triton.language as tl - -# `all_gather_into_tensor` and `reduce_scatter_tensor` are new placeholders for -# `_all_gather_base` and `_reduce_scatter_base`. They require the most recent -# version of PyTorch. The following 2 lines are for backward compatibility with -# older PyTorch. -if "all_gather_into_tensor" not in dir(torch.distributed): - torch.distributed.all_gather_into_tensor = torch.distributed._all_gather_base - - -@triton.heuristics( - { - "HAS_SMOOTHING": lambda args: args["smoothing"] > 0.0, - } -) -@triton.jit -def cross_entropy_fwd_kernel( - loss_ptr, # data ptrs - lse_ptr, - z_loss_ptr, - logits_ptr, - labels_ptr, - smoothing, - logit_scale, - lse_square_scale, - ignore_index, - total_classes, - class_start_idx, # Useful for tensor parallel when each rank only has a subset of classes - n_cols, # shapes - logits_row_stride, # strides - BLOCK_SIZE: tl.constexpr, - HAS_SMOOTHING: tl.constexpr, - # if SPLIT (e.g. tensor parallel), don't include the LSE in the loss since it's not the final LSE - SPLIT: tl.constexpr, - PRECOMPUTED_LSE: tl.constexpr, # If LSE is already computed (also no smoothing and logit_scale == 1.0) -): - row_idx = tl.program_id(0) - logits_ptr = logits_ptr + row_idx * logits_row_stride.to(tl.int64) - sum_logits = 0.0 # For smoothing - if not PRECOMPUTED_LSE: - # Statistics for online softmax - m_i = -float("inf") - l_i = 0.0 - for col_offset in range(0, n_cols, BLOCK_SIZE): - cols = col_offset + tl.arange(0, BLOCK_SIZE) - logits = tl.load(logits_ptr + cols, mask=cols < n_cols, other=-float("inf")).to( - tl.float32 - ) * logit_scale - if HAS_SMOOTHING: - sum_logits += tl.sum(tl.where(cols < n_cols, logits, 0.0)) - m_i_new = tl.maximum(m_i, tl.max(logits)) - l_i = tl.exp(m_i - m_i_new) * l_i + tl.sum(tl.exp(logits - m_i_new)) - m_i = m_i_new - lse = tl.log(l_i) + m_i - tl.store(lse_ptr + row_idx, lse) - else: - lse = tl.load(lse_ptr + row_idx) - label_idx = tl.load(labels_ptr + row_idx) - if label_idx == ignore_index: - loss = 0.0 - z_loss = 0.0 - else: - label_idx -= class_start_idx - if label_idx >= 0 and label_idx < n_cols: - logits_label = tl.load(logits_ptr + label_idx) * logit_scale - if HAS_SMOOTHING: - loss = ( - (lse if not SPLIT else 0.0) - - smoothing * sum_logits / total_classes - - (1 - smoothing) * logits_label - ) - else: - loss = (lse if not SPLIT else 0.0) - logits_label - else: - # If label is out of bounds, we set the CE loss to 0.0. But we still want the smoothing loss - if HAS_SMOOTHING: - loss = smoothing * ((lse if not SPLIT else 0.0) - sum_logits / total_classes) - else: - loss = 0.0 - if not SPLIT: - z_loss = lse_square_scale * lse * lse - loss += z_loss - else: - z_loss = 0.0 - tl.store(loss_ptr + row_idx, loss) - if not SPLIT: - tl.store(z_loss_ptr + row_idx, z_loss) - - -@triton.heuristics( - { - "HAS_SMOOTHING": lambda args: args["smoothing"] > 0.0, - } -) -@triton.jit -def cross_entropy_bwd_kernel( - dlogits_ptr, # data ptrs - dloss_ptr, - logits_ptr, - lse_ptr, - labels_ptr, - smoothing, - logit_scale, - lse_square_scale, - ignore_index, - total_classes, - class_start_idx, # Useful for tensor parallel when each rank only has a subset of classes - n_cols, # shapes - logits_row_stride, # strides - dlogits_row_stride, - dloss_row_stride, - BLOCK_SIZE: tl.constexpr, - HAS_SMOOTHING: tl.constexpr, -): - row_idx = tl.program_id(0) - col_block_idx = tl.program_id(1) - logits_ptr = logits_ptr + row_idx * logits_row_stride.to(tl.int64) - dlogits_ptr = dlogits_ptr + row_idx * dlogits_row_stride.to(tl.int64) - col_offsets = col_block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - label_idx = tl.load(labels_ptr + row_idx) - if label_idx != ignore_index: - dloss = tl.load(dloss_ptr + row_idx * dloss_row_stride) - else: - dloss = 0.0 - logits = tl.load(logits_ptr + col_offsets, mask=col_offsets < n_cols, other=-float("inf")).to( - tl.float32 - ) * logit_scale - lse = tl.load(lse_ptr + row_idx) - probs = tl.exp(logits - lse) - probs += 2.0 * lse_square_scale * lse * probs - label_idx -= class_start_idx - if HAS_SMOOTHING: - smooth_positive = 1.0 - smoothing - smooth_negative = smoothing / total_classes - probs = tl.where(col_offsets == label_idx, probs - smooth_positive, probs) - smooth_negative - else: - probs = tl.where(col_offsets == label_idx, probs - 1.0, probs) - tl.store(dlogits_ptr + col_offsets, (dloss * logit_scale) * probs, mask=col_offsets < n_cols) - - -class CrossEntropyLoss(torch.autograd.Function): - - @staticmethod - def forward( - ctx, - logits, - labels, - precomputed_lse=None, - smoothing=0.0, - logit_scale=1.0, - lse_square_scale=0.0, - ignore_index=-100, - inplace_backward=False, - process_group=None, - ): - # For some reason Triton generates wrong code when labels has dtype long and its address - # is not aligned to 16 bytes. The ld.global.b64 seems to load the wrong label index. - if labels.dtype == torch.long and labels.data_ptr() % 16 != 0: - labels = F.pad(labels, (0, 1))[..., :-1] - assert labels.data_ptr() % 16 == 0 - assert logit_scale > 0.0 - n_rows, n_cols = logits.shape - assert labels.shape == (n_rows,) - world_size = 1 if process_group is None else torch.distributed.get_world_size(process_group) - total_classes = world_size * n_cols - rank = 0 if process_group is None else torch.distributed.get_rank(process_group) - class_start_idx = rank * n_cols - use_precomputed_lse = precomputed_lse is not None and logit_scale == 1.0 and smoothing == 0.0 - - if logits.stride(-1) != 1: - logits = logits.contiguous() - MAX_BLOCK_SIZE = 16 * 1024 - BLOCK_SIZE = min(triton.next_power_of_2(n_cols), MAX_BLOCK_SIZE) - num_warps = ( - 4 - if BLOCK_SIZE < 2048 - else (8 if BLOCK_SIZE < 8192 else (16 if BLOCK_SIZE < 128 * 1024 else 32)) - ) - losses = torch.empty(n_rows, dtype=torch.float, device=logits.device) - if use_precomputed_lse: - assert precomputed_lse.shape == (n_rows,) - lse = precomputed_lse.contiguous() - else: - lse = torch.empty(n_rows, dtype=torch.float, device=logits.device) - z_losses = torch.empty(n_rows, dtype=torch.float, device=logits.device) - # Need this, otherwise Triton tries to launch from cuda:0 and we get - # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?) - with torch.cuda.device(logits.device.index): - cross_entropy_fwd_kernel[(n_rows,)]( - losses, # data ptrs - lse, - z_losses, - logits, - labels, - smoothing, - logit_scale, - lse_square_scale, - ignore_index, - total_classes, - class_start_idx, - n_cols, # shapes - logits.stride(0), # strides - BLOCK_SIZE=BLOCK_SIZE, # constants - SPLIT=world_size > 1, - PRECOMPUTED_LSE=use_precomputed_lse, - num_warps=num_warps, - ) - - if world_size > 1: - # If there's no smoothing, if labels are in the vocab of this partition, losses contains - # - predicted logit, and 0 otherwise. - # If there's smoothing=0.1, for labels in the vocab of this partition, losses contains - # -0.9 * predicted logit - 0.1 * sum logit / total_classes. - # For labels not in the vocab of this partition, losses contains - # -0.1 * sum logit / total_classes. - if world_size > 1: - lse_allgather = torch.empty(world_size, n_rows, dtype=lse.dtype, device=lse.device) - torch.distributed.all_gather_into_tensor(lse_allgather, lse, group=process_group) - handle_losses = torch.distributed.all_reduce( - losses, op=torch.distributed.ReduceOp.SUM, group=process_group, async_op=True - ) - lse = torch.logsumexp(lse_allgather, dim=0) - handle_losses.wait() - # After the allreduce, if there's no smoothing, the total losses are - predicted_logit, - # we just have to add the (global) lse. - # If there's smoothing=0.1, the total losses are - # -0.9 * predicted_logit - 0.1 * sum logit / total_classes. - # Again, we just have to add the (global) lse. - losses += lse - if lse_square_scale != 0.0: - z_losses = lse_square_scale * lse.square() - z_losses.masked_fill_(labels == ignore_index, 0.0) - losses += z_losses - else: - z_losses = torch.zeros_like(losses) - losses.masked_fill_(labels == ignore_index, 0.0) - - ctx.save_for_backward(logits, lse, labels) - ctx.mark_non_differentiable(z_losses) - ctx.smoothing = smoothing - ctx.logit_scale = logit_scale - ctx.lse_square_scale = lse_square_scale - ctx.ignore_index = ignore_index - ctx.total_classes = total_classes - ctx.class_start_idx = class_start_idx - ctx.inplace_backward = inplace_backward - return losses, z_losses - - @staticmethod - def backward(ctx, grad_losses, grad_z_losses): - del grad_z_losses # z_losses are only for logging. - - logits, lse, labels = ctx.saved_tensors - dlogits = logits if ctx.inplace_backward else torch.empty_like(logits) - n_rows, n_cols = logits.shape - BLOCK_SIZE = min(triton.next_power_of_2(n_cols), 4 * 1024) - num_warps = 4 if BLOCK_SIZE < 2048 else (8 if BLOCK_SIZE < 8192 else 16) - grid = lambda META: (n_rows, triton.cdiv(n_cols, META["BLOCK_SIZE"])) # noqa - # Need this, otherwise Triton tries to launch from cuda:0 and we get - # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?) - with torch.cuda.device(logits.device.index): - cross_entropy_bwd_kernel[grid]( - dlogits, # data ptrs - grad_losses, - logits, - lse, - labels, - ctx.smoothing, - ctx.logit_scale, - ctx.lse_square_scale, - ctx.ignore_index, - ctx.total_classes, - ctx.class_start_idx, - n_cols, # shapes - logits.stride(0), # strides - dlogits.stride(0), - grad_losses.stride(0), - BLOCK_SIZE=BLOCK_SIZE, # constants - num_warps=num_warps, - ) - return dlogits, None, None, None, None, None, None, None, None, None - - -def cross_entropy_loss( - logits: torch.Tensor, - labels: torch.Tensor, - precomputed_lse: Optional[torch.Tensor] = None, - label_smoothing: float = 0.0, - logit_scale: float = 1.0, - lse_square_scale: float = 0.0, - ignore_index=-100, - inplace_backward: bool = False, - process_group=None, -) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Arguments: - logits: (batch, vocab_size) - labels: (batch,) - label_smoothing: float - logit_scale: float. Multiply logits by this scale before calculating the loss. - lse_square_scale: float. If > 0, we add lse_square_scale * lse(logits) ^ 2 to the loss. - This is also referred to as "z-loss". - ignore_index: int. If labels == ignore_index, the loss is set to 0.0. - inplace_backward: bool. If True, we do the backward pass in-place by modifying the logits. - This saves memory. - process_group: if not None, we're doing Tensor Parallel: each process is responsible for - one part of the vocab. The loss will be aggregated across processes. - Returns: - losses: (batch,), float - z_losses: (batch,), float - """ - return CrossEntropyLoss.apply( - logits, - labels, - precomputed_lse, - label_smoothing, - logit_scale, - lse_square_scale, - ignore_index, - inplace_backward, - process_group, - ) diff --git a/build/torch29-cxx11-cu126-aarch64-linux/ops/triton/k_activations.py b/build/torch29-cxx11-cu126-aarch64-linux/ops/triton/k_activations.py deleted file mode 100644 index efb83c358eb4a85d069ee340a3c83f418f9a805b..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-aarch64-linux/ops/triton/k_activations.py +++ /dev/null @@ -1,162 +0,0 @@ -# Adapted from https://github.com/facebookresearch/xformers/blob/main/xformers/triton/k_activations.py -# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved. -# -# This source code is licensed under the BSD license found in the -# LICENSE file in the root directory of this source tree. - -import math -from enum import Enum -from typing import Optional - -import triton -import triton.language as tl - -_sqrt2pi = math.sqrt(2.0 / math.pi) -_sqrt1_2 = math.sqrt(1.0 / 2) -_gaussian_pdf_normalization = 1.0 / math.sqrt(2 * math.pi) - - -class Activation(str, Enum): - SquaredReLU = "squared_relu" - GeLU = "gelu" - GeLUApprox = "gelu_approx" - LeakyReLU = "leaky_relu" - ReLU = "relu" - - -def get_triton_activation_kernel(activation: Optional[Activation]): - return ( - { - Activation.ReLU: relu, - Activation.LeakyReLU: leaky_relu, - Activation.GeLU: gelu, - Activation.GeLUApprox: gelu_approx, - Activation.SquaredReLU: squared_relu, - }[activation] - if activation - else None - ) - - -def get_triton_activation_bwd_kernel(activation: Optional[Activation]): - return ( - { - Activation.ReLU: relu_grad, - Activation.LeakyReLU: leaky_relu_grad, - Activation.GeLU: gelu_grad, - Activation.GeLUApprox: gelu_approx_grad, - Activation.SquaredReLU: squared_relu_grad, - }[activation] - if activation - else None - ) - - -@triton.jit -def tanh(x): - # Tanh is just a scaled sigmoid - return 2 * tl.sigmoid(2 * x) - 1 - - -@triton.jit -def cosh(x): - exp_x = tl.exp(x) - return (exp_x + 1.0 / exp_x) * 0.5 - - -# a Triton implementation of the most used activations -# See for instance http://arxiv.org/abs/1606.08415 for an overview - -# ReLU -@triton.jit -def relu(x): - """ - ReLU_ activation function - - .. _ReLU: https://pytorch.org/docs/stable/generated/torch.nn.ReLU.html - """ - zero = 0.0 - return tl.where(x >= 0, x, zero.to(x.dtype)) - - -@triton.jit -def relu_grad(x): - # ReLU is different from other activations - # in that it does not require the input to retrospectively compute its gradient - # here the input is the downstream gradient, and we return the upstream gradient directly - zero = 0.0 - one = 1.0 - return tl.where(x >= 0, one.to(x.dtype), zero.to(x.dtype)) - - -@triton.jit -def squared_relu(x): - """ - Squared ReLU activation, as proposed in the Primer_ paper. - - .. _Primer: https://arxiv.org/abs/2109.08668 - """ - x_ = relu(x) - return (x_ * x_).to(x.dtype) - - -@triton.jit -def squared_relu_grad(x): - return tl.where(x >= 0, 2.0 * x, 0.0) - - -# Leaky ReLU -@triton.jit -def leaky_relu(x): - """ - LeakyReLU_ activation - - .. _LeakyReLU: https://pytorch.org/docs/stable/generated/torch.nn.LeakyReLU.html - """ - scale = 0.01 + 0.0 - scale = scale.to(x.dtype) - return tl.where(x >= 0, x, scale * x) - - -@triton.jit -def leaky_relu_grad(x): - min_grad = 0.01 - max_grad = 1 - - min_grad = min_grad.to(x.dtype) - max_grad = max_grad.to(x.dtype) - - return tl.where(x >= 0, max_grad, min_grad) - - -@triton.jit -def gelu(x): - """Gaussian Error Linear Unit (GELU)""" - return x * 0.5 * (1.0 + tl.libdevice.erf(x * _sqrt1_2)) - - -@triton.jit -def gelu_grad(x): - cdf = 0.5 * (1.0 + tl.libdevice.erf(x * _sqrt1_2)) - pdf = tl.exp(-0.5 * x * x) * _gaussian_pdf_normalization - return cdf + x * pdf - - -@triton.jit -def gelu_approx(x): - """ - GeLU_ activation - Gaussian error linear unit, with tanh approximation - - .. _GeLU: https://arxiv.org/pdf/1606.08415.pdf - """ - return 0.5 * x * (1.0 + tanh(_sqrt2pi * x * (1.0 + 0.044715 * x * x))) - - -@triton.jit -def gelu_approx_grad(x): - # CREDITS: Fast implementation proposed in - # https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/fused_bias_gelu.py#L30 - tanh_out = tanh(0.79788456 * x * (1 + 0.044715 * x * x)) - return 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( - 1 + tanh_out - ) diff --git a/build/torch29-cxx11-cu126-aarch64-linux/ops/triton/linear.py b/build/torch29-cxx11-cu126-aarch64-linux/ops/triton/linear.py deleted file mode 100644 index a8966dbc345ab0e593df0124451ee7be3dae131a..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-aarch64-linux/ops/triton/linear.py +++ /dev/null @@ -1,594 +0,0 @@ -# Adapted from https://github.com/ELS-RD/kernl/blob/main/src/kernl/implementations/linear_layer.py -# and https://github.com/openai/triton/blob/master/python/triton/ops/matmul.py -from typing import Optional - -import torch -import triton -import triton.language as tl -from triton.ops.matmul_perf_model import early_config_prune, estimate_matmul_time - -from flash_attn.ops.triton.k_activations import ( - gelu, - gelu_approx, - gelu_approx_grad, - gelu_grad, - squared_relu, - squared_relu_grad, -) - -# CREDITS: Initially inspired by the Triton tutorial on matrix multiplications - - -def init_to_zero(name): - return lambda nargs: nargs[name].zero_() - - -def get_configs_io_bound(): - configs = [] - for num_stages in [2, 3, 4, 5, 6]: - for block_m in [16, 32]: - for block_k in [32, 64]: - for block_n in [32, 64, 128, 256]: - num_warps = 2 if block_n <= 64 else 4 - configs.append( - triton.Config( - { - "BLOCK_M": block_m, - "BLOCK_N": block_n, - "BLOCK_K": block_k, - "SPLIT_K": 1, - }, - num_stages=num_stages, - num_warps=num_warps, - ) - ) - # split_k not used - # for split_k in [2, 4, 8, 16]: - # configs.append(triton.Config( - # {'BLOCK_M': block_m, 'BLOCK_N': block_n, 'BLOCK_K': block_k, 'SPLIT_K': split_k}, - # num_stages=num_stages, num_warps=num_warps, pre_hook=init_to_zero('C'))) - return configs - - -@triton.autotune( - configs=[ - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - # good for int8 - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=4, - num_warps=4, - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - ] - + get_configs_io_bound(), - key=["CACHE_KEY_M", "CACHE_KEY_N", "CACHE_KEY_K"], - prune_configs_by={ - "early_config_prune": early_config_prune, - "perf_model": estimate_matmul_time, - "top_k": 10, - }, -) -@triton.heuristics( - { - "EVEN_K": lambda args: args["K"] % (args["BLOCK_K"] * args["SPLIT_K"]) == 0, - } -) -@triton.jit -def kernel_fwd( - C, # Pointers to matrices - ACT_INPUT, - A, - B, - bias, - # Matrix dimensions - M, - N, - K, - CACHE_KEY_M, - CACHE_KEY_N, - CACHE_KEY_K, - # The stride variables represent how much to increase the ptr by when moving by 1 - # element in a particular dimension. E.g. stride_am is how much to increase a_ptr - # by to get the element one row down (A has M rows) - stride_cm, - # stride_cn, # Assume that stride_cn == 1 - stride_am, - stride_ak, - stride_bn, - stride_bk, - # Meta-parameters - BLOCK_M: tl.constexpr, - GROUP_M: tl.constexpr, - BLOCK_N: tl.constexpr, - BLOCK_K: tl.constexpr, - # split k not used, not performant with activation, kept because early_config_prune is expecting it - SPLIT_K: tl.constexpr, - EVEN_K: tl.constexpr, - A_ROWMAJOR: tl.constexpr, - B_COLMAJOR: tl.constexpr, - BIAS: tl.constexpr, - SAVE_ACT_INPUT: tl.constexpr, - ACTIVATION: tl.constexpr, -): - - """ - Kernel for computing Out = activation(A x W + C) - - Input has shape (M, K) - - Weight has shape (K, N) - - Bias has shape (N,) - - Output has shape (M, N) - - ActInputs (optional) has shape (M, N) - 'ActInputs' optionally saves the A x W + C intermediate for backward computations - This kernel will consolidate over K - """ - - pid = tl.program_id(axis=0) - - grid_m = (M + BLOCK_M - 1) // BLOCK_M - grid_n = (N + BLOCK_N - 1) // BLOCK_N - # re-order program ID for better L2 performance - width = GROUP_M * grid_n - group_id = pid // width - group_size = min(grid_m - group_id * GROUP_M, GROUP_M) - pid_m = group_id * GROUP_M + (pid % group_size) - pid_n = (pid % width) // (group_size) - - # now compute the block that each program will go through - # rm (resp. rn) denotes a range of indices - # for rows (resp. col) of C - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - # trick to avoid masking on M and N axis - ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M) - rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N) - rk = tl.arange(0, BLOCK_K) - - if A_ROWMAJOR: - A = A + (ram[:, None] * stride_am + rk[None, :]) - else: - A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak) - if B_COLMAJOR: - B = B + (rk[:, None] + rbn[None, :] * stride_bn) - else: - B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn) - - acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) - - for k in range(K, 0, -BLOCK_K): - if EVEN_K: - a = tl.load(A) - b = tl.load(B) - else: - a = tl.load(A, mask=rk[None, :] < k, other=0.0) - b = tl.load(B, mask=rk[:, None] < k, other=0.0) - acc += tl.dot(a, b) - - if A_ROWMAJOR: - A += BLOCK_K - else: - A += BLOCK_K * stride_ak - if B_COLMAJOR: - B += BLOCK_K - else: - B += BLOCK_K * stride_bk - - # Putting bias after the matmul (instead of before) is faster, idk why - if BIAS: - bias = tl.load(bias + rn, mask=rn < N, other=0.0).to(tl.float32) - acc += bias[None, :] - - # optional: save the activation inputs - if SAVE_ACT_INPUT: - # act_in_ptrs = ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :] * stride_cn - act_in_ptrs = ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :] - tl.store(act_in_ptrs, acc) - - # optional: fused activation (while the data is in shared memory) - if ACTIVATION == "gelu": - acc = gelu(acc) - elif ACTIVATION == "gelu_approx": - acc = gelu_approx(acc) - elif ACTIVATION == "squared_relu": - acc = squared_relu(acc) - # rematerialize rm and rn to save registers - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - - # write back result - # C = C + rm[:, None] * stride_cm + rn[None, :] * stride_cn - C = C + rm[:, None] * stride_cm + rn[None, :] - mask = (rm < M)[:, None] & (rn < N)[None, :] - tl.store(C, acc) - - -def triton_linear_act( - x: torch.Tensor, - weight: torch.Tensor, - bias: Optional[torch.Tensor] = None, - activation: str = "id", - save_act_input: bool = False, -) -> torch.Tensor: - """ - Compute e = activation(x @ weight.T + bias). - This wrapper kicks the `kernel_fwd` Triton kernel - :param x: input tensor - :param weight: weight matrix - :param bias: an optional bias tensor - :param activation: Activation name. Needs to be a Triton kernel. - :param act_input: an optional tensor to save the activation inputs (for backward) - :return: result tensor - """ - # if torch.is_autocast_enabled(): - # dtype = torch.get_autocast_gpu_dtype() - # x, weight, bias = [a.to(dtype=dtype) for a in [x, weight, bias]] - - assert activation in ["id", "gelu", "gelu_approx", "squared_relu"] - - batch_shape, n = x.shape[:-1], x.shape[-1] - batch_dim = batch_shape.numel() - x_reshaped = x.reshape(batch_dim, n) - - if x_reshaped.stride(0) > 1 and x_reshaped.stride(1) > 1: - x_reshaped = x_reshaped.contiguous() - if weight.stride(0) > 1 and weight.stride(1) > 1: - weight = weight.contiguous() - bias = bias.contiguous() if bias is not None else None - - assert ( - x.dtype == weight.dtype - ), f"Input and weight must have the same dtype, got {x.dtype} and {weight.dtype}" - if bias is not None: - assert ( - x.dtype == bias.dtype - ), f"Input and bias must have the same dtype, got {x.dtype} and {bias.dtype}" - assert ( - x_reshaped.shape[1] == weight.shape[1] - ), f"Incompatible dimensions: {x_reshaped.shape} - {weight.shape}" - - assert ( - bias is None or bias.shape[0] == weight.shape[0] - ), "Incompatible dimensions in between weight and bias" - - M, K = x_reshaped.shape - N, K = weight.shape - - output = torch.empty((M, N), device=x.device, dtype=x.dtype) - act_input = torch.empty_like(output) if save_act_input else None - - # 1D launch kernel where each block gets its own program. - grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),) # noqa - - kernel_fwd[grid]( - output, - act_input, - x_reshaped, - weight, # data ptrs - bias if bias is not None else x, # auto skip bias if not present - M, # shapes - N, - K, - M // 32, # key for triton cache (limit number of compilations) - N // 32, - K // 32, - stride_cm=output.stride(0), # strides - # stride_cn=output.stride(1), - stride_am=x_reshaped.stride(0), - stride_ak=x_reshaped.stride(1), - stride_bk=weight.stride(1), - stride_bn=weight.stride(0), - BIAS=bias is not None, # optional fused bias - SAVE_ACT_INPUT=save_act_input, # optional save activation inputs - ACTIVATION=activation, # optional fused activation - A_ROWMAJOR=x_reshaped.stride(1) == 1, - B_COLMAJOR=weight.stride(1) == 1, - GROUP_M=8, # speed optimization: group the programs - ) - - if not save_act_input: - return output.reshape(*batch_shape, output.shape[-1]) - else: - return ( - output.reshape(*batch_shape, output.shape[-1]), - act_input.reshape(*batch_shape, act_input.shape[-1]), - ) - - -@triton.autotune( - configs=[ - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - # good for int8 - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=4, - num_warps=4, - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - ] - + get_configs_io_bound(), - key=["CACHE_KEY_M", "CACHE_KEY_N", "CACHE_KEY_K"], - prune_configs_by={ - "early_config_prune": early_config_prune, - "perf_model": estimate_matmul_time, - "top_k": 10, - }, -) -@triton.heuristics( - { - "EVEN_K": lambda args: args["K"] % (args["BLOCK_K"] * args["SPLIT_K"]) == 0, - } -) -@triton.jit -def kernel_bwd( - C, # Pointers to matrices - ACT_INPUT, - A, - B, - # Matrix dimensions - M, - N, - K, - CACHE_KEY_M, - CACHE_KEY_N, - CACHE_KEY_K, - # The stride variables represent how much to increase the ptr by when moving by 1 - # element in a particular dimension. E.g. stride_am is how much to increase a_ptr - # by to get the element one row down (A has M rows) - stride_cm, - # stride_cn, # Assume that stride_cn == 1 - stride_am, - stride_ak, - stride_bk, - stride_bn, - # Meta-parameters - BLOCK_M: tl.constexpr, - GROUP_M: tl.constexpr, - BLOCK_N: tl.constexpr, - BLOCK_K: tl.constexpr, - # split k not used, not performant with activation, kept because early_config_prune is expecting it - SPLIT_K: tl.constexpr, - EVEN_K: tl.constexpr, - ACTIVATION: tl.constexpr, -): - - """ - Kernel for computing Out = activation(A x W + C) - - Input has shape (M, K) - - Weight has shape (K, N) - - Output has shape (M, N) - - ActInputs (optional) has shape (M, N) - 'ActInputs' optionally saves the A x W + C intermediate for backward computations - This kernel will consolidate over K - """ - - pid = tl.program_id(axis=0) - - grid_m = (M + BLOCK_M - 1) // BLOCK_M - grid_n = (N + BLOCK_N - 1) // BLOCK_N - # re-order program ID for better L2 performance - width = GROUP_M * grid_n - group_id = pid // width - group_size = min(grid_m - group_id * GROUP_M, GROUP_M) - pid_m = group_id * GROUP_M + (pid % group_size) - pid_n = (pid % width) // (group_size) - - # now compute the block that each program will go through - # rm (resp. rn) denotes a range of indices - # for rows (resp. col) of C - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - # trick to avoid masking on M and N axis - ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M) - rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N) - rk = tl.arange(0, BLOCK_K) - - A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak) - B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn) - - acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) - - for k in range(K, 0, -BLOCK_K): - if EVEN_K: - a = tl.load(A) - b = tl.load(B) - else: - a = tl.load(A, mask=rk[None, :] < k, other=0.0) - b = tl.load(B, mask=rk[:, None] < k, other=0.0) - acc += tl.dot(a, b) - - A += BLOCK_K * stride_ak - B += BLOCK_K * stride_bk - - # optional: fused activation (while the data is in shared memory) - if ACTIVATION != "id": - act_in_ptrs = ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :] - act_input = tl.load(act_in_ptrs).to(acc.dtype) - if ACTIVATION == "gelu": - acc *= gelu_grad(act_input) - elif ACTIVATION == "gelu_approx": - acc *= gelu_approx_grad(act_input) - elif ACTIVATION == "squared_relu": - acc *= squared_relu_grad(act_input) - - # rematerialize rm and rn to save registers - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - - # write back result - C = C + rm[:, None] * stride_cm + rn[None, :] - mask = (rm < M)[:, None] & (rn < N)[None, :] - tl.store(C, acc, mask=mask) - - -def triton_dgrad_act( - grad_output: torch.Tensor, - weight: torch.Tensor, - activation: str = "id", - act_input: Optional[torch.Tensor] = None, -) -> torch.Tensor: - """ - Compute e = activation(grad_output @ weight + bias). - This wrapper kicks the `kernel_fwd` Triton kernel - :param grad_output: input tensor - :param weight: weight matrix - :param activation: Activation name. Needs to be a Triton kernel. - :param act_input: an optional tensor to save the activation inputs (for backward) - :return: result tensor - """ - assert activation in ["id", "gelu", "gelu_approx", "squared_relu"] - - batch_shape, n = grad_output.shape[:-1], grad_output.shape[-1] - batch_dim = batch_shape.numel() - grad_output_reshaped = grad_output.reshape(batch_dim, n) - - if grad_output_reshaped.stride(0) > 1 and grad_output_reshaped.stride(1) > 1: - grad_output_reshaped = grad_output_reshaped.contiguous() - if weight.stride(0) > 1 and weight.stride(1) > 1: - weight = weight.contiguous() - - assert ( - grad_output.dtype == weight.dtype - ), f"grad_output and weight must have the same dtype, got {grad_output.dtype} and {weight.dtype}" - assert ( - grad_output_reshaped.shape[1] == weight.shape[0] - ), f"Incompatible dimensions: {grad_output_reshaped.shape} - {weight.shape}" - if activation != "id": - assert act_input is not None, f"act_input is required for activation {activation}" - - # M, N, K in bwd are different from M, N, K in fwd - M, K = grad_output_reshaped.shape - K, N = weight.shape - - grad_input = torch.empty((M, N), device=grad_output.device, dtype=grad_output.dtype) - - # 1D launch kernel where each block gets its own program. - grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),) # noqa - - kernel_bwd[grid]( - grad_input, - act_input, - grad_output_reshaped, - weight, # data ptrs - M, # shapes - N, - K, - M // 32, # key for triton cache (limit number of compilations) - N // 32, - K // 32, - stride_cm=grad_input.stride(0), # strides - # stride_cn=grad_input.stride(1), - stride_am=grad_output_reshaped.stride(0), - stride_ak=grad_output_reshaped.stride(1), - stride_bk=weight.stride(0), - stride_bn=weight.stride(1), - ACTIVATION=activation, # optional fused activation - GROUP_M=8, # speed optimization: group the programs - ) - - return grad_input.reshape(*batch_shape, grad_input.shape[-1]) diff --git a/build/torch29-cxx11-cu126-aarch64-linux/ops/triton/mlp.py b/build/torch29-cxx11-cu126-aarch64-linux/ops/triton/mlp.py deleted file mode 100644 index 059f4f8a5e174c1f4824e43d313fca18eaa799b8..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-aarch64-linux/ops/triton/mlp.py +++ /dev/null @@ -1,149 +0,0 @@ -# The triton fused matmul + sqrelu is faster for fp16 but slower for bf16, compared -# to naive implementation. -import fused_dense_lib as fused_dense_cuda -import torch -import torch.nn as nn -import torch.nn.functional as F - -from flash_attn.utils.torch import custom_fwd, custom_bwd -from flash_attn.ops.activations import sqrelu_bwd, sqrelu_fwd -from flash_attn.ops.triton.linear import triton_dgrad_act, triton_linear_act - - -class FusedDenseSqreluDenseFunc(torch.autograd.Function): - @staticmethod - @custom_fwd - def forward(ctx, x, weight1, bias1, weight2, bias2, checkpoint_lvl=0): - """checkpoint_lvl: - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute act_input and gelu_out in the bwd - """ - if torch.is_autocast_enabled(): - dtype = torch.get_autocast_gpu_dtype() - x, weight1, bias1, weight2, bias2 = [ - a.to(dtype=dtype) for a in [x, weight1, bias1, weight2, bias2] - ] - is_bf16 = x.dtype == torch.bfloat16 - assert checkpoint_lvl in [0, 1, 2] - x = x.contiguous() - weight1 = weight1.contiguous() - bias1 = bias1.contiguous() - weight2 = weight2.contiguous() - bias2 = bias2.contiguous() - batch_shape, n = x.shape[:-1], x.shape[-1] - batch_dim = batch_shape.numel() - if is_bf16: - act_input = fused_dense_cuda.linear_bias_forward( - x.reshape(batch_dim, n), weight1, bias1 - ) - output1 = sqrelu_fwd(act_input) - else: - save_act_input = checkpoint_lvl != 2 - result = triton_linear_act( - x.reshape(batch_dim, n), - weight1, - bias1, - activation="squared_relu", - save_act_input=save_act_input, - ) - if save_act_input: - output1, act_input = result - else: - output1 = result - output2 = fused_dense_cuda.linear_bias_forward(output1, weight2, bias2) - ctx.checkpoint_lvl = checkpoint_lvl - if checkpoint_lvl == 0: - ctx.save_for_backward(x, weight1, bias1, weight2, act_input, output1) - elif checkpoint_lvl == 1: - ctx.save_for_backward(x, weight1, bias1, weight2, act_input) - elif checkpoint_lvl == 2: - ctx.save_for_backward(x, weight1, bias1, weight2) - return output2.reshape(*batch_shape, output2.shape[-1]) - - @staticmethod - @custom_bwd - def backward(ctx, grad_output): - grad_output = grad_output.contiguous() - checkpoint_lvl = ctx.checkpoint_lvl - x, weight1, bias1, weight2, *rest = ctx.saved_tensors - batch_shape, n = x.shape[:-1], x.shape[-1] - batch_dim = batch_shape.numel() - is_bf16 = x.dtype == torch.bfloat16 - if checkpoint_lvl == 0: - act_input, output1 = rest - elif checkpoint_lvl == 1: - (act_input,) = rest - output1 = sqrelu_fwd(act_input) - elif checkpoint_lvl == 2: - if is_bf16: - act_input = fused_dense_cuda.linear_bias_forward( - x.reshape(batch_dim, n), weight1, bias1 - ) - output1 = sqrelu_fwd(act_input) - else: - output1, act_input = triton_linear_act( - x.reshape(batch_dim, n), - weight1, - bias1, - activation="squared_relu", - save_act_input=True, - ) - - if is_bf16: - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad(output1, grad_output) - grad_output1 = grad_output @ weight2 - grad_act_input = sqrelu_bwd(grad_output1, act_input) - grad_input, grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_backward( - x.reshape(batch_dim, n), weight1, grad_act_input - ) - else: - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad(output1, grad_output) - grad_act_input = triton_dgrad_act( - grad_output, weight2, activation="squared_relu", act_input=act_input - ) - grad_input, grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_backward( - x.reshape(batch_dim, n), weight1, grad_act_input - ) - return grad_input.reshape_as(x), grad_weight1, grad_bias1, grad_weight2, grad_bias2, None - - -fused_dense_sqrelu_dense_function = FusedDenseSqreluDenseFunc.apply - - -class FusedDenseSqreluDense(nn.Module): - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - bias1=True, - bias2=True, - checkpoint_lvl=0, - device=None, - dtype=None, - ): - """ - checkpoint_lvl (increasing lvl means slower but more memory saving): - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute gelu_in and gelu_out in the bwd - """ - assert checkpoint_lvl in [0, 1, 2] - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features * 4 - assert bias1 == True, "DenseSqreluDense module without bias is currently not supported" - assert bias2 == True, "DenseSqreluDense module without bias is currently not supported" - self.checkpoint_lvl = checkpoint_lvl - self.fc1 = nn.Linear(in_features, hidden_features, bias=bias1, **factory_kwargs) - self.fc2 = nn.Linear(hidden_features, out_features, bias=bias2, **factory_kwargs) - - def forward(self, x): - assert x.is_cuda - return fused_dense_sqrelu_dense_function( - x, self.fc1.weight, self.fc1.bias, self.fc2.weight, self.fc2.bias, self.checkpoint_lvl - ) diff --git a/build/torch29-cxx11-cu126-aarch64-linux/ops/triton/rotary.py b/build/torch29-cxx11-cu126-aarch64-linux/ops/triton/rotary.py deleted file mode 100644 index 6ad3ab412a084270120167ad32b67ad1720f2c37..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-aarch64-linux/ops/triton/rotary.py +++ /dev/null @@ -1,186 +0,0 @@ -# Copyright (c) 2025, Tri Dao. -# As of 2025-04-23, we require triton >= 3.0 - -from typing import Optional, Union - -import torch - -import triton -import triton.language as tl - - -@triton.jit -def rotary_kernel( - OUT, # Pointers to matrices - X, - COS, - SIN, - CU_SEQLENS, - SEQLEN_OFFSETS, # this could be int or a pointer - # Matrix dimensions - seqlen, - nheads, - seqlen_ro, - # strides - stride_out_batch, - stride_out_seqlen, - stride_out_nheads, - stride_out_headdim, - stride_x_batch, - stride_x_seqlen, - stride_x_nheads, - stride_x_headdim, - # Meta-parameters - # We want ROTARY_DIM to be constexpr, otherwise the triton compiler doesn't know that - # the mask is constant every 8 elements, and it will generate LDG.16 instead of LDG.128 - ROTARY_DIM: tl.constexpr, - IS_SEQLEN_OFFSETS_TENSOR: tl.constexpr, - IS_VARLEN: tl.constexpr, - INTERLEAVED: tl.constexpr, - CONJUGATE: tl.constexpr, - BLOCK_H: tl.constexpr, - BLOCK_M: tl.constexpr, -): - BLOCK_K: tl.constexpr = triton.next_power_of_2(ROTARY_DIM) - ROTARY_DIM_HALF = ROTARY_DIM // 2 - pid_head = tl.program_id(axis=0) - pid_m = tl.program_id(axis=1) - pid_batch = tl.program_id(axis=2) - - if not IS_VARLEN: - X = X + pid_batch * stride_x_batch - OUT = OUT + pid_batch * stride_out_batch - else: - start_idx = tl.load(CU_SEQLENS + pid_batch) - seqlen = tl.load(CU_SEQLENS + pid_batch + 1) - start_idx - X = X + start_idx * stride_x_seqlen - OUT = OUT + start_idx * stride_out_seqlen - - if pid_m * BLOCK_M >= seqlen: - return - - rh = pid_head * BLOCK_H + tl.arange(0, BLOCK_H) - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - if not IS_SEQLEN_OFFSETS_TENSOR: - rm_cs = rm + SEQLEN_OFFSETS - else: - rm_cs = rm + tl.load(SEQLEN_OFFSETS + pid_batch) - - rk_half = tl.arange(0, BLOCK_K // 2) - COS = COS + (rm_cs[:, None] * ROTARY_DIM_HALF + rk_half[None, :]) - SIN = SIN + (rm_cs[:, None] * ROTARY_DIM_HALF + rk_half[None, :]) - mask_cs = (rm_cs[:, None] < seqlen_ro) & (rk_half[None, :] < ROTARY_DIM_HALF) - cos = tl.load(COS, mask=mask_cs, other=1.0).to(tl.float32) - sin = tl.load(SIN, mask=mask_cs, other=0.0).to(tl.float32) - if CONJUGATE: - sin = -sin - - if not INTERLEAVED: - # Load the 1st and 2nd halves of X, do calculation, then store to 1st and 2nd halves of OUT - X = X + (rh[:, None, None] * stride_x_nheads + rm[None, :, None] * stride_x_seqlen + rk_half[None, None, :] * stride_x_headdim) - OUT = OUT + (rh[:, None, None] * stride_out_nheads + rm[None, :, None] * stride_out_seqlen + rk_half[None, None, :] * stride_out_headdim) - mask = (rh[:, None, None] < nheads) & (rm[None, :, None] < seqlen) & (rk_half[None, None, :] < ROTARY_DIM_HALF) - x0 = tl.load(X, mask=mask, other=0.0).to(tl.float32) - x1 = tl.load(X + ROTARY_DIM_HALF * stride_x_headdim, mask=mask, other=0.0,).to(tl.float32) - o0 = x0 * cos - x1 * sin - o1 = x0 * sin + x1 * cos - tl.store(OUT, o0, mask=mask) - tl.store(OUT + ROTARY_DIM_HALF * stride_out_headdim, o1, mask=mask) - else: - rk = tl.arange(0, BLOCK_K) - X = X + (rh[:, None, None] * stride_x_nheads + rm[None, :, None] * stride_x_seqlen + rk[None, None, :] * stride_x_headdim) - OUT = OUT + (rh[:, None, None] * stride_out_nheads + rm[None, :, None] * stride_out_seqlen + rk[None, None, :] * stride_out_headdim) - mask = (rh[:, None, None] < nheads) & (rm[None, :, None] < seqlen) & (rk[None, None, :] < ROTARY_DIM) - x = tl.load(X, mask=mask, other=0.0).to(tl.float32) - x0, x1 = tl.split(tl.reshape(x, [BLOCK_H, BLOCK_M, BLOCK_K // 2, 2])) - o0 = x0 * cos - x1 * sin - o1 = x0 * sin + x1 * cos - o = tl.reshape(tl.join(o0, o1), [BLOCK_H, BLOCK_M, BLOCK_K]) - tl.store(OUT, o, mask=mask) - - -def apply_rotary( - x: torch.Tensor, - cos: torch.Tensor, - sin: torch.Tensor, - seqlen_offsets: Union[int, torch.Tensor] = 0, - cu_seqlens: Optional[torch.Tensor] = None, - max_seqlen: Optional[int] = None, - interleaved=False, - inplace=False, - conjugate=False, -) -> torch.Tensor: - """ - Arguments: - x: (batch, seqlen, nheads, headdim) if cu_seqlens is None - else (total_seqlen, nheads, headdim). - cos: (seqlen_ro, rotary_dim / 2) - sin: (seqlen_ro, rotary_dim / 2) - seqlen_offsets: integer or integer tensor of size (batch,) - cu_seqlens: (batch + 1,) or None - max_seqlen: int - Returns: - y: (batch, seqlen, nheads, headdim) - """ - is_varlen = cu_seqlens is not None - if not is_varlen: - batch, seqlen, nheads, headdim = x.shape - else: - assert max_seqlen is not None, "If cu_seqlens is passed in, then max_seqlen must be passed" - total_seqlen, nheads, headdim = x.shape - batch_p_1 = cu_seqlens.shape[0] - batch = batch_p_1 - 1 - seqlen = max_seqlen - seqlen_ro, rotary_dim = cos.shape - assert sin.shape == cos.shape - rotary_dim *= 2 - assert rotary_dim <= headdim, "rotary_dim must be <= headdim" - assert headdim <= 256, "Only support headdim <= 256" - assert seqlen_ro >= seqlen, "seqlen_ro must be >= seqlen" - - cos, sin = cos.contiguous(), sin.contiguous() - if isinstance(seqlen_offsets, torch.Tensor): - assert seqlen_offsets.shape == (batch,) - assert seqlen_offsets.dtype in [torch.int32, torch.int64] - seqlen_offsets = seqlen_offsets.contiguous() - else: - assert seqlen_offsets + seqlen <= seqlen_ro - - output = torch.empty_like(x) if not inplace else x - if rotary_dim < headdim and not inplace: - output[..., rotary_dim:].copy_(x[..., rotary_dim:]) - - grid = lambda META: (triton.cdiv(nheads, META["BLOCK_H"]), triton.cdiv(seqlen, META["BLOCK_M"]), batch) # noqa - BLOCK_M = 8 if rotary_dim <= 128 else 4 - - # Need this, otherwise Triton tries to launch from cuda:0 and we get - # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?) - device_ctx = torch.cuda.device(x.device.index) if x.device.type == 'cuda' else torch.xpu.device(x.device.index) - with device_ctx: - torch.library.wrap_triton(rotary_kernel)[grid]( - output, # data ptrs - x, - cos, - sin, - cu_seqlens, - seqlen_offsets, - seqlen, # shapes - nheads, - seqlen_ro, - output.stride(0) if not is_varlen else 0, # batch_strides if not varlen else 0 - output.stride(-3), # seqlen_stride or total_seqlen_stride - output.stride(-2), # nheads_stride - output.stride(-1), # headdim_stride - x.stride(0) if not is_varlen else 0, # batch_strides if not varlen else 0 - x.stride(-3), # seqlen stride or total_seqlen_stride - x.stride(-2), # nheads stride - x.stride(-1), # headdim stride - rotary_dim, - isinstance(seqlen_offsets, torch.Tensor), - is_varlen, - interleaved, - conjugate, - BLOCK_M=BLOCK_M, - BLOCK_H=2, - ) - return output diff --git a/build/torch29-cxx11-cu126-x86_64-linux/_flash_attn2_588b404.abi3.so b/build/torch29-cxx11-cu126-x86_64-linux/_flash_attn2_588b404.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..99b3eb56b2ddec5b0a2bc242deec3b9f909dea25 --- /dev/null +++ b/build/torch29-cxx11-cu126-x86_64-linux/_flash_attn2_588b404.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:82a33a975de0a2c8e2440d596ecde21e4f3e1e8dcc9df42843e2045edb1e6d47 +size 448648728 diff --git a/build/torch29-cxx11-cu126-x86_64-linux/_flash_attn2_cuda_abda3a0.abi3.so b/build/torch29-cxx11-cu126-x86_64-linux/_flash_attn2_cuda_abda3a0.abi3.so deleted file mode 100644 index 8df52e7c2fda6069b0020287de7a152bc2307f75..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-x86_64-linux/_flash_attn2_cuda_abda3a0.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:b7fa75c66a0e0f26e36e97ca905305ff054ae2936078d1ee1354cb4540bf1cdc -size 448648752 diff --git a/build/torch29-cxx11-cu126-x86_64-linux/_ops.py b/build/torch29-cxx11-cu126-x86_64-linux/_ops.py index fb32434e09e651a658beb38be8c101e5db89374f..03506d8a18d17ca7296a5f1b7a18ec7dc4ad73bb 100644 --- a/build/torch29-cxx11-cu126-x86_64-linux/_ops.py +++ b/build/torch29-cxx11-cu126-x86_64-linux/_ops.py @@ -1,9 +1,9 @@ import torch -from . import _flash_attn2_cuda_abda3a0 -ops = torch.ops._flash_attn2_cuda_abda3a0 +from . import _flash_attn2_588b404 +ops = torch.ops._flash_attn2_588b404 def add_op_namespace_prefix(op_name: str): """ Prefix op by namespace. """ - return f"_flash_attn2_cuda_abda3a0::{op_name}" + return f"_flash_attn2_588b404::{op_name}" \ No newline at end of file diff --git a/build/torch29-cxx11-cu126-x86_64-linux/flash_attn_interface.py b/build/torch29-cxx11-cu126-x86_64-linux/flash_attn_interface.py index 6a0157f25449eb90c6c688bc8ef3f4fcf84d46dc..66e8d3665834b141d703998496b3da6fe1f6fe06 100644 --- a/build/torch29-cxx11-cu126-x86_64-linux/flash_attn_interface.py +++ b/build/torch29-cxx11-cu126-x86_64-linux/flash_attn_interface.py @@ -31,14 +31,12 @@ def _get_device(): else: return "cpu" +_XPU_AVAILABLE = torch.xpu.is_available() if hasattr(torch, "xpu") else False # TODO remove hasattr check when bwd is supported on XPU -def _get_block_size_n(device, head_dim, is_dropout, is_causal): - assert head_dim <= 256 - - if device.type == "xpu": - return 64 +def _get_block_size_n(device, head_dim, is_dropout, is_causal): # This should match the block sizes in the CUDA kernel + assert head_dim <= 256 major, minor = torch.cuda.get_device_capability(device) is_sm8x = major == 8 and minor > 0 # Only include sm86 and sm89, exclude sm80 (A100) is_sm80 = major == 8 and minor == 0 @@ -1066,7 +1064,7 @@ def flash_attn_qkvpacked_func( alibi_slopes, deterministic, return_attn_probs, - torch.is_grad_enabled(), + False if _XPU_AVAILABLE else torch.is_grad_enabled(), ) @@ -1144,7 +1142,7 @@ def flash_attn_kvpacked_func( alibi_slopes, deterministic, return_attn_probs, - torch.is_grad_enabled(), + False if _XPU_AVAILABLE else torch.is_grad_enabled(), ) @@ -1221,7 +1219,7 @@ def flash_attn_func( alibi_slopes, deterministic, return_attn_probs, - torch.is_grad_enabled(), + False if _XPU_AVAILABLE else torch.is_grad_enabled(), ) @@ -1287,7 +1285,7 @@ def flash_attn_varlen_qkvpacked_func( alibi_slopes, deterministic, return_attn_probs, - torch.is_grad_enabled(), + False if _XPU_AVAILABLE else torch.is_grad_enabled(), ) @@ -1379,7 +1377,7 @@ def flash_attn_varlen_kvpacked_func( alibi_slopes, deterministic, return_attn_probs, - torch.is_grad_enabled(), + False if _XPU_AVAILABLE else torch.is_grad_enabled(), ) @@ -1473,7 +1471,7 @@ def flash_attn_varlen_func( deterministic, return_attn_probs, block_table, - False if q.device.type == "cpu" else torch.is_grad_enabled(), + False if _XPU_AVAILABLE or q.device.type == "cpu" else torch.is_grad_enabled(), ) diff --git a/build/torch29-cxx11-cu126-x86_64-linux/metadata.json b/build/torch29-cxx11-cu126-x86_64-linux/metadata.json index 3514c62e2b286f3c5245cd885503179f49edf4f5..9cf5deed9898dce769f4cc73913d3530b92a0bd8 100644 --- a/build/torch29-cxx11-cu126-x86_64-linux/metadata.json +++ b/build/torch29-cxx11-cu126-x86_64-linux/metadata.json @@ -1,12 +1,4 @@ { "version": 1, - "license": "BSD-3-Clause", - "python-depends": [], - "backend": { - "type": "cuda", - "archs": [ - "8.0", - "9.0" - ] - } -} + "python-depends": [] +} \ No newline at end of file diff --git a/build/torch29-cxx11-cu128-aarch64-linux/__init__.py b/build/torch29-cxx11-cu128-aarch64-linux/__init__.py deleted file mode 100644 index ecc2f9d896b6c93f90b0a1499856dc0612177422..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-aarch64-linux/__init__.py +++ /dev/null @@ -1,393 +0,0 @@ -from typing import Optional, List -import torch -from ._ops import ops as flash_attn_ops -from .flash_attn_interface import ( - flash_attn_func, - flash_attn_kvpacked_func, - flash_attn_qkvpacked_func, - flash_attn_varlen_func, - flash_attn_varlen_kvpacked_func, - flash_attn_varlen_qkvpacked_func, - flash_attn_with_kvcache, -) - - -def fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - return_softmax: bool = False, - gen: Optional[torch.Generator] = None, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention. - - Args: - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Optional output tensor, same shape as q - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - return_softmax: Whether to return softmax weights - gen: Optional random number generator - - Returns: - List of tensors: [output, softmax_lse, (softmax if return_softmax)] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.fwd( - q, - k, - v, - out, - alibi_slopes, - p_dropout, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - gen, - ) - - -def varlen_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - out: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - block_table: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - max_seqlen_q: int = 0, - max_seqlen_k: int = 0, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - zero_tensors: bool = False, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - return_softmax: bool = False, - gen: Optional[torch.Generator] = None, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention with variable sequence lengths. - - Args: - q: Query tensor of shape [total_q, num_heads, head_size] - k: Key tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - v: Value tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1] - cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1] - out: Optional output tensor of shape [total_q, num_heads, head_size] - seqused_k: Optional tensor specifying how many keys to use per batch element [batch_size] - leftpad_k: Optional left padding for keys of shape [batch_size] - block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq] - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - max_seqlen_q: Maximum sequence length for queries - max_seqlen_k: Maximum sequence length for keys - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - zero_tensors: Whether to zero tensors before computation - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - return_softmax: Whether to return softmax weights - gen: Optional random number generator - - Returns: - List of tensors: [output, softmax_lse, (softmax if return_softmax)] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.varlen_fwd( - q, - k, - v, - out, - cu_seqlens_q, - cu_seqlens_k, - seqused_k, - leftpad_k, - block_table, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - p_dropout, - softmax_scale, - zero_tensors, - is_causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - gen, - ) - - -def bwd( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - p_dropout: float = 0.0, - 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, - gen: Optional[torch.Generator] = None, - rng_state: Optional[torch.Tensor] = None, -) -> List[torch.Tensor]: - """ - Backward pass for multi-head attention. - - Args: - dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size] - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size] - softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q] - dq: Optional gradient tensor for queries, same shape as q - dk: Optional gradient tensor for keys, same shape as k - dv: Optional gradient tensor for values, same shape as v - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - deterministic: Whether to use deterministic algorithms - gen: Optional random number generator - rng_state: Optional RNG state from forward pass - - Returns: - List of tensors: [dq, dk, dv] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - alibi_slopes, - p_dropout, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - deterministic, - gen, - rng_state, - ) - - -def varlen_bwd( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - max_seqlen_q: int = 0, - max_seqlen_k: int = 0, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - zero_tensors: bool = False, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - deterministic: bool = False, - gen: Optional[torch.Generator] = None, - rng_state: Optional[torch.Tensor] = None, -) -> List[torch.Tensor]: - """ - Backward pass for multi-head attention with variable sequence lengths. - - Args: - dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size] - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size] - softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q] - cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1] - cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1] - dq: Optional gradient tensor for queries, same shape as q - dk: Optional gradient tensor for keys, same shape as k - dv: Optional gradient tensor for values, same shape as v - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - max_seqlen_q: Maximum sequence length for queries - max_seqlen_k: Maximum sequence length for keys - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - zero_tensors: Whether to zero tensors before computation - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - deterministic: Whether to use deterministic algorithms - gen: Optional random number generator - rng_state: Optional RNG state from forward pass - - Returns: - List of tensors: [dq, dk, dv] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.varlen_bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - p_dropout, - softmax_scale, - zero_tensors, - is_causal, - window_size_left, - window_size_right, - softcap, - deterministic, - gen, - rng_state, - ) - - -def fwd_kvcache( - q: torch.Tensor, - kcache: torch.Tensor, - vcache: torch.Tensor, - k: Optional[torch.Tensor] = None, - v: Optional[torch.Tensor] = None, - seqlens_k: Optional[torch.Tensor] = None, - rotary_cos: Optional[torch.Tensor] = None, - rotary_sin: Optional[torch.Tensor] = None, - cache_batch_idx: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - block_table: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - out: 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, - is_rotary_interleaved: bool = False, - num_splits: int = 1, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention with KV cache. - - Args: - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - kcache: Key cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - vcache: Value cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - k: Optional new keys tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size] - v: Optional new values tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size] - seqlens_k: Optional sequence lengths for keys of shape [batch_size] - rotary_cos: Optional rotary cosine tensor of shape [seqlen_ro, rotary_dim/2] - rotary_sin: Optional rotary sine tensor of shape [seqlen_ro, rotary_dim/2] - cache_batch_idx: Optional indices to index into the KV cache - leftpad_k: Optional left padding for keys of shape [batch_size] - block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq] - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - out: Optional output tensor, same shape as q - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - is_rotary_interleaved: Whether rotary embeddings are interleaved - num_splits: Number of splits for computation - - Returns: - List of tensors: [output, softmax_lse] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.fwd_kvcache( - q, - kcache, - vcache, - k, - v, - seqlens_k, - rotary_cos, - rotary_sin, - cache_batch_idx, - leftpad_k, - block_table, - alibi_slopes, - out, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - is_rotary_interleaved, - num_splits, - ) diff --git a/build/torch29-cxx11-cu128-aarch64-linux/_flash_attn2_cuda_abda3a0.abi3.so b/build/torch29-cxx11-cu128-aarch64-linux/_flash_attn2_cuda_abda3a0.abi3.so deleted file mode 100644 index d180a47914c443ec738c52256dbe0eebd5f6cf18..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-aarch64-linux/_flash_attn2_cuda_abda3a0.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:fd01e83103c6287b943a514ab196554c9d70f33a819e231d7286ba27ec4a1db0 -size 1037856008 diff --git a/build/torch29-cxx11-cu128-aarch64-linux/_ops.py b/build/torch29-cxx11-cu128-aarch64-linux/_ops.py deleted file mode 100644 index fb32434e09e651a658beb38be8c101e5db89374f..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-aarch64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _flash_attn2_cuda_abda3a0 -ops = torch.ops._flash_attn2_cuda_abda3a0 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_flash_attn2_cuda_abda3a0::{op_name}" diff --git a/build/torch29-cxx11-cu128-aarch64-linux/bert_padding.py b/build/torch29-cxx11-cu128-aarch64-linux/bert_padding.py deleted file mode 100644 index 3c2d35159a014a9d03aabead9e52e009168696ea..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-aarch64-linux/bert_padding.py +++ /dev/null @@ -1,218 +0,0 @@ -# Adapted from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/padding.py - -import torch -import torch.nn.functional as F -from einops import rearrange, repeat - - -class IndexFirstAxis(torch.autograd.Function): - @staticmethod - def forward(ctx, input, indices): - ctx.save_for_backward(indices) - assert input.ndim >= 2 - ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] - second_dim = other_shape.numel() - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - # return input[indices] - return torch.gather( - rearrange(input, "b ... -> b (...)"), 0, repeat(indices, "z -> z d", d=second_dim) - ).reshape(-1, *other_shape) - - @staticmethod - def backward(ctx, grad_output): - (indices,) = ctx.saved_tensors - assert grad_output.ndim >= 2 - other_shape = grad_output.shape[1:] - grad_output = rearrange(grad_output, "b ... -> b (...)") - grad_input = torch.zeros( - [ctx.first_axis_dim, grad_output.shape[1]], - device=grad_output.device, - dtype=grad_output.dtype, - ) - # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. - # grad_input[indices] = grad_output - grad_input.scatter_(0, repeat(indices, "z -> z d", d=grad_output.shape[1]), grad_output) - return grad_input.reshape(ctx.first_axis_dim, *other_shape), None - - -index_first_axis = IndexFirstAxis.apply - - -class IndexPutFirstAxis(torch.autograd.Function): - @staticmethod - def forward(ctx, values, indices, first_axis_dim): - ctx.save_for_backward(indices) - assert indices.ndim == 1 - assert values.ndim >= 2 - output = torch.zeros( - first_axis_dim, *values.shape[1:], device=values.device, dtype=values.dtype - ) - # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. - output[indices] = values - # output.scatter_(0, repeat(indices, 'z -> z d', d=values.shape[1]), values) - return output - - @staticmethod - def backward(ctx, grad_output): - (indices,) = ctx.saved_tensors - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - grad_values = grad_output[indices] - # grad_values = torch.gather(grad_output, 0, repeat(indices, 'z -> z d', d=grad_output.shape[1])) - return grad_values, None, None - - -index_put_first_axis = IndexPutFirstAxis.apply - - -class IndexFirstAxisResidual(torch.autograd.Function): - @staticmethod - def forward(ctx, input, indices): - ctx.save_for_backward(indices) - assert input.ndim >= 2 - ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] - second_dim = other_shape.numel() - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - output = input[indices] - # We don't want to reshape input (b ... -> b (...)) since it could change the channel_last - # memory format to channel_first. In other words, input might not be contiguous. - # If we don't detach, Pytorch complains about output being a view and is being modified inplace - return output, input.detach() - - @staticmethod - def backward(ctx, grad_output, grad_residual): - (indices,) = ctx.saved_tensors - assert grad_output.ndim >= 2 - other_shape = grad_output.shape[1:] - assert grad_residual.shape[1:] == other_shape - grad_input = grad_residual - # grad_input[indices] += grad_output - indices = indices.reshape(indices.shape[0], *((1,) * (grad_output.ndim - 1))) - indices = indices.expand_as(grad_output) - grad_input.scatter_add_(0, indices, grad_output) - return grad_input.reshape(ctx.first_axis_dim, *other_shape), None - - -index_first_axis_residual = IndexFirstAxisResidual.apply - - -def unpad_input(hidden_states, attention_mask, unused_mask=None): - """ - Arguments: - hidden_states: (batch, seqlen, ...) - attention_mask: (batch, seqlen), bool / int, 1 means valid and 0 means not valid. - unused_mask: (batch, seqlen), bool / int, 1 means the element is allocated but unused. - Return: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens selected in attention_mask + unused_mask. - indices: (total_nnz), the indices of masked tokens from the flattened input sequence. - cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. - max_seqlen_in_batch: int - seqused: (batch), returns the number of tokens selected in attention_mask + unused_mask. - """ - all_masks = (attention_mask + unused_mask) if unused_mask is not None else attention_mask - seqlens_in_batch = all_masks.sum(dim=-1, dtype=torch.int32) - used_seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) - indices = torch.nonzero(all_masks.flatten(), as_tuple=False).flatten() - max_seqlen_in_batch = seqlens_in_batch.max().item() - cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) - # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the - # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim - # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to - # index with integer indices. Moreover, torch's index is a bit slower than it needs to be, - # so we write custom forward and backward to make it a bit faster. - return ( - index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices), - indices, - cu_seqlens, - max_seqlen_in_batch, - used_seqlens_in_batch, - ) - - -def unpad_input_for_concatenated_sequences(hidden_states, attention_mask_in_length): - """ - Supports concatenating short samples in one sequence. The attention_mask_in_length is utilized to mask other short samples. It helps efficient training of variant lengths-based samples (e.g., the supervised fine-tuning task in large language model). - The motivation for this function is explained [here](https://github.com/Dao-AILab/flash-attention/issues/432#issuecomment-1668822286). - - For example, if batch = 3 and seqlen = 6, the attention_mask_in_length is: - ``` - [ - [2, 3, 0, 0, 0, 0], - [3, 2, 0, 0, 0, 0], - [6, 0, 0, 0, 0, 0] - ] - ``` - , which refers to the 3D-attention mask: - ``` - [ - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [0, 0, 1, 0, 0, 0], - [0, 0, 1, 1, 0, 0], - [0, 0, 1, 1, 1, 0], - [0, 0, 0, 0, 0, 1] - ], - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [0, 0, 0, 1, 0, 0], - [0, 0, 0, 1, 1, 0], - [0, 0, 0, 0, 0, 1] - ], - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [1, 1, 1, 1, 0, 0], - [1, 1, 1, 1, 1, 0], - [1, 1, 1, 1, 1, 1] - ] - ] - ```. - - Arguments: - hidden_states: (batch, seqlen, ...) - attention_mask_in_length: (batch, seqlen), int, a nonzero number (e.g., 1, 2, 3, etc.) means length of concatenated sequence in b-th batch, and 0 means none. - Return: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. - indices: (total_nnz), the indices of non-masked tokens from the flattened input sequence. - cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. - max_seqlen_in_batch: int - """ - length = attention_mask_in_length.sum(dim=-1) - seqlen = attention_mask_in_length.size(-1) - attention_mask_2d = torch.arange(seqlen, device=length.device, dtype=length.dtype).expand(len(length), seqlen) < length.unsqueeze(1) - real_indices_idx = torch.nonzero(attention_mask_in_length.flatten(), as_tuple=False).flatten() - seqlens_in_batch = attention_mask_in_length.flatten()[real_indices_idx] - indices = torch.nonzero(attention_mask_2d.flatten(), as_tuple=False).flatten() - max_seqlen_in_batch = seqlens_in_batch.max().item() - cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) - # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the - # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim - # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to - # index with integer indices. Moreover, torch's index is a bit slower than it needs to be, - # so we write custom forward and backward to make it a bit faster. - return ( - index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices), - indices, - cu_seqlens, - max_seqlen_in_batch, - ) - - -def pad_input(hidden_states, indices, batch, seqlen): - """ - Arguments: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. - indices: (total_nnz), the indices that represent the non-masked tokens of the original padded input sequence. - batch: int, batch size for the padded sequence. - seqlen: int, maximum sequence length for the padded sequence. - Return: - hidden_states: (batch, seqlen, ...) - """ - dim = hidden_states.shape[-1] - # output = torch.zeros((batch * seqlen), dim, device=hidden_states.device, dtype=hidden_states.dtype) - # output[indices] = hidden_states - output = index_put_first_axis(hidden_states, indices, batch * seqlen) - return rearrange(output, "(b s) ... -> b s ...", b=batch) diff --git a/build/torch29-cxx11-cu128-aarch64-linux/flash_attn2/__init__.py b/build/torch29-cxx11-cu128-aarch64-linux/flash_attn2/__init__.py deleted file mode 100644 index 03dbc1afe1cf156661a2b1b22003cd5f599a0309..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-aarch64-linux/flash_attn2/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -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-aarch64-linux/flash_attn_interface.py b/build/torch29-cxx11-cu128-aarch64-linux/flash_attn_interface.py deleted file mode 100644 index 6a0157f25449eb90c6c688bc8ef3f4fcf84d46dc..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-aarch64-linux/flash_attn_interface.py +++ /dev/null @@ -1,1622 +0,0 @@ -# Copyright (c) 2023, Tri Dao. - -from typing import Optional, Sequence, Tuple, Union - -import torch -import torch.nn as nn -import os - -# # isort: off -# # We need to import the CUDA kernels after importing torch -# USE_TRITON_ROCM = os.getenv("FLASH_ATTENTION_TRITON_AMD_ENABLE", "FALSE") == "TRUE" -# if USE_TRITON_ROCM: -# from .flash_attn_triton_amd import interface_fa as flash_attn -# else: -# import flash_attn_2_cuda as flash_attn - - -from ._ops import ops as flash_attn - -# # isort: on - -def maybe_contiguous(x): - return x.contiguous() if x is not None and x.stride(-1) != 1 else x - - -def _get_device(): - if torch.xpu.is_available(): - return "xpu" - elif torch.cuda.is_available(): - return "cuda" - else: - return "cpu" - - -def _get_block_size_n(device, head_dim, is_dropout, is_causal): - assert head_dim <= 256 - - if device.type == "xpu": - return 64 - - # This should match the block sizes in the CUDA kernel - major, minor = torch.cuda.get_device_capability(device) - is_sm8x = major == 8 and minor > 0 # Only include sm86 and sm89, exclude sm80 (A100) - is_sm80 = major == 8 and minor == 0 - is_sm90 = major == 9 and minor == 0 - if head_dim <= 32: - return 128 - if head_dim <= 64: - return 128 if not is_dropout else 64 - elif head_dim <= 96: - return 64 - elif head_dim <= 128: - if is_sm8x: - return 64 if (not is_dropout and is_causal) else 32 - else: - return 64 if not is_dropout else 32 - elif head_dim <= 192: - return 64 - elif head_dim <= 224: - return 64 - elif head_dim <= 256: - return 64 - - -def round_multiple(x, m): - return (x + m - 1) // m * m - - -# torch.compile() support is only enabled for pytorch >= 2.4 -# The reason for this is that we are using the new custom_op and register_fake -# APIs, which support inplace modification of inputs in the function itself -if torch.__version__ >= "2.4.0": - _torch_custom_op_wrapper = torch.library.custom_op - _torch_register_fake_wrapper = torch.library.register_fake -else: - def noop_custom_op_wrapper(name, fn=None, /, *, mutates_args, device_types=None, schema=None): - def wrap(func): - return func - if fn is None: - return wrap - return fn - def noop_register_fake_wrapper(op, fn=None, /, *, lib=None, _stacklevel=1): - def wrap(func): - return func - if fn is None: - return wrap - return fn - _torch_custom_op_wrapper = noop_custom_op_wrapper - _torch_register_fake_wrapper = noop_register_fake_wrapper - - -@_torch_custom_op_wrapper("flash_attn::_flash_attn_forward", mutates_args=(), device_types=_get_device()) -def _flash_attn_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - return_softmax: bool -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - out, softmax_lse, S_dmask, rng_state = flash_attn.fwd( - q, - k, - v, - None, - alibi_slopes, - dropout_p, - softmax_scale, - causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - None, - ) - return out, softmax_lse, S_dmask, rng_state - - -@_torch_register_fake_wrapper("flash_attn::_flash_attn_forward") -def _flash_attn_forward_fake( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - return_softmax: bool -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - batch_size, seqlen_q, num_heads, head_size = q.shape - seqlen_k = k.shape[1] - out = torch.empty_like(q) - softmax_lse = torch.empty((batch_size, num_heads, seqlen_q), dtype=torch.float32, device=q.device, layout=q.layout) - p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) - if return_softmax: - p = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128), round_multiple(seqlen_k, 128)), dtype=q.dtype, device=q.device, layout=q.layout) - rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) - - return out, softmax_lse, p, rng_state - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_forward = torch.ops.flash_attn._flash_attn_forward -else: - _wrapped_flash_attn_forward = _flash_attn_forward - - -@_torch_custom_op_wrapper("flash_attn::_flash_attn_varlen_forward", mutates_args=(), device_types=_get_device()) -def _flash_attn_varlen_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - alibi_slopes: Optional[torch.Tensor] = None, - return_softmax: bool = False, - block_table: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - out, softmax_lse, S_dmask, rng_state = flash_attn.varlen_fwd( - q, - k, - v, - None, - cu_seqlens_q, - cu_seqlens_k, - seqused_k, - leftpad_k, - block_table, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - zero_tensors, - causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - None, - ) - # if out.isnan().any() or softmax_lse.isnan().any(): - # breakpoint() - return out, softmax_lse, S_dmask, rng_state - - -@_torch_register_fake_wrapper("flash_attn::_flash_attn_varlen_forward") -def _flash_attn_varlen_forward_fake( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - alibi_slopes: Optional[torch.Tensor] = None, - return_softmax: bool = False, - block_table: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - paged_kv = block_table is not None - batch_size = cu_seqlens_q.numel() - 1 - total_q, num_heads, _ = q.shape - - out = torch.empty_like(q) - softmax_lse = torch.empty((num_heads, total_q), dtype=torch.float32, device=q.device, layout=q.layout) - p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) - seqlen_q_rounded = round_multiple(max_seqlen_q, 128) - seqlen_k_rounded = round_multiple(max_seqlen_k, 128) - if return_softmax: - p = torch.empty((batch_size, num_heads, seqlen_q_rounded, seqlen_k_rounded), dtype=q.dtype, device=q.device, layout=q.layout) - rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) - return out, softmax_lse, p, rng_state - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_varlen_forward = torch.ops.flash_attn._flash_attn_varlen_forward -else: - _wrapped_flash_attn_varlen_forward = _flash_attn_varlen_forward - - -@_torch_custom_op_wrapper("flash_attn::_flash_attn_backward", mutates_args=("dq", "dk", "dv"), device_types=_get_device()) -def _flash_attn_backward( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, -) -> 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)] - ( - dq, - dk, - dv, - softmax_d, - ) = flash_attn.bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - alibi_slopes, - dropout_p, - softmax_scale, - causal, - window_size_left, - window_size_right, - softcap, - deterministic, - None, - rng_state, - ) - return softmax_d - - -@_torch_register_fake_wrapper("flash_attn::_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, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, -) -> torch.Tensor: - dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] - if dq is None: - dq = torch.empty_like(q) - if dk is None: - dk = torch.empty_like(k) - if dv is None: - dv = torch.empty_like(v) - batch_size, seqlen_q, num_heads, _ = q.shape - softmax_d = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128)), device=q.device, dtype=torch.float32) - - return softmax_d - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_backward = torch.ops.flash_attn._flash_attn_backward -else: - _wrapped_flash_attn_backward = _flash_attn_backward - - -@_torch_custom_op_wrapper("flash_attn::_flash_attn_varlen_backward", mutates_args=("dq", "dk", "dv"), device_types=_get_device()) -def _flash_attn_varlen_backward( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> 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)] - ( - dq, - dk, - dv, - softmax_d, - ) = flash_attn.varlen_bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - zero_tensors, - causal, - window_size_left, - window_size_right, - softcap, - deterministic, - None, - rng_state, - ) - # if dk.isnan().any() or dk.isnan().any() or dv.isnan().any() or softmax_d.isnan().any(): - # breakpoint() - return softmax_d - - -@_torch_register_fake_wrapper("flash_attn::_flash_attn_varlen_backward") -def _flash_attn_varlen_backward_fake( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> torch.Tensor: - dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] - batch_size = cu_seqlens_q.numel() - 1 - total_q, num_heads, _ = q.shape - - if dq is None: - dq = torch.empty_like(q) - if dk is None: - dk = torch.empty_like(k) - if dv is None: - dv = torch.empty_like(v) - softmax_d = torch.empty((num_heads, total_q + 128 * batch_size), device=q.device, dtype=torch.float32) - - return softmax_d - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_varlen_backward = torch.ops.flash_attn._flash_attn_varlen_backward -else: - _wrapped_flash_attn_varlen_backward = _flash_attn_varlen_backward - - -class FlashAttnQKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and qkv.requires_grad - if softmax_scale is None: - softmax_scale = qkv.shape[-1] ** (-0.5) - q, k, v = qkv[:, :, 0].detach(), qkv[:, :, 1].detach(), qkv[:, :, 2].detach() - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) - dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dqkv[:, :, 0], - dqkv[:, :, 1], - dqkv[:, :, 2], - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dqkv = dqkv[..., : dout.shape[-1]] # We could have padded the head dimension - return dqkv, None, None, None, None, None, None, None, None, None - - -class FlashAttnVarlenQKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - cu_seqlens, - max_seqlen, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and qkv.requires_grad - if softmax_scale is None: - softmax_scale = qkv.shape[-1] ** (-0.5) - q, k, v = qkv[:, 0].detach(), qkv[:, 1].detach(), qkv[:, 2].detach() - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens, - cu_seqlens, - max_seqlen, - max_seqlen, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=None, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens, rng_state) - ctx.dropout_p = dropout_p - ctx.max_seqlen = max_seqlen - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens, rng_state = ctx.saved_tensors - qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) - dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dqkv[:, 0], - dqkv[:, 1], - dqkv[:, 2], - cu_seqlens, - cu_seqlens, - ctx.max_seqlen, - ctx.max_seqlen, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - 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 FlashAttnKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - kv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, kv] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - k, v = kv[:, :, 0].detach(), kv[:, :, 1].detach() - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - dq = torch.empty_like(q) - kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) - dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dkv[:, :, 0], - dkv[:, :, 1], - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dkv = dkv[..., : dout.shape[-1]] - return dq, dkv, None, None, None, None, None, None, None, None, None - - -class FlashAttnVarlenKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, kv] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - k, v = kv[:, 0].detach(), kv[:, 1].detach() - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=None, - ) - if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state - ) - ctx.dropout_p = dropout_p - 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.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors - dq = torch.empty_like(q) - kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) - dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dkv[:, 0], - dkv[:, 1], - cu_seqlens_q, - cu_seqlens_k, - ctx.max_seqlen_q, - ctx.max_seqlen_k, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dkv = dkv[..., : dout.shape[-1]] - return dq, dkv, None, None, None, None, None, None, None, None, None, None, None, None, None - - -class FlashAttnFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - k, - v, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, k, v] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dk = dk[..., : dout.shape[-1]] - dv = dv[..., : dout.shape[-1]] - return dq, dk, dv, 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, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - block_table, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, k, v] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=block_table, - ) - if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state - ) - ctx.dropout_p = dropout_p - 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.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors - dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - ctx.max_seqlen_q, - ctx.max_seqlen_k, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dk = dk[..., : dout.shape[-1]] - dv = dv[..., : dout.shape[-1]] - return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None - - -def flash_attn_qkvpacked_func( - qkv, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # <=0.0 means deactivate - alibi_slopes=None, - deterministic=False, - 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, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_kvpacked_func( - q, - kv, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, -): - """dropout_p should be set to 0.0 during evaluation - If 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 K, V. - 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) - kv: (batch_size, seqlen, 2, 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. - softcap: float. Anything > 0 activates softcapping 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). - 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 FlashAttnKVPackedFunc.apply( - q, - kv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_func( - q, - k, - v, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - 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). - 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 FlashAttnFunc.apply( - q, - k, - v, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_qkvpacked_func( - qkv, - cu_seqlens, - max_seqlen, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - 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_varlen_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_varlen_kvpacked_func and flash_attn_varlen_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: (total, 3, nheads, headdim), where total = total number of tokens in the batch. - cu_seqlens: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into qkv. - max_seqlen: int. Maximum sequence length in the batch. - 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenQKVPackedFunc.apply( - qkv, - cu_seqlens, - max_seqlen, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_kvpacked_func( - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, -): - """dropout_p should be set to 0.0 during evaluation - If 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 K, V. - 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: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch. - kv: (total_k, 2, nheads_k, headdim), where total_k = total number of key tokens in the batch. - cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into q. - cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into kv. - max_seqlen_q: int. Maximum query sequence length in the batch. - max_seqlen_k: int. Maximum key sequence length in the batch. - 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 + 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenKVPackedFunc.apply( - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_func( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, - block_table=None, -): - """dropout_p should be set to 0.0 during evaluation - Supports multi-query and grouped-query attention (MQA/GQA) by passing in K, V 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: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch. - k: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. - v: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. - cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into q. - cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into kv. - max_seqlen_q: int. Maximum query sequence length in the batch. - max_seqlen_k: int. Maximum key sequence length in the batch. - 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 + 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenFunc.apply( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - block_table, - False if q.device.type == "cpu" else torch.is_grad_enabled(), - ) - - -def flash_attn_with_kvcache( - q, - k_cache, - v_cache, - k=None, - v=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, - block_table: Optional[torch.Tensor] = None, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - rotary_interleaved=True, - alibi_slopes=None, - num_splits=0, - 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 block_table, - or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_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) if there's no block_table, - or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_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). Similar to k. - 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. - block_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). - 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. - 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" - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-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) - cache_batch_idx = maybe_contiguous(cache_batch_idx) - block_table = maybe_contiguous(block_table) - out, softmax_lse = flash_attn.fwd_kvcache( - q, - k_cache, - v_cache, - k, - v, - cache_seqlens, - rotary_cos, - rotary_sin, - cache_batch_idx, - cache_leftpad, - block_table, - alibi_slopes, - None, - softmax_scale, - causal, - window_size[0], - window_size[1], - softcap, - rotary_interleaved, - num_splits, - ) - return (out, softmax_lse) if return_softmax_lse else out diff --git a/build/torch29-cxx11-cu128-aarch64-linux/layers/__init__.py b/build/torch29-cxx11-cu128-aarch64-linux/layers/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/build/torch29-cxx11-cu128-aarch64-linux/layers/patch_embed.py b/build/torch29-cxx11-cu128-aarch64-linux/layers/patch_embed.py deleted file mode 100644 index 05562f8e8bcdb58e947c6f402a49eacd2d031871..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-aarch64-linux/layers/patch_embed.py +++ /dev/null @@ -1,67 +0,0 @@ -# We use the same API as https://github.com/rwightman/pytorch-image-models/blob/v0.6.11/timm/models/layers/patch_embed.py -# But we use nn.Linear instead of Conv2d and it's about 8x faster. - -from functools import partial - -import torch.nn as nn -from einops import rearrange -from torch import _assert -from torch.nn.modules.utils import _pair - -try: - from flash_attn.ops.fused_dense import FusedDense -except ImportError: - FusedDense = None - - -class PatchEmbed(nn.Module): - """2D Image to Patch Embedding""" - - def __init__( - self, - img_size=224, - patch_size=16, - in_chans=3, - embed_dim=768, - norm_layer=None, - flatten=True, - bias=True, - fused_bias_fc=False, - ): - super().__init__() - img_size = _pair(img_size) - patch_size = _pair(patch_size) - self.img_size = img_size - self.patch_size = patch_size - self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1]) - self.num_patches = self.grid_size[0] * self.grid_size[1] - self.flatten = flatten - if fused_bias_fc and FusedDense is None: - raise ImportError("fused_dense is not installed") - - linear_cls = nn.Linear if not fused_bias_fc or not bias else FusedDense - self.proj = linear_cls(in_chans * patch_size[0] * patch_size[1], embed_dim, bias=bias) - self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() - - def forward(self, x): - _, _, H, W = x.shape - _assert( - H == self.img_size[0], - f"Input image height ({H}) doesn't match model ({self.img_size[0]}).", - ) - _assert( - W == self.img_size[1], - f"Input image width ({W}) doesn't match model ({self.img_size[1]}).", - ) - x = self.proj( - rearrange( - x, - "b c (h p1) (w p2) -> b h w (c p1 p2)", - p1=self.patch_size[0], - p2=self.patch_size[1], - ) - ) - if self.flatten: - x = rearrange(x, "b h w c -> b (h w) c") - x = self.norm(x) - return x diff --git a/build/torch29-cxx11-cu128-aarch64-linux/layers/rotary.py b/build/torch29-cxx11-cu128-aarch64-linux/layers/rotary.py deleted file mode 100644 index d1bfc21fc7de1dd287e8f382847b194a48075981..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-aarch64-linux/layers/rotary.py +++ /dev/null @@ -1,483 +0,0 @@ -# Copyright (c) 2025, Tri Dao - -import math -from functools import partial -from typing import Optional, Tuple, Union - -import torch -from torch import Tensor - -from einops import rearrange, repeat -# from flash_attn.ops.triton.rotary import apply_rotary -from ..ops.triton.rotary import apply_rotary - - -def rotate_half(x, interleaved=False): - if not interleaved: - x1, x2 = x.chunk(2, dim=-1) - return torch.cat((-x2, x1), dim=-1) - else: - x1, x2 = x[..., ::2], x[..., 1::2] - return rearrange(torch.stack((-x2, x1), dim=-1), "... d two -> ... (d two)", two=2) - - -def apply_rotary_emb_torch(x, cos, sin, interleaved=False): - """ - x: (batch_size, seqlen, nheads, headdim) - cos, sin: (seqlen, rotary_dim / 2) or (batch_size, seqlen, rotary_dim / 2) - """ - ro_dim = cos.shape[-1] * 2 - assert ro_dim <= x.shape[-1] - cos = repeat(cos, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)") - sin = repeat(sin, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)") - return torch.cat( - [x[..., :ro_dim] * cos + rotate_half(x[..., :ro_dim], interleaved) * sin, x[..., ro_dim:]], - dim=-1, - ) - - -class ApplyRotaryEmb(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x, - cos, - sin, - interleaved=False, - inplace=False, - seqlen_offsets: Union[int, Tensor] = 0, - cu_seqlens: Optional[Tensor] = None, - max_seqlen: Optional[int] = None, - ): - out = apply_rotary( - x, - cos, - sin, - seqlen_offsets=seqlen_offsets, - cu_seqlens=cu_seqlens, - max_seqlen=max_seqlen, - interleaved=interleaved, - inplace=inplace, - ) - if isinstance(seqlen_offsets, int): - ctx.save_for_backward(cos, sin, cu_seqlens) # Can't save int with save_for_backward - ctx.seqlen_offsets = seqlen_offsets - else: - ctx.save_for_backward(cos, sin, cu_seqlens, seqlen_offsets) - ctx.seqlen_offsets = None - ctx.interleaved = interleaved - ctx.inplace = inplace - ctx.max_seqlen = max_seqlen - return out if not inplace else x - - @staticmethod - def backward(ctx, do): - seqlen_offsets = ctx.seqlen_offsets - if seqlen_offsets is None: - cos, sin, cu_seqlens, seqlen_offsets = ctx.saved_tensors - else: - cos, sin, cu_seqlens = ctx.saved_tensors - dx = apply_rotary( - do, - cos, - sin, - seqlen_offsets=seqlen_offsets, - cu_seqlens=cu_seqlens, - max_seqlen=ctx.max_seqlen, - interleaved=ctx.interleaved, - inplace=ctx.inplace, - conjugate=True, - ) - return dx, None, None, None, None, None, None, None - - -def apply_rotary_emb( - x, - cos, - sin, - interleaved=False, - inplace=False, - seqlen_offsets: Union[int, Tensor] = 0, - cu_seqlens: Optional[Tensor] = None, - max_seqlen: Optional[int] = None, -): - """ - Arguments: - x: (batch_size, seqlen, nheads, headdim) if cu_seqlens is None - else (total_seqlen, nheads, headdim) - cos, sin: (seqlen_rotary, rotary_dim / 2) - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead - of 1st half and 2nd half (GPT-NeoX style). - inplace: if True, apply rotary embedding in-place. - seqlen_offsets: (batch_size,) or int. Each sequence in x is shifted by this amount. - Most commonly used in inference when we have KV cache. - cu_seqlens: (batch + 1,) or None - max_seqlen: int - Return: - out: (batch_size, seqlen, nheads, headdim) if cu_seqlens is None - else (total_seqlen, nheads, headdim) - rotary_dim must be <= headdim - Apply rotary embedding to the first rotary_dim of x. - """ - return ApplyRotaryEmb.apply( - x, cos, sin, interleaved, inplace, seqlen_offsets, cu_seqlens, max_seqlen - ) - - -# For backward compatibility -apply_rotary_emb_func = apply_rotary_emb - - -def _apply_rotary_emb_qkv( - qkv, - cos, - sin, - cos_k=None, - sin_k=None, - interleaved=False, - inplace=False, - conjugate=False, - seqlen_offsets: Union[int, Tensor] = 0, - num_heads_q: Optional[int] = None, -): - apply_rotary_fn = partial( - apply_rotary, - interleaved=interleaved, - inplace=inplace, - conjugate=conjugate, - seqlen_offsets=seqlen_offsets - ) - if cos_k is None and sin_k is None and qkv.is_contiguous(): - # Call 1 kernel instead of 2 kernels - # We need qkv to be contiguous so that when we reshape to combine (3, nheads) - # dimensions, we get the same tensor - if qkv.dim() == 5: - batch, seqlen, three, nheads, headdim = qkv.shape - assert three == 3 - # qk = rearrange(qkv[:, :, :2], "b s t h d -> b s (t h) d") - qk = qkv[:, :, :2].reshape(batch, seqlen, -1, headdim) - qk = apply_rotary_fn(qk, cos, sin) - else: - assert qkv.dim() == 4 - assert num_heads_q is not None - num_heads_k = (qkv.shape[2] - num_heads_q) // 2 - assert qkv.shape[2] == num_heads_q + 2 * num_heads_k - qk = qkv[:, :, :num_heads_q + num_heads_k] - qk = apply_rotary_fn(qk, cos, sin) - if not inplace: - if qkv.dim() == 5: - qkv = torch.cat([rearrange(qk, "b s (t h) d -> b s t h d", t=2), qkv[:, :, 2:]], dim=2) - else: - qkv = torch.cat([qk, qkv[:, :, num_heads_q + num_heads_k :]], dim=2) - else: - cos_k = cos if cos_k is None else cos_k - sin_k = sin if sin_k is None else sin_k - if qkv.dim() == 5: - batch, seqlen, three, nheads, headdim = qkv.shape - assert three == 3 - q, k = qkv[:, :, 0], qkv[:, :, 1] - else: - assert qkv.dim() == 4 - assert num_heads_q is not None - num_heads_k = (qkv.shape[2] - num_heads_q) // 2 - assert qkv.shape[2] == num_heads_q + 2 * num_heads_k - q, k = qkv[:, :, :num_heads_q], qkv[:, :, num_heads_q : num_heads_q + num_heads_k] - q = apply_rotary_fn(q, cos, sin) - k = apply_rotary_fn(k, cos_k, sin_k) - if not inplace: - if qkv.dim() == 5: - qkv = torch.stack([q, k, qkv[:, :, 2]], dim=2) - else: - qkv = torch.cat([q, k, qkv[:, :, num_heads_q + num_heads_k:]], dim=2) - return qkv - - -class ApplyRotaryEmbQKV_(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - cos, - sin, - cos_k=None, - sin_k=None, - interleaved=False, - seqlen_offsets: Union[int, torch.Tensor] = 0, - num_heads_q: Optional[int] = None, - ): - # apply_rotary_emb_qkv_inplace( - qkv = _apply_rotary_emb_qkv( - qkv, cos, sin, cos_k, sin_k, interleaved=interleaved, inplace=True, - seqlen_offsets=seqlen_offsets, num_heads_q=num_heads_q, - ) - if isinstance(seqlen_offsets, int): - ctx.save_for_backward(cos, sin, cos_k, sin_k) - ctx.seqlen_offsets = seqlen_offsets - else: - ctx.save_for_backward(cos, sin, cos_k, sin_k, seqlen_offsets) - ctx.seqlen_offsets = None - ctx.interleaved = interleaved - ctx.num_heads_q = num_heads_q - return qkv - - @staticmethod - def backward(ctx, dqkv): - seqlen_offsets = ctx.seqlen_offsets - if seqlen_offsets is None: - cos, sin, cos_k, sin_k, seqlen_offsets = ctx.saved_tensors - else: - cos, sin, cos_k, sin_k = ctx.saved_tensors - dqkv = _apply_rotary_emb_qkv( - dqkv, cos, sin, cos_k, sin_k, interleaved=ctx.interleaved, inplace=True, - seqlen_offsets=seqlen_offsets, num_heads_q=ctx.num_heads_q, conjugate=True, - ) - return dqkv, None, None, None, None, None, None, None - - -def apply_rotary_emb_qkv_( - qkv, - cos, - sin, - cos_k=None, - sin_k=None, - interleaved=False, - seqlen_offsets: Union[int, torch.Tensor] = 0, - num_heads_q: Optional[int] = None, -): - """ - Arguments: - qkv: (batch_size, seqlen, 3, nheads, headdim) or (batch_size, seqlen, num_heads_q + 2 * num_heads_k, headdim). - If qkv has shape (batch_size, seqlen, num_heads_q + 2 * num_heads_k, headdim) (e.g. MQA / GQA), - then num_heads_q must be provided. - cos, sin: (seqlen, rotary_dim / 2) - cos_k, sin_k: (seqlen, rotary_dim / 2), optional - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead of - 1st half and 2nd half (GPT-NeoX style). - seqlen_offsets: (batch_size,) or int. Each sequence in Q and K is shifted by this amount. - Most commonly used in inference when we have KV cache. - Return: - qkv: (batch_size, seqlen, 3, nheads, headdim) or (batch_size, seqlen, num_heads_q + 2 * num_heads_k, headdim) - rotary_dim must be <= headdim - Apply rotary embedding *inplace* to the first rotary_dim of Q and K. - """ - return ApplyRotaryEmbQKV_.apply( - qkv, cos, sin, cos_k, sin_k, interleaved, seqlen_offsets, num_heads_q - ) - - -class ApplyRotaryEmbKV_(torch.autograd.Function): - - @staticmethod - def forward(ctx, kv, cos, sin, interleaved=False, seqlen_offsets: Union[int, torch.Tensor] = 0): - batch, seqlen, two, nheads, headdim = kv.shape - assert two == 2 - k = kv[:, :, 0] - apply_rotary( - k, cos, sin, seqlen_offsets=seqlen_offsets, interleaved=interleaved, inplace=True - ) - if isinstance(seqlen_offsets, int): - ctx.save_for_backward(cos, sin) # Can't save int with save_for_backward - ctx.seqlen_offsets = seqlen_offsets - else: - ctx.save_for_backward(cos, sin, seqlen_offsets) - ctx.seqlen_offsets = None - ctx.interleaved = interleaved - return kv - - @staticmethod - def backward(ctx, dkv): - seqlen_offsets = ctx.seqlen_offsets - if seqlen_offsets is None: - cos, sin, seqlen_offsets = ctx.saved_tensors - else: - cos, sin = ctx.saved_tensors - apply_rotary( - dkv[:, :, 0], - cos, - sin, - seqlen_offsets=seqlen_offsets, - interleaved=ctx.interleaved, - inplace=True, - conjugate=True, - ) - return dkv, None, None, None, None - - -apply_rotary_emb_kv_ = ApplyRotaryEmbKV_.apply - - -def apply_rotary_emb_kv_( - kv, - cos, - sin, - interleaved=False, - seqlen_offsets: Union[int, torch.Tensor] = 0, -): - """ - Arguments: - kv: (batch_size, seqlen, 2, nheads, headdim) - cos, sin: (seqlen, rotary_dim / 2) - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead of - 1st half and 2nd half (GPT-NeoX style). - seqlen_offsets: (batch_size,) or int. Each sequence in Q and K is shifted by this amount. - Most commonly used in inference when we have KV cache. - Return: - kv: (batch_size, seqlen, 2, nheads, headdim) - rotary_dim must be <= headdim - Apply rotary embedding *inplace* to the first rotary_dim of K. - """ - return ApplyRotaryEmbKV_.apply(kv, cos, sin, interleaved, seqlen_offsets) - - -class RotaryEmbedding(torch.nn.Module): - """ - The rotary position embeddings from RoFormer_ (Su et. al). - A crucial insight from the method is that the query and keys are - transformed by rotation matrices which depend on the relative positions. - - Other implementations are available in the Rotary Transformer repo_ and in - GPT-NeoX_, GPT-NeoX was an inspiration - - .. _RoFormer: https://arxiv.org/abs/2104.09864 - .. _repo: https://github.com/ZhuiyiTechnology/roformer - .. _GPT-NeoX: https://github.com/EleutherAI/gpt-neox - - If scale_base is not None, this implements XPos (Sun et al., https://arxiv.org/abs/2212.10554). - A recommended value for scale_base is 512: https://github.com/HazyResearch/flash-attention/issues/96 - Reference: https://github.com/sunyt32/torchscale/blob/main/torchscale/component/xpos_relative_position.py - """ - - def __init__( - self, - dim: int, - base=10000.0, - interleaved=False, - scale_base=None, - device=None, - ): - """ - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead - of 1st half and 2nd half (GPT-NeoX style). - """ - super().__init__() - self.dim = dim - self.base = float(base) - # Generate and save the inverse frequency buffer (non trainable) - inv_freq = self._compute_inv_freq(device) - self.register_buffer("inv_freq", inv_freq, persistent=False) - self.interleaved = interleaved - self.scale_base = scale_base - scale = ( - (torch.arange(0, dim, 2, device=device, dtype=torch.float32) + 0.4 * dim) / (1.4 * dim) - if scale_base is not None - else None - ) - self.register_buffer("scale", scale, persistent=False) - - self._seq_len_cached = 0 - self._cos_cached = None - self._sin_cached = None - self._cos_k_cached = None - self._sin_k_cached = None - - def _compute_inv_freq(self, device=None): - return 1.0 / ( - self.base - ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim) - ) - - def _update_cos_sin_cache(self, seqlen, device=None, dtype=None): - # Reset the tables if the sequence length has changed, - # if we're on a new device (possibly due to tracing for instance), - # or if we're switching from inference mode to training - if ( - seqlen > self._seq_len_cached - or self._cos_cached is None - or self._cos_cached.device != device - or self._cos_cached.dtype != dtype - or (self.training and self._cos_cached.is_inference()) - ): - self._seq_len_cached = seqlen - # We want fp32 here, not self.inv_freq.dtype, since the model could be loaded in bf16 - # And the output of arange can be quite large, so bf16 would lose a lot of precision. - t = torch.arange(seqlen, device=device, dtype=torch.float32) - # We want fp32 here as well since inv_freq will be multiplied with t, and the output - # will be large. Having it in bf16 will lose a lot of precision and cause the - # cos & sin output to change significantly. - # We want to recompute self.inv_freq if it was not loaded in fp32 - if self.inv_freq.dtype != torch.float32: - inv_freq = self._compute_inv_freq(device=device) - else: - inv_freq = self.inv_freq - # Don't do einsum, it converts fp32 to bf16 under AMP - # freqs = torch.einsum("i,j->ij", t, self.inv_freq) - freqs = torch.outer(t, inv_freq) - if self.scale is None: - self._cos_cached = torch.cos(freqs).to(dtype) - self._sin_cached = torch.sin(freqs).to(dtype) - else: - power = ( - torch.arange(seqlen, dtype=self.scale.dtype, device=self.scale.device) - - seqlen // 2 - ) / self.scale_base - scale = self.scale.to(device=power.device) ** rearrange(power, "s -> s 1") - # We want the multiplication by scale to happen in fp32 - self._cos_cached = (torch.cos(freqs) * scale).to(dtype) - self._sin_cached = (torch.sin(freqs) * scale).to(dtype) - self._cos_k_cached = (torch.cos(freqs) / scale).to(dtype) - self._sin_k_cached = (torch.sin(freqs) / scale).to(dtype) - - def forward( - self, - qkv: torch.Tensor, - kv: Optional[torch.Tensor] = None, - seqlen_offset: Union[int, torch.Tensor] = 0, - max_seqlen: Optional[int] = None, - num_heads_q: Optional[int] = None, - ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: - """ - qkv: (batch, seqlen, 3, nheads, headdim) or (batch, seqlen, num_heads_q + 2 * num_heads_k, headdim) - if kv is none, else it's just q of shape (batch, seqlen, nheads, headdim). - If qkv has shape (batch, seqlen, num_heads_q + 2 * num_heads_k, headdim) (e.g. MQA / GQA), - then num_heads_q must be provided. - kv: (batch, seqlen, 2, nheads, headdim) - seqlen_offset: (batch_size,) or int. Each sequence in x is shifted by this amount. - Most commonly used in inference when we have KV cache. - If it's a tensor of shape (batch_size,), then to update the cos / sin cache, one - should pass in max_seqlen, which will update the cos / sin cache up to that length. - Apply rotary embedding *inplace* to qkv and / or kv. - """ - seqlen = qkv.shape[1] - if max_seqlen is not None: - self._update_cos_sin_cache(max_seqlen, device=qkv.device, dtype=qkv.dtype) - elif isinstance(seqlen_offset, int): - self._update_cos_sin_cache(seqlen + seqlen_offset, device=qkv.device, dtype=qkv.dtype) - if kv is None: - return apply_rotary_emb_qkv_( - qkv, - self._cos_cached, - self._sin_cached, - self._cos_k_cached if self.scale is not None else None, - self._sin_k_cached if self.scale is not None else None, - interleaved=self.interleaved, - seqlen_offsets=seqlen_offset, - num_heads_q=num_heads_q, - ) - else: - q = qkv - q = apply_rotary_emb_func( - q, - self._cos_cached, - self._sin_cached, - interleaved=self.interleaved, - inplace=True, - seqlen_offsets=seqlen_offset, - ) - kv = apply_rotary_emb_kv_( - kv, - self._cos_cached if self.scale is None else self._cos_k_cached, - self._sin_cached if self.scale is None else self._sin_k_cached, - interleaved=self.interleaved, - seqlen_offsets=seqlen_offset, - ) - return q, kv diff --git a/build/torch29-cxx11-cu128-aarch64-linux/metadata.json b/build/torch29-cxx11-cu128-aarch64-linux/metadata.json deleted file mode 100644 index 8b00d7fb7685121edf7a97bb007415122444f1d5..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-aarch64-linux/metadata.json +++ /dev/null @@ -1,14 +0,0 @@ -{ - "version": 1, - "license": "BSD-3-Clause", - "python-depends": [], - "backend": { - "type": "cuda", - "archs": [ - "10.0", - "12.0", - "8.0", - "9.0" - ] - } -} diff --git a/build/torch29-cxx11-cu128-aarch64-linux/ops/__init__.py b/build/torch29-cxx11-cu128-aarch64-linux/ops/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/build/torch29-cxx11-cu128-aarch64-linux/ops/activations.py b/build/torch29-cxx11-cu128-aarch64-linux/ops/activations.py deleted file mode 100644 index 7c09649fc41e12d5a360c5672825d8380bc7ec80..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-aarch64-linux/ops/activations.py +++ /dev/null @@ -1,135 +0,0 @@ -# Copied from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/model/layers/activations.py -import math - -import torch -import torch.nn as nn -import torch.nn.functional as F - -# 1/sqrt(2*pi)-> 0.3989423 -# 1/sqrt(2) -> 0.70710678 -# sqrt(2/pi) -> 0.79788456 - -# this function is tanh approximation of gelu -# actual gelu is: -# x * 0.5 * (1.0 + torch.erf(x * 0.70710678)) -@torch.jit.script -def bias_gelu(y, bias): - x = bias + y - return (x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))).to(dtype=y.dtype) - - -# gradient of tanh approximation of gelu -# gradient of actual gelu is: -# 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x) -@torch.jit.script -def bias_gelu_back(g, y, bias): - """Assume that y has shape (B, D) and bias has shape (D)""" - x = bias + y - tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)) - # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243 - ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( - 1 + tanh_out - ) - grad_y = ff * g - return grad_y.to(dtype=y.dtype), grad_y.sum(dim=(0), dtype=bias.dtype) - - -class GeLUFunction(torch.autograd.Function): - @staticmethod - # bias is an optional argument - def forward(ctx, input, bias): - ctx.save_for_backward(input, bias) - return bias_gelu(input, bias) - - @staticmethod - def backward(ctx, grad_output): - input, bias = ctx.saved_tensors - tmp = bias_gelu_back(grad_output, input, bias) - return tmp, tmp - - -bias_gelu_impl = GeLUFunction.apply - -# this function is tanh approximation of gelu -# actual gelu is: -# x * 0.5 * (1.0 + torch.erf(x * 0.70710678)) -@torch.jit.script -def gelu_fwd(x): - return (x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))).to(dtype=x.dtype) - - -# gradient of tanh approximation of gelu -# gradient of actual gelu is: -# 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x) -@torch.jit.script -def gelu_bwd(g, x): - tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)) - # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243 - ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( - 1 + tanh_out - ) - return (ff * g).to(dtype=x.dtype) - - -class FastGeLUFunction(torch.autograd.Function): - @staticmethod - # bias is an optional argument - def forward(ctx, input): - ctx.save_for_backward(input) - return gelu_fwd(input) - - @staticmethod - def backward(ctx, grad_output): - (input,) = ctx.saved_tensors - tmp = gelu_bwd(grad_output, input) - return tmp - - -fast_gelu_impl = FastGeLUFunction.apply - - -@torch.jit.script -def relu_bwd(g, x): - return torch.where(x >= 0, g, 0.0).to(dtype=x.dtype) - - -@torch.jit.script -def sqrelu_fwd(x): - r = F.relu(x) - return (r * r).to(dtype=x.dtype) - - -@torch.jit.script -def sqrelu_bwd(g, x): - return (2.0 * g * F.relu(x)).to(dtype=x.dtype) - - -swiglu_fwd_codestring = """ -template T swiglu_fwd(T x, T y) { - return float(x) * float(y) / (1.0f + ::exp(-float(x))); -} -""" -swiglu_bwd_codestring = """ -template void swiglu_bwd(T x, T y, T g, T& dx, T& dy) { - float x_sigmoid = 1.0f / (1.0f + ::exp(-float(x))); - dx = x_sigmoid * (1 + float(x) * (1.0f - x_sigmoid)) * float(g) * float(y); - dy = float(x) * x_sigmoid * float(g); -} -""" -swiglu_fwd = torch.cuda.jiterator._create_jit_fn(swiglu_fwd_codestring) -swiglu_bwd = torch.cuda.jiterator._create_multi_output_jit_fn(swiglu_bwd_codestring, num_outputs=2) - - -class SwiGLUFunction(torch.autograd.Function): - - @staticmethod - def forward(ctx, x, y): - ctx.save_for_backward(x, y) - return swiglu_fwd(x, y) - - @staticmethod - def backward(ctx, dout): - x, y = ctx.saved_tensors - return swiglu_bwd(x, y, dout) - -swiglu = SwiGLUFunction.apply diff --git a/build/torch29-cxx11-cu128-aarch64-linux/ops/fused_dense.py b/build/torch29-cxx11-cu128-aarch64-linux/ops/fused_dense.py deleted file mode 100644 index 6b4033d134e4093fe278f7b3f8c7d3128ce9f36d..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-aarch64-linux/ops/fused_dense.py +++ /dev/null @@ -1,688 +0,0 @@ -# Copyright (c) 2023, Tri Dao. -# Inspired by https://github.com/NVIDIA/apex/blob/master/apex/fused_dense/fused_dense.py -# We make it work with pytorch amp and with bfloat16. -# The TensorParallel linear modules are inspired by https://github.com/NVIDIA/apex/blob/master/apex/transformer/tensor_parallel/layers.py -from functools import partial -from typing import Optional - -# import fused_dense_cuda # from apex -import fused_dense_lib as fused_dense_cuda -import torch -import torch.nn as nn -import torch.nn.functional as F -from torch import Tensor -from torch.distributed import ProcessGroup - -from flash_attn.utils.torch import custom_fwd, custom_bwd -from flash_attn.ops.activations import gelu_bwd, relu_bwd, sqrelu_bwd, sqrelu_fwd -from flash_attn.utils.distributed import ( - all_gather_raw, - all_reduce, - all_reduce_raw, - reduce_scatter, - reduce_scatter_raw, -) - - -class FusedDenseFunc(torch.autograd.Function): - @staticmethod - @custom_fwd - def forward( - ctx, x, weight, bias, return_residual=False, process_group=None, sequence_parallel=True - ): - """ - If process_group is not None and sequence_parallel=True, we're doing Tensor Parallel - with sequence parallelism: we do an all_gather_raw of x before doing the matmul. - """ - ctx.compute_weight_gradient = weight.requires_grad - ctx.return_residual = return_residual - ctx.process_group = process_group - ctx.sequence_parallel = sequence_parallel - - if torch.is_autocast_enabled(): - x = x.to(dtype=torch.get_autocast_gpu_dtype()) - x = x.contiguous() - if process_group is not None and sequence_parallel: - # We want to kick off the all_gather early, before weight dtype conversion - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - else: - total_x = x - - if torch.is_autocast_enabled(): - weight = weight.to(dtype=torch.get_autocast_gpu_dtype()) - bias = bias.to(dtype=torch.get_autocast_gpu_dtype()) if bias is not None else None - weight = weight.contiguous() - if process_group is not None and sequence_parallel: - handle_x.wait() - batch_shape, n = total_x.shape[:-1], total_x.shape[-1] - batch_dim = batch_shape.numel() - # https://github.com/pytorch/pytorch/blob/5b51849b48a7dbccd297286cc0110def4706f9e7/aten/src/ATen/native/cuda/Blas.cpp#L174 - if min(batch_dim, n, *weight.shape) > 65535 * 32: - raise RuntimeError("fused_dense only supports matrix dims <= 2M") - output = F.linear(total_x, weight, bias) - if ctx.compute_weight_gradient: - ctx.save_for_backward(x, weight) - else: - ctx.save_for_backward(weight) - return output if not return_residual else (output, x) - - @staticmethod - @custom_bwd - def backward(ctx, grad_output, *args): - grad_output = grad_output.contiguous() - if ctx.return_residual: - (grad_input,) = args - grad_input = grad_input.contiguous() - process_group = ctx.process_group - sequence_parallel = ctx.sequence_parallel - if ctx.compute_weight_gradient: - x, weight = ctx.saved_tensors - if process_group is not None and sequence_parallel: - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - else: - total_x = x - else: - (weight,) = ctx.saved_tensors - total_x = None - batch_shape = grad_output.shape[:-1] - batch_dim = batch_shape.numel() - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - if ctx.needs_input_grad[0]: - if not ctx.return_residual: - grad_input = F.linear(grad_output, weight.t()) - else: - grad_input = torch.addmm( - grad_input.reshape(batch_dim, grad_input.shape[-1]), grad_output, weight - ) - grad_input = grad_input.reshape(*batch_shape, grad_input.shape[-1]) - if process_group is not None: - reduce_fn = reduce_scatter_raw if sequence_parallel else all_reduce_raw - grad_input, handle_grad_input = reduce_fn(grad_input, process_group, async_op=True) - else: - grad_input = None - if ctx.needs_input_grad[1]: - assert ctx.compute_weight_gradient - if process_group is not None and sequence_parallel: - handle_x.wait() - grad_weight, grad_bias = fused_dense_cuda.linear_bias_wgrad( - total_x.reshape(batch_dim, total_x.shape[-1]), grad_output, ctx.needs_input_grad[2] - ) - else: - grad_weight = None - grad_bias = grad_output if ctx.needs_input_grad[2] else None - if process_group is not None and ctx.needs_input_grad[0]: - handle_grad_input.wait() - return grad_input, grad_weight, grad_bias, None, None, None - - -def fused_dense_func( - x: Tensor, - weight: Tensor, - bias: Optional[Tensor] = None, - return_residual: bool = False, - process_group: Optional[ProcessGroup] = None, - sequence_parallel: bool = True, -): - dtype_eligible = x.dtype in [torch.float16, torch.bfloat16] or ( - x.dtype == torch.float32 and torch.is_autocast_enabled() - ) - if x.is_cuda and weight.is_cuda and (bias is None or bias.is_cuda) and dtype_eligible: - return FusedDenseFunc.apply( - x, weight, bias, return_residual, process_group, sequence_parallel - ) - else: - assert process_group is None - out = F.linear(x, weight, bias) - return out if not return_residual else (out, x) - - -class FusedDense(nn.Linear): - def __init__( - self, - in_features: int, - out_features: int, - bias: bool = True, - return_residual: bool = False, - device=None, - dtype=None, - ) -> None: - super().__init__(in_features, out_features, bias=bias, device=device, dtype=dtype) - self.return_residual = return_residual - - def forward(self, x, process_group=None): - """ - If process_group is not None, we're doing Tensor Parallel with sequence parallelism: - we do an all_gather of x before doing the matmul. - """ - return fused_dense_func( - x, - self.weight, - self.bias, - return_residual=self.return_residual, - process_group=process_group, - ) - - -class ColumnParallelLinear(nn.Linear): - def __init__( - self, - in_features: int, - out_features: int, - process_group: ProcessGroup, - bias: bool = True, - sequence_parallel=True, - multiple_of=1, - device=None, - dtype=None, - ) -> None: - world_size = torch.distributed.get_world_size(process_group) - if out_features % multiple_of: - raise ValueError(f"out_features ({out_features}) must be a multiple of {multiple_of}") - multiple = out_features // multiple_of - # We want to split @multiple across world_size, but it could be an uneven split - div = multiple // world_size - mod = multiple % world_size - # The first @mod ranks get @div + 1 copies, the rest get @div copies - local_multiple = div + int(torch.distributed.get_rank(process_group) < mod) - super().__init__( - in_features, local_multiple * multiple_of, bias=bias, device=device, dtype=dtype - ) - self.process_group = process_group - self.sequence_parallel = sequence_parallel - - def forward(self, x): - # If self.sequence_parallel is True, we're doing Tensor Parallel with sequence parallelism: - # we do an all_gather of x before doing the matmul. - # If not, then the input is already gathered. - return fused_dense_func( - x, - self.weight, - self.bias, - process_group=self.process_group, - sequence_parallel=self.sequence_parallel, - ) - - -class RowParallelLinear(nn.Linear): - def __init__( - self, - in_features: int, - out_features: int, - process_group: ProcessGroup, - bias: bool = True, - sequence_parallel=True, - multiple_of=1, - device=None, - dtype=None, - ) -> None: - world_size = torch.distributed.get_world_size(process_group) - rank = torch.distributed.get_rank(process_group) - if in_features % multiple_of: - raise ValueError(f"in_features ({in_features}) must be a multiple of {multiple_of}") - multiple = in_features // multiple_of - # We want to split @multiple across world_size, but it could be an uneven split - div = multiple // world_size - mod = multiple % world_size - # The first @mod ranks get @div + 1 copies, the rest get @div copies - local_multiple = div + int(torch.distributed.get_rank(process_group) < mod) - # Only rank 0 will have bias - super().__init__( - local_multiple * multiple_of, - out_features, - bias=bias and rank == 0, - device=device, - dtype=dtype, - ) - self.process_group = process_group - self.sequence_parallel = sequence_parallel - - def forward(self, x): - """ - We're doing Tensor Parallel with sequence parallelism: we do the matmul and then - a reduce_scatter of the result. - """ - out = fused_dense_func(x, self.weight, self.bias) - reduce_fn = reduce_scatter if self.sequence_parallel else all_reduce - return reduce_fn(out, self.process_group) - - -class FusedMLPFunc(torch.autograd.Function): - @staticmethod - @custom_fwd - def forward( - ctx, - x, - weight1, - bias1, - weight2, - bias2, - activation="gelu_approx", - save_pre_act=True, - return_residual=False, - checkpoint_lvl=0, - heuristic=0, - process_group=None, - sequence_parallel=True, - ): - """ - If process_group is not None and sequence_parallel=True, we're doing Tensor Parallel - with sequence parallelism: we do an all_gather of x before doing the matmul. - If sequence_parallel=False, then the input is already gathered. - - checkpoint_lvl: - 0: no recomputation in the bwd - 1: recompute gelu_out / relu_out in the bwd - 2: recompute pre_act and gelu_out / relu_out in the bwd - """ - assert -1 <= heuristic <= 4 - assert activation in ["gelu_approx", "relu", "sqrelu"] - if activation == "sqrelu": - assert heuristic == -1 - if not save_pre_act: - checkpoint_lvl = 2 - assert checkpoint_lvl in [0, 1, 2] - ctx.return_residual = return_residual - ctx.process_group = process_group - ctx.sequence_parallel = sequence_parallel - ctx.checkpoint_lvl = checkpoint_lvl - ctx.activation = activation - ctx.heuristic = heuristic - - if torch.is_autocast_enabled(): - x = x.to(dtype=torch.get_autocast_gpu_dtype()) - x = x.contiguous() - if process_group is not None and sequence_parallel: - # We want to kick off the all_gather early, before weight dtype conversion - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - else: - total_x = x - - if torch.is_autocast_enabled(): - dtype = torch.get_autocast_gpu_dtype() - weight1, weight2 = [a.to(dtype=dtype) for a in [weight1, weight2]] - bias1 = bias1.to(dtype=dtype) if bias1 is not None else None - bias2 = bias2.to(dtype=dtype) if bias2 is not None else None - weight1 = weight1.contiguous() - bias1 = bias1.contiguous() if bias1 is not None else None - weight2 = weight2.contiguous() - bias2 = bias2.contiguous() if bias2 is not None else None - if process_group is not None and sequence_parallel: - handle_x.wait() - batch_shape, n = total_x.shape[:-1], total_x.shape[-1] - batch_dim = batch_shape.numel() - # https://github.com/pytorch/pytorch/blob/5b51849b48a7dbccd297286cc0110def4706f9e7/aten/src/ATen/native/cuda/Blas.cpp#L174 - if min(batch_dim, n, *weight1.shape, *weight2.shape) > 65535 * 32: - raise RuntimeError("fused_dense only supports matrix dims <= 2M") - if heuristic == -1: - pre_act = F.linear(total_x, weight1, bias1) - activation_fn = ( - partial(F.gelu, approximate="tanh") - if activation == "gelu_approx" - else (sqrelu_fwd if activation == "sqrelu" else F.relu) - ) - with torch.jit.fuser("fuser2"): - output1 = activation_fn(pre_act) - # This is before adding bias1 - # pre_act = F.linear(total_x.reshape(batch_dim, n), weight1) - # with torch.jit.fuser('fuser2'): - # output1 = bias_gelu(pre_act, bias1) - else: - is_gelu = activation == "gelu_approx" - output1, *rest = fused_dense_cuda.linear_act_forward( - total_x.reshape(batch_dim, n), weight1, bias1, is_gelu, save_pre_act, heuristic - ) - if save_pre_act: - pre_act = rest[0] - output2 = F.linear(output1, weight2, bias2) - if checkpoint_lvl == 0 or (checkpoint_lvl == 1 and activation == "relu"): - # For RELU the pre_act is very small (just a bit-mask) so we just save it - ctx.save_for_backward(x, weight1, weight2, pre_act, output1) - elif checkpoint_lvl == 1: - ctx.save_for_backward(x, weight1, weight2, pre_act) - elif checkpoint_lvl == 2: - ctx.save_for_backward(x, weight1, weight2, bias1) - output2 = output2.reshape(*batch_shape, output2.shape[-1]) - return output2 if not return_residual else (output2, x) - - @staticmethod - @custom_bwd - def backward(ctx, grad_output, *args): - grad_output = grad_output.contiguous() - checkpoint_lvl = ctx.checkpoint_lvl - activation = ctx.activation - activation_fn = ( - partial(F.gelu, approximate="tanh") - if activation == "gelu_approx" - else (sqrelu_fwd if activation == "sqrelu" else F.relu) - ) - if ctx.return_residual: - (grad_input,) = args - grad_input = grad_input.contiguous() - process_group = ctx.process_group - sequence_parallel = ctx.sequence_parallel - x, weight1, weight2, *rest = ctx.saved_tensors - if process_group is None or not sequence_parallel: - total_x = x - batch_shape = grad_output.shape[:-1] - batch_dim = batch_shape.numel() - if checkpoint_lvl in [0, 1]: - if process_group is not None and sequence_parallel: - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - if checkpoint_lvl == 0 or (checkpoint_lvl == 1 and activation == "relu"): - pre_act, output1 = rest - elif checkpoint_lvl == 1: - (pre_act,) = rest - with torch.jit.fuser("fuser2"): - output1 = activation_fn(pre_act) - elif checkpoint_lvl == 2: - (bias1,) = rest - if process_group is not None and sequence_parallel: - total_x, _ = all_gather_raw(x, process_group) - if ctx.heuristic == -1: - pre_act = F.linear(total_x, weight1, bias1) - with torch.jit.fuser("fuser2"): - output1 = activation_fn(pre_act) - else: - output1, pre_act = fused_dense_cuda.linear_act_forward( - total_x.reshape(batch_dim, total_x.shape[-1]), - weight1, - bias1, - activation == "gelu_approx", - True, - ctx.heuristic, - ) - - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - output1 = output1.reshape(batch_dim, output1.shape[-1]) - pre_act = pre_act.reshape(batch_dim, pre_act.shape[-1]) - if ctx.needs_input_grad[3]: - grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad( - output1, grad_output, ctx.needs_input_grad[4] - ) - else: - grad_weight2 = None - grad_bias2 = grad_output if ctx.needs_input_grad[4] else None - if ctx.heuristic == -1: - # grad_pre_act = matmul_dgelu(grad_output, weight2, pre_act) - grad_output1 = F.linear(grad_output, weight2.t()) - activation_grad_fn = ( - gelu_bwd - if activation == "gelu_approx" - else (sqrelu_bwd if activation == "sqrelu" else relu_bwd) - ) - with torch.jit.fuser("fuser2"): - grad_pre_act = activation_grad_fn(grad_output1, pre_act) - else: - # The cublasLt epilogue has to compute both gelu/relu grad and bias grad, we can't - # just compute gelu/relu grad - grad_pre_act, grad_bias1 = fused_dense_cuda.bias_act_linear_dgrad_bgrad( - weight2, grad_output, pre_act, activation == "gelu_approx", ctx.heuristic - ) - if not ctx.needs_input_grad[2]: - grad_bias1 = None - if ctx.needs_input_grad[0]: - if not ctx.return_residual: - grad_input = F.linear(grad_pre_act, weight1.t()) - else: - grad_input = torch.addmm( - grad_input.reshape(batch_dim, grad_input.shape[-1]), grad_pre_act, weight1 - ) - grad_input = grad_input.reshape(*batch_shape, grad_input.shape[-1]) - if process_group is not None: - reduce_fn = reduce_scatter_raw if sequence_parallel else all_reduce_raw - grad_input, handle_grad_input = reduce_fn(grad_input, process_group, async_op=True) - else: - grad_input = None - if ctx.heuristic == -1: - if ctx.needs_input_grad[1]: - if process_group is not None and sequence_parallel and checkpoint_lvl != 2: - handle_x.wait() - grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_wgrad( - total_x.reshape(batch_dim, total_x.shape[-1]), - grad_pre_act, - ctx.needs_input_grad[2], - ) - else: - grad_weight1 = None - grad_bias1 = grad_pre_act if ctx.needs_input_grad[2] else None - else: - if ctx.needs_input_grad[1]: - if process_group is not None and sequence_parallel and checkpoint_lvl != 2: - handle_x.wait() - grad_weight1 = F.linear( - grad_pre_act.t(), total_x.reshape(batch_dim, total_x.shape[-1]).t() - ) - else: - grad_weight1 = None - if process_group is not None and ctx.needs_input_grad[0]: - handle_grad_input.wait() - return ( - grad_input, - grad_weight1, - grad_bias1, - grad_weight2, - grad_bias2, - None, - None, - None, - None, - None, - None, - None, - ) - - -def fused_mlp_func( - x: Tensor, - weight1: Tensor, - weight2: Tensor, - bias1: Optional[Tensor] = None, - bias2: Optional[Tensor] = None, - activation: str = "gelu_approx", - save_pre_act: bool = True, - return_residual: bool = False, - checkpoint_lvl: int = 0, - heuristic: int = 0, - process_group: Optional[ProcessGroup] = None, - sequence_parallel: bool = True, -): - assert activation in ["gelu_approx", "relu", "sqrelu"] - dtype_eligible = x.dtype in [torch.float16, torch.bfloat16] or ( - x.dtype == torch.float32 and torch.is_autocast_enabled() - ) - # If we save pre-activation, dimension must be divisible by 128 (relu) or 8 (gelu) - dim_eligible = not save_pre_act or (x.shape[-1] % (128 if activation == "relu" else 8) == 0) - if ( - x.is_cuda - and weight1.is_cuda - and weight2.is_cuda - and (bias1 is None or bias1.is_cuda) - and (bias2 is None or bias2.is_cuda) - and dtype_eligible - and dim_eligible - ): - return FusedMLPFunc.apply( - x, - weight1, - bias1, - weight2, - bias2, - activation, - save_pre_act, - return_residual, - checkpoint_lvl, - heuristic, - process_group, - sequence_parallel, - ) - else: - assert process_group is None - pre_act = F.linear(x, weight1, bias1) - activation_fn = ( - partial(F.gelu, approximate="tanh") - if activation == "gelu_approx" - else partial(F.relu, inplace=True) - ) - output1 = activation_fn(pre_act) - output2 = F.linear(output1, weight2, bias2) - return output2 if not return_residual else (output2, x) - - -class FusedMLP(nn.Module): - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - bias1=True, - bias2=True, - activation="gelu_approx", - return_residual=False, - checkpoint_lvl=0, - heuristic="auto", - device=None, - dtype=None, - ): - """ - If process_group is not None, we're doing Tensor Parallel with sequence parallelism: - we do an all_gather of x before doing the matmul, gelu, then matmul. - Finally we do a reduce_scatter of the output. - - checkpoint_lvl (increasing lvl means slower but more memory saving): - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute pre_act and gelu_out in the bwd - heuristic: - -1: don't fuse gemm + gelu (separate kernel) - 0..4: use this heuristic for the algo section in the fused gemm + gelu - 'auto': heuristic will be picked automatically: - For CUDA >= 11.8, we set heuristic=0 for both fp16 and bf16 for best perf. - For CUDA <= 11.7, we set heuristic=1 for fp16 and heuristic=-1 for bf16. - For H100, we set heuristic=-1 for both fp16 and bf16 as the fused cuBlasLt implementation - is slower than the unfused version. - return_residual: whether to return the input x along with the output. This is for - performance reason: for post-norm architecture, returning the input allows us - to fuse the backward of nn.Linear with the residual connection. - """ - assert checkpoint_lvl in [0, 1, 2] - assert activation in ["gelu_approx", "relu", "sqrelu"] - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features * 4 - self.activation = activation - self.return_residual = return_residual - self.checkpoint_lvl = checkpoint_lvl - self.heuristic = heuristic if activation != "sqrelu" else -1 - self.fc1 = nn.Linear(in_features, hidden_features, bias=bias1, **factory_kwargs) - self.fc2 = nn.Linear(hidden_features, out_features, bias=bias2, **factory_kwargs) - - def forward(self, x, process_group=None): - dtype = x.dtype if not torch.is_autocast_enabled() else torch.get_autocast_gpu_dtype() - if self.heuristic == "auto": - if self.activation == "gelu_approx": - if torch.cuda.get_device_capability("cuda") == (9, 0): - heuristic = -1 - else: - cuda_ver = tuple(map(int, torch.version.cuda.split("."))) - heuristic = 0 if cuda_ver >= (11, 8) else (1 if dtype == torch.float16 else -1) - else: - heuristic = 0 - else: - heuristic = self.heuristic - out = fused_mlp_func( - x, - self.fc1.weight, - self.fc2.weight, - self.fc1.bias, - self.fc2.bias, - activation=self.activation, - save_pre_act=self.training, - return_residual=self.return_residual, - checkpoint_lvl=self.checkpoint_lvl, - heuristic=heuristic, - process_group=process_group, - ) - if self.return_residual: - out, x = out - if process_group is not None: - out = reduce_scatter(out, process_group) - return out if not self.return_residual else (out, x) - - -class ParallelFusedMLP(nn.Module): - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - activation="gelu_approx", - process_group: ProcessGroup = None, - bias1=True, - bias2=True, - sequence_parallel=True, - checkpoint_lvl=0, - heuristic="auto", - device=None, - dtype=None, - ): - """ - process_group is required. We're doing Tensor Parallel with sequence parallelism: - we do an all_gather of x before doing the matmul, gelu, then matmul. - Finally we do a reduce_scatter of the output. - - checkpoint_lvl (increasing lvl means slower but more memory saving): - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute pre_act and gelu_out in the bwd - heuristic: - -1: don't fuse gemm + gelu (separate kernel) - 0..4: use this heuristic for the algo section in the fused gemm + gelu - 'auto': heuristic will be picked automatically: - For CUDA >= 11.8, we set heuristic=0 for both fp16 and bf16 for best perf. - For CUDA <= 11.7, we set heuristic=1 for fp16 and heuristic=-1 for bf16. - """ - assert checkpoint_lvl in [0, 1, 2] - assert activation in ["gelu_approx", "relu", "sqrelu"] - assert process_group is not None - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features * 4 - self.activation = activation - self.process_group = process_group - self.sequence_parallel = sequence_parallel - self.checkpoint_lvl = checkpoint_lvl - self.heuristic = heuristic if activation != "sqrelu" else -1 - self.fc1 = ColumnParallelLinear( - in_features, hidden_features, process_group, bias=bias1, **factory_kwargs - ) - self.fc2 = RowParallelLinear( - hidden_features, out_features, process_group, bias=bias2, **factory_kwargs - ) - - def forward(self, x): - dtype = x.dtype if not torch.is_autocast_enabled() else torch.get_autocast_gpu_dtype() - if self.heuristic == "auto": - if self.activation == "gelu_approx": - cuda_ver = tuple(map(int, torch.version.cuda.split("."))) - heuristic = 0 if cuda_ver >= (11, 8) else (1 if dtype == torch.float16 else -1) - else: - heuristic = 0 - else: - heuristic = self.heuristic - out = fused_mlp_func( - x, - self.fc1.weight, - self.fc2.weight, - self.fc1.bias, - self.fc2.bias, - activation=self.activation, - save_pre_act=self.training, - checkpoint_lvl=self.checkpoint_lvl, - heuristic=heuristic, - process_group=self.process_group, - sequence_parallel=self.sequence_parallel, - ) - reduce_fn = reduce_scatter if self.sequence_parallel else all_reduce - return reduce_fn(out, self.process_group) diff --git a/build/torch29-cxx11-cu128-aarch64-linux/ops/layer_norm.py b/build/torch29-cxx11-cu128-aarch64-linux/ops/layer_norm.py deleted file mode 100644 index 4b6cd798fd02844ef9cd3897f8ab95e490e638bf..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-aarch64-linux/ops/layer_norm.py +++ /dev/null @@ -1,800 +0,0 @@ -# Copyright (c) 2022, Tri Dao. -# Adapted from https://github.com/NVIDIA/apex/blob/master/apex/contrib/layer_norm/layer_norm.py - -import dropout_layer_norm -import torch -from torch.nn import init - - -def maybe_align(x, alignment_in_bytes=16): - """Assume that x already has last dim divisible by alignment_in_bytes""" - # TD [2023-07-04] I'm not 100% sure that clone will align the memory - # https://discuss.pytorch.org/t/how-to-ensure-that-tensor-data-ptr-is-aligned-to-16-bytes/183440 - return x if x.data_ptr() % alignment_in_bytes == 0 else x.clone() - - -def _dropout_add_layer_norm_forward( - x0, - residual, - gamma, - beta, - rowscale, - colscale, - dropout_p, - epsilon, - residual_in_fp32=False, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes""" - hidden_size = gamma.numel() - x0mat = x0.view((-1, hidden_size)) - residualmat = residual.view((-1, hidden_size)) if residual is not None else None - rowscale = rowscale.view(-1) if rowscale is not None else None - zmat, xmat, dmask, mu, rsigma = dropout_layer_norm.dropout_add_ln_fwd( - x0mat, - residualmat, - gamma, - beta, - rowscale, - colscale, - None, - None, - dropout_p, - epsilon, - 1.0, - 0, - None, - residual_in_fp32, - is_rms_norm, - ) - # dmask is None if dropout_p == 0.0 - # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype - return zmat, xmat if xmat is not None else x0mat, dmask, mu, rsigma - - -def _dropout_add_layer_norm_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - rowscale, - colscale, - dropout_p, - has_residual, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes - dx == None means that it was a post-norm architecture - (x = drop(x0) + residual was not returned in the fwd). - x0 must not be None if we have colscale. - """ - hidden_size = gamma.numel() - xmat = x.view((-1, hidden_size)) - dzmat = dz.view(xmat.shape) - dxmat = dx.view(xmat.shape) if dx is not None else None - x0mat = x0.view((-1, hidden_size)) if x0 is not None else None - rowscale = rowscale.view(-1) if rowscale is not None else None - if colscale is not None: - assert x0 is not None, "x0 is required to compute the gradient of colscale" - dx0mat, dresidualmat, dgamma, dbeta, _, _, *rest = dropout_layer_norm.dropout_add_ln_bwd( - dzmat, - dxmat, - xmat, - x0mat, - dmask, - mu, - rsigma, - gamma, - rowscale, - colscale, - None, - None, - dropout_p, - 1.0, - 0, - has_residual, - is_rms_norm, - ) - # dresidualmat is None if not has_residual - if colscale is None: - return dx0mat, dresidualmat, dgamma, dbeta - else: - dcolscale = rest[0] - return dx0mat, dresidualmat, dgamma, dbeta, dcolscale - - -def _dropout_add_layer_norm_subset_forward( - x0, - residual, - gamma, - beta, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32=False, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes""" - hidden_size = gamma.numel() - x0mat = x0.view((-1, hidden_size)) - residualmat = residual.view((-1, hidden_size)) if residual is not None else None - x0_subset = x0_subset.view(-1) if x0_subset is not None else None - out_subset = out_subset.view(-1) if out_subset is not None else None - zmat, xmat, dmask, mu, rsigma = dropout_layer_norm.dropout_add_ln_fwd( - x0mat, - residualmat, - gamma, - beta, - None, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - None, - residual_in_fp32, - is_rms_norm, - ) - # dmask is None if dropout_p == 0.0 - # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype - return zmat, xmat if xmat is not None else x0mat, dmask, mu, rsigma - - -def _dropout_add_layer_norm_subset_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - colscale, - x0_subset, - out_subset, - dropout_p, - rowscale_const, - x0_numrows, - has_residual, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes - dx == None means that it was a post-norm architecture - (x = drop(x0) + residual was not returned in the fwd). - x0 must not be None if we have colscale. - """ - hidden_size = gamma.numel() - xmat = x.view((-1, hidden_size)) - dzmat = dz.view(-1, hidden_size) - dxmat = dx.view(xmat.shape) if dx is not None else None - x0mat = x0.view((-1, hidden_size)) if x0 is not None else None - x0_subset = x0_subset.view(-1) if x0_subset is not None else None - out_subset = out_subset.view(-1) if out_subset is not None else None - if colscale is not None: - assert x0 is not None, "x0 is required to compute the gradient of colscale" - dx0mat, dresidualmat, dgamma, dbeta, _, _, *rest = dropout_layer_norm.dropout_add_ln_bwd( - dzmat, - dxmat, - xmat, - x0mat, - dmask, - mu, - rsigma, - gamma, - None, - colscale, - x0_subset, - out_subset, - dropout_p, - rowscale_const, - x0_numrows, - has_residual, - is_rms_norm, - ) - # dresidualmat is None if not has_residual - if colscale is None: - return dx0mat, dresidualmat, dgamma, dbeta - else: - dcolscale = rest[0] - return dx0mat, dresidualmat, dgamma, dbeta, dcolscale - - -def _dropout_add_layer_norm_parallel_residual_forward( - x0, - x1, - residual, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - residual_in_fp32=False, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes""" - hidden_size = gamma0.numel() - x0mat = x0.view((-1, hidden_size)) - x1mat = x1.view((-1, hidden_size)) if x1 is not None else None - residualmat = residual.view((-1, hidden_size)) if residual is not None else None - ( - z0mat, - z1mat, - xmat, - dmask0, - dmask1, - mu, - rsigma, - ) = dropout_layer_norm.dropout_add_ln_parallel_residual_fwd( - x0mat, - x1mat, - residualmat, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - None, - residual_in_fp32, - is_rms_norm, - ) - # dmask0 and dmask1 are None if dropout_p == 0.0 - # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype - return z0mat, z1mat, xmat if xmat is not None else x0mat, dmask0, dmask1, mu, rsigma - - -def _dropout_add_layer_norm_parallel_residual_backward( - dz0, - dz1, - dx, - x, - dmask0, - dmask1, - mu, - rsigma, - gamma0, - gamma1, - dropout_p, - has_x1, - has_residual, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes - dx == None means that it was a post-norm architecture - (x = drop(x0) + residual was not returned in the fwd). - """ - hidden_size = gamma0.numel() - xmat = x.view((-1, hidden_size)) - dz0mat = dz0.view(xmat.shape) - dz1mat = dz1.view(xmat.shape) if dz1 is not None else None - dxmat = dx.view(xmat.shape) if dx is not None else None - ( - dx0mat, - dx1mat, - dresidualmat, - dgamma0, - dbeta0, - dgamma1, - dbeta1, - *rest, - ) = dropout_layer_norm.dropout_add_ln_parallel_residual_bwd( - dz0mat, - dz1mat, - dxmat, - xmat, - dmask0, - dmask1, - mu, - rsigma, - gamma0, - gamma1, - dropout_p, - has_x1, - has_residual, - is_rms_norm, - ) - # dresidualmat is None if not has_residual - return dx0mat, dx1mat, dresidualmat, dgamma0, dbeta0, dgamma1, dbeta1 - - -class DropoutAddLayerNormFn(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x0, - residual, - gamma, - beta, - rowscale, - colscale, - dropout_p, - epsilon, - residual_in_fp32=False, - prenorm=False, - is_rms_norm=False, - return_dmask=False, - ): - x0 = maybe_align(x0.contiguous(), 16) - residual = maybe_align(residual.contiguous(), 16) if residual is not None else None - gamma = maybe_align(gamma.contiguous(), 16) - beta = maybe_align(beta.contiguous(), 16) if beta is not None else None - rowscale = maybe_align(rowscale.contiguous(), 16) if rowscale is not None else None - colscale = maybe_align(colscale.contiguous(), 16) if colscale is not None else None - zmat, xmat, dmask, mu, rsigma = _dropout_add_layer_norm_forward( - x0, - residual, - gamma, - beta, - rowscale, - colscale, - dropout_p, - epsilon, - residual_in_fp32, - is_rms_norm, - ) - # Only need to save x0 if we need to compute gradient wrt colscale - x0_saved = x0 if colscale is not None else None - ctx.save_for_backward( - xmat.view(x0.shape), x0_saved, dmask, gamma, mu, rsigma, rowscale, colscale - ) - ctx.prenorm = prenorm - ctx.dropout_p = dropout_p - ctx.has_residual = residual is not None - ctx.is_rms_norm = is_rms_norm - ctx.has_beta = beta is not None - if not return_dmask: - return ( - zmat.view(x0.shape) if not prenorm else (zmat.view(x0.shape), xmat.view(x0.shape)) - ) - else: - dmask = ( - dmask.view(x0.shape) - if dropout_p > 0.0 - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - ctx.mark_non_differentiable(dmask) - return ( - (zmat.view(x0.shape), dmask) - if not prenorm - else (zmat.view(x0.shape), xmat.view(x0.shape), dmask) - ) - - @staticmethod - def backward(ctx, dz, *args): - # assert dz.is_contiguous() - dz = maybe_align(dz.contiguous(), 16) # this happens! - dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None - x, x0, dmask, gamma, mu, rsigma, rowscale, colscale = ctx.saved_tensors - # x0 is None if colscale is None - dropout_p = ctx.dropout_p - has_residual = ctx.has_residual - dx0mat, dresidualmat, dgamma, dbeta, *rest = _dropout_add_layer_norm_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - rowscale, - colscale, - dropout_p, - has_residual, - ctx.is_rms_norm, - ) - dx0 = dx0mat.view(x.shape) - dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None - dcolscale = rest[0] if colscale is not None else None - return ( - dx0, - dresidual, - dgamma, - dbeta if ctx.has_beta else None, - None, - dcolscale, - None, - None, - None, - None, - None, - None, - ) - - -class DropoutAddLayerNormSubsetFn(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x0, - residual, - gamma, - beta, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32=False, - prenorm=False, - is_rms_norm=False, - return_dmask=False, - ): - x0 = maybe_align(x0.contiguous(), 16) - residual = maybe_align(residual.contiguous(), 16) if residual is not None else None - gamma = maybe_align(gamma.contiguous(), 16) - beta = maybe_align(beta.contiguous(), 16) if beta is not None else None - colscale = maybe_align(colscale.contiguous(), 16) if colscale is not None else None - zmat, xmat, dmask, mu, rsigma = _dropout_add_layer_norm_subset_forward( - x0, - residual, - gamma, - beta, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32, - is_rms_norm, - ) - # Only need to save x0 if we need to compute gradient wrt colscale - x0_saved = x0 if colscale is not None else None - x_shape = (-1, *x0.shape[1:]) - ctx.save_for_backward( - xmat.view(x_shape), x0_saved, dmask, gamma, mu, rsigma, colscale, x0_subset, out_subset - ) - ctx.prenorm = prenorm - ctx.dropout_p = dropout_p - ctx.rowscale_const = rowscale_const - ctx.x0_numrows = x0.shape[:-1].numel() - ctx.has_residual = residual is not None - ctx.is_rms_norm = is_rms_norm - ctx.has_beta = beta is not None - z_shape = (-1, *x0.shape[1:]) - if not return_dmask: - return zmat.view(z_shape) if not prenorm else (zmat.view(z_shape), xmat.view(x0.shape)) - else: - z = zmat.view(z_shape) - dmask = ( - dmask.view(x0.shape) - if dropout_p > 0.0 - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - ctx.mark_non_differentiable(dmask) - return (z, dmask) if not prenorm else (z, xmat.view(x_shape), dmask) - - @staticmethod - def backward(ctx, dz, *args): - # assert dz.is_contiguous() - dz = maybe_align(dz.contiguous(), 16) # this happens! - dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None - x, x0, dmask, gamma, mu, rsigma, colscale, x0_subset, out_subset = ctx.saved_tensors - # x0 is None if colscale is None - dropout_p = ctx.dropout_p - has_residual = ctx.has_residual - dx0mat, dresidualmat, dgamma, dbeta, *rest = _dropout_add_layer_norm_subset_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - colscale, - x0_subset, - out_subset, - dropout_p, - ctx.rowscale_const, - ctx.x0_numrows, - has_residual, - ctx.is_rms_norm, - ) - dx0 = dx0mat.view(-1, *x.shape[1:]) - dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None - dcolscale = rest[0] if colscale is not None else None - return ( - dx0, - dresidual, - dgamma, - dbeta if ctx.has_beta else None, - dcolscale, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) - - -class DropoutAddLayerNormParallelResidualFn(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x0, - x1, - residual, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - residual_in_fp32=False, - prenorm=False, - is_rms_norm=False, - return_dmask=False, - ): - x0 = maybe_align(x0.contiguous(), 16) - x1 = maybe_align(x1.contiguous(), 16) if x1 is not None else None - residual = maybe_align(residual.contiguous(), 16) if residual is not None else None - gamma0 = maybe_align(gamma0.contiguous(), 16) - beta0 = maybe_align(beta0.contiguous(), 16) if beta0 is not None else None - gamma1 = maybe_align(gamma1.contiguous(), 16) if gamma1 is not None else None - beta1 = maybe_align(beta1.contiguous(), 16) if beta1 is not None else None - ( - z0mat, - z1mat, - xmat, - dmask0, - dmask1, - mu, - rsigma, - ) = _dropout_add_layer_norm_parallel_residual_forward( - x0, - x1, - residual, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - residual_in_fp32, - is_rms_norm, - ) - ctx.save_for_backward(xmat.view(x0.shape), dmask0, dmask1, gamma0, gamma1, mu, rsigma) - ctx.prenorm = prenorm - ctx.dropout_p = dropout_p - ctx.has_x1 = x1 is not None - ctx.has_residual = residual is not None - ctx.is_rms_norm = is_rms_norm - ctx.has_beta = beta0 is not None - z = (z0mat.view(x0.shape), z1mat.view(x0.shape) if z1mat is not None else None) - if not return_dmask: - return z if not prenorm else (*z, xmat.view(x0.shape)) - else: - dmask0 = ( - dmask0.view(x0.shape) - if dropout_p > 0.0 - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - dmask1 = ( - dmask1.view(x0.shape) - if dropout_p > 0.0 and x1 is not None - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - ctx.mark_non_differentiable(dmask0) - ctx.mark_non_differentiable(dmask1) - return ( - (*z, dmask0, dmask1) if not prenorm else (*z, xmat.view(x0.shape), dmask0, dmask1) - ) - - @staticmethod - def backward(ctx, dz0, dz1, *args): - dz0 = maybe_align(dz0.contiguous(), 16) # this happens! - dz1 = maybe_align(dz1.contiguous(), 16) if dz1 is not None else None - dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None - x, dmask0, dmask1, gamma0, gamma1, mu, rsigma = ctx.saved_tensors - dropout_p = ctx.dropout_p - has_x1 = ctx.has_x1 - has_residual = ctx.has_residual - ( - dx0mat, - dx1mat, - dresidualmat, - dgamma0, - dbeta0, - dgamma1, - dbeta1, - ) = _dropout_add_layer_norm_parallel_residual_backward( - dz0, - dz1, - dx, - x, - dmask0, - dmask1, - mu, - rsigma, - gamma0, - gamma1, - dropout_p, - has_x1, - has_residual, - ctx.is_rms_norm, - ) - dx0 = dx0mat.view(x.shape) - dx1 = dx1mat.view(x.shape) if dx1mat is not None else None - dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None - return ( - dx0, - dx1, - dresidual, - dgamma0, - dbeta0 if ctx.has_beta else None, - dgamma1, - dbeta1 if ctx.has_beta else None, - None, - None, - None, - None, - None, - None, - ) - - -def layer_norm(x, weight, bias, epsilon): - return DropoutAddLayerNormFn.apply(x, None, weight, bias, None, None, 0.0, epsilon, False) - - -def dropout_add_layer_norm( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - rowscale=None, - layerscale=None, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormFn.apply( - x0, - residual, - weight, - bias, - rowscale, - layerscale, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - False, - return_dropout_mask, - ) - - -def dropout_add_layer_norm_subset( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - layerscale=None, - x0_subset=None, - out_subset=None, - rowscale_const=1.0, - out_numrows=0, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormSubsetFn.apply( - x0, - residual, - weight, - bias, - layerscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32, - prenorm, - False, - return_dropout_mask, - ) - - -def dropout_add_layer_norm_parallel_residual( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormParallelResidualFn.apply( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - False, - return_dropout_mask, - ) - - -class DropoutAddLayerNorm(torch.nn.Module): - def __init__( - self, - hidden_size, - prenorm=False, - p=0.0, - eps=1e-5, - residual_in_fp32=False, - device=None, - dtype=None, - ): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.prenorm = prenorm - self.p = p - self.eps = eps - self.residual_in_fp32 = residual_in_fp32 - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.bias = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.reset_parameters() - - def reset_parameters(self): - init.ones_(self.weight) - init.zeros_(self.bias) - - def forward(self, x0, residual=None): - return dropout_add_layer_norm( - x0, - residual, - self.weight, - self.bias, - self.p if self.training else 0.0, - self.eps, - prenorm=self.prenorm, - residual_in_fp32=self.residual_in_fp32, - ) diff --git a/build/torch29-cxx11-cu128-aarch64-linux/ops/rms_norm.py b/build/torch29-cxx11-cu128-aarch64-linux/ops/rms_norm.py deleted file mode 100644 index 068348d61290e3839dd082b540d898578ba1e8e2..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-aarch64-linux/ops/rms_norm.py +++ /dev/null @@ -1,174 +0,0 @@ -# Copyright (c) 2022, Tri Dao. -# Adapted from https://github.com/NVIDIA/apex/blob/master/apex/contrib/layer_norm/layer_norm.py - -import torch -from torch.nn import init - -from flash_attn.ops.layer_norm import ( - DropoutAddLayerNormFn, - DropoutAddLayerNormParallelResidualFn, - DropoutAddLayerNormSubsetFn, -) - - -def rms_norm(x, weight, epsilon): - return DropoutAddLayerNormFn.apply( - x, None, weight, None, None, None, 0.0, epsilon, False, False, True - ) - - -def dropout_add_rms_norm( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - rowscale=None, - layerscale=None, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormFn.apply( - x0, - residual, - weight, - bias, - rowscale, - layerscale, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - True, - return_dropout_mask, - ) - - -def dropout_add_rms_norm_subset( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - layerscale=None, - x0_subset=None, - out_subset=None, - rowscale_const=1.0, - out_numrows=0, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormSubsetFn.apply( - x0, - residual, - weight, - bias, - layerscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32, - prenorm, - True, - return_dropout_mask, - ) - - -def dropout_add_rms_norm_parallel_residual( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormParallelResidualFn.apply( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - True, - return_dropout_mask, - ) - - -class RMSNorm(torch.nn.Module): - def __init__(self, hidden_size, eps=1e-5, device=None, dtype=None): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.eps = eps - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.register_parameter("bias", None) - self.reset_parameters() - - def reset_parameters(self): - init.ones_(self.weight) - - def forward(self, x): - return rms_norm(x, self.weight, self.eps) - - -class DropoutAddRMSNorm(torch.nn.Module): - def __init__( - self, - hidden_size, - prenorm=False, - p=0.0, - eps=1e-5, - residual_in_fp32=False, - device=None, - dtype=None, - ): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.prenorm = prenorm - self.p = p - self.eps = eps - self.residual_in_fp32 = residual_in_fp32 - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.register_parameter("bias", None) - self.reset_parameters() - - def reset_parameters(self): - init.ones_(self.weight) - - def forward(self, x0, residual=None): - return dropout_add_rms_norm( - x0, - residual, - self.weight, - None, - self.p if self.training else 0.0, - self.eps, - prenorm=self.prenorm, - residual_in_fp32=self.residual_in_fp32, - ) diff --git a/build/torch29-cxx11-cu128-aarch64-linux/ops/triton/__init__.py b/build/torch29-cxx11-cu128-aarch64-linux/ops/triton/__init__.py deleted file mode 100644 index 8b137891791fe96927ad78e64b0aad7bded08bdc..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-aarch64-linux/ops/triton/__init__.py +++ /dev/null @@ -1 +0,0 @@ - diff --git a/build/torch29-cxx11-cu128-aarch64-linux/ops/triton/cross_entropy.py b/build/torch29-cxx11-cu128-aarch64-linux/ops/triton/cross_entropy.py deleted file mode 100644 index 1b5a415b73f236f3e05fb14b9141959559e18526..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-aarch64-linux/ops/triton/cross_entropy.py +++ /dev/null @@ -1,330 +0,0 @@ -# Copyright (c) 2023, Tri Dao. - -from typing import Tuple, Optional, Union - -import torch -import torch.nn.functional as F - -import triton -import triton.language as tl - -# `all_gather_into_tensor` and `reduce_scatter_tensor` are new placeholders for -# `_all_gather_base` and `_reduce_scatter_base`. They require the most recent -# version of PyTorch. The following 2 lines are for backward compatibility with -# older PyTorch. -if "all_gather_into_tensor" not in dir(torch.distributed): - torch.distributed.all_gather_into_tensor = torch.distributed._all_gather_base - - -@triton.heuristics( - { - "HAS_SMOOTHING": lambda args: args["smoothing"] > 0.0, - } -) -@triton.jit -def cross_entropy_fwd_kernel( - loss_ptr, # data ptrs - lse_ptr, - z_loss_ptr, - logits_ptr, - labels_ptr, - smoothing, - logit_scale, - lse_square_scale, - ignore_index, - total_classes, - class_start_idx, # Useful for tensor parallel when each rank only has a subset of classes - n_cols, # shapes - logits_row_stride, # strides - BLOCK_SIZE: tl.constexpr, - HAS_SMOOTHING: tl.constexpr, - # if SPLIT (e.g. tensor parallel), don't include the LSE in the loss since it's not the final LSE - SPLIT: tl.constexpr, - PRECOMPUTED_LSE: tl.constexpr, # If LSE is already computed (also no smoothing and logit_scale == 1.0) -): - row_idx = tl.program_id(0) - logits_ptr = logits_ptr + row_idx * logits_row_stride.to(tl.int64) - sum_logits = 0.0 # For smoothing - if not PRECOMPUTED_LSE: - # Statistics for online softmax - m_i = -float("inf") - l_i = 0.0 - for col_offset in range(0, n_cols, BLOCK_SIZE): - cols = col_offset + tl.arange(0, BLOCK_SIZE) - logits = tl.load(logits_ptr + cols, mask=cols < n_cols, other=-float("inf")).to( - tl.float32 - ) * logit_scale - if HAS_SMOOTHING: - sum_logits += tl.sum(tl.where(cols < n_cols, logits, 0.0)) - m_i_new = tl.maximum(m_i, tl.max(logits)) - l_i = tl.exp(m_i - m_i_new) * l_i + tl.sum(tl.exp(logits - m_i_new)) - m_i = m_i_new - lse = tl.log(l_i) + m_i - tl.store(lse_ptr + row_idx, lse) - else: - lse = tl.load(lse_ptr + row_idx) - label_idx = tl.load(labels_ptr + row_idx) - if label_idx == ignore_index: - loss = 0.0 - z_loss = 0.0 - else: - label_idx -= class_start_idx - if label_idx >= 0 and label_idx < n_cols: - logits_label = tl.load(logits_ptr + label_idx) * logit_scale - if HAS_SMOOTHING: - loss = ( - (lse if not SPLIT else 0.0) - - smoothing * sum_logits / total_classes - - (1 - smoothing) * logits_label - ) - else: - loss = (lse if not SPLIT else 0.0) - logits_label - else: - # If label is out of bounds, we set the CE loss to 0.0. But we still want the smoothing loss - if HAS_SMOOTHING: - loss = smoothing * ((lse if not SPLIT else 0.0) - sum_logits / total_classes) - else: - loss = 0.0 - if not SPLIT: - z_loss = lse_square_scale * lse * lse - loss += z_loss - else: - z_loss = 0.0 - tl.store(loss_ptr + row_idx, loss) - if not SPLIT: - tl.store(z_loss_ptr + row_idx, z_loss) - - -@triton.heuristics( - { - "HAS_SMOOTHING": lambda args: args["smoothing"] > 0.0, - } -) -@triton.jit -def cross_entropy_bwd_kernel( - dlogits_ptr, # data ptrs - dloss_ptr, - logits_ptr, - lse_ptr, - labels_ptr, - smoothing, - logit_scale, - lse_square_scale, - ignore_index, - total_classes, - class_start_idx, # Useful for tensor parallel when each rank only has a subset of classes - n_cols, # shapes - logits_row_stride, # strides - dlogits_row_stride, - dloss_row_stride, - BLOCK_SIZE: tl.constexpr, - HAS_SMOOTHING: tl.constexpr, -): - row_idx = tl.program_id(0) - col_block_idx = tl.program_id(1) - logits_ptr = logits_ptr + row_idx * logits_row_stride.to(tl.int64) - dlogits_ptr = dlogits_ptr + row_idx * dlogits_row_stride.to(tl.int64) - col_offsets = col_block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - label_idx = tl.load(labels_ptr + row_idx) - if label_idx != ignore_index: - dloss = tl.load(dloss_ptr + row_idx * dloss_row_stride) - else: - dloss = 0.0 - logits = tl.load(logits_ptr + col_offsets, mask=col_offsets < n_cols, other=-float("inf")).to( - tl.float32 - ) * logit_scale - lse = tl.load(lse_ptr + row_idx) - probs = tl.exp(logits - lse) - probs += 2.0 * lse_square_scale * lse * probs - label_idx -= class_start_idx - if HAS_SMOOTHING: - smooth_positive = 1.0 - smoothing - smooth_negative = smoothing / total_classes - probs = tl.where(col_offsets == label_idx, probs - smooth_positive, probs) - smooth_negative - else: - probs = tl.where(col_offsets == label_idx, probs - 1.0, probs) - tl.store(dlogits_ptr + col_offsets, (dloss * logit_scale) * probs, mask=col_offsets < n_cols) - - -class CrossEntropyLoss(torch.autograd.Function): - - @staticmethod - def forward( - ctx, - logits, - labels, - precomputed_lse=None, - smoothing=0.0, - logit_scale=1.0, - lse_square_scale=0.0, - ignore_index=-100, - inplace_backward=False, - process_group=None, - ): - # For some reason Triton generates wrong code when labels has dtype long and its address - # is not aligned to 16 bytes. The ld.global.b64 seems to load the wrong label index. - if labels.dtype == torch.long and labels.data_ptr() % 16 != 0: - labels = F.pad(labels, (0, 1))[..., :-1] - assert labels.data_ptr() % 16 == 0 - assert logit_scale > 0.0 - n_rows, n_cols = logits.shape - assert labels.shape == (n_rows,) - world_size = 1 if process_group is None else torch.distributed.get_world_size(process_group) - total_classes = world_size * n_cols - rank = 0 if process_group is None else torch.distributed.get_rank(process_group) - class_start_idx = rank * n_cols - use_precomputed_lse = precomputed_lse is not None and logit_scale == 1.0 and smoothing == 0.0 - - if logits.stride(-1) != 1: - logits = logits.contiguous() - MAX_BLOCK_SIZE = 16 * 1024 - BLOCK_SIZE = min(triton.next_power_of_2(n_cols), MAX_BLOCK_SIZE) - num_warps = ( - 4 - if BLOCK_SIZE < 2048 - else (8 if BLOCK_SIZE < 8192 else (16 if BLOCK_SIZE < 128 * 1024 else 32)) - ) - losses = torch.empty(n_rows, dtype=torch.float, device=logits.device) - if use_precomputed_lse: - assert precomputed_lse.shape == (n_rows,) - lse = precomputed_lse.contiguous() - else: - lse = torch.empty(n_rows, dtype=torch.float, device=logits.device) - z_losses = torch.empty(n_rows, dtype=torch.float, device=logits.device) - # Need this, otherwise Triton tries to launch from cuda:0 and we get - # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?) - with torch.cuda.device(logits.device.index): - cross_entropy_fwd_kernel[(n_rows,)]( - losses, # data ptrs - lse, - z_losses, - logits, - labels, - smoothing, - logit_scale, - lse_square_scale, - ignore_index, - total_classes, - class_start_idx, - n_cols, # shapes - logits.stride(0), # strides - BLOCK_SIZE=BLOCK_SIZE, # constants - SPLIT=world_size > 1, - PRECOMPUTED_LSE=use_precomputed_lse, - num_warps=num_warps, - ) - - if world_size > 1: - # If there's no smoothing, if labels are in the vocab of this partition, losses contains - # - predicted logit, and 0 otherwise. - # If there's smoothing=0.1, for labels in the vocab of this partition, losses contains - # -0.9 * predicted logit - 0.1 * sum logit / total_classes. - # For labels not in the vocab of this partition, losses contains - # -0.1 * sum logit / total_classes. - if world_size > 1: - lse_allgather = torch.empty(world_size, n_rows, dtype=lse.dtype, device=lse.device) - torch.distributed.all_gather_into_tensor(lse_allgather, lse, group=process_group) - handle_losses = torch.distributed.all_reduce( - losses, op=torch.distributed.ReduceOp.SUM, group=process_group, async_op=True - ) - lse = torch.logsumexp(lse_allgather, dim=0) - handle_losses.wait() - # After the allreduce, if there's no smoothing, the total losses are - predicted_logit, - # we just have to add the (global) lse. - # If there's smoothing=0.1, the total losses are - # -0.9 * predicted_logit - 0.1 * sum logit / total_classes. - # Again, we just have to add the (global) lse. - losses += lse - if lse_square_scale != 0.0: - z_losses = lse_square_scale * lse.square() - z_losses.masked_fill_(labels == ignore_index, 0.0) - losses += z_losses - else: - z_losses = torch.zeros_like(losses) - losses.masked_fill_(labels == ignore_index, 0.0) - - ctx.save_for_backward(logits, lse, labels) - ctx.mark_non_differentiable(z_losses) - ctx.smoothing = smoothing - ctx.logit_scale = logit_scale - ctx.lse_square_scale = lse_square_scale - ctx.ignore_index = ignore_index - ctx.total_classes = total_classes - ctx.class_start_idx = class_start_idx - ctx.inplace_backward = inplace_backward - return losses, z_losses - - @staticmethod - def backward(ctx, grad_losses, grad_z_losses): - del grad_z_losses # z_losses are only for logging. - - logits, lse, labels = ctx.saved_tensors - dlogits = logits if ctx.inplace_backward else torch.empty_like(logits) - n_rows, n_cols = logits.shape - BLOCK_SIZE = min(triton.next_power_of_2(n_cols), 4 * 1024) - num_warps = 4 if BLOCK_SIZE < 2048 else (8 if BLOCK_SIZE < 8192 else 16) - grid = lambda META: (n_rows, triton.cdiv(n_cols, META["BLOCK_SIZE"])) # noqa - # Need this, otherwise Triton tries to launch from cuda:0 and we get - # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?) - with torch.cuda.device(logits.device.index): - cross_entropy_bwd_kernel[grid]( - dlogits, # data ptrs - grad_losses, - logits, - lse, - labels, - ctx.smoothing, - ctx.logit_scale, - ctx.lse_square_scale, - ctx.ignore_index, - ctx.total_classes, - ctx.class_start_idx, - n_cols, # shapes - logits.stride(0), # strides - dlogits.stride(0), - grad_losses.stride(0), - BLOCK_SIZE=BLOCK_SIZE, # constants - num_warps=num_warps, - ) - return dlogits, None, None, None, None, None, None, None, None, None - - -def cross_entropy_loss( - logits: torch.Tensor, - labels: torch.Tensor, - precomputed_lse: Optional[torch.Tensor] = None, - label_smoothing: float = 0.0, - logit_scale: float = 1.0, - lse_square_scale: float = 0.0, - ignore_index=-100, - inplace_backward: bool = False, - process_group=None, -) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Arguments: - logits: (batch, vocab_size) - labels: (batch,) - label_smoothing: float - logit_scale: float. Multiply logits by this scale before calculating the loss. - lse_square_scale: float. If > 0, we add lse_square_scale * lse(logits) ^ 2 to the loss. - This is also referred to as "z-loss". - ignore_index: int. If labels == ignore_index, the loss is set to 0.0. - inplace_backward: bool. If True, we do the backward pass in-place by modifying the logits. - This saves memory. - process_group: if not None, we're doing Tensor Parallel: each process is responsible for - one part of the vocab. The loss will be aggregated across processes. - Returns: - losses: (batch,), float - z_losses: (batch,), float - """ - return CrossEntropyLoss.apply( - logits, - labels, - precomputed_lse, - label_smoothing, - logit_scale, - lse_square_scale, - ignore_index, - inplace_backward, - process_group, - ) diff --git a/build/torch29-cxx11-cu128-aarch64-linux/ops/triton/k_activations.py b/build/torch29-cxx11-cu128-aarch64-linux/ops/triton/k_activations.py deleted file mode 100644 index efb83c358eb4a85d069ee340a3c83f418f9a805b..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-aarch64-linux/ops/triton/k_activations.py +++ /dev/null @@ -1,162 +0,0 @@ -# Adapted from https://github.com/facebookresearch/xformers/blob/main/xformers/triton/k_activations.py -# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved. -# -# This source code is licensed under the BSD license found in the -# LICENSE file in the root directory of this source tree. - -import math -from enum import Enum -from typing import Optional - -import triton -import triton.language as tl - -_sqrt2pi = math.sqrt(2.0 / math.pi) -_sqrt1_2 = math.sqrt(1.0 / 2) -_gaussian_pdf_normalization = 1.0 / math.sqrt(2 * math.pi) - - -class Activation(str, Enum): - SquaredReLU = "squared_relu" - GeLU = "gelu" - GeLUApprox = "gelu_approx" - LeakyReLU = "leaky_relu" - ReLU = "relu" - - -def get_triton_activation_kernel(activation: Optional[Activation]): - return ( - { - Activation.ReLU: relu, - Activation.LeakyReLU: leaky_relu, - Activation.GeLU: gelu, - Activation.GeLUApprox: gelu_approx, - Activation.SquaredReLU: squared_relu, - }[activation] - if activation - else None - ) - - -def get_triton_activation_bwd_kernel(activation: Optional[Activation]): - return ( - { - Activation.ReLU: relu_grad, - Activation.LeakyReLU: leaky_relu_grad, - Activation.GeLU: gelu_grad, - Activation.GeLUApprox: gelu_approx_grad, - Activation.SquaredReLU: squared_relu_grad, - }[activation] - if activation - else None - ) - - -@triton.jit -def tanh(x): - # Tanh is just a scaled sigmoid - return 2 * tl.sigmoid(2 * x) - 1 - - -@triton.jit -def cosh(x): - exp_x = tl.exp(x) - return (exp_x + 1.0 / exp_x) * 0.5 - - -# a Triton implementation of the most used activations -# See for instance http://arxiv.org/abs/1606.08415 for an overview - -# ReLU -@triton.jit -def relu(x): - """ - ReLU_ activation function - - .. _ReLU: https://pytorch.org/docs/stable/generated/torch.nn.ReLU.html - """ - zero = 0.0 - return tl.where(x >= 0, x, zero.to(x.dtype)) - - -@triton.jit -def relu_grad(x): - # ReLU is different from other activations - # in that it does not require the input to retrospectively compute its gradient - # here the input is the downstream gradient, and we return the upstream gradient directly - zero = 0.0 - one = 1.0 - return tl.where(x >= 0, one.to(x.dtype), zero.to(x.dtype)) - - -@triton.jit -def squared_relu(x): - """ - Squared ReLU activation, as proposed in the Primer_ paper. - - .. _Primer: https://arxiv.org/abs/2109.08668 - """ - x_ = relu(x) - return (x_ * x_).to(x.dtype) - - -@triton.jit -def squared_relu_grad(x): - return tl.where(x >= 0, 2.0 * x, 0.0) - - -# Leaky ReLU -@triton.jit -def leaky_relu(x): - """ - LeakyReLU_ activation - - .. _LeakyReLU: https://pytorch.org/docs/stable/generated/torch.nn.LeakyReLU.html - """ - scale = 0.01 + 0.0 - scale = scale.to(x.dtype) - return tl.where(x >= 0, x, scale * x) - - -@triton.jit -def leaky_relu_grad(x): - min_grad = 0.01 - max_grad = 1 - - min_grad = min_grad.to(x.dtype) - max_grad = max_grad.to(x.dtype) - - return tl.where(x >= 0, max_grad, min_grad) - - -@triton.jit -def gelu(x): - """Gaussian Error Linear Unit (GELU)""" - return x * 0.5 * (1.0 + tl.libdevice.erf(x * _sqrt1_2)) - - -@triton.jit -def gelu_grad(x): - cdf = 0.5 * (1.0 + tl.libdevice.erf(x * _sqrt1_2)) - pdf = tl.exp(-0.5 * x * x) * _gaussian_pdf_normalization - return cdf + x * pdf - - -@triton.jit -def gelu_approx(x): - """ - GeLU_ activation - Gaussian error linear unit, with tanh approximation - - .. _GeLU: https://arxiv.org/pdf/1606.08415.pdf - """ - return 0.5 * x * (1.0 + tanh(_sqrt2pi * x * (1.0 + 0.044715 * x * x))) - - -@triton.jit -def gelu_approx_grad(x): - # CREDITS: Fast implementation proposed in - # https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/fused_bias_gelu.py#L30 - tanh_out = tanh(0.79788456 * x * (1 + 0.044715 * x * x)) - return 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( - 1 + tanh_out - ) diff --git a/build/torch29-cxx11-cu128-aarch64-linux/ops/triton/linear.py b/build/torch29-cxx11-cu128-aarch64-linux/ops/triton/linear.py deleted file mode 100644 index a8966dbc345ab0e593df0124451ee7be3dae131a..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-aarch64-linux/ops/triton/linear.py +++ /dev/null @@ -1,594 +0,0 @@ -# Adapted from https://github.com/ELS-RD/kernl/blob/main/src/kernl/implementations/linear_layer.py -# and https://github.com/openai/triton/blob/master/python/triton/ops/matmul.py -from typing import Optional - -import torch -import triton -import triton.language as tl -from triton.ops.matmul_perf_model import early_config_prune, estimate_matmul_time - -from flash_attn.ops.triton.k_activations import ( - gelu, - gelu_approx, - gelu_approx_grad, - gelu_grad, - squared_relu, - squared_relu_grad, -) - -# CREDITS: Initially inspired by the Triton tutorial on matrix multiplications - - -def init_to_zero(name): - return lambda nargs: nargs[name].zero_() - - -def get_configs_io_bound(): - configs = [] - for num_stages in [2, 3, 4, 5, 6]: - for block_m in [16, 32]: - for block_k in [32, 64]: - for block_n in [32, 64, 128, 256]: - num_warps = 2 if block_n <= 64 else 4 - configs.append( - triton.Config( - { - "BLOCK_M": block_m, - "BLOCK_N": block_n, - "BLOCK_K": block_k, - "SPLIT_K": 1, - }, - num_stages=num_stages, - num_warps=num_warps, - ) - ) - # split_k not used - # for split_k in [2, 4, 8, 16]: - # configs.append(triton.Config( - # {'BLOCK_M': block_m, 'BLOCK_N': block_n, 'BLOCK_K': block_k, 'SPLIT_K': split_k}, - # num_stages=num_stages, num_warps=num_warps, pre_hook=init_to_zero('C'))) - return configs - - -@triton.autotune( - configs=[ - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - # good for int8 - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=4, - num_warps=4, - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - ] - + get_configs_io_bound(), - key=["CACHE_KEY_M", "CACHE_KEY_N", "CACHE_KEY_K"], - prune_configs_by={ - "early_config_prune": early_config_prune, - "perf_model": estimate_matmul_time, - "top_k": 10, - }, -) -@triton.heuristics( - { - "EVEN_K": lambda args: args["K"] % (args["BLOCK_K"] * args["SPLIT_K"]) == 0, - } -) -@triton.jit -def kernel_fwd( - C, # Pointers to matrices - ACT_INPUT, - A, - B, - bias, - # Matrix dimensions - M, - N, - K, - CACHE_KEY_M, - CACHE_KEY_N, - CACHE_KEY_K, - # The stride variables represent how much to increase the ptr by when moving by 1 - # element in a particular dimension. E.g. stride_am is how much to increase a_ptr - # by to get the element one row down (A has M rows) - stride_cm, - # stride_cn, # Assume that stride_cn == 1 - stride_am, - stride_ak, - stride_bn, - stride_bk, - # Meta-parameters - BLOCK_M: tl.constexpr, - GROUP_M: tl.constexpr, - BLOCK_N: tl.constexpr, - BLOCK_K: tl.constexpr, - # split k not used, not performant with activation, kept because early_config_prune is expecting it - SPLIT_K: tl.constexpr, - EVEN_K: tl.constexpr, - A_ROWMAJOR: tl.constexpr, - B_COLMAJOR: tl.constexpr, - BIAS: tl.constexpr, - SAVE_ACT_INPUT: tl.constexpr, - ACTIVATION: tl.constexpr, -): - - """ - Kernel for computing Out = activation(A x W + C) - - Input has shape (M, K) - - Weight has shape (K, N) - - Bias has shape (N,) - - Output has shape (M, N) - - ActInputs (optional) has shape (M, N) - 'ActInputs' optionally saves the A x W + C intermediate for backward computations - This kernel will consolidate over K - """ - - pid = tl.program_id(axis=0) - - grid_m = (M + BLOCK_M - 1) // BLOCK_M - grid_n = (N + BLOCK_N - 1) // BLOCK_N - # re-order program ID for better L2 performance - width = GROUP_M * grid_n - group_id = pid // width - group_size = min(grid_m - group_id * GROUP_M, GROUP_M) - pid_m = group_id * GROUP_M + (pid % group_size) - pid_n = (pid % width) // (group_size) - - # now compute the block that each program will go through - # rm (resp. rn) denotes a range of indices - # for rows (resp. col) of C - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - # trick to avoid masking on M and N axis - ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M) - rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N) - rk = tl.arange(0, BLOCK_K) - - if A_ROWMAJOR: - A = A + (ram[:, None] * stride_am + rk[None, :]) - else: - A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak) - if B_COLMAJOR: - B = B + (rk[:, None] + rbn[None, :] * stride_bn) - else: - B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn) - - acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) - - for k in range(K, 0, -BLOCK_K): - if EVEN_K: - a = tl.load(A) - b = tl.load(B) - else: - a = tl.load(A, mask=rk[None, :] < k, other=0.0) - b = tl.load(B, mask=rk[:, None] < k, other=0.0) - acc += tl.dot(a, b) - - if A_ROWMAJOR: - A += BLOCK_K - else: - A += BLOCK_K * stride_ak - if B_COLMAJOR: - B += BLOCK_K - else: - B += BLOCK_K * stride_bk - - # Putting bias after the matmul (instead of before) is faster, idk why - if BIAS: - bias = tl.load(bias + rn, mask=rn < N, other=0.0).to(tl.float32) - acc += bias[None, :] - - # optional: save the activation inputs - if SAVE_ACT_INPUT: - # act_in_ptrs = ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :] * stride_cn - act_in_ptrs = ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :] - tl.store(act_in_ptrs, acc) - - # optional: fused activation (while the data is in shared memory) - if ACTIVATION == "gelu": - acc = gelu(acc) - elif ACTIVATION == "gelu_approx": - acc = gelu_approx(acc) - elif ACTIVATION == "squared_relu": - acc = squared_relu(acc) - # rematerialize rm and rn to save registers - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - - # write back result - # C = C + rm[:, None] * stride_cm + rn[None, :] * stride_cn - C = C + rm[:, None] * stride_cm + rn[None, :] - mask = (rm < M)[:, None] & (rn < N)[None, :] - tl.store(C, acc) - - -def triton_linear_act( - x: torch.Tensor, - weight: torch.Tensor, - bias: Optional[torch.Tensor] = None, - activation: str = "id", - save_act_input: bool = False, -) -> torch.Tensor: - """ - Compute e = activation(x @ weight.T + bias). - This wrapper kicks the `kernel_fwd` Triton kernel - :param x: input tensor - :param weight: weight matrix - :param bias: an optional bias tensor - :param activation: Activation name. Needs to be a Triton kernel. - :param act_input: an optional tensor to save the activation inputs (for backward) - :return: result tensor - """ - # if torch.is_autocast_enabled(): - # dtype = torch.get_autocast_gpu_dtype() - # x, weight, bias = [a.to(dtype=dtype) for a in [x, weight, bias]] - - assert activation in ["id", "gelu", "gelu_approx", "squared_relu"] - - batch_shape, n = x.shape[:-1], x.shape[-1] - batch_dim = batch_shape.numel() - x_reshaped = x.reshape(batch_dim, n) - - if x_reshaped.stride(0) > 1 and x_reshaped.stride(1) > 1: - x_reshaped = x_reshaped.contiguous() - if weight.stride(0) > 1 and weight.stride(1) > 1: - weight = weight.contiguous() - bias = bias.contiguous() if bias is not None else None - - assert ( - x.dtype == weight.dtype - ), f"Input and weight must have the same dtype, got {x.dtype} and {weight.dtype}" - if bias is not None: - assert ( - x.dtype == bias.dtype - ), f"Input and bias must have the same dtype, got {x.dtype} and {bias.dtype}" - assert ( - x_reshaped.shape[1] == weight.shape[1] - ), f"Incompatible dimensions: {x_reshaped.shape} - {weight.shape}" - - assert ( - bias is None or bias.shape[0] == weight.shape[0] - ), "Incompatible dimensions in between weight and bias" - - M, K = x_reshaped.shape - N, K = weight.shape - - output = torch.empty((M, N), device=x.device, dtype=x.dtype) - act_input = torch.empty_like(output) if save_act_input else None - - # 1D launch kernel where each block gets its own program. - grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),) # noqa - - kernel_fwd[grid]( - output, - act_input, - x_reshaped, - weight, # data ptrs - bias if bias is not None else x, # auto skip bias if not present - M, # shapes - N, - K, - M // 32, # key for triton cache (limit number of compilations) - N // 32, - K // 32, - stride_cm=output.stride(0), # strides - # stride_cn=output.stride(1), - stride_am=x_reshaped.stride(0), - stride_ak=x_reshaped.stride(1), - stride_bk=weight.stride(1), - stride_bn=weight.stride(0), - BIAS=bias is not None, # optional fused bias - SAVE_ACT_INPUT=save_act_input, # optional save activation inputs - ACTIVATION=activation, # optional fused activation - A_ROWMAJOR=x_reshaped.stride(1) == 1, - B_COLMAJOR=weight.stride(1) == 1, - GROUP_M=8, # speed optimization: group the programs - ) - - if not save_act_input: - return output.reshape(*batch_shape, output.shape[-1]) - else: - return ( - output.reshape(*batch_shape, output.shape[-1]), - act_input.reshape(*batch_shape, act_input.shape[-1]), - ) - - -@triton.autotune( - configs=[ - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - # good for int8 - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=4, - num_warps=4, - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - ] - + get_configs_io_bound(), - key=["CACHE_KEY_M", "CACHE_KEY_N", "CACHE_KEY_K"], - prune_configs_by={ - "early_config_prune": early_config_prune, - "perf_model": estimate_matmul_time, - "top_k": 10, - }, -) -@triton.heuristics( - { - "EVEN_K": lambda args: args["K"] % (args["BLOCK_K"] * args["SPLIT_K"]) == 0, - } -) -@triton.jit -def kernel_bwd( - C, # Pointers to matrices - ACT_INPUT, - A, - B, - # Matrix dimensions - M, - N, - K, - CACHE_KEY_M, - CACHE_KEY_N, - CACHE_KEY_K, - # The stride variables represent how much to increase the ptr by when moving by 1 - # element in a particular dimension. E.g. stride_am is how much to increase a_ptr - # by to get the element one row down (A has M rows) - stride_cm, - # stride_cn, # Assume that stride_cn == 1 - stride_am, - stride_ak, - stride_bk, - stride_bn, - # Meta-parameters - BLOCK_M: tl.constexpr, - GROUP_M: tl.constexpr, - BLOCK_N: tl.constexpr, - BLOCK_K: tl.constexpr, - # split k not used, not performant with activation, kept because early_config_prune is expecting it - SPLIT_K: tl.constexpr, - EVEN_K: tl.constexpr, - ACTIVATION: tl.constexpr, -): - - """ - Kernel for computing Out = activation(A x W + C) - - Input has shape (M, K) - - Weight has shape (K, N) - - Output has shape (M, N) - - ActInputs (optional) has shape (M, N) - 'ActInputs' optionally saves the A x W + C intermediate for backward computations - This kernel will consolidate over K - """ - - pid = tl.program_id(axis=0) - - grid_m = (M + BLOCK_M - 1) // BLOCK_M - grid_n = (N + BLOCK_N - 1) // BLOCK_N - # re-order program ID for better L2 performance - width = GROUP_M * grid_n - group_id = pid // width - group_size = min(grid_m - group_id * GROUP_M, GROUP_M) - pid_m = group_id * GROUP_M + (pid % group_size) - pid_n = (pid % width) // (group_size) - - # now compute the block that each program will go through - # rm (resp. rn) denotes a range of indices - # for rows (resp. col) of C - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - # trick to avoid masking on M and N axis - ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M) - rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N) - rk = tl.arange(0, BLOCK_K) - - A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak) - B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn) - - acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) - - for k in range(K, 0, -BLOCK_K): - if EVEN_K: - a = tl.load(A) - b = tl.load(B) - else: - a = tl.load(A, mask=rk[None, :] < k, other=0.0) - b = tl.load(B, mask=rk[:, None] < k, other=0.0) - acc += tl.dot(a, b) - - A += BLOCK_K * stride_ak - B += BLOCK_K * stride_bk - - # optional: fused activation (while the data is in shared memory) - if ACTIVATION != "id": - act_in_ptrs = ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :] - act_input = tl.load(act_in_ptrs).to(acc.dtype) - if ACTIVATION == "gelu": - acc *= gelu_grad(act_input) - elif ACTIVATION == "gelu_approx": - acc *= gelu_approx_grad(act_input) - elif ACTIVATION == "squared_relu": - acc *= squared_relu_grad(act_input) - - # rematerialize rm and rn to save registers - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - - # write back result - C = C + rm[:, None] * stride_cm + rn[None, :] - mask = (rm < M)[:, None] & (rn < N)[None, :] - tl.store(C, acc, mask=mask) - - -def triton_dgrad_act( - grad_output: torch.Tensor, - weight: torch.Tensor, - activation: str = "id", - act_input: Optional[torch.Tensor] = None, -) -> torch.Tensor: - """ - Compute e = activation(grad_output @ weight + bias). - This wrapper kicks the `kernel_fwd` Triton kernel - :param grad_output: input tensor - :param weight: weight matrix - :param activation: Activation name. Needs to be a Triton kernel. - :param act_input: an optional tensor to save the activation inputs (for backward) - :return: result tensor - """ - assert activation in ["id", "gelu", "gelu_approx", "squared_relu"] - - batch_shape, n = grad_output.shape[:-1], grad_output.shape[-1] - batch_dim = batch_shape.numel() - grad_output_reshaped = grad_output.reshape(batch_dim, n) - - if grad_output_reshaped.stride(0) > 1 and grad_output_reshaped.stride(1) > 1: - grad_output_reshaped = grad_output_reshaped.contiguous() - if weight.stride(0) > 1 and weight.stride(1) > 1: - weight = weight.contiguous() - - assert ( - grad_output.dtype == weight.dtype - ), f"grad_output and weight must have the same dtype, got {grad_output.dtype} and {weight.dtype}" - assert ( - grad_output_reshaped.shape[1] == weight.shape[0] - ), f"Incompatible dimensions: {grad_output_reshaped.shape} - {weight.shape}" - if activation != "id": - assert act_input is not None, f"act_input is required for activation {activation}" - - # M, N, K in bwd are different from M, N, K in fwd - M, K = grad_output_reshaped.shape - K, N = weight.shape - - grad_input = torch.empty((M, N), device=grad_output.device, dtype=grad_output.dtype) - - # 1D launch kernel where each block gets its own program. - grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),) # noqa - - kernel_bwd[grid]( - grad_input, - act_input, - grad_output_reshaped, - weight, # data ptrs - M, # shapes - N, - K, - M // 32, # key for triton cache (limit number of compilations) - N // 32, - K // 32, - stride_cm=grad_input.stride(0), # strides - # stride_cn=grad_input.stride(1), - stride_am=grad_output_reshaped.stride(0), - stride_ak=grad_output_reshaped.stride(1), - stride_bk=weight.stride(0), - stride_bn=weight.stride(1), - ACTIVATION=activation, # optional fused activation - GROUP_M=8, # speed optimization: group the programs - ) - - return grad_input.reshape(*batch_shape, grad_input.shape[-1]) diff --git a/build/torch29-cxx11-cu128-aarch64-linux/ops/triton/mlp.py b/build/torch29-cxx11-cu128-aarch64-linux/ops/triton/mlp.py deleted file mode 100644 index 059f4f8a5e174c1f4824e43d313fca18eaa799b8..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-aarch64-linux/ops/triton/mlp.py +++ /dev/null @@ -1,149 +0,0 @@ -# The triton fused matmul + sqrelu is faster for fp16 but slower for bf16, compared -# to naive implementation. -import fused_dense_lib as fused_dense_cuda -import torch -import torch.nn as nn -import torch.nn.functional as F - -from flash_attn.utils.torch import custom_fwd, custom_bwd -from flash_attn.ops.activations import sqrelu_bwd, sqrelu_fwd -from flash_attn.ops.triton.linear import triton_dgrad_act, triton_linear_act - - -class FusedDenseSqreluDenseFunc(torch.autograd.Function): - @staticmethod - @custom_fwd - def forward(ctx, x, weight1, bias1, weight2, bias2, checkpoint_lvl=0): - """checkpoint_lvl: - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute act_input and gelu_out in the bwd - """ - if torch.is_autocast_enabled(): - dtype = torch.get_autocast_gpu_dtype() - x, weight1, bias1, weight2, bias2 = [ - a.to(dtype=dtype) for a in [x, weight1, bias1, weight2, bias2] - ] - is_bf16 = x.dtype == torch.bfloat16 - assert checkpoint_lvl in [0, 1, 2] - x = x.contiguous() - weight1 = weight1.contiguous() - bias1 = bias1.contiguous() - weight2 = weight2.contiguous() - bias2 = bias2.contiguous() - batch_shape, n = x.shape[:-1], x.shape[-1] - batch_dim = batch_shape.numel() - if is_bf16: - act_input = fused_dense_cuda.linear_bias_forward( - x.reshape(batch_dim, n), weight1, bias1 - ) - output1 = sqrelu_fwd(act_input) - else: - save_act_input = checkpoint_lvl != 2 - result = triton_linear_act( - x.reshape(batch_dim, n), - weight1, - bias1, - activation="squared_relu", - save_act_input=save_act_input, - ) - if save_act_input: - output1, act_input = result - else: - output1 = result - output2 = fused_dense_cuda.linear_bias_forward(output1, weight2, bias2) - ctx.checkpoint_lvl = checkpoint_lvl - if checkpoint_lvl == 0: - ctx.save_for_backward(x, weight1, bias1, weight2, act_input, output1) - elif checkpoint_lvl == 1: - ctx.save_for_backward(x, weight1, bias1, weight2, act_input) - elif checkpoint_lvl == 2: - ctx.save_for_backward(x, weight1, bias1, weight2) - return output2.reshape(*batch_shape, output2.shape[-1]) - - @staticmethod - @custom_bwd - def backward(ctx, grad_output): - grad_output = grad_output.contiguous() - checkpoint_lvl = ctx.checkpoint_lvl - x, weight1, bias1, weight2, *rest = ctx.saved_tensors - batch_shape, n = x.shape[:-1], x.shape[-1] - batch_dim = batch_shape.numel() - is_bf16 = x.dtype == torch.bfloat16 - if checkpoint_lvl == 0: - act_input, output1 = rest - elif checkpoint_lvl == 1: - (act_input,) = rest - output1 = sqrelu_fwd(act_input) - elif checkpoint_lvl == 2: - if is_bf16: - act_input = fused_dense_cuda.linear_bias_forward( - x.reshape(batch_dim, n), weight1, bias1 - ) - output1 = sqrelu_fwd(act_input) - else: - output1, act_input = triton_linear_act( - x.reshape(batch_dim, n), - weight1, - bias1, - activation="squared_relu", - save_act_input=True, - ) - - if is_bf16: - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad(output1, grad_output) - grad_output1 = grad_output @ weight2 - grad_act_input = sqrelu_bwd(grad_output1, act_input) - grad_input, grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_backward( - x.reshape(batch_dim, n), weight1, grad_act_input - ) - else: - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad(output1, grad_output) - grad_act_input = triton_dgrad_act( - grad_output, weight2, activation="squared_relu", act_input=act_input - ) - grad_input, grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_backward( - x.reshape(batch_dim, n), weight1, grad_act_input - ) - return grad_input.reshape_as(x), grad_weight1, grad_bias1, grad_weight2, grad_bias2, None - - -fused_dense_sqrelu_dense_function = FusedDenseSqreluDenseFunc.apply - - -class FusedDenseSqreluDense(nn.Module): - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - bias1=True, - bias2=True, - checkpoint_lvl=0, - device=None, - dtype=None, - ): - """ - checkpoint_lvl (increasing lvl means slower but more memory saving): - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute gelu_in and gelu_out in the bwd - """ - assert checkpoint_lvl in [0, 1, 2] - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features * 4 - assert bias1 == True, "DenseSqreluDense module without bias is currently not supported" - assert bias2 == True, "DenseSqreluDense module without bias is currently not supported" - self.checkpoint_lvl = checkpoint_lvl - self.fc1 = nn.Linear(in_features, hidden_features, bias=bias1, **factory_kwargs) - self.fc2 = nn.Linear(hidden_features, out_features, bias=bias2, **factory_kwargs) - - def forward(self, x): - assert x.is_cuda - return fused_dense_sqrelu_dense_function( - x, self.fc1.weight, self.fc1.bias, self.fc2.weight, self.fc2.bias, self.checkpoint_lvl - ) diff --git a/build/torch29-cxx11-cu128-aarch64-linux/ops/triton/rotary.py b/build/torch29-cxx11-cu128-aarch64-linux/ops/triton/rotary.py deleted file mode 100644 index 6ad3ab412a084270120167ad32b67ad1720f2c37..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-aarch64-linux/ops/triton/rotary.py +++ /dev/null @@ -1,186 +0,0 @@ -# Copyright (c) 2025, Tri Dao. -# As of 2025-04-23, we require triton >= 3.0 - -from typing import Optional, Union - -import torch - -import triton -import triton.language as tl - - -@triton.jit -def rotary_kernel( - OUT, # Pointers to matrices - X, - COS, - SIN, - CU_SEQLENS, - SEQLEN_OFFSETS, # this could be int or a pointer - # Matrix dimensions - seqlen, - nheads, - seqlen_ro, - # strides - stride_out_batch, - stride_out_seqlen, - stride_out_nheads, - stride_out_headdim, - stride_x_batch, - stride_x_seqlen, - stride_x_nheads, - stride_x_headdim, - # Meta-parameters - # We want ROTARY_DIM to be constexpr, otherwise the triton compiler doesn't know that - # the mask is constant every 8 elements, and it will generate LDG.16 instead of LDG.128 - ROTARY_DIM: tl.constexpr, - IS_SEQLEN_OFFSETS_TENSOR: tl.constexpr, - IS_VARLEN: tl.constexpr, - INTERLEAVED: tl.constexpr, - CONJUGATE: tl.constexpr, - BLOCK_H: tl.constexpr, - BLOCK_M: tl.constexpr, -): - BLOCK_K: tl.constexpr = triton.next_power_of_2(ROTARY_DIM) - ROTARY_DIM_HALF = ROTARY_DIM // 2 - pid_head = tl.program_id(axis=0) - pid_m = tl.program_id(axis=1) - pid_batch = tl.program_id(axis=2) - - if not IS_VARLEN: - X = X + pid_batch * stride_x_batch - OUT = OUT + pid_batch * stride_out_batch - else: - start_idx = tl.load(CU_SEQLENS + pid_batch) - seqlen = tl.load(CU_SEQLENS + pid_batch + 1) - start_idx - X = X + start_idx * stride_x_seqlen - OUT = OUT + start_idx * stride_out_seqlen - - if pid_m * BLOCK_M >= seqlen: - return - - rh = pid_head * BLOCK_H + tl.arange(0, BLOCK_H) - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - if not IS_SEQLEN_OFFSETS_TENSOR: - rm_cs = rm + SEQLEN_OFFSETS - else: - rm_cs = rm + tl.load(SEQLEN_OFFSETS + pid_batch) - - rk_half = tl.arange(0, BLOCK_K // 2) - COS = COS + (rm_cs[:, None] * ROTARY_DIM_HALF + rk_half[None, :]) - SIN = SIN + (rm_cs[:, None] * ROTARY_DIM_HALF + rk_half[None, :]) - mask_cs = (rm_cs[:, None] < seqlen_ro) & (rk_half[None, :] < ROTARY_DIM_HALF) - cos = tl.load(COS, mask=mask_cs, other=1.0).to(tl.float32) - sin = tl.load(SIN, mask=mask_cs, other=0.0).to(tl.float32) - if CONJUGATE: - sin = -sin - - if not INTERLEAVED: - # Load the 1st and 2nd halves of X, do calculation, then store to 1st and 2nd halves of OUT - X = X + (rh[:, None, None] * stride_x_nheads + rm[None, :, None] * stride_x_seqlen + rk_half[None, None, :] * stride_x_headdim) - OUT = OUT + (rh[:, None, None] * stride_out_nheads + rm[None, :, None] * stride_out_seqlen + rk_half[None, None, :] * stride_out_headdim) - mask = (rh[:, None, None] < nheads) & (rm[None, :, None] < seqlen) & (rk_half[None, None, :] < ROTARY_DIM_HALF) - x0 = tl.load(X, mask=mask, other=0.0).to(tl.float32) - x1 = tl.load(X + ROTARY_DIM_HALF * stride_x_headdim, mask=mask, other=0.0,).to(tl.float32) - o0 = x0 * cos - x1 * sin - o1 = x0 * sin + x1 * cos - tl.store(OUT, o0, mask=mask) - tl.store(OUT + ROTARY_DIM_HALF * stride_out_headdim, o1, mask=mask) - else: - rk = tl.arange(0, BLOCK_K) - X = X + (rh[:, None, None] * stride_x_nheads + rm[None, :, None] * stride_x_seqlen + rk[None, None, :] * stride_x_headdim) - OUT = OUT + (rh[:, None, None] * stride_out_nheads + rm[None, :, None] * stride_out_seqlen + rk[None, None, :] * stride_out_headdim) - mask = (rh[:, None, None] < nheads) & (rm[None, :, None] < seqlen) & (rk[None, None, :] < ROTARY_DIM) - x = tl.load(X, mask=mask, other=0.0).to(tl.float32) - x0, x1 = tl.split(tl.reshape(x, [BLOCK_H, BLOCK_M, BLOCK_K // 2, 2])) - o0 = x0 * cos - x1 * sin - o1 = x0 * sin + x1 * cos - o = tl.reshape(tl.join(o0, o1), [BLOCK_H, BLOCK_M, BLOCK_K]) - tl.store(OUT, o, mask=mask) - - -def apply_rotary( - x: torch.Tensor, - cos: torch.Tensor, - sin: torch.Tensor, - seqlen_offsets: Union[int, torch.Tensor] = 0, - cu_seqlens: Optional[torch.Tensor] = None, - max_seqlen: Optional[int] = None, - interleaved=False, - inplace=False, - conjugate=False, -) -> torch.Tensor: - """ - Arguments: - x: (batch, seqlen, nheads, headdim) if cu_seqlens is None - else (total_seqlen, nheads, headdim). - cos: (seqlen_ro, rotary_dim / 2) - sin: (seqlen_ro, rotary_dim / 2) - seqlen_offsets: integer or integer tensor of size (batch,) - cu_seqlens: (batch + 1,) or None - max_seqlen: int - Returns: - y: (batch, seqlen, nheads, headdim) - """ - is_varlen = cu_seqlens is not None - if not is_varlen: - batch, seqlen, nheads, headdim = x.shape - else: - assert max_seqlen is not None, "If cu_seqlens is passed in, then max_seqlen must be passed" - total_seqlen, nheads, headdim = x.shape - batch_p_1 = cu_seqlens.shape[0] - batch = batch_p_1 - 1 - seqlen = max_seqlen - seqlen_ro, rotary_dim = cos.shape - assert sin.shape == cos.shape - rotary_dim *= 2 - assert rotary_dim <= headdim, "rotary_dim must be <= headdim" - assert headdim <= 256, "Only support headdim <= 256" - assert seqlen_ro >= seqlen, "seqlen_ro must be >= seqlen" - - cos, sin = cos.contiguous(), sin.contiguous() - if isinstance(seqlen_offsets, torch.Tensor): - assert seqlen_offsets.shape == (batch,) - assert seqlen_offsets.dtype in [torch.int32, torch.int64] - seqlen_offsets = seqlen_offsets.contiguous() - else: - assert seqlen_offsets + seqlen <= seqlen_ro - - output = torch.empty_like(x) if not inplace else x - if rotary_dim < headdim and not inplace: - output[..., rotary_dim:].copy_(x[..., rotary_dim:]) - - grid = lambda META: (triton.cdiv(nheads, META["BLOCK_H"]), triton.cdiv(seqlen, META["BLOCK_M"]), batch) # noqa - BLOCK_M = 8 if rotary_dim <= 128 else 4 - - # Need this, otherwise Triton tries to launch from cuda:0 and we get - # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?) - device_ctx = torch.cuda.device(x.device.index) if x.device.type == 'cuda' else torch.xpu.device(x.device.index) - with device_ctx: - torch.library.wrap_triton(rotary_kernel)[grid]( - output, # data ptrs - x, - cos, - sin, - cu_seqlens, - seqlen_offsets, - seqlen, # shapes - nheads, - seqlen_ro, - output.stride(0) if not is_varlen else 0, # batch_strides if not varlen else 0 - output.stride(-3), # seqlen_stride or total_seqlen_stride - output.stride(-2), # nheads_stride - output.stride(-1), # headdim_stride - x.stride(0) if not is_varlen else 0, # batch_strides if not varlen else 0 - x.stride(-3), # seqlen stride or total_seqlen_stride - x.stride(-2), # nheads stride - x.stride(-1), # headdim stride - rotary_dim, - isinstance(seqlen_offsets, torch.Tensor), - is_varlen, - interleaved, - conjugate, - BLOCK_M=BLOCK_M, - BLOCK_H=2, - ) - return output diff --git a/build/torch29-cxx11-cu128-x86_64-linux/_flash_attn2_588b404.abi3.so b/build/torch29-cxx11-cu128-x86_64-linux/_flash_attn2_588b404.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..294c145205f0adc893f74d8a0ff5b626a3c2476c --- /dev/null +++ b/build/torch29-cxx11-cu128-x86_64-linux/_flash_attn2_588b404.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:7b3afc2eda58b3649ac67513c775ce1cb124e5498f8dbbbe4ef07db6857d56d3 +size 1037644608 diff --git a/build/torch29-cxx11-cu128-x86_64-linux/_flash_attn2_cuda_abda3a0.abi3.so b/build/torch29-cxx11-cu128-x86_64-linux/_flash_attn2_cuda_abda3a0.abi3.so deleted file mode 100644 index 9c5a036f0030070eac6bda9b7d7a4cbcb26d60a2..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-x86_64-linux/_flash_attn2_cuda_abda3a0.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:54202fb14acadd754654a7822be7943e513822cbe0f07d4ddee5f9ee99d6fe50 -size 1037644632 diff --git a/build/torch29-cxx11-cu128-x86_64-linux/_ops.py b/build/torch29-cxx11-cu128-x86_64-linux/_ops.py index fb32434e09e651a658beb38be8c101e5db89374f..03506d8a18d17ca7296a5f1b7a18ec7dc4ad73bb 100644 --- a/build/torch29-cxx11-cu128-x86_64-linux/_ops.py +++ b/build/torch29-cxx11-cu128-x86_64-linux/_ops.py @@ -1,9 +1,9 @@ import torch -from . import _flash_attn2_cuda_abda3a0 -ops = torch.ops._flash_attn2_cuda_abda3a0 +from . import _flash_attn2_588b404 +ops = torch.ops._flash_attn2_588b404 def add_op_namespace_prefix(op_name: str): """ Prefix op by namespace. """ - return f"_flash_attn2_cuda_abda3a0::{op_name}" + return f"_flash_attn2_588b404::{op_name}" \ No newline at end of file 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 index 6a0157f25449eb90c6c688bc8ef3f4fcf84d46dc..66e8d3665834b141d703998496b3da6fe1f6fe06 100644 --- a/build/torch29-cxx11-cu128-x86_64-linux/flash_attn_interface.py +++ b/build/torch29-cxx11-cu128-x86_64-linux/flash_attn_interface.py @@ -31,14 +31,12 @@ def _get_device(): else: return "cpu" +_XPU_AVAILABLE = torch.xpu.is_available() if hasattr(torch, "xpu") else False # TODO remove hasattr check when bwd is supported on XPU -def _get_block_size_n(device, head_dim, is_dropout, is_causal): - assert head_dim <= 256 - - if device.type == "xpu": - return 64 +def _get_block_size_n(device, head_dim, is_dropout, is_causal): # This should match the block sizes in the CUDA kernel + assert head_dim <= 256 major, minor = torch.cuda.get_device_capability(device) is_sm8x = major == 8 and minor > 0 # Only include sm86 and sm89, exclude sm80 (A100) is_sm80 = major == 8 and minor == 0 @@ -1066,7 +1064,7 @@ def flash_attn_qkvpacked_func( alibi_slopes, deterministic, return_attn_probs, - torch.is_grad_enabled(), + False if _XPU_AVAILABLE else torch.is_grad_enabled(), ) @@ -1144,7 +1142,7 @@ def flash_attn_kvpacked_func( alibi_slopes, deterministic, return_attn_probs, - torch.is_grad_enabled(), + False if _XPU_AVAILABLE else torch.is_grad_enabled(), ) @@ -1221,7 +1219,7 @@ def flash_attn_func( alibi_slopes, deterministic, return_attn_probs, - torch.is_grad_enabled(), + False if _XPU_AVAILABLE else torch.is_grad_enabled(), ) @@ -1287,7 +1285,7 @@ def flash_attn_varlen_qkvpacked_func( alibi_slopes, deterministic, return_attn_probs, - torch.is_grad_enabled(), + False if _XPU_AVAILABLE else torch.is_grad_enabled(), ) @@ -1379,7 +1377,7 @@ def flash_attn_varlen_kvpacked_func( alibi_slopes, deterministic, return_attn_probs, - torch.is_grad_enabled(), + False if _XPU_AVAILABLE else torch.is_grad_enabled(), ) @@ -1473,7 +1471,7 @@ def flash_attn_varlen_func( deterministic, return_attn_probs, block_table, - False if q.device.type == "cpu" else torch.is_grad_enabled(), + False if _XPU_AVAILABLE or q.device.type == "cpu" else torch.is_grad_enabled(), ) diff --git a/build/torch29-cxx11-cu128-x86_64-linux/metadata.json b/build/torch29-cxx11-cu128-x86_64-linux/metadata.json index 8b00d7fb7685121edf7a97bb007415122444f1d5..9cf5deed9898dce769f4cc73913d3530b92a0bd8 100644 --- a/build/torch29-cxx11-cu128-x86_64-linux/metadata.json +++ b/build/torch29-cxx11-cu128-x86_64-linux/metadata.json @@ -1,14 +1,4 @@ { "version": 1, - "license": "BSD-3-Clause", - "python-depends": [], - "backend": { - "type": "cuda", - "archs": [ - "10.0", - "12.0", - "8.0", - "9.0" - ] - } -} + "python-depends": [] +} \ No newline at end of file diff --git a/build/torch29-cxx11-cu129-aarch64-linux/__init__.py b/build/torch29-cxx11-cu129-aarch64-linux/__init__.py deleted file mode 100644 index ecc2f9d896b6c93f90b0a1499856dc0612177422..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-aarch64-linux/__init__.py +++ /dev/null @@ -1,393 +0,0 @@ -from typing import Optional, List -import torch -from ._ops import ops as flash_attn_ops -from .flash_attn_interface import ( - flash_attn_func, - flash_attn_kvpacked_func, - flash_attn_qkvpacked_func, - flash_attn_varlen_func, - flash_attn_varlen_kvpacked_func, - flash_attn_varlen_qkvpacked_func, - flash_attn_with_kvcache, -) - - -def fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - return_softmax: bool = False, - gen: Optional[torch.Generator] = None, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention. - - Args: - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Optional output tensor, same shape as q - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - return_softmax: Whether to return softmax weights - gen: Optional random number generator - - Returns: - List of tensors: [output, softmax_lse, (softmax if return_softmax)] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.fwd( - q, - k, - v, - out, - alibi_slopes, - p_dropout, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - gen, - ) - - -def varlen_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - out: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - block_table: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - max_seqlen_q: int = 0, - max_seqlen_k: int = 0, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - zero_tensors: bool = False, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - return_softmax: bool = False, - gen: Optional[torch.Generator] = None, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention with variable sequence lengths. - - Args: - q: Query tensor of shape [total_q, num_heads, head_size] - k: Key tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - v: Value tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1] - cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1] - out: Optional output tensor of shape [total_q, num_heads, head_size] - seqused_k: Optional tensor specifying how many keys to use per batch element [batch_size] - leftpad_k: Optional left padding for keys of shape [batch_size] - block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq] - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - max_seqlen_q: Maximum sequence length for queries - max_seqlen_k: Maximum sequence length for keys - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - zero_tensors: Whether to zero tensors before computation - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - return_softmax: Whether to return softmax weights - gen: Optional random number generator - - Returns: - List of tensors: [output, softmax_lse, (softmax if return_softmax)] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.varlen_fwd( - q, - k, - v, - out, - cu_seqlens_q, - cu_seqlens_k, - seqused_k, - leftpad_k, - block_table, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - p_dropout, - softmax_scale, - zero_tensors, - is_causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - gen, - ) - - -def bwd( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - p_dropout: float = 0.0, - 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, - gen: Optional[torch.Generator] = None, - rng_state: Optional[torch.Tensor] = None, -) -> List[torch.Tensor]: - """ - Backward pass for multi-head attention. - - Args: - dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size] - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size] - softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q] - dq: Optional gradient tensor for queries, same shape as q - dk: Optional gradient tensor for keys, same shape as k - dv: Optional gradient tensor for values, same shape as v - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - deterministic: Whether to use deterministic algorithms - gen: Optional random number generator - rng_state: Optional RNG state from forward pass - - Returns: - List of tensors: [dq, dk, dv] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - alibi_slopes, - p_dropout, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - deterministic, - gen, - rng_state, - ) - - -def varlen_bwd( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - max_seqlen_q: int = 0, - max_seqlen_k: int = 0, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - zero_tensors: bool = False, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - deterministic: bool = False, - gen: Optional[torch.Generator] = None, - rng_state: Optional[torch.Tensor] = None, -) -> List[torch.Tensor]: - """ - Backward pass for multi-head attention with variable sequence lengths. - - Args: - dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size] - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size] - softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q] - cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1] - cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1] - dq: Optional gradient tensor for queries, same shape as q - dk: Optional gradient tensor for keys, same shape as k - dv: Optional gradient tensor for values, same shape as v - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - max_seqlen_q: Maximum sequence length for queries - max_seqlen_k: Maximum sequence length for keys - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - zero_tensors: Whether to zero tensors before computation - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - deterministic: Whether to use deterministic algorithms - gen: Optional random number generator - rng_state: Optional RNG state from forward pass - - Returns: - List of tensors: [dq, dk, dv] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.varlen_bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - p_dropout, - softmax_scale, - zero_tensors, - is_causal, - window_size_left, - window_size_right, - softcap, - deterministic, - gen, - rng_state, - ) - - -def fwd_kvcache( - q: torch.Tensor, - kcache: torch.Tensor, - vcache: torch.Tensor, - k: Optional[torch.Tensor] = None, - v: Optional[torch.Tensor] = None, - seqlens_k: Optional[torch.Tensor] = None, - rotary_cos: Optional[torch.Tensor] = None, - rotary_sin: Optional[torch.Tensor] = None, - cache_batch_idx: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - block_table: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - out: 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, - is_rotary_interleaved: bool = False, - num_splits: int = 1, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention with KV cache. - - Args: - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - kcache: Key cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - vcache: Value cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - k: Optional new keys tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size] - v: Optional new values tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size] - seqlens_k: Optional sequence lengths for keys of shape [batch_size] - rotary_cos: Optional rotary cosine tensor of shape [seqlen_ro, rotary_dim/2] - rotary_sin: Optional rotary sine tensor of shape [seqlen_ro, rotary_dim/2] - cache_batch_idx: Optional indices to index into the KV cache - leftpad_k: Optional left padding for keys of shape [batch_size] - block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq] - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - out: Optional output tensor, same shape as q - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - is_rotary_interleaved: Whether rotary embeddings are interleaved - num_splits: Number of splits for computation - - Returns: - List of tensors: [output, softmax_lse] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.fwd_kvcache( - q, - kcache, - vcache, - k, - v, - seqlens_k, - rotary_cos, - rotary_sin, - cache_batch_idx, - leftpad_k, - block_table, - alibi_slopes, - out, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - is_rotary_interleaved, - num_splits, - ) diff --git a/build/torch29-cxx11-cu129-aarch64-linux/_flash_attn2_cuda_f1a742f.abi3.so b/build/torch29-cxx11-cu129-aarch64-linux/_flash_attn2_cuda_f1a742f.abi3.so deleted file mode 100644 index f367c821b16db32c89cf38aa29b9a5e5d62906fa..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-aarch64-linux/_flash_attn2_cuda_f1a742f.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:f6d2481a77b64ab1a0048a1f45333bf306d3bb985319758a31f57859a2211c20 -size 1043164184 diff --git a/build/torch29-cxx11-cu129-aarch64-linux/_ops.py b/build/torch29-cxx11-cu129-aarch64-linux/_ops.py deleted file mode 100644 index 659b9bb637b9f59c68d4987f781194173595c596..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-aarch64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _flash_attn2_cuda_f1a742f -ops = torch.ops._flash_attn2_cuda_f1a742f - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_flash_attn2_cuda_f1a742f::{op_name}" diff --git a/build/torch29-cxx11-cu129-aarch64-linux/bert_padding.py b/build/torch29-cxx11-cu129-aarch64-linux/bert_padding.py deleted file mode 100644 index 3c2d35159a014a9d03aabead9e52e009168696ea..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-aarch64-linux/bert_padding.py +++ /dev/null @@ -1,218 +0,0 @@ -# Adapted from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/padding.py - -import torch -import torch.nn.functional as F -from einops import rearrange, repeat - - -class IndexFirstAxis(torch.autograd.Function): - @staticmethod - def forward(ctx, input, indices): - ctx.save_for_backward(indices) - assert input.ndim >= 2 - ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] - second_dim = other_shape.numel() - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - # return input[indices] - return torch.gather( - rearrange(input, "b ... -> b (...)"), 0, repeat(indices, "z -> z d", d=second_dim) - ).reshape(-1, *other_shape) - - @staticmethod - def backward(ctx, grad_output): - (indices,) = ctx.saved_tensors - assert grad_output.ndim >= 2 - other_shape = grad_output.shape[1:] - grad_output = rearrange(grad_output, "b ... -> b (...)") - grad_input = torch.zeros( - [ctx.first_axis_dim, grad_output.shape[1]], - device=grad_output.device, - dtype=grad_output.dtype, - ) - # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. - # grad_input[indices] = grad_output - grad_input.scatter_(0, repeat(indices, "z -> z d", d=grad_output.shape[1]), grad_output) - return grad_input.reshape(ctx.first_axis_dim, *other_shape), None - - -index_first_axis = IndexFirstAxis.apply - - -class IndexPutFirstAxis(torch.autograd.Function): - @staticmethod - def forward(ctx, values, indices, first_axis_dim): - ctx.save_for_backward(indices) - assert indices.ndim == 1 - assert values.ndim >= 2 - output = torch.zeros( - first_axis_dim, *values.shape[1:], device=values.device, dtype=values.dtype - ) - # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. - output[indices] = values - # output.scatter_(0, repeat(indices, 'z -> z d', d=values.shape[1]), values) - return output - - @staticmethod - def backward(ctx, grad_output): - (indices,) = ctx.saved_tensors - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - grad_values = grad_output[indices] - # grad_values = torch.gather(grad_output, 0, repeat(indices, 'z -> z d', d=grad_output.shape[1])) - return grad_values, None, None - - -index_put_first_axis = IndexPutFirstAxis.apply - - -class IndexFirstAxisResidual(torch.autograd.Function): - @staticmethod - def forward(ctx, input, indices): - ctx.save_for_backward(indices) - assert input.ndim >= 2 - ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] - second_dim = other_shape.numel() - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - output = input[indices] - # We don't want to reshape input (b ... -> b (...)) since it could change the channel_last - # memory format to channel_first. In other words, input might not be contiguous. - # If we don't detach, Pytorch complains about output being a view and is being modified inplace - return output, input.detach() - - @staticmethod - def backward(ctx, grad_output, grad_residual): - (indices,) = ctx.saved_tensors - assert grad_output.ndim >= 2 - other_shape = grad_output.shape[1:] - assert grad_residual.shape[1:] == other_shape - grad_input = grad_residual - # grad_input[indices] += grad_output - indices = indices.reshape(indices.shape[0], *((1,) * (grad_output.ndim - 1))) - indices = indices.expand_as(grad_output) - grad_input.scatter_add_(0, indices, grad_output) - return grad_input.reshape(ctx.first_axis_dim, *other_shape), None - - -index_first_axis_residual = IndexFirstAxisResidual.apply - - -def unpad_input(hidden_states, attention_mask, unused_mask=None): - """ - Arguments: - hidden_states: (batch, seqlen, ...) - attention_mask: (batch, seqlen), bool / int, 1 means valid and 0 means not valid. - unused_mask: (batch, seqlen), bool / int, 1 means the element is allocated but unused. - Return: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens selected in attention_mask + unused_mask. - indices: (total_nnz), the indices of masked tokens from the flattened input sequence. - cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. - max_seqlen_in_batch: int - seqused: (batch), returns the number of tokens selected in attention_mask + unused_mask. - """ - all_masks = (attention_mask + unused_mask) if unused_mask is not None else attention_mask - seqlens_in_batch = all_masks.sum(dim=-1, dtype=torch.int32) - used_seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) - indices = torch.nonzero(all_masks.flatten(), as_tuple=False).flatten() - max_seqlen_in_batch = seqlens_in_batch.max().item() - cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) - # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the - # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim - # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to - # index with integer indices. Moreover, torch's index is a bit slower than it needs to be, - # so we write custom forward and backward to make it a bit faster. - return ( - index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices), - indices, - cu_seqlens, - max_seqlen_in_batch, - used_seqlens_in_batch, - ) - - -def unpad_input_for_concatenated_sequences(hidden_states, attention_mask_in_length): - """ - Supports concatenating short samples in one sequence. The attention_mask_in_length is utilized to mask other short samples. It helps efficient training of variant lengths-based samples (e.g., the supervised fine-tuning task in large language model). - The motivation for this function is explained [here](https://github.com/Dao-AILab/flash-attention/issues/432#issuecomment-1668822286). - - For example, if batch = 3 and seqlen = 6, the attention_mask_in_length is: - ``` - [ - [2, 3, 0, 0, 0, 0], - [3, 2, 0, 0, 0, 0], - [6, 0, 0, 0, 0, 0] - ] - ``` - , which refers to the 3D-attention mask: - ``` - [ - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [0, 0, 1, 0, 0, 0], - [0, 0, 1, 1, 0, 0], - [0, 0, 1, 1, 1, 0], - [0, 0, 0, 0, 0, 1] - ], - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [0, 0, 0, 1, 0, 0], - [0, 0, 0, 1, 1, 0], - [0, 0, 0, 0, 0, 1] - ], - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [1, 1, 1, 1, 0, 0], - [1, 1, 1, 1, 1, 0], - [1, 1, 1, 1, 1, 1] - ] - ] - ```. - - Arguments: - hidden_states: (batch, seqlen, ...) - attention_mask_in_length: (batch, seqlen), int, a nonzero number (e.g., 1, 2, 3, etc.) means length of concatenated sequence in b-th batch, and 0 means none. - Return: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. - indices: (total_nnz), the indices of non-masked tokens from the flattened input sequence. - cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. - max_seqlen_in_batch: int - """ - length = attention_mask_in_length.sum(dim=-1) - seqlen = attention_mask_in_length.size(-1) - attention_mask_2d = torch.arange(seqlen, device=length.device, dtype=length.dtype).expand(len(length), seqlen) < length.unsqueeze(1) - real_indices_idx = torch.nonzero(attention_mask_in_length.flatten(), as_tuple=False).flatten() - seqlens_in_batch = attention_mask_in_length.flatten()[real_indices_idx] - indices = torch.nonzero(attention_mask_2d.flatten(), as_tuple=False).flatten() - max_seqlen_in_batch = seqlens_in_batch.max().item() - cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) - # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the - # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim - # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to - # index with integer indices. Moreover, torch's index is a bit slower than it needs to be, - # so we write custom forward and backward to make it a bit faster. - return ( - index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices), - indices, - cu_seqlens, - max_seqlen_in_batch, - ) - - -def pad_input(hidden_states, indices, batch, seqlen): - """ - Arguments: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. - indices: (total_nnz), the indices that represent the non-masked tokens of the original padded input sequence. - batch: int, batch size for the padded sequence. - seqlen: int, maximum sequence length for the padded sequence. - Return: - hidden_states: (batch, seqlen, ...) - """ - dim = hidden_states.shape[-1] - # output = torch.zeros((batch * seqlen), dim, device=hidden_states.device, dtype=hidden_states.dtype) - # output[indices] = hidden_states - output = index_put_first_axis(hidden_states, indices, batch * seqlen) - return rearrange(output, "(b s) ... -> b s ...", b=batch) diff --git a/build/torch29-cxx11-cu129-aarch64-linux/flash_attn2/__init__.py b/build/torch29-cxx11-cu129-aarch64-linux/flash_attn2/__init__.py deleted file mode 100644 index a9b2672c1cd85b74c1b3ded0fc0b2100e1aeac23..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-aarch64-linux/flash_attn2/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -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-aarch64-linux/flash_attn_interface.py b/build/torch29-cxx11-cu129-aarch64-linux/flash_attn_interface.py deleted file mode 100644 index 6a0157f25449eb90c6c688bc8ef3f4fcf84d46dc..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-aarch64-linux/flash_attn_interface.py +++ /dev/null @@ -1,1622 +0,0 @@ -# Copyright (c) 2023, Tri Dao. - -from typing import Optional, Sequence, Tuple, Union - -import torch -import torch.nn as nn -import os - -# # isort: off -# # We need to import the CUDA kernels after importing torch -# USE_TRITON_ROCM = os.getenv("FLASH_ATTENTION_TRITON_AMD_ENABLE", "FALSE") == "TRUE" -# if USE_TRITON_ROCM: -# from .flash_attn_triton_amd import interface_fa as flash_attn -# else: -# import flash_attn_2_cuda as flash_attn - - -from ._ops import ops as flash_attn - -# # isort: on - -def maybe_contiguous(x): - return x.contiguous() if x is not None and x.stride(-1) != 1 else x - - -def _get_device(): - if torch.xpu.is_available(): - return "xpu" - elif torch.cuda.is_available(): - return "cuda" - else: - return "cpu" - - -def _get_block_size_n(device, head_dim, is_dropout, is_causal): - assert head_dim <= 256 - - if device.type == "xpu": - return 64 - - # This should match the block sizes in the CUDA kernel - major, minor = torch.cuda.get_device_capability(device) - is_sm8x = major == 8 and minor > 0 # Only include sm86 and sm89, exclude sm80 (A100) - is_sm80 = major == 8 and minor == 0 - is_sm90 = major == 9 and minor == 0 - if head_dim <= 32: - return 128 - if head_dim <= 64: - return 128 if not is_dropout else 64 - elif head_dim <= 96: - return 64 - elif head_dim <= 128: - if is_sm8x: - return 64 if (not is_dropout and is_causal) else 32 - else: - return 64 if not is_dropout else 32 - elif head_dim <= 192: - return 64 - elif head_dim <= 224: - return 64 - elif head_dim <= 256: - return 64 - - -def round_multiple(x, m): - return (x + m - 1) // m * m - - -# torch.compile() support is only enabled for pytorch >= 2.4 -# The reason for this is that we are using the new custom_op and register_fake -# APIs, which support inplace modification of inputs in the function itself -if torch.__version__ >= "2.4.0": - _torch_custom_op_wrapper = torch.library.custom_op - _torch_register_fake_wrapper = torch.library.register_fake -else: - def noop_custom_op_wrapper(name, fn=None, /, *, mutates_args, device_types=None, schema=None): - def wrap(func): - return func - if fn is None: - return wrap - return fn - def noop_register_fake_wrapper(op, fn=None, /, *, lib=None, _stacklevel=1): - def wrap(func): - return func - if fn is None: - return wrap - return fn - _torch_custom_op_wrapper = noop_custom_op_wrapper - _torch_register_fake_wrapper = noop_register_fake_wrapper - - -@_torch_custom_op_wrapper("flash_attn::_flash_attn_forward", mutates_args=(), device_types=_get_device()) -def _flash_attn_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - return_softmax: bool -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - out, softmax_lse, S_dmask, rng_state = flash_attn.fwd( - q, - k, - v, - None, - alibi_slopes, - dropout_p, - softmax_scale, - causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - None, - ) - return out, softmax_lse, S_dmask, rng_state - - -@_torch_register_fake_wrapper("flash_attn::_flash_attn_forward") -def _flash_attn_forward_fake( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - return_softmax: bool -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - batch_size, seqlen_q, num_heads, head_size = q.shape - seqlen_k = k.shape[1] - out = torch.empty_like(q) - softmax_lse = torch.empty((batch_size, num_heads, seqlen_q), dtype=torch.float32, device=q.device, layout=q.layout) - p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) - if return_softmax: - p = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128), round_multiple(seqlen_k, 128)), dtype=q.dtype, device=q.device, layout=q.layout) - rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) - - return out, softmax_lse, p, rng_state - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_forward = torch.ops.flash_attn._flash_attn_forward -else: - _wrapped_flash_attn_forward = _flash_attn_forward - - -@_torch_custom_op_wrapper("flash_attn::_flash_attn_varlen_forward", mutates_args=(), device_types=_get_device()) -def _flash_attn_varlen_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - alibi_slopes: Optional[torch.Tensor] = None, - return_softmax: bool = False, - block_table: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - out, softmax_lse, S_dmask, rng_state = flash_attn.varlen_fwd( - q, - k, - v, - None, - cu_seqlens_q, - cu_seqlens_k, - seqused_k, - leftpad_k, - block_table, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - zero_tensors, - causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - None, - ) - # if out.isnan().any() or softmax_lse.isnan().any(): - # breakpoint() - return out, softmax_lse, S_dmask, rng_state - - -@_torch_register_fake_wrapper("flash_attn::_flash_attn_varlen_forward") -def _flash_attn_varlen_forward_fake( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - alibi_slopes: Optional[torch.Tensor] = None, - return_softmax: bool = False, - block_table: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - paged_kv = block_table is not None - batch_size = cu_seqlens_q.numel() - 1 - total_q, num_heads, _ = q.shape - - out = torch.empty_like(q) - softmax_lse = torch.empty((num_heads, total_q), dtype=torch.float32, device=q.device, layout=q.layout) - p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) - seqlen_q_rounded = round_multiple(max_seqlen_q, 128) - seqlen_k_rounded = round_multiple(max_seqlen_k, 128) - if return_softmax: - p = torch.empty((batch_size, num_heads, seqlen_q_rounded, seqlen_k_rounded), dtype=q.dtype, device=q.device, layout=q.layout) - rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) - return out, softmax_lse, p, rng_state - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_varlen_forward = torch.ops.flash_attn._flash_attn_varlen_forward -else: - _wrapped_flash_attn_varlen_forward = _flash_attn_varlen_forward - - -@_torch_custom_op_wrapper("flash_attn::_flash_attn_backward", mutates_args=("dq", "dk", "dv"), device_types=_get_device()) -def _flash_attn_backward( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, -) -> 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)] - ( - dq, - dk, - dv, - softmax_d, - ) = flash_attn.bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - alibi_slopes, - dropout_p, - softmax_scale, - causal, - window_size_left, - window_size_right, - softcap, - deterministic, - None, - rng_state, - ) - return softmax_d - - -@_torch_register_fake_wrapper("flash_attn::_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, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, -) -> torch.Tensor: - dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] - if dq is None: - dq = torch.empty_like(q) - if dk is None: - dk = torch.empty_like(k) - if dv is None: - dv = torch.empty_like(v) - batch_size, seqlen_q, num_heads, _ = q.shape - softmax_d = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128)), device=q.device, dtype=torch.float32) - - return softmax_d - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_backward = torch.ops.flash_attn._flash_attn_backward -else: - _wrapped_flash_attn_backward = _flash_attn_backward - - -@_torch_custom_op_wrapper("flash_attn::_flash_attn_varlen_backward", mutates_args=("dq", "dk", "dv"), device_types=_get_device()) -def _flash_attn_varlen_backward( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> 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)] - ( - dq, - dk, - dv, - softmax_d, - ) = flash_attn.varlen_bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - zero_tensors, - causal, - window_size_left, - window_size_right, - softcap, - deterministic, - None, - rng_state, - ) - # if dk.isnan().any() or dk.isnan().any() or dv.isnan().any() or softmax_d.isnan().any(): - # breakpoint() - return softmax_d - - -@_torch_register_fake_wrapper("flash_attn::_flash_attn_varlen_backward") -def _flash_attn_varlen_backward_fake( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> torch.Tensor: - dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] - batch_size = cu_seqlens_q.numel() - 1 - total_q, num_heads, _ = q.shape - - if dq is None: - dq = torch.empty_like(q) - if dk is None: - dk = torch.empty_like(k) - if dv is None: - dv = torch.empty_like(v) - softmax_d = torch.empty((num_heads, total_q + 128 * batch_size), device=q.device, dtype=torch.float32) - - return softmax_d - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_varlen_backward = torch.ops.flash_attn._flash_attn_varlen_backward -else: - _wrapped_flash_attn_varlen_backward = _flash_attn_varlen_backward - - -class FlashAttnQKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and qkv.requires_grad - if softmax_scale is None: - softmax_scale = qkv.shape[-1] ** (-0.5) - q, k, v = qkv[:, :, 0].detach(), qkv[:, :, 1].detach(), qkv[:, :, 2].detach() - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) - dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dqkv[:, :, 0], - dqkv[:, :, 1], - dqkv[:, :, 2], - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dqkv = dqkv[..., : dout.shape[-1]] # We could have padded the head dimension - return dqkv, None, None, None, None, None, None, None, None, None - - -class FlashAttnVarlenQKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - cu_seqlens, - max_seqlen, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and qkv.requires_grad - if softmax_scale is None: - softmax_scale = qkv.shape[-1] ** (-0.5) - q, k, v = qkv[:, 0].detach(), qkv[:, 1].detach(), qkv[:, 2].detach() - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens, - cu_seqlens, - max_seqlen, - max_seqlen, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=None, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens, rng_state) - ctx.dropout_p = dropout_p - ctx.max_seqlen = max_seqlen - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens, rng_state = ctx.saved_tensors - qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) - dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dqkv[:, 0], - dqkv[:, 1], - dqkv[:, 2], - cu_seqlens, - cu_seqlens, - ctx.max_seqlen, - ctx.max_seqlen, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - 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 FlashAttnKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - kv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, kv] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - k, v = kv[:, :, 0].detach(), kv[:, :, 1].detach() - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - dq = torch.empty_like(q) - kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) - dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dkv[:, :, 0], - dkv[:, :, 1], - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dkv = dkv[..., : dout.shape[-1]] - return dq, dkv, None, None, None, None, None, None, None, None, None - - -class FlashAttnVarlenKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, kv] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - k, v = kv[:, 0].detach(), kv[:, 1].detach() - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=None, - ) - if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state - ) - ctx.dropout_p = dropout_p - 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.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors - dq = torch.empty_like(q) - kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) - dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dkv[:, 0], - dkv[:, 1], - cu_seqlens_q, - cu_seqlens_k, - ctx.max_seqlen_q, - ctx.max_seqlen_k, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dkv = dkv[..., : dout.shape[-1]] - return dq, dkv, None, None, None, None, None, None, None, None, None, None, None, None, None - - -class FlashAttnFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - k, - v, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, k, v] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dk = dk[..., : dout.shape[-1]] - dv = dv[..., : dout.shape[-1]] - return dq, dk, dv, 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, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - block_table, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, k, v] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=block_table, - ) - if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state - ) - ctx.dropout_p = dropout_p - 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.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors - dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - ctx.max_seqlen_q, - ctx.max_seqlen_k, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dk = dk[..., : dout.shape[-1]] - dv = dv[..., : dout.shape[-1]] - return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None - - -def flash_attn_qkvpacked_func( - qkv, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # <=0.0 means deactivate - alibi_slopes=None, - deterministic=False, - 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, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_kvpacked_func( - q, - kv, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, -): - """dropout_p should be set to 0.0 during evaluation - If 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 K, V. - 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) - kv: (batch_size, seqlen, 2, 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. - softcap: float. Anything > 0 activates softcapping 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). - 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 FlashAttnKVPackedFunc.apply( - q, - kv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_func( - q, - k, - v, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - 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). - 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 FlashAttnFunc.apply( - q, - k, - v, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_qkvpacked_func( - qkv, - cu_seqlens, - max_seqlen, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - 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_varlen_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_varlen_kvpacked_func and flash_attn_varlen_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: (total, 3, nheads, headdim), where total = total number of tokens in the batch. - cu_seqlens: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into qkv. - max_seqlen: int. Maximum sequence length in the batch. - 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenQKVPackedFunc.apply( - qkv, - cu_seqlens, - max_seqlen, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_kvpacked_func( - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, -): - """dropout_p should be set to 0.0 during evaluation - If 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 K, V. - 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: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch. - kv: (total_k, 2, nheads_k, headdim), where total_k = total number of key tokens in the batch. - cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into q. - cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into kv. - max_seqlen_q: int. Maximum query sequence length in the batch. - max_seqlen_k: int. Maximum key sequence length in the batch. - 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 + 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenKVPackedFunc.apply( - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_func( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, - block_table=None, -): - """dropout_p should be set to 0.0 during evaluation - Supports multi-query and grouped-query attention (MQA/GQA) by passing in K, V 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: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch. - k: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. - v: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. - cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into q. - cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into kv. - max_seqlen_q: int. Maximum query sequence length in the batch. - max_seqlen_k: int. Maximum key sequence length in the batch. - 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 + 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenFunc.apply( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - block_table, - False if q.device.type == "cpu" else torch.is_grad_enabled(), - ) - - -def flash_attn_with_kvcache( - q, - k_cache, - v_cache, - k=None, - v=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, - block_table: Optional[torch.Tensor] = None, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - rotary_interleaved=True, - alibi_slopes=None, - num_splits=0, - 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 block_table, - or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_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) if there's no block_table, - or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_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). Similar to k. - 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. - block_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). - 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. - 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" - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-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) - cache_batch_idx = maybe_contiguous(cache_batch_idx) - block_table = maybe_contiguous(block_table) - out, softmax_lse = flash_attn.fwd_kvcache( - q, - k_cache, - v_cache, - k, - v, - cache_seqlens, - rotary_cos, - rotary_sin, - cache_batch_idx, - cache_leftpad, - block_table, - alibi_slopes, - None, - softmax_scale, - causal, - window_size[0], - window_size[1], - softcap, - rotary_interleaved, - num_splits, - ) - return (out, softmax_lse) if return_softmax_lse else out diff --git a/build/torch29-cxx11-cu129-aarch64-linux/layers/__init__.py b/build/torch29-cxx11-cu129-aarch64-linux/layers/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/build/torch29-cxx11-cu129-aarch64-linux/layers/patch_embed.py b/build/torch29-cxx11-cu129-aarch64-linux/layers/patch_embed.py deleted file mode 100644 index 05562f8e8bcdb58e947c6f402a49eacd2d031871..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-aarch64-linux/layers/patch_embed.py +++ /dev/null @@ -1,67 +0,0 @@ -# We use the same API as https://github.com/rwightman/pytorch-image-models/blob/v0.6.11/timm/models/layers/patch_embed.py -# But we use nn.Linear instead of Conv2d and it's about 8x faster. - -from functools import partial - -import torch.nn as nn -from einops import rearrange -from torch import _assert -from torch.nn.modules.utils import _pair - -try: - from flash_attn.ops.fused_dense import FusedDense -except ImportError: - FusedDense = None - - -class PatchEmbed(nn.Module): - """2D Image to Patch Embedding""" - - def __init__( - self, - img_size=224, - patch_size=16, - in_chans=3, - embed_dim=768, - norm_layer=None, - flatten=True, - bias=True, - fused_bias_fc=False, - ): - super().__init__() - img_size = _pair(img_size) - patch_size = _pair(patch_size) - self.img_size = img_size - self.patch_size = patch_size - self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1]) - self.num_patches = self.grid_size[0] * self.grid_size[1] - self.flatten = flatten - if fused_bias_fc and FusedDense is None: - raise ImportError("fused_dense is not installed") - - linear_cls = nn.Linear if not fused_bias_fc or not bias else FusedDense - self.proj = linear_cls(in_chans * patch_size[0] * patch_size[1], embed_dim, bias=bias) - self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() - - def forward(self, x): - _, _, H, W = x.shape - _assert( - H == self.img_size[0], - f"Input image height ({H}) doesn't match model ({self.img_size[0]}).", - ) - _assert( - W == self.img_size[1], - f"Input image width ({W}) doesn't match model ({self.img_size[1]}).", - ) - x = self.proj( - rearrange( - x, - "b c (h p1) (w p2) -> b h w (c p1 p2)", - p1=self.patch_size[0], - p2=self.patch_size[1], - ) - ) - if self.flatten: - x = rearrange(x, "b h w c -> b (h w) c") - x = self.norm(x) - return x diff --git a/build/torch29-cxx11-cu129-aarch64-linux/layers/rotary.py b/build/torch29-cxx11-cu129-aarch64-linux/layers/rotary.py deleted file mode 100644 index d1bfc21fc7de1dd287e8f382847b194a48075981..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-aarch64-linux/layers/rotary.py +++ /dev/null @@ -1,483 +0,0 @@ -# Copyright (c) 2025, Tri Dao - -import math -from functools import partial -from typing import Optional, Tuple, Union - -import torch -from torch import Tensor - -from einops import rearrange, repeat -# from flash_attn.ops.triton.rotary import apply_rotary -from ..ops.triton.rotary import apply_rotary - - -def rotate_half(x, interleaved=False): - if not interleaved: - x1, x2 = x.chunk(2, dim=-1) - return torch.cat((-x2, x1), dim=-1) - else: - x1, x2 = x[..., ::2], x[..., 1::2] - return rearrange(torch.stack((-x2, x1), dim=-1), "... d two -> ... (d two)", two=2) - - -def apply_rotary_emb_torch(x, cos, sin, interleaved=False): - """ - x: (batch_size, seqlen, nheads, headdim) - cos, sin: (seqlen, rotary_dim / 2) or (batch_size, seqlen, rotary_dim / 2) - """ - ro_dim = cos.shape[-1] * 2 - assert ro_dim <= x.shape[-1] - cos = repeat(cos, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)") - sin = repeat(sin, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)") - return torch.cat( - [x[..., :ro_dim] * cos + rotate_half(x[..., :ro_dim], interleaved) * sin, x[..., ro_dim:]], - dim=-1, - ) - - -class ApplyRotaryEmb(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x, - cos, - sin, - interleaved=False, - inplace=False, - seqlen_offsets: Union[int, Tensor] = 0, - cu_seqlens: Optional[Tensor] = None, - max_seqlen: Optional[int] = None, - ): - out = apply_rotary( - x, - cos, - sin, - seqlen_offsets=seqlen_offsets, - cu_seqlens=cu_seqlens, - max_seqlen=max_seqlen, - interleaved=interleaved, - inplace=inplace, - ) - if isinstance(seqlen_offsets, int): - ctx.save_for_backward(cos, sin, cu_seqlens) # Can't save int with save_for_backward - ctx.seqlen_offsets = seqlen_offsets - else: - ctx.save_for_backward(cos, sin, cu_seqlens, seqlen_offsets) - ctx.seqlen_offsets = None - ctx.interleaved = interleaved - ctx.inplace = inplace - ctx.max_seqlen = max_seqlen - return out if not inplace else x - - @staticmethod - def backward(ctx, do): - seqlen_offsets = ctx.seqlen_offsets - if seqlen_offsets is None: - cos, sin, cu_seqlens, seqlen_offsets = ctx.saved_tensors - else: - cos, sin, cu_seqlens = ctx.saved_tensors - dx = apply_rotary( - do, - cos, - sin, - seqlen_offsets=seqlen_offsets, - cu_seqlens=cu_seqlens, - max_seqlen=ctx.max_seqlen, - interleaved=ctx.interleaved, - inplace=ctx.inplace, - conjugate=True, - ) - return dx, None, None, None, None, None, None, None - - -def apply_rotary_emb( - x, - cos, - sin, - interleaved=False, - inplace=False, - seqlen_offsets: Union[int, Tensor] = 0, - cu_seqlens: Optional[Tensor] = None, - max_seqlen: Optional[int] = None, -): - """ - Arguments: - x: (batch_size, seqlen, nheads, headdim) if cu_seqlens is None - else (total_seqlen, nheads, headdim) - cos, sin: (seqlen_rotary, rotary_dim / 2) - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead - of 1st half and 2nd half (GPT-NeoX style). - inplace: if True, apply rotary embedding in-place. - seqlen_offsets: (batch_size,) or int. Each sequence in x is shifted by this amount. - Most commonly used in inference when we have KV cache. - cu_seqlens: (batch + 1,) or None - max_seqlen: int - Return: - out: (batch_size, seqlen, nheads, headdim) if cu_seqlens is None - else (total_seqlen, nheads, headdim) - rotary_dim must be <= headdim - Apply rotary embedding to the first rotary_dim of x. - """ - return ApplyRotaryEmb.apply( - x, cos, sin, interleaved, inplace, seqlen_offsets, cu_seqlens, max_seqlen - ) - - -# For backward compatibility -apply_rotary_emb_func = apply_rotary_emb - - -def _apply_rotary_emb_qkv( - qkv, - cos, - sin, - cos_k=None, - sin_k=None, - interleaved=False, - inplace=False, - conjugate=False, - seqlen_offsets: Union[int, Tensor] = 0, - num_heads_q: Optional[int] = None, -): - apply_rotary_fn = partial( - apply_rotary, - interleaved=interleaved, - inplace=inplace, - conjugate=conjugate, - seqlen_offsets=seqlen_offsets - ) - if cos_k is None and sin_k is None and qkv.is_contiguous(): - # Call 1 kernel instead of 2 kernels - # We need qkv to be contiguous so that when we reshape to combine (3, nheads) - # dimensions, we get the same tensor - if qkv.dim() == 5: - batch, seqlen, three, nheads, headdim = qkv.shape - assert three == 3 - # qk = rearrange(qkv[:, :, :2], "b s t h d -> b s (t h) d") - qk = qkv[:, :, :2].reshape(batch, seqlen, -1, headdim) - qk = apply_rotary_fn(qk, cos, sin) - else: - assert qkv.dim() == 4 - assert num_heads_q is not None - num_heads_k = (qkv.shape[2] - num_heads_q) // 2 - assert qkv.shape[2] == num_heads_q + 2 * num_heads_k - qk = qkv[:, :, :num_heads_q + num_heads_k] - qk = apply_rotary_fn(qk, cos, sin) - if not inplace: - if qkv.dim() == 5: - qkv = torch.cat([rearrange(qk, "b s (t h) d -> b s t h d", t=2), qkv[:, :, 2:]], dim=2) - else: - qkv = torch.cat([qk, qkv[:, :, num_heads_q + num_heads_k :]], dim=2) - else: - cos_k = cos if cos_k is None else cos_k - sin_k = sin if sin_k is None else sin_k - if qkv.dim() == 5: - batch, seqlen, three, nheads, headdim = qkv.shape - assert three == 3 - q, k = qkv[:, :, 0], qkv[:, :, 1] - else: - assert qkv.dim() == 4 - assert num_heads_q is not None - num_heads_k = (qkv.shape[2] - num_heads_q) // 2 - assert qkv.shape[2] == num_heads_q + 2 * num_heads_k - q, k = qkv[:, :, :num_heads_q], qkv[:, :, num_heads_q : num_heads_q + num_heads_k] - q = apply_rotary_fn(q, cos, sin) - k = apply_rotary_fn(k, cos_k, sin_k) - if not inplace: - if qkv.dim() == 5: - qkv = torch.stack([q, k, qkv[:, :, 2]], dim=2) - else: - qkv = torch.cat([q, k, qkv[:, :, num_heads_q + num_heads_k:]], dim=2) - return qkv - - -class ApplyRotaryEmbQKV_(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - cos, - sin, - cos_k=None, - sin_k=None, - interleaved=False, - seqlen_offsets: Union[int, torch.Tensor] = 0, - num_heads_q: Optional[int] = None, - ): - # apply_rotary_emb_qkv_inplace( - qkv = _apply_rotary_emb_qkv( - qkv, cos, sin, cos_k, sin_k, interleaved=interleaved, inplace=True, - seqlen_offsets=seqlen_offsets, num_heads_q=num_heads_q, - ) - if isinstance(seqlen_offsets, int): - ctx.save_for_backward(cos, sin, cos_k, sin_k) - ctx.seqlen_offsets = seqlen_offsets - else: - ctx.save_for_backward(cos, sin, cos_k, sin_k, seqlen_offsets) - ctx.seqlen_offsets = None - ctx.interleaved = interleaved - ctx.num_heads_q = num_heads_q - return qkv - - @staticmethod - def backward(ctx, dqkv): - seqlen_offsets = ctx.seqlen_offsets - if seqlen_offsets is None: - cos, sin, cos_k, sin_k, seqlen_offsets = ctx.saved_tensors - else: - cos, sin, cos_k, sin_k = ctx.saved_tensors - dqkv = _apply_rotary_emb_qkv( - dqkv, cos, sin, cos_k, sin_k, interleaved=ctx.interleaved, inplace=True, - seqlen_offsets=seqlen_offsets, num_heads_q=ctx.num_heads_q, conjugate=True, - ) - return dqkv, None, None, None, None, None, None, None - - -def apply_rotary_emb_qkv_( - qkv, - cos, - sin, - cos_k=None, - sin_k=None, - interleaved=False, - seqlen_offsets: Union[int, torch.Tensor] = 0, - num_heads_q: Optional[int] = None, -): - """ - Arguments: - qkv: (batch_size, seqlen, 3, nheads, headdim) or (batch_size, seqlen, num_heads_q + 2 * num_heads_k, headdim). - If qkv has shape (batch_size, seqlen, num_heads_q + 2 * num_heads_k, headdim) (e.g. MQA / GQA), - then num_heads_q must be provided. - cos, sin: (seqlen, rotary_dim / 2) - cos_k, sin_k: (seqlen, rotary_dim / 2), optional - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead of - 1st half and 2nd half (GPT-NeoX style). - seqlen_offsets: (batch_size,) or int. Each sequence in Q and K is shifted by this amount. - Most commonly used in inference when we have KV cache. - Return: - qkv: (batch_size, seqlen, 3, nheads, headdim) or (batch_size, seqlen, num_heads_q + 2 * num_heads_k, headdim) - rotary_dim must be <= headdim - Apply rotary embedding *inplace* to the first rotary_dim of Q and K. - """ - return ApplyRotaryEmbQKV_.apply( - qkv, cos, sin, cos_k, sin_k, interleaved, seqlen_offsets, num_heads_q - ) - - -class ApplyRotaryEmbKV_(torch.autograd.Function): - - @staticmethod - def forward(ctx, kv, cos, sin, interleaved=False, seqlen_offsets: Union[int, torch.Tensor] = 0): - batch, seqlen, two, nheads, headdim = kv.shape - assert two == 2 - k = kv[:, :, 0] - apply_rotary( - k, cos, sin, seqlen_offsets=seqlen_offsets, interleaved=interleaved, inplace=True - ) - if isinstance(seqlen_offsets, int): - ctx.save_for_backward(cos, sin) # Can't save int with save_for_backward - ctx.seqlen_offsets = seqlen_offsets - else: - ctx.save_for_backward(cos, sin, seqlen_offsets) - ctx.seqlen_offsets = None - ctx.interleaved = interleaved - return kv - - @staticmethod - def backward(ctx, dkv): - seqlen_offsets = ctx.seqlen_offsets - if seqlen_offsets is None: - cos, sin, seqlen_offsets = ctx.saved_tensors - else: - cos, sin = ctx.saved_tensors - apply_rotary( - dkv[:, :, 0], - cos, - sin, - seqlen_offsets=seqlen_offsets, - interleaved=ctx.interleaved, - inplace=True, - conjugate=True, - ) - return dkv, None, None, None, None - - -apply_rotary_emb_kv_ = ApplyRotaryEmbKV_.apply - - -def apply_rotary_emb_kv_( - kv, - cos, - sin, - interleaved=False, - seqlen_offsets: Union[int, torch.Tensor] = 0, -): - """ - Arguments: - kv: (batch_size, seqlen, 2, nheads, headdim) - cos, sin: (seqlen, rotary_dim / 2) - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead of - 1st half and 2nd half (GPT-NeoX style). - seqlen_offsets: (batch_size,) or int. Each sequence in Q and K is shifted by this amount. - Most commonly used in inference when we have KV cache. - Return: - kv: (batch_size, seqlen, 2, nheads, headdim) - rotary_dim must be <= headdim - Apply rotary embedding *inplace* to the first rotary_dim of K. - """ - return ApplyRotaryEmbKV_.apply(kv, cos, sin, interleaved, seqlen_offsets) - - -class RotaryEmbedding(torch.nn.Module): - """ - The rotary position embeddings from RoFormer_ (Su et. al). - A crucial insight from the method is that the query and keys are - transformed by rotation matrices which depend on the relative positions. - - Other implementations are available in the Rotary Transformer repo_ and in - GPT-NeoX_, GPT-NeoX was an inspiration - - .. _RoFormer: https://arxiv.org/abs/2104.09864 - .. _repo: https://github.com/ZhuiyiTechnology/roformer - .. _GPT-NeoX: https://github.com/EleutherAI/gpt-neox - - If scale_base is not None, this implements XPos (Sun et al., https://arxiv.org/abs/2212.10554). - A recommended value for scale_base is 512: https://github.com/HazyResearch/flash-attention/issues/96 - Reference: https://github.com/sunyt32/torchscale/blob/main/torchscale/component/xpos_relative_position.py - """ - - def __init__( - self, - dim: int, - base=10000.0, - interleaved=False, - scale_base=None, - device=None, - ): - """ - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead - of 1st half and 2nd half (GPT-NeoX style). - """ - super().__init__() - self.dim = dim - self.base = float(base) - # Generate and save the inverse frequency buffer (non trainable) - inv_freq = self._compute_inv_freq(device) - self.register_buffer("inv_freq", inv_freq, persistent=False) - self.interleaved = interleaved - self.scale_base = scale_base - scale = ( - (torch.arange(0, dim, 2, device=device, dtype=torch.float32) + 0.4 * dim) / (1.4 * dim) - if scale_base is not None - else None - ) - self.register_buffer("scale", scale, persistent=False) - - self._seq_len_cached = 0 - self._cos_cached = None - self._sin_cached = None - self._cos_k_cached = None - self._sin_k_cached = None - - def _compute_inv_freq(self, device=None): - return 1.0 / ( - self.base - ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim) - ) - - def _update_cos_sin_cache(self, seqlen, device=None, dtype=None): - # Reset the tables if the sequence length has changed, - # if we're on a new device (possibly due to tracing for instance), - # or if we're switching from inference mode to training - if ( - seqlen > self._seq_len_cached - or self._cos_cached is None - or self._cos_cached.device != device - or self._cos_cached.dtype != dtype - or (self.training and self._cos_cached.is_inference()) - ): - self._seq_len_cached = seqlen - # We want fp32 here, not self.inv_freq.dtype, since the model could be loaded in bf16 - # And the output of arange can be quite large, so bf16 would lose a lot of precision. - t = torch.arange(seqlen, device=device, dtype=torch.float32) - # We want fp32 here as well since inv_freq will be multiplied with t, and the output - # will be large. Having it in bf16 will lose a lot of precision and cause the - # cos & sin output to change significantly. - # We want to recompute self.inv_freq if it was not loaded in fp32 - if self.inv_freq.dtype != torch.float32: - inv_freq = self._compute_inv_freq(device=device) - else: - inv_freq = self.inv_freq - # Don't do einsum, it converts fp32 to bf16 under AMP - # freqs = torch.einsum("i,j->ij", t, self.inv_freq) - freqs = torch.outer(t, inv_freq) - if self.scale is None: - self._cos_cached = torch.cos(freqs).to(dtype) - self._sin_cached = torch.sin(freqs).to(dtype) - else: - power = ( - torch.arange(seqlen, dtype=self.scale.dtype, device=self.scale.device) - - seqlen // 2 - ) / self.scale_base - scale = self.scale.to(device=power.device) ** rearrange(power, "s -> s 1") - # We want the multiplication by scale to happen in fp32 - self._cos_cached = (torch.cos(freqs) * scale).to(dtype) - self._sin_cached = (torch.sin(freqs) * scale).to(dtype) - self._cos_k_cached = (torch.cos(freqs) / scale).to(dtype) - self._sin_k_cached = (torch.sin(freqs) / scale).to(dtype) - - def forward( - self, - qkv: torch.Tensor, - kv: Optional[torch.Tensor] = None, - seqlen_offset: Union[int, torch.Tensor] = 0, - max_seqlen: Optional[int] = None, - num_heads_q: Optional[int] = None, - ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: - """ - qkv: (batch, seqlen, 3, nheads, headdim) or (batch, seqlen, num_heads_q + 2 * num_heads_k, headdim) - if kv is none, else it's just q of shape (batch, seqlen, nheads, headdim). - If qkv has shape (batch, seqlen, num_heads_q + 2 * num_heads_k, headdim) (e.g. MQA / GQA), - then num_heads_q must be provided. - kv: (batch, seqlen, 2, nheads, headdim) - seqlen_offset: (batch_size,) or int. Each sequence in x is shifted by this amount. - Most commonly used in inference when we have KV cache. - If it's a tensor of shape (batch_size,), then to update the cos / sin cache, one - should pass in max_seqlen, which will update the cos / sin cache up to that length. - Apply rotary embedding *inplace* to qkv and / or kv. - """ - seqlen = qkv.shape[1] - if max_seqlen is not None: - self._update_cos_sin_cache(max_seqlen, device=qkv.device, dtype=qkv.dtype) - elif isinstance(seqlen_offset, int): - self._update_cos_sin_cache(seqlen + seqlen_offset, device=qkv.device, dtype=qkv.dtype) - if kv is None: - return apply_rotary_emb_qkv_( - qkv, - self._cos_cached, - self._sin_cached, - self._cos_k_cached if self.scale is not None else None, - self._sin_k_cached if self.scale is not None else None, - interleaved=self.interleaved, - seqlen_offsets=seqlen_offset, - num_heads_q=num_heads_q, - ) - else: - q = qkv - q = apply_rotary_emb_func( - q, - self._cos_cached, - self._sin_cached, - interleaved=self.interleaved, - inplace=True, - seqlen_offsets=seqlen_offset, - ) - kv = apply_rotary_emb_kv_( - kv, - self._cos_cached if self.scale is None else self._cos_k_cached, - self._sin_cached if self.scale is None else self._sin_k_cached, - interleaved=self.interleaved, - seqlen_offsets=seqlen_offset, - ) - return q, kv diff --git a/build/torch29-cxx11-cu129-aarch64-linux/metadata.json b/build/torch29-cxx11-cu129-aarch64-linux/metadata.json deleted file mode 100644 index 8b00d7fb7685121edf7a97bb007415122444f1d5..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-aarch64-linux/metadata.json +++ /dev/null @@ -1,14 +0,0 @@ -{ - "version": 1, - "license": "BSD-3-Clause", - "python-depends": [], - "backend": { - "type": "cuda", - "archs": [ - "10.0", - "12.0", - "8.0", - "9.0" - ] - } -} diff --git a/build/torch29-cxx11-cu129-aarch64-linux/ops/__init__.py b/build/torch29-cxx11-cu129-aarch64-linux/ops/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/build/torch29-cxx11-cu129-aarch64-linux/ops/activations.py b/build/torch29-cxx11-cu129-aarch64-linux/ops/activations.py deleted file mode 100644 index 7c09649fc41e12d5a360c5672825d8380bc7ec80..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-aarch64-linux/ops/activations.py +++ /dev/null @@ -1,135 +0,0 @@ -# Copied from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/model/layers/activations.py -import math - -import torch -import torch.nn as nn -import torch.nn.functional as F - -# 1/sqrt(2*pi)-> 0.3989423 -# 1/sqrt(2) -> 0.70710678 -# sqrt(2/pi) -> 0.79788456 - -# this function is tanh approximation of gelu -# actual gelu is: -# x * 0.5 * (1.0 + torch.erf(x * 0.70710678)) -@torch.jit.script -def bias_gelu(y, bias): - x = bias + y - return (x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))).to(dtype=y.dtype) - - -# gradient of tanh approximation of gelu -# gradient of actual gelu is: -# 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x) -@torch.jit.script -def bias_gelu_back(g, y, bias): - """Assume that y has shape (B, D) and bias has shape (D)""" - x = bias + y - tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)) - # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243 - ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( - 1 + tanh_out - ) - grad_y = ff * g - return grad_y.to(dtype=y.dtype), grad_y.sum(dim=(0), dtype=bias.dtype) - - -class GeLUFunction(torch.autograd.Function): - @staticmethod - # bias is an optional argument - def forward(ctx, input, bias): - ctx.save_for_backward(input, bias) - return bias_gelu(input, bias) - - @staticmethod - def backward(ctx, grad_output): - input, bias = ctx.saved_tensors - tmp = bias_gelu_back(grad_output, input, bias) - return tmp, tmp - - -bias_gelu_impl = GeLUFunction.apply - -# this function is tanh approximation of gelu -# actual gelu is: -# x * 0.5 * (1.0 + torch.erf(x * 0.70710678)) -@torch.jit.script -def gelu_fwd(x): - return (x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))).to(dtype=x.dtype) - - -# gradient of tanh approximation of gelu -# gradient of actual gelu is: -# 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x) -@torch.jit.script -def gelu_bwd(g, x): - tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)) - # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243 - ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( - 1 + tanh_out - ) - return (ff * g).to(dtype=x.dtype) - - -class FastGeLUFunction(torch.autograd.Function): - @staticmethod - # bias is an optional argument - def forward(ctx, input): - ctx.save_for_backward(input) - return gelu_fwd(input) - - @staticmethod - def backward(ctx, grad_output): - (input,) = ctx.saved_tensors - tmp = gelu_bwd(grad_output, input) - return tmp - - -fast_gelu_impl = FastGeLUFunction.apply - - -@torch.jit.script -def relu_bwd(g, x): - return torch.where(x >= 0, g, 0.0).to(dtype=x.dtype) - - -@torch.jit.script -def sqrelu_fwd(x): - r = F.relu(x) - return (r * r).to(dtype=x.dtype) - - -@torch.jit.script -def sqrelu_bwd(g, x): - return (2.0 * g * F.relu(x)).to(dtype=x.dtype) - - -swiglu_fwd_codestring = """ -template T swiglu_fwd(T x, T y) { - return float(x) * float(y) / (1.0f + ::exp(-float(x))); -} -""" -swiglu_bwd_codestring = """ -template void swiglu_bwd(T x, T y, T g, T& dx, T& dy) { - float x_sigmoid = 1.0f / (1.0f + ::exp(-float(x))); - dx = x_sigmoid * (1 + float(x) * (1.0f - x_sigmoid)) * float(g) * float(y); - dy = float(x) * x_sigmoid * float(g); -} -""" -swiglu_fwd = torch.cuda.jiterator._create_jit_fn(swiglu_fwd_codestring) -swiglu_bwd = torch.cuda.jiterator._create_multi_output_jit_fn(swiglu_bwd_codestring, num_outputs=2) - - -class SwiGLUFunction(torch.autograd.Function): - - @staticmethod - def forward(ctx, x, y): - ctx.save_for_backward(x, y) - return swiglu_fwd(x, y) - - @staticmethod - def backward(ctx, dout): - x, y = ctx.saved_tensors - return swiglu_bwd(x, y, dout) - -swiglu = SwiGLUFunction.apply diff --git a/build/torch29-cxx11-cu129-aarch64-linux/ops/fused_dense.py b/build/torch29-cxx11-cu129-aarch64-linux/ops/fused_dense.py deleted file mode 100644 index 6b4033d134e4093fe278f7b3f8c7d3128ce9f36d..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-aarch64-linux/ops/fused_dense.py +++ /dev/null @@ -1,688 +0,0 @@ -# Copyright (c) 2023, Tri Dao. -# Inspired by https://github.com/NVIDIA/apex/blob/master/apex/fused_dense/fused_dense.py -# We make it work with pytorch amp and with bfloat16. -# The TensorParallel linear modules are inspired by https://github.com/NVIDIA/apex/blob/master/apex/transformer/tensor_parallel/layers.py -from functools import partial -from typing import Optional - -# import fused_dense_cuda # from apex -import fused_dense_lib as fused_dense_cuda -import torch -import torch.nn as nn -import torch.nn.functional as F -from torch import Tensor -from torch.distributed import ProcessGroup - -from flash_attn.utils.torch import custom_fwd, custom_bwd -from flash_attn.ops.activations import gelu_bwd, relu_bwd, sqrelu_bwd, sqrelu_fwd -from flash_attn.utils.distributed import ( - all_gather_raw, - all_reduce, - all_reduce_raw, - reduce_scatter, - reduce_scatter_raw, -) - - -class FusedDenseFunc(torch.autograd.Function): - @staticmethod - @custom_fwd - def forward( - ctx, x, weight, bias, return_residual=False, process_group=None, sequence_parallel=True - ): - """ - If process_group is not None and sequence_parallel=True, we're doing Tensor Parallel - with sequence parallelism: we do an all_gather_raw of x before doing the matmul. - """ - ctx.compute_weight_gradient = weight.requires_grad - ctx.return_residual = return_residual - ctx.process_group = process_group - ctx.sequence_parallel = sequence_parallel - - if torch.is_autocast_enabled(): - x = x.to(dtype=torch.get_autocast_gpu_dtype()) - x = x.contiguous() - if process_group is not None and sequence_parallel: - # We want to kick off the all_gather early, before weight dtype conversion - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - else: - total_x = x - - if torch.is_autocast_enabled(): - weight = weight.to(dtype=torch.get_autocast_gpu_dtype()) - bias = bias.to(dtype=torch.get_autocast_gpu_dtype()) if bias is not None else None - weight = weight.contiguous() - if process_group is not None and sequence_parallel: - handle_x.wait() - batch_shape, n = total_x.shape[:-1], total_x.shape[-1] - batch_dim = batch_shape.numel() - # https://github.com/pytorch/pytorch/blob/5b51849b48a7dbccd297286cc0110def4706f9e7/aten/src/ATen/native/cuda/Blas.cpp#L174 - if min(batch_dim, n, *weight.shape) > 65535 * 32: - raise RuntimeError("fused_dense only supports matrix dims <= 2M") - output = F.linear(total_x, weight, bias) - if ctx.compute_weight_gradient: - ctx.save_for_backward(x, weight) - else: - ctx.save_for_backward(weight) - return output if not return_residual else (output, x) - - @staticmethod - @custom_bwd - def backward(ctx, grad_output, *args): - grad_output = grad_output.contiguous() - if ctx.return_residual: - (grad_input,) = args - grad_input = grad_input.contiguous() - process_group = ctx.process_group - sequence_parallel = ctx.sequence_parallel - if ctx.compute_weight_gradient: - x, weight = ctx.saved_tensors - if process_group is not None and sequence_parallel: - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - else: - total_x = x - else: - (weight,) = ctx.saved_tensors - total_x = None - batch_shape = grad_output.shape[:-1] - batch_dim = batch_shape.numel() - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - if ctx.needs_input_grad[0]: - if not ctx.return_residual: - grad_input = F.linear(grad_output, weight.t()) - else: - grad_input = torch.addmm( - grad_input.reshape(batch_dim, grad_input.shape[-1]), grad_output, weight - ) - grad_input = grad_input.reshape(*batch_shape, grad_input.shape[-1]) - if process_group is not None: - reduce_fn = reduce_scatter_raw if sequence_parallel else all_reduce_raw - grad_input, handle_grad_input = reduce_fn(grad_input, process_group, async_op=True) - else: - grad_input = None - if ctx.needs_input_grad[1]: - assert ctx.compute_weight_gradient - if process_group is not None and sequence_parallel: - handle_x.wait() - grad_weight, grad_bias = fused_dense_cuda.linear_bias_wgrad( - total_x.reshape(batch_dim, total_x.shape[-1]), grad_output, ctx.needs_input_grad[2] - ) - else: - grad_weight = None - grad_bias = grad_output if ctx.needs_input_grad[2] else None - if process_group is not None and ctx.needs_input_grad[0]: - handle_grad_input.wait() - return grad_input, grad_weight, grad_bias, None, None, None - - -def fused_dense_func( - x: Tensor, - weight: Tensor, - bias: Optional[Tensor] = None, - return_residual: bool = False, - process_group: Optional[ProcessGroup] = None, - sequence_parallel: bool = True, -): - dtype_eligible = x.dtype in [torch.float16, torch.bfloat16] or ( - x.dtype == torch.float32 and torch.is_autocast_enabled() - ) - if x.is_cuda and weight.is_cuda and (bias is None or bias.is_cuda) and dtype_eligible: - return FusedDenseFunc.apply( - x, weight, bias, return_residual, process_group, sequence_parallel - ) - else: - assert process_group is None - out = F.linear(x, weight, bias) - return out if not return_residual else (out, x) - - -class FusedDense(nn.Linear): - def __init__( - self, - in_features: int, - out_features: int, - bias: bool = True, - return_residual: bool = False, - device=None, - dtype=None, - ) -> None: - super().__init__(in_features, out_features, bias=bias, device=device, dtype=dtype) - self.return_residual = return_residual - - def forward(self, x, process_group=None): - """ - If process_group is not None, we're doing Tensor Parallel with sequence parallelism: - we do an all_gather of x before doing the matmul. - """ - return fused_dense_func( - x, - self.weight, - self.bias, - return_residual=self.return_residual, - process_group=process_group, - ) - - -class ColumnParallelLinear(nn.Linear): - def __init__( - self, - in_features: int, - out_features: int, - process_group: ProcessGroup, - bias: bool = True, - sequence_parallel=True, - multiple_of=1, - device=None, - dtype=None, - ) -> None: - world_size = torch.distributed.get_world_size(process_group) - if out_features % multiple_of: - raise ValueError(f"out_features ({out_features}) must be a multiple of {multiple_of}") - multiple = out_features // multiple_of - # We want to split @multiple across world_size, but it could be an uneven split - div = multiple // world_size - mod = multiple % world_size - # The first @mod ranks get @div + 1 copies, the rest get @div copies - local_multiple = div + int(torch.distributed.get_rank(process_group) < mod) - super().__init__( - in_features, local_multiple * multiple_of, bias=bias, device=device, dtype=dtype - ) - self.process_group = process_group - self.sequence_parallel = sequence_parallel - - def forward(self, x): - # If self.sequence_parallel is True, we're doing Tensor Parallel with sequence parallelism: - # we do an all_gather of x before doing the matmul. - # If not, then the input is already gathered. - return fused_dense_func( - x, - self.weight, - self.bias, - process_group=self.process_group, - sequence_parallel=self.sequence_parallel, - ) - - -class RowParallelLinear(nn.Linear): - def __init__( - self, - in_features: int, - out_features: int, - process_group: ProcessGroup, - bias: bool = True, - sequence_parallel=True, - multiple_of=1, - device=None, - dtype=None, - ) -> None: - world_size = torch.distributed.get_world_size(process_group) - rank = torch.distributed.get_rank(process_group) - if in_features % multiple_of: - raise ValueError(f"in_features ({in_features}) must be a multiple of {multiple_of}") - multiple = in_features // multiple_of - # We want to split @multiple across world_size, but it could be an uneven split - div = multiple // world_size - mod = multiple % world_size - # The first @mod ranks get @div + 1 copies, the rest get @div copies - local_multiple = div + int(torch.distributed.get_rank(process_group) < mod) - # Only rank 0 will have bias - super().__init__( - local_multiple * multiple_of, - out_features, - bias=bias and rank == 0, - device=device, - dtype=dtype, - ) - self.process_group = process_group - self.sequence_parallel = sequence_parallel - - def forward(self, x): - """ - We're doing Tensor Parallel with sequence parallelism: we do the matmul and then - a reduce_scatter of the result. - """ - out = fused_dense_func(x, self.weight, self.bias) - reduce_fn = reduce_scatter if self.sequence_parallel else all_reduce - return reduce_fn(out, self.process_group) - - -class FusedMLPFunc(torch.autograd.Function): - @staticmethod - @custom_fwd - def forward( - ctx, - x, - weight1, - bias1, - weight2, - bias2, - activation="gelu_approx", - save_pre_act=True, - return_residual=False, - checkpoint_lvl=0, - heuristic=0, - process_group=None, - sequence_parallel=True, - ): - """ - If process_group is not None and sequence_parallel=True, we're doing Tensor Parallel - with sequence parallelism: we do an all_gather of x before doing the matmul. - If sequence_parallel=False, then the input is already gathered. - - checkpoint_lvl: - 0: no recomputation in the bwd - 1: recompute gelu_out / relu_out in the bwd - 2: recompute pre_act and gelu_out / relu_out in the bwd - """ - assert -1 <= heuristic <= 4 - assert activation in ["gelu_approx", "relu", "sqrelu"] - if activation == "sqrelu": - assert heuristic == -1 - if not save_pre_act: - checkpoint_lvl = 2 - assert checkpoint_lvl in [0, 1, 2] - ctx.return_residual = return_residual - ctx.process_group = process_group - ctx.sequence_parallel = sequence_parallel - ctx.checkpoint_lvl = checkpoint_lvl - ctx.activation = activation - ctx.heuristic = heuristic - - if torch.is_autocast_enabled(): - x = x.to(dtype=torch.get_autocast_gpu_dtype()) - x = x.contiguous() - if process_group is not None and sequence_parallel: - # We want to kick off the all_gather early, before weight dtype conversion - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - else: - total_x = x - - if torch.is_autocast_enabled(): - dtype = torch.get_autocast_gpu_dtype() - weight1, weight2 = [a.to(dtype=dtype) for a in [weight1, weight2]] - bias1 = bias1.to(dtype=dtype) if bias1 is not None else None - bias2 = bias2.to(dtype=dtype) if bias2 is not None else None - weight1 = weight1.contiguous() - bias1 = bias1.contiguous() if bias1 is not None else None - weight2 = weight2.contiguous() - bias2 = bias2.contiguous() if bias2 is not None else None - if process_group is not None and sequence_parallel: - handle_x.wait() - batch_shape, n = total_x.shape[:-1], total_x.shape[-1] - batch_dim = batch_shape.numel() - # https://github.com/pytorch/pytorch/blob/5b51849b48a7dbccd297286cc0110def4706f9e7/aten/src/ATen/native/cuda/Blas.cpp#L174 - if min(batch_dim, n, *weight1.shape, *weight2.shape) > 65535 * 32: - raise RuntimeError("fused_dense only supports matrix dims <= 2M") - if heuristic == -1: - pre_act = F.linear(total_x, weight1, bias1) - activation_fn = ( - partial(F.gelu, approximate="tanh") - if activation == "gelu_approx" - else (sqrelu_fwd if activation == "sqrelu" else F.relu) - ) - with torch.jit.fuser("fuser2"): - output1 = activation_fn(pre_act) - # This is before adding bias1 - # pre_act = F.linear(total_x.reshape(batch_dim, n), weight1) - # with torch.jit.fuser('fuser2'): - # output1 = bias_gelu(pre_act, bias1) - else: - is_gelu = activation == "gelu_approx" - output1, *rest = fused_dense_cuda.linear_act_forward( - total_x.reshape(batch_dim, n), weight1, bias1, is_gelu, save_pre_act, heuristic - ) - if save_pre_act: - pre_act = rest[0] - output2 = F.linear(output1, weight2, bias2) - if checkpoint_lvl == 0 or (checkpoint_lvl == 1 and activation == "relu"): - # For RELU the pre_act is very small (just a bit-mask) so we just save it - ctx.save_for_backward(x, weight1, weight2, pre_act, output1) - elif checkpoint_lvl == 1: - ctx.save_for_backward(x, weight1, weight2, pre_act) - elif checkpoint_lvl == 2: - ctx.save_for_backward(x, weight1, weight2, bias1) - output2 = output2.reshape(*batch_shape, output2.shape[-1]) - return output2 if not return_residual else (output2, x) - - @staticmethod - @custom_bwd - def backward(ctx, grad_output, *args): - grad_output = grad_output.contiguous() - checkpoint_lvl = ctx.checkpoint_lvl - activation = ctx.activation - activation_fn = ( - partial(F.gelu, approximate="tanh") - if activation == "gelu_approx" - else (sqrelu_fwd if activation == "sqrelu" else F.relu) - ) - if ctx.return_residual: - (grad_input,) = args - grad_input = grad_input.contiguous() - process_group = ctx.process_group - sequence_parallel = ctx.sequence_parallel - x, weight1, weight2, *rest = ctx.saved_tensors - if process_group is None or not sequence_parallel: - total_x = x - batch_shape = grad_output.shape[:-1] - batch_dim = batch_shape.numel() - if checkpoint_lvl in [0, 1]: - if process_group is not None and sequence_parallel: - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - if checkpoint_lvl == 0 or (checkpoint_lvl == 1 and activation == "relu"): - pre_act, output1 = rest - elif checkpoint_lvl == 1: - (pre_act,) = rest - with torch.jit.fuser("fuser2"): - output1 = activation_fn(pre_act) - elif checkpoint_lvl == 2: - (bias1,) = rest - if process_group is not None and sequence_parallel: - total_x, _ = all_gather_raw(x, process_group) - if ctx.heuristic == -1: - pre_act = F.linear(total_x, weight1, bias1) - with torch.jit.fuser("fuser2"): - output1 = activation_fn(pre_act) - else: - output1, pre_act = fused_dense_cuda.linear_act_forward( - total_x.reshape(batch_dim, total_x.shape[-1]), - weight1, - bias1, - activation == "gelu_approx", - True, - ctx.heuristic, - ) - - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - output1 = output1.reshape(batch_dim, output1.shape[-1]) - pre_act = pre_act.reshape(batch_dim, pre_act.shape[-1]) - if ctx.needs_input_grad[3]: - grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad( - output1, grad_output, ctx.needs_input_grad[4] - ) - else: - grad_weight2 = None - grad_bias2 = grad_output if ctx.needs_input_grad[4] else None - if ctx.heuristic == -1: - # grad_pre_act = matmul_dgelu(grad_output, weight2, pre_act) - grad_output1 = F.linear(grad_output, weight2.t()) - activation_grad_fn = ( - gelu_bwd - if activation == "gelu_approx" - else (sqrelu_bwd if activation == "sqrelu" else relu_bwd) - ) - with torch.jit.fuser("fuser2"): - grad_pre_act = activation_grad_fn(grad_output1, pre_act) - else: - # The cublasLt epilogue has to compute both gelu/relu grad and bias grad, we can't - # just compute gelu/relu grad - grad_pre_act, grad_bias1 = fused_dense_cuda.bias_act_linear_dgrad_bgrad( - weight2, grad_output, pre_act, activation == "gelu_approx", ctx.heuristic - ) - if not ctx.needs_input_grad[2]: - grad_bias1 = None - if ctx.needs_input_grad[0]: - if not ctx.return_residual: - grad_input = F.linear(grad_pre_act, weight1.t()) - else: - grad_input = torch.addmm( - grad_input.reshape(batch_dim, grad_input.shape[-1]), grad_pre_act, weight1 - ) - grad_input = grad_input.reshape(*batch_shape, grad_input.shape[-1]) - if process_group is not None: - reduce_fn = reduce_scatter_raw if sequence_parallel else all_reduce_raw - grad_input, handle_grad_input = reduce_fn(grad_input, process_group, async_op=True) - else: - grad_input = None - if ctx.heuristic == -1: - if ctx.needs_input_grad[1]: - if process_group is not None and sequence_parallel and checkpoint_lvl != 2: - handle_x.wait() - grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_wgrad( - total_x.reshape(batch_dim, total_x.shape[-1]), - grad_pre_act, - ctx.needs_input_grad[2], - ) - else: - grad_weight1 = None - grad_bias1 = grad_pre_act if ctx.needs_input_grad[2] else None - else: - if ctx.needs_input_grad[1]: - if process_group is not None and sequence_parallel and checkpoint_lvl != 2: - handle_x.wait() - grad_weight1 = F.linear( - grad_pre_act.t(), total_x.reshape(batch_dim, total_x.shape[-1]).t() - ) - else: - grad_weight1 = None - if process_group is not None and ctx.needs_input_grad[0]: - handle_grad_input.wait() - return ( - grad_input, - grad_weight1, - grad_bias1, - grad_weight2, - grad_bias2, - None, - None, - None, - None, - None, - None, - None, - ) - - -def fused_mlp_func( - x: Tensor, - weight1: Tensor, - weight2: Tensor, - bias1: Optional[Tensor] = None, - bias2: Optional[Tensor] = None, - activation: str = "gelu_approx", - save_pre_act: bool = True, - return_residual: bool = False, - checkpoint_lvl: int = 0, - heuristic: int = 0, - process_group: Optional[ProcessGroup] = None, - sequence_parallel: bool = True, -): - assert activation in ["gelu_approx", "relu", "sqrelu"] - dtype_eligible = x.dtype in [torch.float16, torch.bfloat16] or ( - x.dtype == torch.float32 and torch.is_autocast_enabled() - ) - # If we save pre-activation, dimension must be divisible by 128 (relu) or 8 (gelu) - dim_eligible = not save_pre_act or (x.shape[-1] % (128 if activation == "relu" else 8) == 0) - if ( - x.is_cuda - and weight1.is_cuda - and weight2.is_cuda - and (bias1 is None or bias1.is_cuda) - and (bias2 is None or bias2.is_cuda) - and dtype_eligible - and dim_eligible - ): - return FusedMLPFunc.apply( - x, - weight1, - bias1, - weight2, - bias2, - activation, - save_pre_act, - return_residual, - checkpoint_lvl, - heuristic, - process_group, - sequence_parallel, - ) - else: - assert process_group is None - pre_act = F.linear(x, weight1, bias1) - activation_fn = ( - partial(F.gelu, approximate="tanh") - if activation == "gelu_approx" - else partial(F.relu, inplace=True) - ) - output1 = activation_fn(pre_act) - output2 = F.linear(output1, weight2, bias2) - return output2 if not return_residual else (output2, x) - - -class FusedMLP(nn.Module): - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - bias1=True, - bias2=True, - activation="gelu_approx", - return_residual=False, - checkpoint_lvl=0, - heuristic="auto", - device=None, - dtype=None, - ): - """ - If process_group is not None, we're doing Tensor Parallel with sequence parallelism: - we do an all_gather of x before doing the matmul, gelu, then matmul. - Finally we do a reduce_scatter of the output. - - checkpoint_lvl (increasing lvl means slower but more memory saving): - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute pre_act and gelu_out in the bwd - heuristic: - -1: don't fuse gemm + gelu (separate kernel) - 0..4: use this heuristic for the algo section in the fused gemm + gelu - 'auto': heuristic will be picked automatically: - For CUDA >= 11.8, we set heuristic=0 for both fp16 and bf16 for best perf. - For CUDA <= 11.7, we set heuristic=1 for fp16 and heuristic=-1 for bf16. - For H100, we set heuristic=-1 for both fp16 and bf16 as the fused cuBlasLt implementation - is slower than the unfused version. - return_residual: whether to return the input x along with the output. This is for - performance reason: for post-norm architecture, returning the input allows us - to fuse the backward of nn.Linear with the residual connection. - """ - assert checkpoint_lvl in [0, 1, 2] - assert activation in ["gelu_approx", "relu", "sqrelu"] - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features * 4 - self.activation = activation - self.return_residual = return_residual - self.checkpoint_lvl = checkpoint_lvl - self.heuristic = heuristic if activation != "sqrelu" else -1 - self.fc1 = nn.Linear(in_features, hidden_features, bias=bias1, **factory_kwargs) - self.fc2 = nn.Linear(hidden_features, out_features, bias=bias2, **factory_kwargs) - - def forward(self, x, process_group=None): - dtype = x.dtype if not torch.is_autocast_enabled() else torch.get_autocast_gpu_dtype() - if self.heuristic == "auto": - if self.activation == "gelu_approx": - if torch.cuda.get_device_capability("cuda") == (9, 0): - heuristic = -1 - else: - cuda_ver = tuple(map(int, torch.version.cuda.split("."))) - heuristic = 0 if cuda_ver >= (11, 8) else (1 if dtype == torch.float16 else -1) - else: - heuristic = 0 - else: - heuristic = self.heuristic - out = fused_mlp_func( - x, - self.fc1.weight, - self.fc2.weight, - self.fc1.bias, - self.fc2.bias, - activation=self.activation, - save_pre_act=self.training, - return_residual=self.return_residual, - checkpoint_lvl=self.checkpoint_lvl, - heuristic=heuristic, - process_group=process_group, - ) - if self.return_residual: - out, x = out - if process_group is not None: - out = reduce_scatter(out, process_group) - return out if not self.return_residual else (out, x) - - -class ParallelFusedMLP(nn.Module): - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - activation="gelu_approx", - process_group: ProcessGroup = None, - bias1=True, - bias2=True, - sequence_parallel=True, - checkpoint_lvl=0, - heuristic="auto", - device=None, - dtype=None, - ): - """ - process_group is required. We're doing Tensor Parallel with sequence parallelism: - we do an all_gather of x before doing the matmul, gelu, then matmul. - Finally we do a reduce_scatter of the output. - - checkpoint_lvl (increasing lvl means slower but more memory saving): - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute pre_act and gelu_out in the bwd - heuristic: - -1: don't fuse gemm + gelu (separate kernel) - 0..4: use this heuristic for the algo section in the fused gemm + gelu - 'auto': heuristic will be picked automatically: - For CUDA >= 11.8, we set heuristic=0 for both fp16 and bf16 for best perf. - For CUDA <= 11.7, we set heuristic=1 for fp16 and heuristic=-1 for bf16. - """ - assert checkpoint_lvl in [0, 1, 2] - assert activation in ["gelu_approx", "relu", "sqrelu"] - assert process_group is not None - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features * 4 - self.activation = activation - self.process_group = process_group - self.sequence_parallel = sequence_parallel - self.checkpoint_lvl = checkpoint_lvl - self.heuristic = heuristic if activation != "sqrelu" else -1 - self.fc1 = ColumnParallelLinear( - in_features, hidden_features, process_group, bias=bias1, **factory_kwargs - ) - self.fc2 = RowParallelLinear( - hidden_features, out_features, process_group, bias=bias2, **factory_kwargs - ) - - def forward(self, x): - dtype = x.dtype if not torch.is_autocast_enabled() else torch.get_autocast_gpu_dtype() - if self.heuristic == "auto": - if self.activation == "gelu_approx": - cuda_ver = tuple(map(int, torch.version.cuda.split("."))) - heuristic = 0 if cuda_ver >= (11, 8) else (1 if dtype == torch.float16 else -1) - else: - heuristic = 0 - else: - heuristic = self.heuristic - out = fused_mlp_func( - x, - self.fc1.weight, - self.fc2.weight, - self.fc1.bias, - self.fc2.bias, - activation=self.activation, - save_pre_act=self.training, - checkpoint_lvl=self.checkpoint_lvl, - heuristic=heuristic, - process_group=self.process_group, - sequence_parallel=self.sequence_parallel, - ) - reduce_fn = reduce_scatter if self.sequence_parallel else all_reduce - return reduce_fn(out, self.process_group) diff --git a/build/torch29-cxx11-cu129-aarch64-linux/ops/layer_norm.py b/build/torch29-cxx11-cu129-aarch64-linux/ops/layer_norm.py deleted file mode 100644 index 4b6cd798fd02844ef9cd3897f8ab95e490e638bf..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-aarch64-linux/ops/layer_norm.py +++ /dev/null @@ -1,800 +0,0 @@ -# Copyright (c) 2022, Tri Dao. -# Adapted from https://github.com/NVIDIA/apex/blob/master/apex/contrib/layer_norm/layer_norm.py - -import dropout_layer_norm -import torch -from torch.nn import init - - -def maybe_align(x, alignment_in_bytes=16): - """Assume that x already has last dim divisible by alignment_in_bytes""" - # TD [2023-07-04] I'm not 100% sure that clone will align the memory - # https://discuss.pytorch.org/t/how-to-ensure-that-tensor-data-ptr-is-aligned-to-16-bytes/183440 - return x if x.data_ptr() % alignment_in_bytes == 0 else x.clone() - - -def _dropout_add_layer_norm_forward( - x0, - residual, - gamma, - beta, - rowscale, - colscale, - dropout_p, - epsilon, - residual_in_fp32=False, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes""" - hidden_size = gamma.numel() - x0mat = x0.view((-1, hidden_size)) - residualmat = residual.view((-1, hidden_size)) if residual is not None else None - rowscale = rowscale.view(-1) if rowscale is not None else None - zmat, xmat, dmask, mu, rsigma = dropout_layer_norm.dropout_add_ln_fwd( - x0mat, - residualmat, - gamma, - beta, - rowscale, - colscale, - None, - None, - dropout_p, - epsilon, - 1.0, - 0, - None, - residual_in_fp32, - is_rms_norm, - ) - # dmask is None if dropout_p == 0.0 - # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype - return zmat, xmat if xmat is not None else x0mat, dmask, mu, rsigma - - -def _dropout_add_layer_norm_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - rowscale, - colscale, - dropout_p, - has_residual, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes - dx == None means that it was a post-norm architecture - (x = drop(x0) + residual was not returned in the fwd). - x0 must not be None if we have colscale. - """ - hidden_size = gamma.numel() - xmat = x.view((-1, hidden_size)) - dzmat = dz.view(xmat.shape) - dxmat = dx.view(xmat.shape) if dx is not None else None - x0mat = x0.view((-1, hidden_size)) if x0 is not None else None - rowscale = rowscale.view(-1) if rowscale is not None else None - if colscale is not None: - assert x0 is not None, "x0 is required to compute the gradient of colscale" - dx0mat, dresidualmat, dgamma, dbeta, _, _, *rest = dropout_layer_norm.dropout_add_ln_bwd( - dzmat, - dxmat, - xmat, - x0mat, - dmask, - mu, - rsigma, - gamma, - rowscale, - colscale, - None, - None, - dropout_p, - 1.0, - 0, - has_residual, - is_rms_norm, - ) - # dresidualmat is None if not has_residual - if colscale is None: - return dx0mat, dresidualmat, dgamma, dbeta - else: - dcolscale = rest[0] - return dx0mat, dresidualmat, dgamma, dbeta, dcolscale - - -def _dropout_add_layer_norm_subset_forward( - x0, - residual, - gamma, - beta, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32=False, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes""" - hidden_size = gamma.numel() - x0mat = x0.view((-1, hidden_size)) - residualmat = residual.view((-1, hidden_size)) if residual is not None else None - x0_subset = x0_subset.view(-1) if x0_subset is not None else None - out_subset = out_subset.view(-1) if out_subset is not None else None - zmat, xmat, dmask, mu, rsigma = dropout_layer_norm.dropout_add_ln_fwd( - x0mat, - residualmat, - gamma, - beta, - None, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - None, - residual_in_fp32, - is_rms_norm, - ) - # dmask is None if dropout_p == 0.0 - # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype - return zmat, xmat if xmat is not None else x0mat, dmask, mu, rsigma - - -def _dropout_add_layer_norm_subset_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - colscale, - x0_subset, - out_subset, - dropout_p, - rowscale_const, - x0_numrows, - has_residual, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes - dx == None means that it was a post-norm architecture - (x = drop(x0) + residual was not returned in the fwd). - x0 must not be None if we have colscale. - """ - hidden_size = gamma.numel() - xmat = x.view((-1, hidden_size)) - dzmat = dz.view(-1, hidden_size) - dxmat = dx.view(xmat.shape) if dx is not None else None - x0mat = x0.view((-1, hidden_size)) if x0 is not None else None - x0_subset = x0_subset.view(-1) if x0_subset is not None else None - out_subset = out_subset.view(-1) if out_subset is not None else None - if colscale is not None: - assert x0 is not None, "x0 is required to compute the gradient of colscale" - dx0mat, dresidualmat, dgamma, dbeta, _, _, *rest = dropout_layer_norm.dropout_add_ln_bwd( - dzmat, - dxmat, - xmat, - x0mat, - dmask, - mu, - rsigma, - gamma, - None, - colscale, - x0_subset, - out_subset, - dropout_p, - rowscale_const, - x0_numrows, - has_residual, - is_rms_norm, - ) - # dresidualmat is None if not has_residual - if colscale is None: - return dx0mat, dresidualmat, dgamma, dbeta - else: - dcolscale = rest[0] - return dx0mat, dresidualmat, dgamma, dbeta, dcolscale - - -def _dropout_add_layer_norm_parallel_residual_forward( - x0, - x1, - residual, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - residual_in_fp32=False, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes""" - hidden_size = gamma0.numel() - x0mat = x0.view((-1, hidden_size)) - x1mat = x1.view((-1, hidden_size)) if x1 is not None else None - residualmat = residual.view((-1, hidden_size)) if residual is not None else None - ( - z0mat, - z1mat, - xmat, - dmask0, - dmask1, - mu, - rsigma, - ) = dropout_layer_norm.dropout_add_ln_parallel_residual_fwd( - x0mat, - x1mat, - residualmat, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - None, - residual_in_fp32, - is_rms_norm, - ) - # dmask0 and dmask1 are None if dropout_p == 0.0 - # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype - return z0mat, z1mat, xmat if xmat is not None else x0mat, dmask0, dmask1, mu, rsigma - - -def _dropout_add_layer_norm_parallel_residual_backward( - dz0, - dz1, - dx, - x, - dmask0, - dmask1, - mu, - rsigma, - gamma0, - gamma1, - dropout_p, - has_x1, - has_residual, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes - dx == None means that it was a post-norm architecture - (x = drop(x0) + residual was not returned in the fwd). - """ - hidden_size = gamma0.numel() - xmat = x.view((-1, hidden_size)) - dz0mat = dz0.view(xmat.shape) - dz1mat = dz1.view(xmat.shape) if dz1 is not None else None - dxmat = dx.view(xmat.shape) if dx is not None else None - ( - dx0mat, - dx1mat, - dresidualmat, - dgamma0, - dbeta0, - dgamma1, - dbeta1, - *rest, - ) = dropout_layer_norm.dropout_add_ln_parallel_residual_bwd( - dz0mat, - dz1mat, - dxmat, - xmat, - dmask0, - dmask1, - mu, - rsigma, - gamma0, - gamma1, - dropout_p, - has_x1, - has_residual, - is_rms_norm, - ) - # dresidualmat is None if not has_residual - return dx0mat, dx1mat, dresidualmat, dgamma0, dbeta0, dgamma1, dbeta1 - - -class DropoutAddLayerNormFn(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x0, - residual, - gamma, - beta, - rowscale, - colscale, - dropout_p, - epsilon, - residual_in_fp32=False, - prenorm=False, - is_rms_norm=False, - return_dmask=False, - ): - x0 = maybe_align(x0.contiguous(), 16) - residual = maybe_align(residual.contiguous(), 16) if residual is not None else None - gamma = maybe_align(gamma.contiguous(), 16) - beta = maybe_align(beta.contiguous(), 16) if beta is not None else None - rowscale = maybe_align(rowscale.contiguous(), 16) if rowscale is not None else None - colscale = maybe_align(colscale.contiguous(), 16) if colscale is not None else None - zmat, xmat, dmask, mu, rsigma = _dropout_add_layer_norm_forward( - x0, - residual, - gamma, - beta, - rowscale, - colscale, - dropout_p, - epsilon, - residual_in_fp32, - is_rms_norm, - ) - # Only need to save x0 if we need to compute gradient wrt colscale - x0_saved = x0 if colscale is not None else None - ctx.save_for_backward( - xmat.view(x0.shape), x0_saved, dmask, gamma, mu, rsigma, rowscale, colscale - ) - ctx.prenorm = prenorm - ctx.dropout_p = dropout_p - ctx.has_residual = residual is not None - ctx.is_rms_norm = is_rms_norm - ctx.has_beta = beta is not None - if not return_dmask: - return ( - zmat.view(x0.shape) if not prenorm else (zmat.view(x0.shape), xmat.view(x0.shape)) - ) - else: - dmask = ( - dmask.view(x0.shape) - if dropout_p > 0.0 - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - ctx.mark_non_differentiable(dmask) - return ( - (zmat.view(x0.shape), dmask) - if not prenorm - else (zmat.view(x0.shape), xmat.view(x0.shape), dmask) - ) - - @staticmethod - def backward(ctx, dz, *args): - # assert dz.is_contiguous() - dz = maybe_align(dz.contiguous(), 16) # this happens! - dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None - x, x0, dmask, gamma, mu, rsigma, rowscale, colscale = ctx.saved_tensors - # x0 is None if colscale is None - dropout_p = ctx.dropout_p - has_residual = ctx.has_residual - dx0mat, dresidualmat, dgamma, dbeta, *rest = _dropout_add_layer_norm_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - rowscale, - colscale, - dropout_p, - has_residual, - ctx.is_rms_norm, - ) - dx0 = dx0mat.view(x.shape) - dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None - dcolscale = rest[0] if colscale is not None else None - return ( - dx0, - dresidual, - dgamma, - dbeta if ctx.has_beta else None, - None, - dcolscale, - None, - None, - None, - None, - None, - None, - ) - - -class DropoutAddLayerNormSubsetFn(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x0, - residual, - gamma, - beta, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32=False, - prenorm=False, - is_rms_norm=False, - return_dmask=False, - ): - x0 = maybe_align(x0.contiguous(), 16) - residual = maybe_align(residual.contiguous(), 16) if residual is not None else None - gamma = maybe_align(gamma.contiguous(), 16) - beta = maybe_align(beta.contiguous(), 16) if beta is not None else None - colscale = maybe_align(colscale.contiguous(), 16) if colscale is not None else None - zmat, xmat, dmask, mu, rsigma = _dropout_add_layer_norm_subset_forward( - x0, - residual, - gamma, - beta, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32, - is_rms_norm, - ) - # Only need to save x0 if we need to compute gradient wrt colscale - x0_saved = x0 if colscale is not None else None - x_shape = (-1, *x0.shape[1:]) - ctx.save_for_backward( - xmat.view(x_shape), x0_saved, dmask, gamma, mu, rsigma, colscale, x0_subset, out_subset - ) - ctx.prenorm = prenorm - ctx.dropout_p = dropout_p - ctx.rowscale_const = rowscale_const - ctx.x0_numrows = x0.shape[:-1].numel() - ctx.has_residual = residual is not None - ctx.is_rms_norm = is_rms_norm - ctx.has_beta = beta is not None - z_shape = (-1, *x0.shape[1:]) - if not return_dmask: - return zmat.view(z_shape) if not prenorm else (zmat.view(z_shape), xmat.view(x0.shape)) - else: - z = zmat.view(z_shape) - dmask = ( - dmask.view(x0.shape) - if dropout_p > 0.0 - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - ctx.mark_non_differentiable(dmask) - return (z, dmask) if not prenorm else (z, xmat.view(x_shape), dmask) - - @staticmethod - def backward(ctx, dz, *args): - # assert dz.is_contiguous() - dz = maybe_align(dz.contiguous(), 16) # this happens! - dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None - x, x0, dmask, gamma, mu, rsigma, colscale, x0_subset, out_subset = ctx.saved_tensors - # x0 is None if colscale is None - dropout_p = ctx.dropout_p - has_residual = ctx.has_residual - dx0mat, dresidualmat, dgamma, dbeta, *rest = _dropout_add_layer_norm_subset_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - colscale, - x0_subset, - out_subset, - dropout_p, - ctx.rowscale_const, - ctx.x0_numrows, - has_residual, - ctx.is_rms_norm, - ) - dx0 = dx0mat.view(-1, *x.shape[1:]) - dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None - dcolscale = rest[0] if colscale is not None else None - return ( - dx0, - dresidual, - dgamma, - dbeta if ctx.has_beta else None, - dcolscale, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) - - -class DropoutAddLayerNormParallelResidualFn(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x0, - x1, - residual, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - residual_in_fp32=False, - prenorm=False, - is_rms_norm=False, - return_dmask=False, - ): - x0 = maybe_align(x0.contiguous(), 16) - x1 = maybe_align(x1.contiguous(), 16) if x1 is not None else None - residual = maybe_align(residual.contiguous(), 16) if residual is not None else None - gamma0 = maybe_align(gamma0.contiguous(), 16) - beta0 = maybe_align(beta0.contiguous(), 16) if beta0 is not None else None - gamma1 = maybe_align(gamma1.contiguous(), 16) if gamma1 is not None else None - beta1 = maybe_align(beta1.contiguous(), 16) if beta1 is not None else None - ( - z0mat, - z1mat, - xmat, - dmask0, - dmask1, - mu, - rsigma, - ) = _dropout_add_layer_norm_parallel_residual_forward( - x0, - x1, - residual, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - residual_in_fp32, - is_rms_norm, - ) - ctx.save_for_backward(xmat.view(x0.shape), dmask0, dmask1, gamma0, gamma1, mu, rsigma) - ctx.prenorm = prenorm - ctx.dropout_p = dropout_p - ctx.has_x1 = x1 is not None - ctx.has_residual = residual is not None - ctx.is_rms_norm = is_rms_norm - ctx.has_beta = beta0 is not None - z = (z0mat.view(x0.shape), z1mat.view(x0.shape) if z1mat is not None else None) - if not return_dmask: - return z if not prenorm else (*z, xmat.view(x0.shape)) - else: - dmask0 = ( - dmask0.view(x0.shape) - if dropout_p > 0.0 - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - dmask1 = ( - dmask1.view(x0.shape) - if dropout_p > 0.0 and x1 is not None - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - ctx.mark_non_differentiable(dmask0) - ctx.mark_non_differentiable(dmask1) - return ( - (*z, dmask0, dmask1) if not prenorm else (*z, xmat.view(x0.shape), dmask0, dmask1) - ) - - @staticmethod - def backward(ctx, dz0, dz1, *args): - dz0 = maybe_align(dz0.contiguous(), 16) # this happens! - dz1 = maybe_align(dz1.contiguous(), 16) if dz1 is not None else None - dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None - x, dmask0, dmask1, gamma0, gamma1, mu, rsigma = ctx.saved_tensors - dropout_p = ctx.dropout_p - has_x1 = ctx.has_x1 - has_residual = ctx.has_residual - ( - dx0mat, - dx1mat, - dresidualmat, - dgamma0, - dbeta0, - dgamma1, - dbeta1, - ) = _dropout_add_layer_norm_parallel_residual_backward( - dz0, - dz1, - dx, - x, - dmask0, - dmask1, - mu, - rsigma, - gamma0, - gamma1, - dropout_p, - has_x1, - has_residual, - ctx.is_rms_norm, - ) - dx0 = dx0mat.view(x.shape) - dx1 = dx1mat.view(x.shape) if dx1mat is not None else None - dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None - return ( - dx0, - dx1, - dresidual, - dgamma0, - dbeta0 if ctx.has_beta else None, - dgamma1, - dbeta1 if ctx.has_beta else None, - None, - None, - None, - None, - None, - None, - ) - - -def layer_norm(x, weight, bias, epsilon): - return DropoutAddLayerNormFn.apply(x, None, weight, bias, None, None, 0.0, epsilon, False) - - -def dropout_add_layer_norm( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - rowscale=None, - layerscale=None, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormFn.apply( - x0, - residual, - weight, - bias, - rowscale, - layerscale, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - False, - return_dropout_mask, - ) - - -def dropout_add_layer_norm_subset( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - layerscale=None, - x0_subset=None, - out_subset=None, - rowscale_const=1.0, - out_numrows=0, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormSubsetFn.apply( - x0, - residual, - weight, - bias, - layerscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32, - prenorm, - False, - return_dropout_mask, - ) - - -def dropout_add_layer_norm_parallel_residual( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormParallelResidualFn.apply( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - False, - return_dropout_mask, - ) - - -class DropoutAddLayerNorm(torch.nn.Module): - def __init__( - self, - hidden_size, - prenorm=False, - p=0.0, - eps=1e-5, - residual_in_fp32=False, - device=None, - dtype=None, - ): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.prenorm = prenorm - self.p = p - self.eps = eps - self.residual_in_fp32 = residual_in_fp32 - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.bias = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.reset_parameters() - - def reset_parameters(self): - init.ones_(self.weight) - init.zeros_(self.bias) - - def forward(self, x0, residual=None): - return dropout_add_layer_norm( - x0, - residual, - self.weight, - self.bias, - self.p if self.training else 0.0, - self.eps, - prenorm=self.prenorm, - residual_in_fp32=self.residual_in_fp32, - ) diff --git a/build/torch29-cxx11-cu129-aarch64-linux/ops/rms_norm.py b/build/torch29-cxx11-cu129-aarch64-linux/ops/rms_norm.py deleted file mode 100644 index 068348d61290e3839dd082b540d898578ba1e8e2..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-aarch64-linux/ops/rms_norm.py +++ /dev/null @@ -1,174 +0,0 @@ -# Copyright (c) 2022, Tri Dao. -# Adapted from https://github.com/NVIDIA/apex/blob/master/apex/contrib/layer_norm/layer_norm.py - -import torch -from torch.nn import init - -from flash_attn.ops.layer_norm import ( - DropoutAddLayerNormFn, - DropoutAddLayerNormParallelResidualFn, - DropoutAddLayerNormSubsetFn, -) - - -def rms_norm(x, weight, epsilon): - return DropoutAddLayerNormFn.apply( - x, None, weight, None, None, None, 0.0, epsilon, False, False, True - ) - - -def dropout_add_rms_norm( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - rowscale=None, - layerscale=None, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormFn.apply( - x0, - residual, - weight, - bias, - rowscale, - layerscale, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - True, - return_dropout_mask, - ) - - -def dropout_add_rms_norm_subset( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - layerscale=None, - x0_subset=None, - out_subset=None, - rowscale_const=1.0, - out_numrows=0, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormSubsetFn.apply( - x0, - residual, - weight, - bias, - layerscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32, - prenorm, - True, - return_dropout_mask, - ) - - -def dropout_add_rms_norm_parallel_residual( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormParallelResidualFn.apply( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - True, - return_dropout_mask, - ) - - -class RMSNorm(torch.nn.Module): - def __init__(self, hidden_size, eps=1e-5, device=None, dtype=None): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.eps = eps - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.register_parameter("bias", None) - self.reset_parameters() - - def reset_parameters(self): - init.ones_(self.weight) - - def forward(self, x): - return rms_norm(x, self.weight, self.eps) - - -class DropoutAddRMSNorm(torch.nn.Module): - def __init__( - self, - hidden_size, - prenorm=False, - p=0.0, - eps=1e-5, - residual_in_fp32=False, - device=None, - dtype=None, - ): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.prenorm = prenorm - self.p = p - self.eps = eps - self.residual_in_fp32 = residual_in_fp32 - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.register_parameter("bias", None) - self.reset_parameters() - - def reset_parameters(self): - init.ones_(self.weight) - - def forward(self, x0, residual=None): - return dropout_add_rms_norm( - x0, - residual, - self.weight, - None, - self.p if self.training else 0.0, - self.eps, - prenorm=self.prenorm, - residual_in_fp32=self.residual_in_fp32, - ) diff --git a/build/torch29-cxx11-cu129-aarch64-linux/ops/triton/__init__.py b/build/torch29-cxx11-cu129-aarch64-linux/ops/triton/__init__.py deleted file mode 100644 index 8b137891791fe96927ad78e64b0aad7bded08bdc..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-aarch64-linux/ops/triton/__init__.py +++ /dev/null @@ -1 +0,0 @@ - diff --git a/build/torch29-cxx11-cu129-aarch64-linux/ops/triton/cross_entropy.py b/build/torch29-cxx11-cu129-aarch64-linux/ops/triton/cross_entropy.py deleted file mode 100644 index 1b5a415b73f236f3e05fb14b9141959559e18526..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-aarch64-linux/ops/triton/cross_entropy.py +++ /dev/null @@ -1,330 +0,0 @@ -# Copyright (c) 2023, Tri Dao. - -from typing import Tuple, Optional, Union - -import torch -import torch.nn.functional as F - -import triton -import triton.language as tl - -# `all_gather_into_tensor` and `reduce_scatter_tensor` are new placeholders for -# `_all_gather_base` and `_reduce_scatter_base`. They require the most recent -# version of PyTorch. The following 2 lines are for backward compatibility with -# older PyTorch. -if "all_gather_into_tensor" not in dir(torch.distributed): - torch.distributed.all_gather_into_tensor = torch.distributed._all_gather_base - - -@triton.heuristics( - { - "HAS_SMOOTHING": lambda args: args["smoothing"] > 0.0, - } -) -@triton.jit -def cross_entropy_fwd_kernel( - loss_ptr, # data ptrs - lse_ptr, - z_loss_ptr, - logits_ptr, - labels_ptr, - smoothing, - logit_scale, - lse_square_scale, - ignore_index, - total_classes, - class_start_idx, # Useful for tensor parallel when each rank only has a subset of classes - n_cols, # shapes - logits_row_stride, # strides - BLOCK_SIZE: tl.constexpr, - HAS_SMOOTHING: tl.constexpr, - # if SPLIT (e.g. tensor parallel), don't include the LSE in the loss since it's not the final LSE - SPLIT: tl.constexpr, - PRECOMPUTED_LSE: tl.constexpr, # If LSE is already computed (also no smoothing and logit_scale == 1.0) -): - row_idx = tl.program_id(0) - logits_ptr = logits_ptr + row_idx * logits_row_stride.to(tl.int64) - sum_logits = 0.0 # For smoothing - if not PRECOMPUTED_LSE: - # Statistics for online softmax - m_i = -float("inf") - l_i = 0.0 - for col_offset in range(0, n_cols, BLOCK_SIZE): - cols = col_offset + tl.arange(0, BLOCK_SIZE) - logits = tl.load(logits_ptr + cols, mask=cols < n_cols, other=-float("inf")).to( - tl.float32 - ) * logit_scale - if HAS_SMOOTHING: - sum_logits += tl.sum(tl.where(cols < n_cols, logits, 0.0)) - m_i_new = tl.maximum(m_i, tl.max(logits)) - l_i = tl.exp(m_i - m_i_new) * l_i + tl.sum(tl.exp(logits - m_i_new)) - m_i = m_i_new - lse = tl.log(l_i) + m_i - tl.store(lse_ptr + row_idx, lse) - else: - lse = tl.load(lse_ptr + row_idx) - label_idx = tl.load(labels_ptr + row_idx) - if label_idx == ignore_index: - loss = 0.0 - z_loss = 0.0 - else: - label_idx -= class_start_idx - if label_idx >= 0 and label_idx < n_cols: - logits_label = tl.load(logits_ptr + label_idx) * logit_scale - if HAS_SMOOTHING: - loss = ( - (lse if not SPLIT else 0.0) - - smoothing * sum_logits / total_classes - - (1 - smoothing) * logits_label - ) - else: - loss = (lse if not SPLIT else 0.0) - logits_label - else: - # If label is out of bounds, we set the CE loss to 0.0. But we still want the smoothing loss - if HAS_SMOOTHING: - loss = smoothing * ((lse if not SPLIT else 0.0) - sum_logits / total_classes) - else: - loss = 0.0 - if not SPLIT: - z_loss = lse_square_scale * lse * lse - loss += z_loss - else: - z_loss = 0.0 - tl.store(loss_ptr + row_idx, loss) - if not SPLIT: - tl.store(z_loss_ptr + row_idx, z_loss) - - -@triton.heuristics( - { - "HAS_SMOOTHING": lambda args: args["smoothing"] > 0.0, - } -) -@triton.jit -def cross_entropy_bwd_kernel( - dlogits_ptr, # data ptrs - dloss_ptr, - logits_ptr, - lse_ptr, - labels_ptr, - smoothing, - logit_scale, - lse_square_scale, - ignore_index, - total_classes, - class_start_idx, # Useful for tensor parallel when each rank only has a subset of classes - n_cols, # shapes - logits_row_stride, # strides - dlogits_row_stride, - dloss_row_stride, - BLOCK_SIZE: tl.constexpr, - HAS_SMOOTHING: tl.constexpr, -): - row_idx = tl.program_id(0) - col_block_idx = tl.program_id(1) - logits_ptr = logits_ptr + row_idx * logits_row_stride.to(tl.int64) - dlogits_ptr = dlogits_ptr + row_idx * dlogits_row_stride.to(tl.int64) - col_offsets = col_block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - label_idx = tl.load(labels_ptr + row_idx) - if label_idx != ignore_index: - dloss = tl.load(dloss_ptr + row_idx * dloss_row_stride) - else: - dloss = 0.0 - logits = tl.load(logits_ptr + col_offsets, mask=col_offsets < n_cols, other=-float("inf")).to( - tl.float32 - ) * logit_scale - lse = tl.load(lse_ptr + row_idx) - probs = tl.exp(logits - lse) - probs += 2.0 * lse_square_scale * lse * probs - label_idx -= class_start_idx - if HAS_SMOOTHING: - smooth_positive = 1.0 - smoothing - smooth_negative = smoothing / total_classes - probs = tl.where(col_offsets == label_idx, probs - smooth_positive, probs) - smooth_negative - else: - probs = tl.where(col_offsets == label_idx, probs - 1.0, probs) - tl.store(dlogits_ptr + col_offsets, (dloss * logit_scale) * probs, mask=col_offsets < n_cols) - - -class CrossEntropyLoss(torch.autograd.Function): - - @staticmethod - def forward( - ctx, - logits, - labels, - precomputed_lse=None, - smoothing=0.0, - logit_scale=1.0, - lse_square_scale=0.0, - ignore_index=-100, - inplace_backward=False, - process_group=None, - ): - # For some reason Triton generates wrong code when labels has dtype long and its address - # is not aligned to 16 bytes. The ld.global.b64 seems to load the wrong label index. - if labels.dtype == torch.long and labels.data_ptr() % 16 != 0: - labels = F.pad(labels, (0, 1))[..., :-1] - assert labels.data_ptr() % 16 == 0 - assert logit_scale > 0.0 - n_rows, n_cols = logits.shape - assert labels.shape == (n_rows,) - world_size = 1 if process_group is None else torch.distributed.get_world_size(process_group) - total_classes = world_size * n_cols - rank = 0 if process_group is None else torch.distributed.get_rank(process_group) - class_start_idx = rank * n_cols - use_precomputed_lse = precomputed_lse is not None and logit_scale == 1.0 and smoothing == 0.0 - - if logits.stride(-1) != 1: - logits = logits.contiguous() - MAX_BLOCK_SIZE = 16 * 1024 - BLOCK_SIZE = min(triton.next_power_of_2(n_cols), MAX_BLOCK_SIZE) - num_warps = ( - 4 - if BLOCK_SIZE < 2048 - else (8 if BLOCK_SIZE < 8192 else (16 if BLOCK_SIZE < 128 * 1024 else 32)) - ) - losses = torch.empty(n_rows, dtype=torch.float, device=logits.device) - if use_precomputed_lse: - assert precomputed_lse.shape == (n_rows,) - lse = precomputed_lse.contiguous() - else: - lse = torch.empty(n_rows, dtype=torch.float, device=logits.device) - z_losses = torch.empty(n_rows, dtype=torch.float, device=logits.device) - # Need this, otherwise Triton tries to launch from cuda:0 and we get - # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?) - with torch.cuda.device(logits.device.index): - cross_entropy_fwd_kernel[(n_rows,)]( - losses, # data ptrs - lse, - z_losses, - logits, - labels, - smoothing, - logit_scale, - lse_square_scale, - ignore_index, - total_classes, - class_start_idx, - n_cols, # shapes - logits.stride(0), # strides - BLOCK_SIZE=BLOCK_SIZE, # constants - SPLIT=world_size > 1, - PRECOMPUTED_LSE=use_precomputed_lse, - num_warps=num_warps, - ) - - if world_size > 1: - # If there's no smoothing, if labels are in the vocab of this partition, losses contains - # - predicted logit, and 0 otherwise. - # If there's smoothing=0.1, for labels in the vocab of this partition, losses contains - # -0.9 * predicted logit - 0.1 * sum logit / total_classes. - # For labels not in the vocab of this partition, losses contains - # -0.1 * sum logit / total_classes. - if world_size > 1: - lse_allgather = torch.empty(world_size, n_rows, dtype=lse.dtype, device=lse.device) - torch.distributed.all_gather_into_tensor(lse_allgather, lse, group=process_group) - handle_losses = torch.distributed.all_reduce( - losses, op=torch.distributed.ReduceOp.SUM, group=process_group, async_op=True - ) - lse = torch.logsumexp(lse_allgather, dim=0) - handle_losses.wait() - # After the allreduce, if there's no smoothing, the total losses are - predicted_logit, - # we just have to add the (global) lse. - # If there's smoothing=0.1, the total losses are - # -0.9 * predicted_logit - 0.1 * sum logit / total_classes. - # Again, we just have to add the (global) lse. - losses += lse - if lse_square_scale != 0.0: - z_losses = lse_square_scale * lse.square() - z_losses.masked_fill_(labels == ignore_index, 0.0) - losses += z_losses - else: - z_losses = torch.zeros_like(losses) - losses.masked_fill_(labels == ignore_index, 0.0) - - ctx.save_for_backward(logits, lse, labels) - ctx.mark_non_differentiable(z_losses) - ctx.smoothing = smoothing - ctx.logit_scale = logit_scale - ctx.lse_square_scale = lse_square_scale - ctx.ignore_index = ignore_index - ctx.total_classes = total_classes - ctx.class_start_idx = class_start_idx - ctx.inplace_backward = inplace_backward - return losses, z_losses - - @staticmethod - def backward(ctx, grad_losses, grad_z_losses): - del grad_z_losses # z_losses are only for logging. - - logits, lse, labels = ctx.saved_tensors - dlogits = logits if ctx.inplace_backward else torch.empty_like(logits) - n_rows, n_cols = logits.shape - BLOCK_SIZE = min(triton.next_power_of_2(n_cols), 4 * 1024) - num_warps = 4 if BLOCK_SIZE < 2048 else (8 if BLOCK_SIZE < 8192 else 16) - grid = lambda META: (n_rows, triton.cdiv(n_cols, META["BLOCK_SIZE"])) # noqa - # Need this, otherwise Triton tries to launch from cuda:0 and we get - # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?) - with torch.cuda.device(logits.device.index): - cross_entropy_bwd_kernel[grid]( - dlogits, # data ptrs - grad_losses, - logits, - lse, - labels, - ctx.smoothing, - ctx.logit_scale, - ctx.lse_square_scale, - ctx.ignore_index, - ctx.total_classes, - ctx.class_start_idx, - n_cols, # shapes - logits.stride(0), # strides - dlogits.stride(0), - grad_losses.stride(0), - BLOCK_SIZE=BLOCK_SIZE, # constants - num_warps=num_warps, - ) - return dlogits, None, None, None, None, None, None, None, None, None - - -def cross_entropy_loss( - logits: torch.Tensor, - labels: torch.Tensor, - precomputed_lse: Optional[torch.Tensor] = None, - label_smoothing: float = 0.0, - logit_scale: float = 1.0, - lse_square_scale: float = 0.0, - ignore_index=-100, - inplace_backward: bool = False, - process_group=None, -) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Arguments: - logits: (batch, vocab_size) - labels: (batch,) - label_smoothing: float - logit_scale: float. Multiply logits by this scale before calculating the loss. - lse_square_scale: float. If > 0, we add lse_square_scale * lse(logits) ^ 2 to the loss. - This is also referred to as "z-loss". - ignore_index: int. If labels == ignore_index, the loss is set to 0.0. - inplace_backward: bool. If True, we do the backward pass in-place by modifying the logits. - This saves memory. - process_group: if not None, we're doing Tensor Parallel: each process is responsible for - one part of the vocab. The loss will be aggregated across processes. - Returns: - losses: (batch,), float - z_losses: (batch,), float - """ - return CrossEntropyLoss.apply( - logits, - labels, - precomputed_lse, - label_smoothing, - logit_scale, - lse_square_scale, - ignore_index, - inplace_backward, - process_group, - ) diff --git a/build/torch29-cxx11-cu129-aarch64-linux/ops/triton/k_activations.py b/build/torch29-cxx11-cu129-aarch64-linux/ops/triton/k_activations.py deleted file mode 100644 index efb83c358eb4a85d069ee340a3c83f418f9a805b..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-aarch64-linux/ops/triton/k_activations.py +++ /dev/null @@ -1,162 +0,0 @@ -# Adapted from https://github.com/facebookresearch/xformers/blob/main/xformers/triton/k_activations.py -# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved. -# -# This source code is licensed under the BSD license found in the -# LICENSE file in the root directory of this source tree. - -import math -from enum import Enum -from typing import Optional - -import triton -import triton.language as tl - -_sqrt2pi = math.sqrt(2.0 / math.pi) -_sqrt1_2 = math.sqrt(1.0 / 2) -_gaussian_pdf_normalization = 1.0 / math.sqrt(2 * math.pi) - - -class Activation(str, Enum): - SquaredReLU = "squared_relu" - GeLU = "gelu" - GeLUApprox = "gelu_approx" - LeakyReLU = "leaky_relu" - ReLU = "relu" - - -def get_triton_activation_kernel(activation: Optional[Activation]): - return ( - { - Activation.ReLU: relu, - Activation.LeakyReLU: leaky_relu, - Activation.GeLU: gelu, - Activation.GeLUApprox: gelu_approx, - Activation.SquaredReLU: squared_relu, - }[activation] - if activation - else None - ) - - -def get_triton_activation_bwd_kernel(activation: Optional[Activation]): - return ( - { - Activation.ReLU: relu_grad, - Activation.LeakyReLU: leaky_relu_grad, - Activation.GeLU: gelu_grad, - Activation.GeLUApprox: gelu_approx_grad, - Activation.SquaredReLU: squared_relu_grad, - }[activation] - if activation - else None - ) - - -@triton.jit -def tanh(x): - # Tanh is just a scaled sigmoid - return 2 * tl.sigmoid(2 * x) - 1 - - -@triton.jit -def cosh(x): - exp_x = tl.exp(x) - return (exp_x + 1.0 / exp_x) * 0.5 - - -# a Triton implementation of the most used activations -# See for instance http://arxiv.org/abs/1606.08415 for an overview - -# ReLU -@triton.jit -def relu(x): - """ - ReLU_ activation function - - .. _ReLU: https://pytorch.org/docs/stable/generated/torch.nn.ReLU.html - """ - zero = 0.0 - return tl.where(x >= 0, x, zero.to(x.dtype)) - - -@triton.jit -def relu_grad(x): - # ReLU is different from other activations - # in that it does not require the input to retrospectively compute its gradient - # here the input is the downstream gradient, and we return the upstream gradient directly - zero = 0.0 - one = 1.0 - return tl.where(x >= 0, one.to(x.dtype), zero.to(x.dtype)) - - -@triton.jit -def squared_relu(x): - """ - Squared ReLU activation, as proposed in the Primer_ paper. - - .. _Primer: https://arxiv.org/abs/2109.08668 - """ - x_ = relu(x) - return (x_ * x_).to(x.dtype) - - -@triton.jit -def squared_relu_grad(x): - return tl.where(x >= 0, 2.0 * x, 0.0) - - -# Leaky ReLU -@triton.jit -def leaky_relu(x): - """ - LeakyReLU_ activation - - .. _LeakyReLU: https://pytorch.org/docs/stable/generated/torch.nn.LeakyReLU.html - """ - scale = 0.01 + 0.0 - scale = scale.to(x.dtype) - return tl.where(x >= 0, x, scale * x) - - -@triton.jit -def leaky_relu_grad(x): - min_grad = 0.01 - max_grad = 1 - - min_grad = min_grad.to(x.dtype) - max_grad = max_grad.to(x.dtype) - - return tl.where(x >= 0, max_grad, min_grad) - - -@triton.jit -def gelu(x): - """Gaussian Error Linear Unit (GELU)""" - return x * 0.5 * (1.0 + tl.libdevice.erf(x * _sqrt1_2)) - - -@triton.jit -def gelu_grad(x): - cdf = 0.5 * (1.0 + tl.libdevice.erf(x * _sqrt1_2)) - pdf = tl.exp(-0.5 * x * x) * _gaussian_pdf_normalization - return cdf + x * pdf - - -@triton.jit -def gelu_approx(x): - """ - GeLU_ activation - Gaussian error linear unit, with tanh approximation - - .. _GeLU: https://arxiv.org/pdf/1606.08415.pdf - """ - return 0.5 * x * (1.0 + tanh(_sqrt2pi * x * (1.0 + 0.044715 * x * x))) - - -@triton.jit -def gelu_approx_grad(x): - # CREDITS: Fast implementation proposed in - # https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/fused_bias_gelu.py#L30 - tanh_out = tanh(0.79788456 * x * (1 + 0.044715 * x * x)) - return 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( - 1 + tanh_out - ) diff --git a/build/torch29-cxx11-cu129-aarch64-linux/ops/triton/linear.py b/build/torch29-cxx11-cu129-aarch64-linux/ops/triton/linear.py deleted file mode 100644 index a8966dbc345ab0e593df0124451ee7be3dae131a..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-aarch64-linux/ops/triton/linear.py +++ /dev/null @@ -1,594 +0,0 @@ -# Adapted from https://github.com/ELS-RD/kernl/blob/main/src/kernl/implementations/linear_layer.py -# and https://github.com/openai/triton/blob/master/python/triton/ops/matmul.py -from typing import Optional - -import torch -import triton -import triton.language as tl -from triton.ops.matmul_perf_model import early_config_prune, estimate_matmul_time - -from flash_attn.ops.triton.k_activations import ( - gelu, - gelu_approx, - gelu_approx_grad, - gelu_grad, - squared_relu, - squared_relu_grad, -) - -# CREDITS: Initially inspired by the Triton tutorial on matrix multiplications - - -def init_to_zero(name): - return lambda nargs: nargs[name].zero_() - - -def get_configs_io_bound(): - configs = [] - for num_stages in [2, 3, 4, 5, 6]: - for block_m in [16, 32]: - for block_k in [32, 64]: - for block_n in [32, 64, 128, 256]: - num_warps = 2 if block_n <= 64 else 4 - configs.append( - triton.Config( - { - "BLOCK_M": block_m, - "BLOCK_N": block_n, - "BLOCK_K": block_k, - "SPLIT_K": 1, - }, - num_stages=num_stages, - num_warps=num_warps, - ) - ) - # split_k not used - # for split_k in [2, 4, 8, 16]: - # configs.append(triton.Config( - # {'BLOCK_M': block_m, 'BLOCK_N': block_n, 'BLOCK_K': block_k, 'SPLIT_K': split_k}, - # num_stages=num_stages, num_warps=num_warps, pre_hook=init_to_zero('C'))) - return configs - - -@triton.autotune( - configs=[ - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - # good for int8 - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=4, - num_warps=4, - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - ] - + get_configs_io_bound(), - key=["CACHE_KEY_M", "CACHE_KEY_N", "CACHE_KEY_K"], - prune_configs_by={ - "early_config_prune": early_config_prune, - "perf_model": estimate_matmul_time, - "top_k": 10, - }, -) -@triton.heuristics( - { - "EVEN_K": lambda args: args["K"] % (args["BLOCK_K"] * args["SPLIT_K"]) == 0, - } -) -@triton.jit -def kernel_fwd( - C, # Pointers to matrices - ACT_INPUT, - A, - B, - bias, - # Matrix dimensions - M, - N, - K, - CACHE_KEY_M, - CACHE_KEY_N, - CACHE_KEY_K, - # The stride variables represent how much to increase the ptr by when moving by 1 - # element in a particular dimension. E.g. stride_am is how much to increase a_ptr - # by to get the element one row down (A has M rows) - stride_cm, - # stride_cn, # Assume that stride_cn == 1 - stride_am, - stride_ak, - stride_bn, - stride_bk, - # Meta-parameters - BLOCK_M: tl.constexpr, - GROUP_M: tl.constexpr, - BLOCK_N: tl.constexpr, - BLOCK_K: tl.constexpr, - # split k not used, not performant with activation, kept because early_config_prune is expecting it - SPLIT_K: tl.constexpr, - EVEN_K: tl.constexpr, - A_ROWMAJOR: tl.constexpr, - B_COLMAJOR: tl.constexpr, - BIAS: tl.constexpr, - SAVE_ACT_INPUT: tl.constexpr, - ACTIVATION: tl.constexpr, -): - - """ - Kernel for computing Out = activation(A x W + C) - - Input has shape (M, K) - - Weight has shape (K, N) - - Bias has shape (N,) - - Output has shape (M, N) - - ActInputs (optional) has shape (M, N) - 'ActInputs' optionally saves the A x W + C intermediate for backward computations - This kernel will consolidate over K - """ - - pid = tl.program_id(axis=0) - - grid_m = (M + BLOCK_M - 1) // BLOCK_M - grid_n = (N + BLOCK_N - 1) // BLOCK_N - # re-order program ID for better L2 performance - width = GROUP_M * grid_n - group_id = pid // width - group_size = min(grid_m - group_id * GROUP_M, GROUP_M) - pid_m = group_id * GROUP_M + (pid % group_size) - pid_n = (pid % width) // (group_size) - - # now compute the block that each program will go through - # rm (resp. rn) denotes a range of indices - # for rows (resp. col) of C - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - # trick to avoid masking on M and N axis - ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M) - rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N) - rk = tl.arange(0, BLOCK_K) - - if A_ROWMAJOR: - A = A + (ram[:, None] * stride_am + rk[None, :]) - else: - A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak) - if B_COLMAJOR: - B = B + (rk[:, None] + rbn[None, :] * stride_bn) - else: - B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn) - - acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) - - for k in range(K, 0, -BLOCK_K): - if EVEN_K: - a = tl.load(A) - b = tl.load(B) - else: - a = tl.load(A, mask=rk[None, :] < k, other=0.0) - b = tl.load(B, mask=rk[:, None] < k, other=0.0) - acc += tl.dot(a, b) - - if A_ROWMAJOR: - A += BLOCK_K - else: - A += BLOCK_K * stride_ak - if B_COLMAJOR: - B += BLOCK_K - else: - B += BLOCK_K * stride_bk - - # Putting bias after the matmul (instead of before) is faster, idk why - if BIAS: - bias = tl.load(bias + rn, mask=rn < N, other=0.0).to(tl.float32) - acc += bias[None, :] - - # optional: save the activation inputs - if SAVE_ACT_INPUT: - # act_in_ptrs = ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :] * stride_cn - act_in_ptrs = ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :] - tl.store(act_in_ptrs, acc) - - # optional: fused activation (while the data is in shared memory) - if ACTIVATION == "gelu": - acc = gelu(acc) - elif ACTIVATION == "gelu_approx": - acc = gelu_approx(acc) - elif ACTIVATION == "squared_relu": - acc = squared_relu(acc) - # rematerialize rm and rn to save registers - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - - # write back result - # C = C + rm[:, None] * stride_cm + rn[None, :] * stride_cn - C = C + rm[:, None] * stride_cm + rn[None, :] - mask = (rm < M)[:, None] & (rn < N)[None, :] - tl.store(C, acc) - - -def triton_linear_act( - x: torch.Tensor, - weight: torch.Tensor, - bias: Optional[torch.Tensor] = None, - activation: str = "id", - save_act_input: bool = False, -) -> torch.Tensor: - """ - Compute e = activation(x @ weight.T + bias). - This wrapper kicks the `kernel_fwd` Triton kernel - :param x: input tensor - :param weight: weight matrix - :param bias: an optional bias tensor - :param activation: Activation name. Needs to be a Triton kernel. - :param act_input: an optional tensor to save the activation inputs (for backward) - :return: result tensor - """ - # if torch.is_autocast_enabled(): - # dtype = torch.get_autocast_gpu_dtype() - # x, weight, bias = [a.to(dtype=dtype) for a in [x, weight, bias]] - - assert activation in ["id", "gelu", "gelu_approx", "squared_relu"] - - batch_shape, n = x.shape[:-1], x.shape[-1] - batch_dim = batch_shape.numel() - x_reshaped = x.reshape(batch_dim, n) - - if x_reshaped.stride(0) > 1 and x_reshaped.stride(1) > 1: - x_reshaped = x_reshaped.contiguous() - if weight.stride(0) > 1 and weight.stride(1) > 1: - weight = weight.contiguous() - bias = bias.contiguous() if bias is not None else None - - assert ( - x.dtype == weight.dtype - ), f"Input and weight must have the same dtype, got {x.dtype} and {weight.dtype}" - if bias is not None: - assert ( - x.dtype == bias.dtype - ), f"Input and bias must have the same dtype, got {x.dtype} and {bias.dtype}" - assert ( - x_reshaped.shape[1] == weight.shape[1] - ), f"Incompatible dimensions: {x_reshaped.shape} - {weight.shape}" - - assert ( - bias is None or bias.shape[0] == weight.shape[0] - ), "Incompatible dimensions in between weight and bias" - - M, K = x_reshaped.shape - N, K = weight.shape - - output = torch.empty((M, N), device=x.device, dtype=x.dtype) - act_input = torch.empty_like(output) if save_act_input else None - - # 1D launch kernel where each block gets its own program. - grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),) # noqa - - kernel_fwd[grid]( - output, - act_input, - x_reshaped, - weight, # data ptrs - bias if bias is not None else x, # auto skip bias if not present - M, # shapes - N, - K, - M // 32, # key for triton cache (limit number of compilations) - N // 32, - K // 32, - stride_cm=output.stride(0), # strides - # stride_cn=output.stride(1), - stride_am=x_reshaped.stride(0), - stride_ak=x_reshaped.stride(1), - stride_bk=weight.stride(1), - stride_bn=weight.stride(0), - BIAS=bias is not None, # optional fused bias - SAVE_ACT_INPUT=save_act_input, # optional save activation inputs - ACTIVATION=activation, # optional fused activation - A_ROWMAJOR=x_reshaped.stride(1) == 1, - B_COLMAJOR=weight.stride(1) == 1, - GROUP_M=8, # speed optimization: group the programs - ) - - if not save_act_input: - return output.reshape(*batch_shape, output.shape[-1]) - else: - return ( - output.reshape(*batch_shape, output.shape[-1]), - act_input.reshape(*batch_shape, act_input.shape[-1]), - ) - - -@triton.autotune( - configs=[ - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - # good for int8 - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=4, - num_warps=4, - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - ] - + get_configs_io_bound(), - key=["CACHE_KEY_M", "CACHE_KEY_N", "CACHE_KEY_K"], - prune_configs_by={ - "early_config_prune": early_config_prune, - "perf_model": estimate_matmul_time, - "top_k": 10, - }, -) -@triton.heuristics( - { - "EVEN_K": lambda args: args["K"] % (args["BLOCK_K"] * args["SPLIT_K"]) == 0, - } -) -@triton.jit -def kernel_bwd( - C, # Pointers to matrices - ACT_INPUT, - A, - B, - # Matrix dimensions - M, - N, - K, - CACHE_KEY_M, - CACHE_KEY_N, - CACHE_KEY_K, - # The stride variables represent how much to increase the ptr by when moving by 1 - # element in a particular dimension. E.g. stride_am is how much to increase a_ptr - # by to get the element one row down (A has M rows) - stride_cm, - # stride_cn, # Assume that stride_cn == 1 - stride_am, - stride_ak, - stride_bk, - stride_bn, - # Meta-parameters - BLOCK_M: tl.constexpr, - GROUP_M: tl.constexpr, - BLOCK_N: tl.constexpr, - BLOCK_K: tl.constexpr, - # split k not used, not performant with activation, kept because early_config_prune is expecting it - SPLIT_K: tl.constexpr, - EVEN_K: tl.constexpr, - ACTIVATION: tl.constexpr, -): - - """ - Kernel for computing Out = activation(A x W + C) - - Input has shape (M, K) - - Weight has shape (K, N) - - Output has shape (M, N) - - ActInputs (optional) has shape (M, N) - 'ActInputs' optionally saves the A x W + C intermediate for backward computations - This kernel will consolidate over K - """ - - pid = tl.program_id(axis=0) - - grid_m = (M + BLOCK_M - 1) // BLOCK_M - grid_n = (N + BLOCK_N - 1) // BLOCK_N - # re-order program ID for better L2 performance - width = GROUP_M * grid_n - group_id = pid // width - group_size = min(grid_m - group_id * GROUP_M, GROUP_M) - pid_m = group_id * GROUP_M + (pid % group_size) - pid_n = (pid % width) // (group_size) - - # now compute the block that each program will go through - # rm (resp. rn) denotes a range of indices - # for rows (resp. col) of C - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - # trick to avoid masking on M and N axis - ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M) - rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N) - rk = tl.arange(0, BLOCK_K) - - A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak) - B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn) - - acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) - - for k in range(K, 0, -BLOCK_K): - if EVEN_K: - a = tl.load(A) - b = tl.load(B) - else: - a = tl.load(A, mask=rk[None, :] < k, other=0.0) - b = tl.load(B, mask=rk[:, None] < k, other=0.0) - acc += tl.dot(a, b) - - A += BLOCK_K * stride_ak - B += BLOCK_K * stride_bk - - # optional: fused activation (while the data is in shared memory) - if ACTIVATION != "id": - act_in_ptrs = ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :] - act_input = tl.load(act_in_ptrs).to(acc.dtype) - if ACTIVATION == "gelu": - acc *= gelu_grad(act_input) - elif ACTIVATION == "gelu_approx": - acc *= gelu_approx_grad(act_input) - elif ACTIVATION == "squared_relu": - acc *= squared_relu_grad(act_input) - - # rematerialize rm and rn to save registers - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - - # write back result - C = C + rm[:, None] * stride_cm + rn[None, :] - mask = (rm < M)[:, None] & (rn < N)[None, :] - tl.store(C, acc, mask=mask) - - -def triton_dgrad_act( - grad_output: torch.Tensor, - weight: torch.Tensor, - activation: str = "id", - act_input: Optional[torch.Tensor] = None, -) -> torch.Tensor: - """ - Compute e = activation(grad_output @ weight + bias). - This wrapper kicks the `kernel_fwd` Triton kernel - :param grad_output: input tensor - :param weight: weight matrix - :param activation: Activation name. Needs to be a Triton kernel. - :param act_input: an optional tensor to save the activation inputs (for backward) - :return: result tensor - """ - assert activation in ["id", "gelu", "gelu_approx", "squared_relu"] - - batch_shape, n = grad_output.shape[:-1], grad_output.shape[-1] - batch_dim = batch_shape.numel() - grad_output_reshaped = grad_output.reshape(batch_dim, n) - - if grad_output_reshaped.stride(0) > 1 and grad_output_reshaped.stride(1) > 1: - grad_output_reshaped = grad_output_reshaped.contiguous() - if weight.stride(0) > 1 and weight.stride(1) > 1: - weight = weight.contiguous() - - assert ( - grad_output.dtype == weight.dtype - ), f"grad_output and weight must have the same dtype, got {grad_output.dtype} and {weight.dtype}" - assert ( - grad_output_reshaped.shape[1] == weight.shape[0] - ), f"Incompatible dimensions: {grad_output_reshaped.shape} - {weight.shape}" - if activation != "id": - assert act_input is not None, f"act_input is required for activation {activation}" - - # M, N, K in bwd are different from M, N, K in fwd - M, K = grad_output_reshaped.shape - K, N = weight.shape - - grad_input = torch.empty((M, N), device=grad_output.device, dtype=grad_output.dtype) - - # 1D launch kernel where each block gets its own program. - grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),) # noqa - - kernel_bwd[grid]( - grad_input, - act_input, - grad_output_reshaped, - weight, # data ptrs - M, # shapes - N, - K, - M // 32, # key for triton cache (limit number of compilations) - N // 32, - K // 32, - stride_cm=grad_input.stride(0), # strides - # stride_cn=grad_input.stride(1), - stride_am=grad_output_reshaped.stride(0), - stride_ak=grad_output_reshaped.stride(1), - stride_bk=weight.stride(0), - stride_bn=weight.stride(1), - ACTIVATION=activation, # optional fused activation - GROUP_M=8, # speed optimization: group the programs - ) - - return grad_input.reshape(*batch_shape, grad_input.shape[-1]) diff --git a/build/torch29-cxx11-cu129-aarch64-linux/ops/triton/mlp.py b/build/torch29-cxx11-cu129-aarch64-linux/ops/triton/mlp.py deleted file mode 100644 index 059f4f8a5e174c1f4824e43d313fca18eaa799b8..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-aarch64-linux/ops/triton/mlp.py +++ /dev/null @@ -1,149 +0,0 @@ -# The triton fused matmul + sqrelu is faster for fp16 but slower for bf16, compared -# to naive implementation. -import fused_dense_lib as fused_dense_cuda -import torch -import torch.nn as nn -import torch.nn.functional as F - -from flash_attn.utils.torch import custom_fwd, custom_bwd -from flash_attn.ops.activations import sqrelu_bwd, sqrelu_fwd -from flash_attn.ops.triton.linear import triton_dgrad_act, triton_linear_act - - -class FusedDenseSqreluDenseFunc(torch.autograd.Function): - @staticmethod - @custom_fwd - def forward(ctx, x, weight1, bias1, weight2, bias2, checkpoint_lvl=0): - """checkpoint_lvl: - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute act_input and gelu_out in the bwd - """ - if torch.is_autocast_enabled(): - dtype = torch.get_autocast_gpu_dtype() - x, weight1, bias1, weight2, bias2 = [ - a.to(dtype=dtype) for a in [x, weight1, bias1, weight2, bias2] - ] - is_bf16 = x.dtype == torch.bfloat16 - assert checkpoint_lvl in [0, 1, 2] - x = x.contiguous() - weight1 = weight1.contiguous() - bias1 = bias1.contiguous() - weight2 = weight2.contiguous() - bias2 = bias2.contiguous() - batch_shape, n = x.shape[:-1], x.shape[-1] - batch_dim = batch_shape.numel() - if is_bf16: - act_input = fused_dense_cuda.linear_bias_forward( - x.reshape(batch_dim, n), weight1, bias1 - ) - output1 = sqrelu_fwd(act_input) - else: - save_act_input = checkpoint_lvl != 2 - result = triton_linear_act( - x.reshape(batch_dim, n), - weight1, - bias1, - activation="squared_relu", - save_act_input=save_act_input, - ) - if save_act_input: - output1, act_input = result - else: - output1 = result - output2 = fused_dense_cuda.linear_bias_forward(output1, weight2, bias2) - ctx.checkpoint_lvl = checkpoint_lvl - if checkpoint_lvl == 0: - ctx.save_for_backward(x, weight1, bias1, weight2, act_input, output1) - elif checkpoint_lvl == 1: - ctx.save_for_backward(x, weight1, bias1, weight2, act_input) - elif checkpoint_lvl == 2: - ctx.save_for_backward(x, weight1, bias1, weight2) - return output2.reshape(*batch_shape, output2.shape[-1]) - - @staticmethod - @custom_bwd - def backward(ctx, grad_output): - grad_output = grad_output.contiguous() - checkpoint_lvl = ctx.checkpoint_lvl - x, weight1, bias1, weight2, *rest = ctx.saved_tensors - batch_shape, n = x.shape[:-1], x.shape[-1] - batch_dim = batch_shape.numel() - is_bf16 = x.dtype == torch.bfloat16 - if checkpoint_lvl == 0: - act_input, output1 = rest - elif checkpoint_lvl == 1: - (act_input,) = rest - output1 = sqrelu_fwd(act_input) - elif checkpoint_lvl == 2: - if is_bf16: - act_input = fused_dense_cuda.linear_bias_forward( - x.reshape(batch_dim, n), weight1, bias1 - ) - output1 = sqrelu_fwd(act_input) - else: - output1, act_input = triton_linear_act( - x.reshape(batch_dim, n), - weight1, - bias1, - activation="squared_relu", - save_act_input=True, - ) - - if is_bf16: - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad(output1, grad_output) - grad_output1 = grad_output @ weight2 - grad_act_input = sqrelu_bwd(grad_output1, act_input) - grad_input, grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_backward( - x.reshape(batch_dim, n), weight1, grad_act_input - ) - else: - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad(output1, grad_output) - grad_act_input = triton_dgrad_act( - grad_output, weight2, activation="squared_relu", act_input=act_input - ) - grad_input, grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_backward( - x.reshape(batch_dim, n), weight1, grad_act_input - ) - return grad_input.reshape_as(x), grad_weight1, grad_bias1, grad_weight2, grad_bias2, None - - -fused_dense_sqrelu_dense_function = FusedDenseSqreluDenseFunc.apply - - -class FusedDenseSqreluDense(nn.Module): - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - bias1=True, - bias2=True, - checkpoint_lvl=0, - device=None, - dtype=None, - ): - """ - checkpoint_lvl (increasing lvl means slower but more memory saving): - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute gelu_in and gelu_out in the bwd - """ - assert checkpoint_lvl in [0, 1, 2] - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features * 4 - assert bias1 == True, "DenseSqreluDense module without bias is currently not supported" - assert bias2 == True, "DenseSqreluDense module without bias is currently not supported" - self.checkpoint_lvl = checkpoint_lvl - self.fc1 = nn.Linear(in_features, hidden_features, bias=bias1, **factory_kwargs) - self.fc2 = nn.Linear(hidden_features, out_features, bias=bias2, **factory_kwargs) - - def forward(self, x): - assert x.is_cuda - return fused_dense_sqrelu_dense_function( - x, self.fc1.weight, self.fc1.bias, self.fc2.weight, self.fc2.bias, self.checkpoint_lvl - ) diff --git a/build/torch29-cxx11-cu129-aarch64-linux/ops/triton/rotary.py b/build/torch29-cxx11-cu129-aarch64-linux/ops/triton/rotary.py deleted file mode 100644 index 6ad3ab412a084270120167ad32b67ad1720f2c37..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-aarch64-linux/ops/triton/rotary.py +++ /dev/null @@ -1,186 +0,0 @@ -# Copyright (c) 2025, Tri Dao. -# As of 2025-04-23, we require triton >= 3.0 - -from typing import Optional, Union - -import torch - -import triton -import triton.language as tl - - -@triton.jit -def rotary_kernel( - OUT, # Pointers to matrices - X, - COS, - SIN, - CU_SEQLENS, - SEQLEN_OFFSETS, # this could be int or a pointer - # Matrix dimensions - seqlen, - nheads, - seqlen_ro, - # strides - stride_out_batch, - stride_out_seqlen, - stride_out_nheads, - stride_out_headdim, - stride_x_batch, - stride_x_seqlen, - stride_x_nheads, - stride_x_headdim, - # Meta-parameters - # We want ROTARY_DIM to be constexpr, otherwise the triton compiler doesn't know that - # the mask is constant every 8 elements, and it will generate LDG.16 instead of LDG.128 - ROTARY_DIM: tl.constexpr, - IS_SEQLEN_OFFSETS_TENSOR: tl.constexpr, - IS_VARLEN: tl.constexpr, - INTERLEAVED: tl.constexpr, - CONJUGATE: tl.constexpr, - BLOCK_H: tl.constexpr, - BLOCK_M: tl.constexpr, -): - BLOCK_K: tl.constexpr = triton.next_power_of_2(ROTARY_DIM) - ROTARY_DIM_HALF = ROTARY_DIM // 2 - pid_head = tl.program_id(axis=0) - pid_m = tl.program_id(axis=1) - pid_batch = tl.program_id(axis=2) - - if not IS_VARLEN: - X = X + pid_batch * stride_x_batch - OUT = OUT + pid_batch * stride_out_batch - else: - start_idx = tl.load(CU_SEQLENS + pid_batch) - seqlen = tl.load(CU_SEQLENS + pid_batch + 1) - start_idx - X = X + start_idx * stride_x_seqlen - OUT = OUT + start_idx * stride_out_seqlen - - if pid_m * BLOCK_M >= seqlen: - return - - rh = pid_head * BLOCK_H + tl.arange(0, BLOCK_H) - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - if not IS_SEQLEN_OFFSETS_TENSOR: - rm_cs = rm + SEQLEN_OFFSETS - else: - rm_cs = rm + tl.load(SEQLEN_OFFSETS + pid_batch) - - rk_half = tl.arange(0, BLOCK_K // 2) - COS = COS + (rm_cs[:, None] * ROTARY_DIM_HALF + rk_half[None, :]) - SIN = SIN + (rm_cs[:, None] * ROTARY_DIM_HALF + rk_half[None, :]) - mask_cs = (rm_cs[:, None] < seqlen_ro) & (rk_half[None, :] < ROTARY_DIM_HALF) - cos = tl.load(COS, mask=mask_cs, other=1.0).to(tl.float32) - sin = tl.load(SIN, mask=mask_cs, other=0.0).to(tl.float32) - if CONJUGATE: - sin = -sin - - if not INTERLEAVED: - # Load the 1st and 2nd halves of X, do calculation, then store to 1st and 2nd halves of OUT - X = X + (rh[:, None, None] * stride_x_nheads + rm[None, :, None] * stride_x_seqlen + rk_half[None, None, :] * stride_x_headdim) - OUT = OUT + (rh[:, None, None] * stride_out_nheads + rm[None, :, None] * stride_out_seqlen + rk_half[None, None, :] * stride_out_headdim) - mask = (rh[:, None, None] < nheads) & (rm[None, :, None] < seqlen) & (rk_half[None, None, :] < ROTARY_DIM_HALF) - x0 = tl.load(X, mask=mask, other=0.0).to(tl.float32) - x1 = tl.load(X + ROTARY_DIM_HALF * stride_x_headdim, mask=mask, other=0.0,).to(tl.float32) - o0 = x0 * cos - x1 * sin - o1 = x0 * sin + x1 * cos - tl.store(OUT, o0, mask=mask) - tl.store(OUT + ROTARY_DIM_HALF * stride_out_headdim, o1, mask=mask) - else: - rk = tl.arange(0, BLOCK_K) - X = X + (rh[:, None, None] * stride_x_nheads + rm[None, :, None] * stride_x_seqlen + rk[None, None, :] * stride_x_headdim) - OUT = OUT + (rh[:, None, None] * stride_out_nheads + rm[None, :, None] * stride_out_seqlen + rk[None, None, :] * stride_out_headdim) - mask = (rh[:, None, None] < nheads) & (rm[None, :, None] < seqlen) & (rk[None, None, :] < ROTARY_DIM) - x = tl.load(X, mask=mask, other=0.0).to(tl.float32) - x0, x1 = tl.split(tl.reshape(x, [BLOCK_H, BLOCK_M, BLOCK_K // 2, 2])) - o0 = x0 * cos - x1 * sin - o1 = x0 * sin + x1 * cos - o = tl.reshape(tl.join(o0, o1), [BLOCK_H, BLOCK_M, BLOCK_K]) - tl.store(OUT, o, mask=mask) - - -def apply_rotary( - x: torch.Tensor, - cos: torch.Tensor, - sin: torch.Tensor, - seqlen_offsets: Union[int, torch.Tensor] = 0, - cu_seqlens: Optional[torch.Tensor] = None, - max_seqlen: Optional[int] = None, - interleaved=False, - inplace=False, - conjugate=False, -) -> torch.Tensor: - """ - Arguments: - x: (batch, seqlen, nheads, headdim) if cu_seqlens is None - else (total_seqlen, nheads, headdim). - cos: (seqlen_ro, rotary_dim / 2) - sin: (seqlen_ro, rotary_dim / 2) - seqlen_offsets: integer or integer tensor of size (batch,) - cu_seqlens: (batch + 1,) or None - max_seqlen: int - Returns: - y: (batch, seqlen, nheads, headdim) - """ - is_varlen = cu_seqlens is not None - if not is_varlen: - batch, seqlen, nheads, headdim = x.shape - else: - assert max_seqlen is not None, "If cu_seqlens is passed in, then max_seqlen must be passed" - total_seqlen, nheads, headdim = x.shape - batch_p_1 = cu_seqlens.shape[0] - batch = batch_p_1 - 1 - seqlen = max_seqlen - seqlen_ro, rotary_dim = cos.shape - assert sin.shape == cos.shape - rotary_dim *= 2 - assert rotary_dim <= headdim, "rotary_dim must be <= headdim" - assert headdim <= 256, "Only support headdim <= 256" - assert seqlen_ro >= seqlen, "seqlen_ro must be >= seqlen" - - cos, sin = cos.contiguous(), sin.contiguous() - if isinstance(seqlen_offsets, torch.Tensor): - assert seqlen_offsets.shape == (batch,) - assert seqlen_offsets.dtype in [torch.int32, torch.int64] - seqlen_offsets = seqlen_offsets.contiguous() - else: - assert seqlen_offsets + seqlen <= seqlen_ro - - output = torch.empty_like(x) if not inplace else x - if rotary_dim < headdim and not inplace: - output[..., rotary_dim:].copy_(x[..., rotary_dim:]) - - grid = lambda META: (triton.cdiv(nheads, META["BLOCK_H"]), triton.cdiv(seqlen, META["BLOCK_M"]), batch) # noqa - BLOCK_M = 8 if rotary_dim <= 128 else 4 - - # Need this, otherwise Triton tries to launch from cuda:0 and we get - # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?) - device_ctx = torch.cuda.device(x.device.index) if x.device.type == 'cuda' else torch.xpu.device(x.device.index) - with device_ctx: - torch.library.wrap_triton(rotary_kernel)[grid]( - output, # data ptrs - x, - cos, - sin, - cu_seqlens, - seqlen_offsets, - seqlen, # shapes - nheads, - seqlen_ro, - output.stride(0) if not is_varlen else 0, # batch_strides if not varlen else 0 - output.stride(-3), # seqlen_stride or total_seqlen_stride - output.stride(-2), # nheads_stride - output.stride(-1), # headdim_stride - x.stride(0) if not is_varlen else 0, # batch_strides if not varlen else 0 - x.stride(-3), # seqlen stride or total_seqlen_stride - x.stride(-2), # nheads stride - x.stride(-1), # headdim stride - rotary_dim, - isinstance(seqlen_offsets, torch.Tensor), - is_varlen, - interleaved, - conjugate, - BLOCK_M=BLOCK_M, - BLOCK_H=2, - ) - return output diff --git a/build/torch29-cxx11-cu129-x86_64-linux/__init__.py b/build/torch29-cxx11-cu129-x86_64-linux/__init__.py deleted file mode 100644 index ecc2f9d896b6c93f90b0a1499856dc0612177422..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-x86_64-linux/__init__.py +++ /dev/null @@ -1,393 +0,0 @@ -from typing import Optional, List -import torch -from ._ops import ops as flash_attn_ops -from .flash_attn_interface import ( - flash_attn_func, - flash_attn_kvpacked_func, - flash_attn_qkvpacked_func, - flash_attn_varlen_func, - flash_attn_varlen_kvpacked_func, - flash_attn_varlen_qkvpacked_func, - flash_attn_with_kvcache, -) - - -def fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - return_softmax: bool = False, - gen: Optional[torch.Generator] = None, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention. - - Args: - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Optional output tensor, same shape as q - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - return_softmax: Whether to return softmax weights - gen: Optional random number generator - - Returns: - List of tensors: [output, softmax_lse, (softmax if return_softmax)] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.fwd( - q, - k, - v, - out, - alibi_slopes, - p_dropout, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - gen, - ) - - -def varlen_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - out: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - block_table: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - max_seqlen_q: int = 0, - max_seqlen_k: int = 0, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - zero_tensors: bool = False, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - return_softmax: bool = False, - gen: Optional[torch.Generator] = None, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention with variable sequence lengths. - - Args: - q: Query tensor of shape [total_q, num_heads, head_size] - k: Key tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - v: Value tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1] - cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1] - out: Optional output tensor of shape [total_q, num_heads, head_size] - seqused_k: Optional tensor specifying how many keys to use per batch element [batch_size] - leftpad_k: Optional left padding for keys of shape [batch_size] - block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq] - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - max_seqlen_q: Maximum sequence length for queries - max_seqlen_k: Maximum sequence length for keys - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - zero_tensors: Whether to zero tensors before computation - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - return_softmax: Whether to return softmax weights - gen: Optional random number generator - - Returns: - List of tensors: [output, softmax_lse, (softmax if return_softmax)] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.varlen_fwd( - q, - k, - v, - out, - cu_seqlens_q, - cu_seqlens_k, - seqused_k, - leftpad_k, - block_table, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - p_dropout, - softmax_scale, - zero_tensors, - is_causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - gen, - ) - - -def bwd( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - p_dropout: float = 0.0, - 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, - gen: Optional[torch.Generator] = None, - rng_state: Optional[torch.Tensor] = None, -) -> List[torch.Tensor]: - """ - Backward pass for multi-head attention. - - Args: - dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size] - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size] - softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q] - dq: Optional gradient tensor for queries, same shape as q - dk: Optional gradient tensor for keys, same shape as k - dv: Optional gradient tensor for values, same shape as v - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - deterministic: Whether to use deterministic algorithms - gen: Optional random number generator - rng_state: Optional RNG state from forward pass - - Returns: - List of tensors: [dq, dk, dv] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - alibi_slopes, - p_dropout, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - deterministic, - gen, - rng_state, - ) - - -def varlen_bwd( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - max_seqlen_q: int = 0, - max_seqlen_k: int = 0, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - zero_tensors: bool = False, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - deterministic: bool = False, - gen: Optional[torch.Generator] = None, - rng_state: Optional[torch.Tensor] = None, -) -> List[torch.Tensor]: - """ - Backward pass for multi-head attention with variable sequence lengths. - - Args: - dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size] - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size] - softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q] - cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1] - cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1] - dq: Optional gradient tensor for queries, same shape as q - dk: Optional gradient tensor for keys, same shape as k - dv: Optional gradient tensor for values, same shape as v - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - max_seqlen_q: Maximum sequence length for queries - max_seqlen_k: Maximum sequence length for keys - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - zero_tensors: Whether to zero tensors before computation - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - deterministic: Whether to use deterministic algorithms - gen: Optional random number generator - rng_state: Optional RNG state from forward pass - - Returns: - List of tensors: [dq, dk, dv] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.varlen_bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - p_dropout, - softmax_scale, - zero_tensors, - is_causal, - window_size_left, - window_size_right, - softcap, - deterministic, - gen, - rng_state, - ) - - -def fwd_kvcache( - q: torch.Tensor, - kcache: torch.Tensor, - vcache: torch.Tensor, - k: Optional[torch.Tensor] = None, - v: Optional[torch.Tensor] = None, - seqlens_k: Optional[torch.Tensor] = None, - rotary_cos: Optional[torch.Tensor] = None, - rotary_sin: Optional[torch.Tensor] = None, - cache_batch_idx: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - block_table: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - out: 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, - is_rotary_interleaved: bool = False, - num_splits: int = 1, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention with KV cache. - - Args: - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - kcache: Key cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - vcache: Value cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - k: Optional new keys tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size] - v: Optional new values tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size] - seqlens_k: Optional sequence lengths for keys of shape [batch_size] - rotary_cos: Optional rotary cosine tensor of shape [seqlen_ro, rotary_dim/2] - rotary_sin: Optional rotary sine tensor of shape [seqlen_ro, rotary_dim/2] - cache_batch_idx: Optional indices to index into the KV cache - leftpad_k: Optional left padding for keys of shape [batch_size] - block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq] - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - out: Optional output tensor, same shape as q - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - is_rotary_interleaved: Whether rotary embeddings are interleaved - num_splits: Number of splits for computation - - Returns: - List of tensors: [output, softmax_lse] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.fwd_kvcache( - q, - kcache, - vcache, - k, - v, - seqlens_k, - rotary_cos, - rotary_sin, - cache_batch_idx, - leftpad_k, - block_table, - alibi_slopes, - out, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - is_rotary_interleaved, - num_splits, - ) diff --git a/build/torch29-cxx11-cu129-x86_64-linux/_flash_attn2_cuda_f1a742f.abi3.so b/build/torch29-cxx11-cu129-x86_64-linux/_flash_attn2_cuda_f1a742f.abi3.so deleted file mode 100644 index 08fb5f61c5758fe5302404e15e1aafcd40fd63e9..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-x86_64-linux/_flash_attn2_cuda_f1a742f.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:dec1424f566a194f2d51147e6b75e69a66d0107c9e2e3ac82fffe90bf4554786 -size 1042952368 diff --git a/build/torch29-cxx11-cu129-x86_64-linux/_ops.py b/build/torch29-cxx11-cu129-x86_64-linux/_ops.py deleted file mode 100644 index 659b9bb637b9f59c68d4987f781194173595c596..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-x86_64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _flash_attn2_cuda_f1a742f -ops = torch.ops._flash_attn2_cuda_f1a742f - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_flash_attn2_cuda_f1a742f::{op_name}" diff --git a/build/torch29-cxx11-cu129-x86_64-linux/bert_padding.py b/build/torch29-cxx11-cu129-x86_64-linux/bert_padding.py deleted file mode 100644 index 3c2d35159a014a9d03aabead9e52e009168696ea..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-x86_64-linux/bert_padding.py +++ /dev/null @@ -1,218 +0,0 @@ -# Adapted from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/padding.py - -import torch -import torch.nn.functional as F -from einops import rearrange, repeat - - -class IndexFirstAxis(torch.autograd.Function): - @staticmethod - def forward(ctx, input, indices): - ctx.save_for_backward(indices) - assert input.ndim >= 2 - ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] - second_dim = other_shape.numel() - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - # return input[indices] - return torch.gather( - rearrange(input, "b ... -> b (...)"), 0, repeat(indices, "z -> z d", d=second_dim) - ).reshape(-1, *other_shape) - - @staticmethod - def backward(ctx, grad_output): - (indices,) = ctx.saved_tensors - assert grad_output.ndim >= 2 - other_shape = grad_output.shape[1:] - grad_output = rearrange(grad_output, "b ... -> b (...)") - grad_input = torch.zeros( - [ctx.first_axis_dim, grad_output.shape[1]], - device=grad_output.device, - dtype=grad_output.dtype, - ) - # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. - # grad_input[indices] = grad_output - grad_input.scatter_(0, repeat(indices, "z -> z d", d=grad_output.shape[1]), grad_output) - return grad_input.reshape(ctx.first_axis_dim, *other_shape), None - - -index_first_axis = IndexFirstAxis.apply - - -class IndexPutFirstAxis(torch.autograd.Function): - @staticmethod - def forward(ctx, values, indices, first_axis_dim): - ctx.save_for_backward(indices) - assert indices.ndim == 1 - assert values.ndim >= 2 - output = torch.zeros( - first_axis_dim, *values.shape[1:], device=values.device, dtype=values.dtype - ) - # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. - output[indices] = values - # output.scatter_(0, repeat(indices, 'z -> z d', d=values.shape[1]), values) - return output - - @staticmethod - def backward(ctx, grad_output): - (indices,) = ctx.saved_tensors - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - grad_values = grad_output[indices] - # grad_values = torch.gather(grad_output, 0, repeat(indices, 'z -> z d', d=grad_output.shape[1])) - return grad_values, None, None - - -index_put_first_axis = IndexPutFirstAxis.apply - - -class IndexFirstAxisResidual(torch.autograd.Function): - @staticmethod - def forward(ctx, input, indices): - ctx.save_for_backward(indices) - assert input.ndim >= 2 - ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] - second_dim = other_shape.numel() - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - output = input[indices] - # We don't want to reshape input (b ... -> b (...)) since it could change the channel_last - # memory format to channel_first. In other words, input might not be contiguous. - # If we don't detach, Pytorch complains about output being a view and is being modified inplace - return output, input.detach() - - @staticmethod - def backward(ctx, grad_output, grad_residual): - (indices,) = ctx.saved_tensors - assert grad_output.ndim >= 2 - other_shape = grad_output.shape[1:] - assert grad_residual.shape[1:] == other_shape - grad_input = grad_residual - # grad_input[indices] += grad_output - indices = indices.reshape(indices.shape[0], *((1,) * (grad_output.ndim - 1))) - indices = indices.expand_as(grad_output) - grad_input.scatter_add_(0, indices, grad_output) - return grad_input.reshape(ctx.first_axis_dim, *other_shape), None - - -index_first_axis_residual = IndexFirstAxisResidual.apply - - -def unpad_input(hidden_states, attention_mask, unused_mask=None): - """ - Arguments: - hidden_states: (batch, seqlen, ...) - attention_mask: (batch, seqlen), bool / int, 1 means valid and 0 means not valid. - unused_mask: (batch, seqlen), bool / int, 1 means the element is allocated but unused. - Return: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens selected in attention_mask + unused_mask. - indices: (total_nnz), the indices of masked tokens from the flattened input sequence. - cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. - max_seqlen_in_batch: int - seqused: (batch), returns the number of tokens selected in attention_mask + unused_mask. - """ - all_masks = (attention_mask + unused_mask) if unused_mask is not None else attention_mask - seqlens_in_batch = all_masks.sum(dim=-1, dtype=torch.int32) - used_seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) - indices = torch.nonzero(all_masks.flatten(), as_tuple=False).flatten() - max_seqlen_in_batch = seqlens_in_batch.max().item() - cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) - # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the - # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim - # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to - # index with integer indices. Moreover, torch's index is a bit slower than it needs to be, - # so we write custom forward and backward to make it a bit faster. - return ( - index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices), - indices, - cu_seqlens, - max_seqlen_in_batch, - used_seqlens_in_batch, - ) - - -def unpad_input_for_concatenated_sequences(hidden_states, attention_mask_in_length): - """ - Supports concatenating short samples in one sequence. The attention_mask_in_length is utilized to mask other short samples. It helps efficient training of variant lengths-based samples (e.g., the supervised fine-tuning task in large language model). - The motivation for this function is explained [here](https://github.com/Dao-AILab/flash-attention/issues/432#issuecomment-1668822286). - - For example, if batch = 3 and seqlen = 6, the attention_mask_in_length is: - ``` - [ - [2, 3, 0, 0, 0, 0], - [3, 2, 0, 0, 0, 0], - [6, 0, 0, 0, 0, 0] - ] - ``` - , which refers to the 3D-attention mask: - ``` - [ - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [0, 0, 1, 0, 0, 0], - [0, 0, 1, 1, 0, 0], - [0, 0, 1, 1, 1, 0], - [0, 0, 0, 0, 0, 1] - ], - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [0, 0, 0, 1, 0, 0], - [0, 0, 0, 1, 1, 0], - [0, 0, 0, 0, 0, 1] - ], - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [1, 1, 1, 1, 0, 0], - [1, 1, 1, 1, 1, 0], - [1, 1, 1, 1, 1, 1] - ] - ] - ```. - - Arguments: - hidden_states: (batch, seqlen, ...) - attention_mask_in_length: (batch, seqlen), int, a nonzero number (e.g., 1, 2, 3, etc.) means length of concatenated sequence in b-th batch, and 0 means none. - Return: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. - indices: (total_nnz), the indices of non-masked tokens from the flattened input sequence. - cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. - max_seqlen_in_batch: int - """ - length = attention_mask_in_length.sum(dim=-1) - seqlen = attention_mask_in_length.size(-1) - attention_mask_2d = torch.arange(seqlen, device=length.device, dtype=length.dtype).expand(len(length), seqlen) < length.unsqueeze(1) - real_indices_idx = torch.nonzero(attention_mask_in_length.flatten(), as_tuple=False).flatten() - seqlens_in_batch = attention_mask_in_length.flatten()[real_indices_idx] - indices = torch.nonzero(attention_mask_2d.flatten(), as_tuple=False).flatten() - max_seqlen_in_batch = seqlens_in_batch.max().item() - cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) - # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the - # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim - # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to - # index with integer indices. Moreover, torch's index is a bit slower than it needs to be, - # so we write custom forward and backward to make it a bit faster. - return ( - index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices), - indices, - cu_seqlens, - max_seqlen_in_batch, - ) - - -def pad_input(hidden_states, indices, batch, seqlen): - """ - Arguments: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. - indices: (total_nnz), the indices that represent the non-masked tokens of the original padded input sequence. - batch: int, batch size for the padded sequence. - seqlen: int, maximum sequence length for the padded sequence. - Return: - hidden_states: (batch, seqlen, ...) - """ - dim = hidden_states.shape[-1] - # output = torch.zeros((batch * seqlen), dim, device=hidden_states.device, dtype=hidden_states.dtype) - # output[indices] = hidden_states - output = index_put_first_axis(hidden_states, indices, batch * seqlen) - return rearrange(output, "(b s) ... -> b s ...", b=batch) diff --git a/build/torch29-cxx11-cu129-x86_64-linux/flash_attn2/__init__.py b/build/torch29-cxx11-cu129-x86_64-linux/flash_attn2/__init__.py deleted file mode 100644 index a9b2672c1cd85b74c1b3ded0fc0b2100e1aeac23..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-x86_64-linux/flash_attn2/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -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_interface.py b/build/torch29-cxx11-cu129-x86_64-linux/flash_attn_interface.py deleted file mode 100644 index 6a0157f25449eb90c6c688bc8ef3f4fcf84d46dc..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-x86_64-linux/flash_attn_interface.py +++ /dev/null @@ -1,1622 +0,0 @@ -# Copyright (c) 2023, Tri Dao. - -from typing import Optional, Sequence, Tuple, Union - -import torch -import torch.nn as nn -import os - -# # isort: off -# # We need to import the CUDA kernels after importing torch -# USE_TRITON_ROCM = os.getenv("FLASH_ATTENTION_TRITON_AMD_ENABLE", "FALSE") == "TRUE" -# if USE_TRITON_ROCM: -# from .flash_attn_triton_amd import interface_fa as flash_attn -# else: -# import flash_attn_2_cuda as flash_attn - - -from ._ops import ops as flash_attn - -# # isort: on - -def maybe_contiguous(x): - return x.contiguous() if x is not None and x.stride(-1) != 1 else x - - -def _get_device(): - if torch.xpu.is_available(): - return "xpu" - elif torch.cuda.is_available(): - return "cuda" - else: - return "cpu" - - -def _get_block_size_n(device, head_dim, is_dropout, is_causal): - assert head_dim <= 256 - - if device.type == "xpu": - return 64 - - # This should match the block sizes in the CUDA kernel - major, minor = torch.cuda.get_device_capability(device) - is_sm8x = major == 8 and minor > 0 # Only include sm86 and sm89, exclude sm80 (A100) - is_sm80 = major == 8 and minor == 0 - is_sm90 = major == 9 and minor == 0 - if head_dim <= 32: - return 128 - if head_dim <= 64: - return 128 if not is_dropout else 64 - elif head_dim <= 96: - return 64 - elif head_dim <= 128: - if is_sm8x: - return 64 if (not is_dropout and is_causal) else 32 - else: - return 64 if not is_dropout else 32 - elif head_dim <= 192: - return 64 - elif head_dim <= 224: - return 64 - elif head_dim <= 256: - return 64 - - -def round_multiple(x, m): - return (x + m - 1) // m * m - - -# torch.compile() support is only enabled for pytorch >= 2.4 -# The reason for this is that we are using the new custom_op and register_fake -# APIs, which support inplace modification of inputs in the function itself -if torch.__version__ >= "2.4.0": - _torch_custom_op_wrapper = torch.library.custom_op - _torch_register_fake_wrapper = torch.library.register_fake -else: - def noop_custom_op_wrapper(name, fn=None, /, *, mutates_args, device_types=None, schema=None): - def wrap(func): - return func - if fn is None: - return wrap - return fn - def noop_register_fake_wrapper(op, fn=None, /, *, lib=None, _stacklevel=1): - def wrap(func): - return func - if fn is None: - return wrap - return fn - _torch_custom_op_wrapper = noop_custom_op_wrapper - _torch_register_fake_wrapper = noop_register_fake_wrapper - - -@_torch_custom_op_wrapper("flash_attn::_flash_attn_forward", mutates_args=(), device_types=_get_device()) -def _flash_attn_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - return_softmax: bool -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - out, softmax_lse, S_dmask, rng_state = flash_attn.fwd( - q, - k, - v, - None, - alibi_slopes, - dropout_p, - softmax_scale, - causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - None, - ) - return out, softmax_lse, S_dmask, rng_state - - -@_torch_register_fake_wrapper("flash_attn::_flash_attn_forward") -def _flash_attn_forward_fake( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - return_softmax: bool -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - batch_size, seqlen_q, num_heads, head_size = q.shape - seqlen_k = k.shape[1] - out = torch.empty_like(q) - softmax_lse = torch.empty((batch_size, num_heads, seqlen_q), dtype=torch.float32, device=q.device, layout=q.layout) - p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) - if return_softmax: - p = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128), round_multiple(seqlen_k, 128)), dtype=q.dtype, device=q.device, layout=q.layout) - rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) - - return out, softmax_lse, p, rng_state - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_forward = torch.ops.flash_attn._flash_attn_forward -else: - _wrapped_flash_attn_forward = _flash_attn_forward - - -@_torch_custom_op_wrapper("flash_attn::_flash_attn_varlen_forward", mutates_args=(), device_types=_get_device()) -def _flash_attn_varlen_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - alibi_slopes: Optional[torch.Tensor] = None, - return_softmax: bool = False, - block_table: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - out, softmax_lse, S_dmask, rng_state = flash_attn.varlen_fwd( - q, - k, - v, - None, - cu_seqlens_q, - cu_seqlens_k, - seqused_k, - leftpad_k, - block_table, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - zero_tensors, - causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - None, - ) - # if out.isnan().any() or softmax_lse.isnan().any(): - # breakpoint() - return out, softmax_lse, S_dmask, rng_state - - -@_torch_register_fake_wrapper("flash_attn::_flash_attn_varlen_forward") -def _flash_attn_varlen_forward_fake( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - alibi_slopes: Optional[torch.Tensor] = None, - return_softmax: bool = False, - block_table: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - paged_kv = block_table is not None - batch_size = cu_seqlens_q.numel() - 1 - total_q, num_heads, _ = q.shape - - out = torch.empty_like(q) - softmax_lse = torch.empty((num_heads, total_q), dtype=torch.float32, device=q.device, layout=q.layout) - p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) - seqlen_q_rounded = round_multiple(max_seqlen_q, 128) - seqlen_k_rounded = round_multiple(max_seqlen_k, 128) - if return_softmax: - p = torch.empty((batch_size, num_heads, seqlen_q_rounded, seqlen_k_rounded), dtype=q.dtype, device=q.device, layout=q.layout) - rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) - return out, softmax_lse, p, rng_state - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_varlen_forward = torch.ops.flash_attn._flash_attn_varlen_forward -else: - _wrapped_flash_attn_varlen_forward = _flash_attn_varlen_forward - - -@_torch_custom_op_wrapper("flash_attn::_flash_attn_backward", mutates_args=("dq", "dk", "dv"), device_types=_get_device()) -def _flash_attn_backward( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, -) -> 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)] - ( - dq, - dk, - dv, - softmax_d, - ) = flash_attn.bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - alibi_slopes, - dropout_p, - softmax_scale, - causal, - window_size_left, - window_size_right, - softcap, - deterministic, - None, - rng_state, - ) - return softmax_d - - -@_torch_register_fake_wrapper("flash_attn::_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, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, -) -> torch.Tensor: - dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] - if dq is None: - dq = torch.empty_like(q) - if dk is None: - dk = torch.empty_like(k) - if dv is None: - dv = torch.empty_like(v) - batch_size, seqlen_q, num_heads, _ = q.shape - softmax_d = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128)), device=q.device, dtype=torch.float32) - - return softmax_d - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_backward = torch.ops.flash_attn._flash_attn_backward -else: - _wrapped_flash_attn_backward = _flash_attn_backward - - -@_torch_custom_op_wrapper("flash_attn::_flash_attn_varlen_backward", mutates_args=("dq", "dk", "dv"), device_types=_get_device()) -def _flash_attn_varlen_backward( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> 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)] - ( - dq, - dk, - dv, - softmax_d, - ) = flash_attn.varlen_bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - zero_tensors, - causal, - window_size_left, - window_size_right, - softcap, - deterministic, - None, - rng_state, - ) - # if dk.isnan().any() or dk.isnan().any() or dv.isnan().any() or softmax_d.isnan().any(): - # breakpoint() - return softmax_d - - -@_torch_register_fake_wrapper("flash_attn::_flash_attn_varlen_backward") -def _flash_attn_varlen_backward_fake( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> torch.Tensor: - dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] - batch_size = cu_seqlens_q.numel() - 1 - total_q, num_heads, _ = q.shape - - if dq is None: - dq = torch.empty_like(q) - if dk is None: - dk = torch.empty_like(k) - if dv is None: - dv = torch.empty_like(v) - softmax_d = torch.empty((num_heads, total_q + 128 * batch_size), device=q.device, dtype=torch.float32) - - return softmax_d - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_varlen_backward = torch.ops.flash_attn._flash_attn_varlen_backward -else: - _wrapped_flash_attn_varlen_backward = _flash_attn_varlen_backward - - -class FlashAttnQKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and qkv.requires_grad - if softmax_scale is None: - softmax_scale = qkv.shape[-1] ** (-0.5) - q, k, v = qkv[:, :, 0].detach(), qkv[:, :, 1].detach(), qkv[:, :, 2].detach() - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) - dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dqkv[:, :, 0], - dqkv[:, :, 1], - dqkv[:, :, 2], - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dqkv = dqkv[..., : dout.shape[-1]] # We could have padded the head dimension - return dqkv, None, None, None, None, None, None, None, None, None - - -class FlashAttnVarlenQKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - cu_seqlens, - max_seqlen, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and qkv.requires_grad - if softmax_scale is None: - softmax_scale = qkv.shape[-1] ** (-0.5) - q, k, v = qkv[:, 0].detach(), qkv[:, 1].detach(), qkv[:, 2].detach() - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens, - cu_seqlens, - max_seqlen, - max_seqlen, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=None, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens, rng_state) - ctx.dropout_p = dropout_p - ctx.max_seqlen = max_seqlen - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens, rng_state = ctx.saved_tensors - qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) - dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dqkv[:, 0], - dqkv[:, 1], - dqkv[:, 2], - cu_seqlens, - cu_seqlens, - ctx.max_seqlen, - ctx.max_seqlen, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - 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 FlashAttnKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - kv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, kv] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - k, v = kv[:, :, 0].detach(), kv[:, :, 1].detach() - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - dq = torch.empty_like(q) - kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) - dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dkv[:, :, 0], - dkv[:, :, 1], - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dkv = dkv[..., : dout.shape[-1]] - return dq, dkv, None, None, None, None, None, None, None, None, None - - -class FlashAttnVarlenKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, kv] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - k, v = kv[:, 0].detach(), kv[:, 1].detach() - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=None, - ) - if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state - ) - ctx.dropout_p = dropout_p - 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.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors - dq = torch.empty_like(q) - kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) - dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dkv[:, 0], - dkv[:, 1], - cu_seqlens_q, - cu_seqlens_k, - ctx.max_seqlen_q, - ctx.max_seqlen_k, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dkv = dkv[..., : dout.shape[-1]] - return dq, dkv, None, None, None, None, None, None, None, None, None, None, None, None, None - - -class FlashAttnFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - k, - v, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, k, v] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dk = dk[..., : dout.shape[-1]] - dv = dv[..., : dout.shape[-1]] - return dq, dk, dv, 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, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - block_table, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, k, v] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=block_table, - ) - if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state - ) - ctx.dropout_p = dropout_p - 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.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors - dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - ctx.max_seqlen_q, - ctx.max_seqlen_k, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dk = dk[..., : dout.shape[-1]] - dv = dv[..., : dout.shape[-1]] - return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None - - -def flash_attn_qkvpacked_func( - qkv, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # <=0.0 means deactivate - alibi_slopes=None, - deterministic=False, - 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, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_kvpacked_func( - q, - kv, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, -): - """dropout_p should be set to 0.0 during evaluation - If 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 K, V. - 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) - kv: (batch_size, seqlen, 2, 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. - softcap: float. Anything > 0 activates softcapping 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). - 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 FlashAttnKVPackedFunc.apply( - q, - kv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_func( - q, - k, - v, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - 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). - 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 FlashAttnFunc.apply( - q, - k, - v, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_qkvpacked_func( - qkv, - cu_seqlens, - max_seqlen, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - 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_varlen_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_varlen_kvpacked_func and flash_attn_varlen_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: (total, 3, nheads, headdim), where total = total number of tokens in the batch. - cu_seqlens: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into qkv. - max_seqlen: int. Maximum sequence length in the batch. - 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenQKVPackedFunc.apply( - qkv, - cu_seqlens, - max_seqlen, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_kvpacked_func( - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, -): - """dropout_p should be set to 0.0 during evaluation - If 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 K, V. - 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: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch. - kv: (total_k, 2, nheads_k, headdim), where total_k = total number of key tokens in the batch. - cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into q. - cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into kv. - max_seqlen_q: int. Maximum query sequence length in the batch. - max_seqlen_k: int. Maximum key sequence length in the batch. - 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 + 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenKVPackedFunc.apply( - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_func( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, - block_table=None, -): - """dropout_p should be set to 0.0 during evaluation - Supports multi-query and grouped-query attention (MQA/GQA) by passing in K, V 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: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch. - k: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. - v: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. - cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into q. - cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into kv. - max_seqlen_q: int. Maximum query sequence length in the batch. - max_seqlen_k: int. Maximum key sequence length in the batch. - 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 + 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenFunc.apply( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - block_table, - False if q.device.type == "cpu" else torch.is_grad_enabled(), - ) - - -def flash_attn_with_kvcache( - q, - k_cache, - v_cache, - k=None, - v=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, - block_table: Optional[torch.Tensor] = None, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - rotary_interleaved=True, - alibi_slopes=None, - num_splits=0, - 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 block_table, - or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_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) if there's no block_table, - or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_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). Similar to k. - 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. - block_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). - 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. - 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" - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-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) - cache_batch_idx = maybe_contiguous(cache_batch_idx) - block_table = maybe_contiguous(block_table) - out, softmax_lse = flash_attn.fwd_kvcache( - q, - k_cache, - v_cache, - k, - v, - cache_seqlens, - rotary_cos, - rotary_sin, - cache_batch_idx, - cache_leftpad, - block_table, - alibi_slopes, - None, - softmax_scale, - causal, - window_size[0], - window_size[1], - softcap, - rotary_interleaved, - num_splits, - ) - return (out, softmax_lse) if return_softmax_lse else out diff --git a/build/torch29-cxx11-cu129-x86_64-linux/layers/__init__.py b/build/torch29-cxx11-cu129-x86_64-linux/layers/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/build/torch29-cxx11-cu129-x86_64-linux/layers/patch_embed.py b/build/torch29-cxx11-cu129-x86_64-linux/layers/patch_embed.py deleted file mode 100644 index 05562f8e8bcdb58e947c6f402a49eacd2d031871..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-x86_64-linux/layers/patch_embed.py +++ /dev/null @@ -1,67 +0,0 @@ -# We use the same API as https://github.com/rwightman/pytorch-image-models/blob/v0.6.11/timm/models/layers/patch_embed.py -# But we use nn.Linear instead of Conv2d and it's about 8x faster. - -from functools import partial - -import torch.nn as nn -from einops import rearrange -from torch import _assert -from torch.nn.modules.utils import _pair - -try: - from flash_attn.ops.fused_dense import FusedDense -except ImportError: - FusedDense = None - - -class PatchEmbed(nn.Module): - """2D Image to Patch Embedding""" - - def __init__( - self, - img_size=224, - patch_size=16, - in_chans=3, - embed_dim=768, - norm_layer=None, - flatten=True, - bias=True, - fused_bias_fc=False, - ): - super().__init__() - img_size = _pair(img_size) - patch_size = _pair(patch_size) - self.img_size = img_size - self.patch_size = patch_size - self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1]) - self.num_patches = self.grid_size[0] * self.grid_size[1] - self.flatten = flatten - if fused_bias_fc and FusedDense is None: - raise ImportError("fused_dense is not installed") - - linear_cls = nn.Linear if not fused_bias_fc or not bias else FusedDense - self.proj = linear_cls(in_chans * patch_size[0] * patch_size[1], embed_dim, bias=bias) - self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() - - def forward(self, x): - _, _, H, W = x.shape - _assert( - H == self.img_size[0], - f"Input image height ({H}) doesn't match model ({self.img_size[0]}).", - ) - _assert( - W == self.img_size[1], - f"Input image width ({W}) doesn't match model ({self.img_size[1]}).", - ) - x = self.proj( - rearrange( - x, - "b c (h p1) (w p2) -> b h w (c p1 p2)", - p1=self.patch_size[0], - p2=self.patch_size[1], - ) - ) - if self.flatten: - x = rearrange(x, "b h w c -> b (h w) c") - x = self.norm(x) - return x diff --git a/build/torch29-cxx11-cu129-x86_64-linux/layers/rotary.py b/build/torch29-cxx11-cu129-x86_64-linux/layers/rotary.py deleted file mode 100644 index d1bfc21fc7de1dd287e8f382847b194a48075981..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-x86_64-linux/layers/rotary.py +++ /dev/null @@ -1,483 +0,0 @@ -# Copyright (c) 2025, Tri Dao - -import math -from functools import partial -from typing import Optional, Tuple, Union - -import torch -from torch import Tensor - -from einops import rearrange, repeat -# from flash_attn.ops.triton.rotary import apply_rotary -from ..ops.triton.rotary import apply_rotary - - -def rotate_half(x, interleaved=False): - if not interleaved: - x1, x2 = x.chunk(2, dim=-1) - return torch.cat((-x2, x1), dim=-1) - else: - x1, x2 = x[..., ::2], x[..., 1::2] - return rearrange(torch.stack((-x2, x1), dim=-1), "... d two -> ... (d two)", two=2) - - -def apply_rotary_emb_torch(x, cos, sin, interleaved=False): - """ - x: (batch_size, seqlen, nheads, headdim) - cos, sin: (seqlen, rotary_dim / 2) or (batch_size, seqlen, rotary_dim / 2) - """ - ro_dim = cos.shape[-1] * 2 - assert ro_dim <= x.shape[-1] - cos = repeat(cos, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)") - sin = repeat(sin, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)") - return torch.cat( - [x[..., :ro_dim] * cos + rotate_half(x[..., :ro_dim], interleaved) * sin, x[..., ro_dim:]], - dim=-1, - ) - - -class ApplyRotaryEmb(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x, - cos, - sin, - interleaved=False, - inplace=False, - seqlen_offsets: Union[int, Tensor] = 0, - cu_seqlens: Optional[Tensor] = None, - max_seqlen: Optional[int] = None, - ): - out = apply_rotary( - x, - cos, - sin, - seqlen_offsets=seqlen_offsets, - cu_seqlens=cu_seqlens, - max_seqlen=max_seqlen, - interleaved=interleaved, - inplace=inplace, - ) - if isinstance(seqlen_offsets, int): - ctx.save_for_backward(cos, sin, cu_seqlens) # Can't save int with save_for_backward - ctx.seqlen_offsets = seqlen_offsets - else: - ctx.save_for_backward(cos, sin, cu_seqlens, seqlen_offsets) - ctx.seqlen_offsets = None - ctx.interleaved = interleaved - ctx.inplace = inplace - ctx.max_seqlen = max_seqlen - return out if not inplace else x - - @staticmethod - def backward(ctx, do): - seqlen_offsets = ctx.seqlen_offsets - if seqlen_offsets is None: - cos, sin, cu_seqlens, seqlen_offsets = ctx.saved_tensors - else: - cos, sin, cu_seqlens = ctx.saved_tensors - dx = apply_rotary( - do, - cos, - sin, - seqlen_offsets=seqlen_offsets, - cu_seqlens=cu_seqlens, - max_seqlen=ctx.max_seqlen, - interleaved=ctx.interleaved, - inplace=ctx.inplace, - conjugate=True, - ) - return dx, None, None, None, None, None, None, None - - -def apply_rotary_emb( - x, - cos, - sin, - interleaved=False, - inplace=False, - seqlen_offsets: Union[int, Tensor] = 0, - cu_seqlens: Optional[Tensor] = None, - max_seqlen: Optional[int] = None, -): - """ - Arguments: - x: (batch_size, seqlen, nheads, headdim) if cu_seqlens is None - else (total_seqlen, nheads, headdim) - cos, sin: (seqlen_rotary, rotary_dim / 2) - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead - of 1st half and 2nd half (GPT-NeoX style). - inplace: if True, apply rotary embedding in-place. - seqlen_offsets: (batch_size,) or int. Each sequence in x is shifted by this amount. - Most commonly used in inference when we have KV cache. - cu_seqlens: (batch + 1,) or None - max_seqlen: int - Return: - out: (batch_size, seqlen, nheads, headdim) if cu_seqlens is None - else (total_seqlen, nheads, headdim) - rotary_dim must be <= headdim - Apply rotary embedding to the first rotary_dim of x. - """ - return ApplyRotaryEmb.apply( - x, cos, sin, interleaved, inplace, seqlen_offsets, cu_seqlens, max_seqlen - ) - - -# For backward compatibility -apply_rotary_emb_func = apply_rotary_emb - - -def _apply_rotary_emb_qkv( - qkv, - cos, - sin, - cos_k=None, - sin_k=None, - interleaved=False, - inplace=False, - conjugate=False, - seqlen_offsets: Union[int, Tensor] = 0, - num_heads_q: Optional[int] = None, -): - apply_rotary_fn = partial( - apply_rotary, - interleaved=interleaved, - inplace=inplace, - conjugate=conjugate, - seqlen_offsets=seqlen_offsets - ) - if cos_k is None and sin_k is None and qkv.is_contiguous(): - # Call 1 kernel instead of 2 kernels - # We need qkv to be contiguous so that when we reshape to combine (3, nheads) - # dimensions, we get the same tensor - if qkv.dim() == 5: - batch, seqlen, three, nheads, headdim = qkv.shape - assert three == 3 - # qk = rearrange(qkv[:, :, :2], "b s t h d -> b s (t h) d") - qk = qkv[:, :, :2].reshape(batch, seqlen, -1, headdim) - qk = apply_rotary_fn(qk, cos, sin) - else: - assert qkv.dim() == 4 - assert num_heads_q is not None - num_heads_k = (qkv.shape[2] - num_heads_q) // 2 - assert qkv.shape[2] == num_heads_q + 2 * num_heads_k - qk = qkv[:, :, :num_heads_q + num_heads_k] - qk = apply_rotary_fn(qk, cos, sin) - if not inplace: - if qkv.dim() == 5: - qkv = torch.cat([rearrange(qk, "b s (t h) d -> b s t h d", t=2), qkv[:, :, 2:]], dim=2) - else: - qkv = torch.cat([qk, qkv[:, :, num_heads_q + num_heads_k :]], dim=2) - else: - cos_k = cos if cos_k is None else cos_k - sin_k = sin if sin_k is None else sin_k - if qkv.dim() == 5: - batch, seqlen, three, nheads, headdim = qkv.shape - assert three == 3 - q, k = qkv[:, :, 0], qkv[:, :, 1] - else: - assert qkv.dim() == 4 - assert num_heads_q is not None - num_heads_k = (qkv.shape[2] - num_heads_q) // 2 - assert qkv.shape[2] == num_heads_q + 2 * num_heads_k - q, k = qkv[:, :, :num_heads_q], qkv[:, :, num_heads_q : num_heads_q + num_heads_k] - q = apply_rotary_fn(q, cos, sin) - k = apply_rotary_fn(k, cos_k, sin_k) - if not inplace: - if qkv.dim() == 5: - qkv = torch.stack([q, k, qkv[:, :, 2]], dim=2) - else: - qkv = torch.cat([q, k, qkv[:, :, num_heads_q + num_heads_k:]], dim=2) - return qkv - - -class ApplyRotaryEmbQKV_(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - cos, - sin, - cos_k=None, - sin_k=None, - interleaved=False, - seqlen_offsets: Union[int, torch.Tensor] = 0, - num_heads_q: Optional[int] = None, - ): - # apply_rotary_emb_qkv_inplace( - qkv = _apply_rotary_emb_qkv( - qkv, cos, sin, cos_k, sin_k, interleaved=interleaved, inplace=True, - seqlen_offsets=seqlen_offsets, num_heads_q=num_heads_q, - ) - if isinstance(seqlen_offsets, int): - ctx.save_for_backward(cos, sin, cos_k, sin_k) - ctx.seqlen_offsets = seqlen_offsets - else: - ctx.save_for_backward(cos, sin, cos_k, sin_k, seqlen_offsets) - ctx.seqlen_offsets = None - ctx.interleaved = interleaved - ctx.num_heads_q = num_heads_q - return qkv - - @staticmethod - def backward(ctx, dqkv): - seqlen_offsets = ctx.seqlen_offsets - if seqlen_offsets is None: - cos, sin, cos_k, sin_k, seqlen_offsets = ctx.saved_tensors - else: - cos, sin, cos_k, sin_k = ctx.saved_tensors - dqkv = _apply_rotary_emb_qkv( - dqkv, cos, sin, cos_k, sin_k, interleaved=ctx.interleaved, inplace=True, - seqlen_offsets=seqlen_offsets, num_heads_q=ctx.num_heads_q, conjugate=True, - ) - return dqkv, None, None, None, None, None, None, None - - -def apply_rotary_emb_qkv_( - qkv, - cos, - sin, - cos_k=None, - sin_k=None, - interleaved=False, - seqlen_offsets: Union[int, torch.Tensor] = 0, - num_heads_q: Optional[int] = None, -): - """ - Arguments: - qkv: (batch_size, seqlen, 3, nheads, headdim) or (batch_size, seqlen, num_heads_q + 2 * num_heads_k, headdim). - If qkv has shape (batch_size, seqlen, num_heads_q + 2 * num_heads_k, headdim) (e.g. MQA / GQA), - then num_heads_q must be provided. - cos, sin: (seqlen, rotary_dim / 2) - cos_k, sin_k: (seqlen, rotary_dim / 2), optional - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead of - 1st half and 2nd half (GPT-NeoX style). - seqlen_offsets: (batch_size,) or int. Each sequence in Q and K is shifted by this amount. - Most commonly used in inference when we have KV cache. - Return: - qkv: (batch_size, seqlen, 3, nheads, headdim) or (batch_size, seqlen, num_heads_q + 2 * num_heads_k, headdim) - rotary_dim must be <= headdim - Apply rotary embedding *inplace* to the first rotary_dim of Q and K. - """ - return ApplyRotaryEmbQKV_.apply( - qkv, cos, sin, cos_k, sin_k, interleaved, seqlen_offsets, num_heads_q - ) - - -class ApplyRotaryEmbKV_(torch.autograd.Function): - - @staticmethod - def forward(ctx, kv, cos, sin, interleaved=False, seqlen_offsets: Union[int, torch.Tensor] = 0): - batch, seqlen, two, nheads, headdim = kv.shape - assert two == 2 - k = kv[:, :, 0] - apply_rotary( - k, cos, sin, seqlen_offsets=seqlen_offsets, interleaved=interleaved, inplace=True - ) - if isinstance(seqlen_offsets, int): - ctx.save_for_backward(cos, sin) # Can't save int with save_for_backward - ctx.seqlen_offsets = seqlen_offsets - else: - ctx.save_for_backward(cos, sin, seqlen_offsets) - ctx.seqlen_offsets = None - ctx.interleaved = interleaved - return kv - - @staticmethod - def backward(ctx, dkv): - seqlen_offsets = ctx.seqlen_offsets - if seqlen_offsets is None: - cos, sin, seqlen_offsets = ctx.saved_tensors - else: - cos, sin = ctx.saved_tensors - apply_rotary( - dkv[:, :, 0], - cos, - sin, - seqlen_offsets=seqlen_offsets, - interleaved=ctx.interleaved, - inplace=True, - conjugate=True, - ) - return dkv, None, None, None, None - - -apply_rotary_emb_kv_ = ApplyRotaryEmbKV_.apply - - -def apply_rotary_emb_kv_( - kv, - cos, - sin, - interleaved=False, - seqlen_offsets: Union[int, torch.Tensor] = 0, -): - """ - Arguments: - kv: (batch_size, seqlen, 2, nheads, headdim) - cos, sin: (seqlen, rotary_dim / 2) - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead of - 1st half and 2nd half (GPT-NeoX style). - seqlen_offsets: (batch_size,) or int. Each sequence in Q and K is shifted by this amount. - Most commonly used in inference when we have KV cache. - Return: - kv: (batch_size, seqlen, 2, nheads, headdim) - rotary_dim must be <= headdim - Apply rotary embedding *inplace* to the first rotary_dim of K. - """ - return ApplyRotaryEmbKV_.apply(kv, cos, sin, interleaved, seqlen_offsets) - - -class RotaryEmbedding(torch.nn.Module): - """ - The rotary position embeddings from RoFormer_ (Su et. al). - A crucial insight from the method is that the query and keys are - transformed by rotation matrices which depend on the relative positions. - - Other implementations are available in the Rotary Transformer repo_ and in - GPT-NeoX_, GPT-NeoX was an inspiration - - .. _RoFormer: https://arxiv.org/abs/2104.09864 - .. _repo: https://github.com/ZhuiyiTechnology/roformer - .. _GPT-NeoX: https://github.com/EleutherAI/gpt-neox - - If scale_base is not None, this implements XPos (Sun et al., https://arxiv.org/abs/2212.10554). - A recommended value for scale_base is 512: https://github.com/HazyResearch/flash-attention/issues/96 - Reference: https://github.com/sunyt32/torchscale/blob/main/torchscale/component/xpos_relative_position.py - """ - - def __init__( - self, - dim: int, - base=10000.0, - interleaved=False, - scale_base=None, - device=None, - ): - """ - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead - of 1st half and 2nd half (GPT-NeoX style). - """ - super().__init__() - self.dim = dim - self.base = float(base) - # Generate and save the inverse frequency buffer (non trainable) - inv_freq = self._compute_inv_freq(device) - self.register_buffer("inv_freq", inv_freq, persistent=False) - self.interleaved = interleaved - self.scale_base = scale_base - scale = ( - (torch.arange(0, dim, 2, device=device, dtype=torch.float32) + 0.4 * dim) / (1.4 * dim) - if scale_base is not None - else None - ) - self.register_buffer("scale", scale, persistent=False) - - self._seq_len_cached = 0 - self._cos_cached = None - self._sin_cached = None - self._cos_k_cached = None - self._sin_k_cached = None - - def _compute_inv_freq(self, device=None): - return 1.0 / ( - self.base - ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim) - ) - - def _update_cos_sin_cache(self, seqlen, device=None, dtype=None): - # Reset the tables if the sequence length has changed, - # if we're on a new device (possibly due to tracing for instance), - # or if we're switching from inference mode to training - if ( - seqlen > self._seq_len_cached - or self._cos_cached is None - or self._cos_cached.device != device - or self._cos_cached.dtype != dtype - or (self.training and self._cos_cached.is_inference()) - ): - self._seq_len_cached = seqlen - # We want fp32 here, not self.inv_freq.dtype, since the model could be loaded in bf16 - # And the output of arange can be quite large, so bf16 would lose a lot of precision. - t = torch.arange(seqlen, device=device, dtype=torch.float32) - # We want fp32 here as well since inv_freq will be multiplied with t, and the output - # will be large. Having it in bf16 will lose a lot of precision and cause the - # cos & sin output to change significantly. - # We want to recompute self.inv_freq if it was not loaded in fp32 - if self.inv_freq.dtype != torch.float32: - inv_freq = self._compute_inv_freq(device=device) - else: - inv_freq = self.inv_freq - # Don't do einsum, it converts fp32 to bf16 under AMP - # freqs = torch.einsum("i,j->ij", t, self.inv_freq) - freqs = torch.outer(t, inv_freq) - if self.scale is None: - self._cos_cached = torch.cos(freqs).to(dtype) - self._sin_cached = torch.sin(freqs).to(dtype) - else: - power = ( - torch.arange(seqlen, dtype=self.scale.dtype, device=self.scale.device) - - seqlen // 2 - ) / self.scale_base - scale = self.scale.to(device=power.device) ** rearrange(power, "s -> s 1") - # We want the multiplication by scale to happen in fp32 - self._cos_cached = (torch.cos(freqs) * scale).to(dtype) - self._sin_cached = (torch.sin(freqs) * scale).to(dtype) - self._cos_k_cached = (torch.cos(freqs) / scale).to(dtype) - self._sin_k_cached = (torch.sin(freqs) / scale).to(dtype) - - def forward( - self, - qkv: torch.Tensor, - kv: Optional[torch.Tensor] = None, - seqlen_offset: Union[int, torch.Tensor] = 0, - max_seqlen: Optional[int] = None, - num_heads_q: Optional[int] = None, - ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: - """ - qkv: (batch, seqlen, 3, nheads, headdim) or (batch, seqlen, num_heads_q + 2 * num_heads_k, headdim) - if kv is none, else it's just q of shape (batch, seqlen, nheads, headdim). - If qkv has shape (batch, seqlen, num_heads_q + 2 * num_heads_k, headdim) (e.g. MQA / GQA), - then num_heads_q must be provided. - kv: (batch, seqlen, 2, nheads, headdim) - seqlen_offset: (batch_size,) or int. Each sequence in x is shifted by this amount. - Most commonly used in inference when we have KV cache. - If it's a tensor of shape (batch_size,), then to update the cos / sin cache, one - should pass in max_seqlen, which will update the cos / sin cache up to that length. - Apply rotary embedding *inplace* to qkv and / or kv. - """ - seqlen = qkv.shape[1] - if max_seqlen is not None: - self._update_cos_sin_cache(max_seqlen, device=qkv.device, dtype=qkv.dtype) - elif isinstance(seqlen_offset, int): - self._update_cos_sin_cache(seqlen + seqlen_offset, device=qkv.device, dtype=qkv.dtype) - if kv is None: - return apply_rotary_emb_qkv_( - qkv, - self._cos_cached, - self._sin_cached, - self._cos_k_cached if self.scale is not None else None, - self._sin_k_cached if self.scale is not None else None, - interleaved=self.interleaved, - seqlen_offsets=seqlen_offset, - num_heads_q=num_heads_q, - ) - else: - q = qkv - q = apply_rotary_emb_func( - q, - self._cos_cached, - self._sin_cached, - interleaved=self.interleaved, - inplace=True, - seqlen_offsets=seqlen_offset, - ) - kv = apply_rotary_emb_kv_( - kv, - self._cos_cached if self.scale is None else self._cos_k_cached, - self._sin_cached if self.scale is None else self._sin_k_cached, - interleaved=self.interleaved, - seqlen_offsets=seqlen_offset, - ) - return q, kv diff --git a/build/torch29-cxx11-cu129-x86_64-linux/metadata.json b/build/torch29-cxx11-cu129-x86_64-linux/metadata.json deleted file mode 100644 index 8b00d7fb7685121edf7a97bb007415122444f1d5..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-x86_64-linux/metadata.json +++ /dev/null @@ -1,14 +0,0 @@ -{ - "version": 1, - "license": "BSD-3-Clause", - "python-depends": [], - "backend": { - "type": "cuda", - "archs": [ - "10.0", - "12.0", - "8.0", - "9.0" - ] - } -} diff --git a/build/torch29-cxx11-cu129-x86_64-linux/ops/__init__.py b/build/torch29-cxx11-cu129-x86_64-linux/ops/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/build/torch29-cxx11-cu129-x86_64-linux/ops/activations.py b/build/torch29-cxx11-cu129-x86_64-linux/ops/activations.py deleted file mode 100644 index 7c09649fc41e12d5a360c5672825d8380bc7ec80..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-x86_64-linux/ops/activations.py +++ /dev/null @@ -1,135 +0,0 @@ -# Copied from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/model/layers/activations.py -import math - -import torch -import torch.nn as nn -import torch.nn.functional as F - -# 1/sqrt(2*pi)-> 0.3989423 -# 1/sqrt(2) -> 0.70710678 -# sqrt(2/pi) -> 0.79788456 - -# this function is tanh approximation of gelu -# actual gelu is: -# x * 0.5 * (1.0 + torch.erf(x * 0.70710678)) -@torch.jit.script -def bias_gelu(y, bias): - x = bias + y - return (x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))).to(dtype=y.dtype) - - -# gradient of tanh approximation of gelu -# gradient of actual gelu is: -# 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x) -@torch.jit.script -def bias_gelu_back(g, y, bias): - """Assume that y has shape (B, D) and bias has shape (D)""" - x = bias + y - tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)) - # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243 - ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( - 1 + tanh_out - ) - grad_y = ff * g - return grad_y.to(dtype=y.dtype), grad_y.sum(dim=(0), dtype=bias.dtype) - - -class GeLUFunction(torch.autograd.Function): - @staticmethod - # bias is an optional argument - def forward(ctx, input, bias): - ctx.save_for_backward(input, bias) - return bias_gelu(input, bias) - - @staticmethod - def backward(ctx, grad_output): - input, bias = ctx.saved_tensors - tmp = bias_gelu_back(grad_output, input, bias) - return tmp, tmp - - -bias_gelu_impl = GeLUFunction.apply - -# this function is tanh approximation of gelu -# actual gelu is: -# x * 0.5 * (1.0 + torch.erf(x * 0.70710678)) -@torch.jit.script -def gelu_fwd(x): - return (x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))).to(dtype=x.dtype) - - -# gradient of tanh approximation of gelu -# gradient of actual gelu is: -# 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x) -@torch.jit.script -def gelu_bwd(g, x): - tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)) - # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243 - ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( - 1 + tanh_out - ) - return (ff * g).to(dtype=x.dtype) - - -class FastGeLUFunction(torch.autograd.Function): - @staticmethod - # bias is an optional argument - def forward(ctx, input): - ctx.save_for_backward(input) - return gelu_fwd(input) - - @staticmethod - def backward(ctx, grad_output): - (input,) = ctx.saved_tensors - tmp = gelu_bwd(grad_output, input) - return tmp - - -fast_gelu_impl = FastGeLUFunction.apply - - -@torch.jit.script -def relu_bwd(g, x): - return torch.where(x >= 0, g, 0.0).to(dtype=x.dtype) - - -@torch.jit.script -def sqrelu_fwd(x): - r = F.relu(x) - return (r * r).to(dtype=x.dtype) - - -@torch.jit.script -def sqrelu_bwd(g, x): - return (2.0 * g * F.relu(x)).to(dtype=x.dtype) - - -swiglu_fwd_codestring = """ -template T swiglu_fwd(T x, T y) { - return float(x) * float(y) / (1.0f + ::exp(-float(x))); -} -""" -swiglu_bwd_codestring = """ -template void swiglu_bwd(T x, T y, T g, T& dx, T& dy) { - float x_sigmoid = 1.0f / (1.0f + ::exp(-float(x))); - dx = x_sigmoid * (1 + float(x) * (1.0f - x_sigmoid)) * float(g) * float(y); - dy = float(x) * x_sigmoid * float(g); -} -""" -swiglu_fwd = torch.cuda.jiterator._create_jit_fn(swiglu_fwd_codestring) -swiglu_bwd = torch.cuda.jiterator._create_multi_output_jit_fn(swiglu_bwd_codestring, num_outputs=2) - - -class SwiGLUFunction(torch.autograd.Function): - - @staticmethod - def forward(ctx, x, y): - ctx.save_for_backward(x, y) - return swiglu_fwd(x, y) - - @staticmethod - def backward(ctx, dout): - x, y = ctx.saved_tensors - return swiglu_bwd(x, y, dout) - -swiglu = SwiGLUFunction.apply diff --git a/build/torch29-cxx11-cu129-x86_64-linux/ops/fused_dense.py b/build/torch29-cxx11-cu129-x86_64-linux/ops/fused_dense.py deleted file mode 100644 index 6b4033d134e4093fe278f7b3f8c7d3128ce9f36d..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-x86_64-linux/ops/fused_dense.py +++ /dev/null @@ -1,688 +0,0 @@ -# Copyright (c) 2023, Tri Dao. -# Inspired by https://github.com/NVIDIA/apex/blob/master/apex/fused_dense/fused_dense.py -# We make it work with pytorch amp and with bfloat16. -# The TensorParallel linear modules are inspired by https://github.com/NVIDIA/apex/blob/master/apex/transformer/tensor_parallel/layers.py -from functools import partial -from typing import Optional - -# import fused_dense_cuda # from apex -import fused_dense_lib as fused_dense_cuda -import torch -import torch.nn as nn -import torch.nn.functional as F -from torch import Tensor -from torch.distributed import ProcessGroup - -from flash_attn.utils.torch import custom_fwd, custom_bwd -from flash_attn.ops.activations import gelu_bwd, relu_bwd, sqrelu_bwd, sqrelu_fwd -from flash_attn.utils.distributed import ( - all_gather_raw, - all_reduce, - all_reduce_raw, - reduce_scatter, - reduce_scatter_raw, -) - - -class FusedDenseFunc(torch.autograd.Function): - @staticmethod - @custom_fwd - def forward( - ctx, x, weight, bias, return_residual=False, process_group=None, sequence_parallel=True - ): - """ - If process_group is not None and sequence_parallel=True, we're doing Tensor Parallel - with sequence parallelism: we do an all_gather_raw of x before doing the matmul. - """ - ctx.compute_weight_gradient = weight.requires_grad - ctx.return_residual = return_residual - ctx.process_group = process_group - ctx.sequence_parallel = sequence_parallel - - if torch.is_autocast_enabled(): - x = x.to(dtype=torch.get_autocast_gpu_dtype()) - x = x.contiguous() - if process_group is not None and sequence_parallel: - # We want to kick off the all_gather early, before weight dtype conversion - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - else: - total_x = x - - if torch.is_autocast_enabled(): - weight = weight.to(dtype=torch.get_autocast_gpu_dtype()) - bias = bias.to(dtype=torch.get_autocast_gpu_dtype()) if bias is not None else None - weight = weight.contiguous() - if process_group is not None and sequence_parallel: - handle_x.wait() - batch_shape, n = total_x.shape[:-1], total_x.shape[-1] - batch_dim = batch_shape.numel() - # https://github.com/pytorch/pytorch/blob/5b51849b48a7dbccd297286cc0110def4706f9e7/aten/src/ATen/native/cuda/Blas.cpp#L174 - if min(batch_dim, n, *weight.shape) > 65535 * 32: - raise RuntimeError("fused_dense only supports matrix dims <= 2M") - output = F.linear(total_x, weight, bias) - if ctx.compute_weight_gradient: - ctx.save_for_backward(x, weight) - else: - ctx.save_for_backward(weight) - return output if not return_residual else (output, x) - - @staticmethod - @custom_bwd - def backward(ctx, grad_output, *args): - grad_output = grad_output.contiguous() - if ctx.return_residual: - (grad_input,) = args - grad_input = grad_input.contiguous() - process_group = ctx.process_group - sequence_parallel = ctx.sequence_parallel - if ctx.compute_weight_gradient: - x, weight = ctx.saved_tensors - if process_group is not None and sequence_parallel: - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - else: - total_x = x - else: - (weight,) = ctx.saved_tensors - total_x = None - batch_shape = grad_output.shape[:-1] - batch_dim = batch_shape.numel() - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - if ctx.needs_input_grad[0]: - if not ctx.return_residual: - grad_input = F.linear(grad_output, weight.t()) - else: - grad_input = torch.addmm( - grad_input.reshape(batch_dim, grad_input.shape[-1]), grad_output, weight - ) - grad_input = grad_input.reshape(*batch_shape, grad_input.shape[-1]) - if process_group is not None: - reduce_fn = reduce_scatter_raw if sequence_parallel else all_reduce_raw - grad_input, handle_grad_input = reduce_fn(grad_input, process_group, async_op=True) - else: - grad_input = None - if ctx.needs_input_grad[1]: - assert ctx.compute_weight_gradient - if process_group is not None and sequence_parallel: - handle_x.wait() - grad_weight, grad_bias = fused_dense_cuda.linear_bias_wgrad( - total_x.reshape(batch_dim, total_x.shape[-1]), grad_output, ctx.needs_input_grad[2] - ) - else: - grad_weight = None - grad_bias = grad_output if ctx.needs_input_grad[2] else None - if process_group is not None and ctx.needs_input_grad[0]: - handle_grad_input.wait() - return grad_input, grad_weight, grad_bias, None, None, None - - -def fused_dense_func( - x: Tensor, - weight: Tensor, - bias: Optional[Tensor] = None, - return_residual: bool = False, - process_group: Optional[ProcessGroup] = None, - sequence_parallel: bool = True, -): - dtype_eligible = x.dtype in [torch.float16, torch.bfloat16] or ( - x.dtype == torch.float32 and torch.is_autocast_enabled() - ) - if x.is_cuda and weight.is_cuda and (bias is None or bias.is_cuda) and dtype_eligible: - return FusedDenseFunc.apply( - x, weight, bias, return_residual, process_group, sequence_parallel - ) - else: - assert process_group is None - out = F.linear(x, weight, bias) - return out if not return_residual else (out, x) - - -class FusedDense(nn.Linear): - def __init__( - self, - in_features: int, - out_features: int, - bias: bool = True, - return_residual: bool = False, - device=None, - dtype=None, - ) -> None: - super().__init__(in_features, out_features, bias=bias, device=device, dtype=dtype) - self.return_residual = return_residual - - def forward(self, x, process_group=None): - """ - If process_group is not None, we're doing Tensor Parallel with sequence parallelism: - we do an all_gather of x before doing the matmul. - """ - return fused_dense_func( - x, - self.weight, - self.bias, - return_residual=self.return_residual, - process_group=process_group, - ) - - -class ColumnParallelLinear(nn.Linear): - def __init__( - self, - in_features: int, - out_features: int, - process_group: ProcessGroup, - bias: bool = True, - sequence_parallel=True, - multiple_of=1, - device=None, - dtype=None, - ) -> None: - world_size = torch.distributed.get_world_size(process_group) - if out_features % multiple_of: - raise ValueError(f"out_features ({out_features}) must be a multiple of {multiple_of}") - multiple = out_features // multiple_of - # We want to split @multiple across world_size, but it could be an uneven split - div = multiple // world_size - mod = multiple % world_size - # The first @mod ranks get @div + 1 copies, the rest get @div copies - local_multiple = div + int(torch.distributed.get_rank(process_group) < mod) - super().__init__( - in_features, local_multiple * multiple_of, bias=bias, device=device, dtype=dtype - ) - self.process_group = process_group - self.sequence_parallel = sequence_parallel - - def forward(self, x): - # If self.sequence_parallel is True, we're doing Tensor Parallel with sequence parallelism: - # we do an all_gather of x before doing the matmul. - # If not, then the input is already gathered. - return fused_dense_func( - x, - self.weight, - self.bias, - process_group=self.process_group, - sequence_parallel=self.sequence_parallel, - ) - - -class RowParallelLinear(nn.Linear): - def __init__( - self, - in_features: int, - out_features: int, - process_group: ProcessGroup, - bias: bool = True, - sequence_parallel=True, - multiple_of=1, - device=None, - dtype=None, - ) -> None: - world_size = torch.distributed.get_world_size(process_group) - rank = torch.distributed.get_rank(process_group) - if in_features % multiple_of: - raise ValueError(f"in_features ({in_features}) must be a multiple of {multiple_of}") - multiple = in_features // multiple_of - # We want to split @multiple across world_size, but it could be an uneven split - div = multiple // world_size - mod = multiple % world_size - # The first @mod ranks get @div + 1 copies, the rest get @div copies - local_multiple = div + int(torch.distributed.get_rank(process_group) < mod) - # Only rank 0 will have bias - super().__init__( - local_multiple * multiple_of, - out_features, - bias=bias and rank == 0, - device=device, - dtype=dtype, - ) - self.process_group = process_group - self.sequence_parallel = sequence_parallel - - def forward(self, x): - """ - We're doing Tensor Parallel with sequence parallelism: we do the matmul and then - a reduce_scatter of the result. - """ - out = fused_dense_func(x, self.weight, self.bias) - reduce_fn = reduce_scatter if self.sequence_parallel else all_reduce - return reduce_fn(out, self.process_group) - - -class FusedMLPFunc(torch.autograd.Function): - @staticmethod - @custom_fwd - def forward( - ctx, - x, - weight1, - bias1, - weight2, - bias2, - activation="gelu_approx", - save_pre_act=True, - return_residual=False, - checkpoint_lvl=0, - heuristic=0, - process_group=None, - sequence_parallel=True, - ): - """ - If process_group is not None and sequence_parallel=True, we're doing Tensor Parallel - with sequence parallelism: we do an all_gather of x before doing the matmul. - If sequence_parallel=False, then the input is already gathered. - - checkpoint_lvl: - 0: no recomputation in the bwd - 1: recompute gelu_out / relu_out in the bwd - 2: recompute pre_act and gelu_out / relu_out in the bwd - """ - assert -1 <= heuristic <= 4 - assert activation in ["gelu_approx", "relu", "sqrelu"] - if activation == "sqrelu": - assert heuristic == -1 - if not save_pre_act: - checkpoint_lvl = 2 - assert checkpoint_lvl in [0, 1, 2] - ctx.return_residual = return_residual - ctx.process_group = process_group - ctx.sequence_parallel = sequence_parallel - ctx.checkpoint_lvl = checkpoint_lvl - ctx.activation = activation - ctx.heuristic = heuristic - - if torch.is_autocast_enabled(): - x = x.to(dtype=torch.get_autocast_gpu_dtype()) - x = x.contiguous() - if process_group is not None and sequence_parallel: - # We want to kick off the all_gather early, before weight dtype conversion - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - else: - total_x = x - - if torch.is_autocast_enabled(): - dtype = torch.get_autocast_gpu_dtype() - weight1, weight2 = [a.to(dtype=dtype) for a in [weight1, weight2]] - bias1 = bias1.to(dtype=dtype) if bias1 is not None else None - bias2 = bias2.to(dtype=dtype) if bias2 is not None else None - weight1 = weight1.contiguous() - bias1 = bias1.contiguous() if bias1 is not None else None - weight2 = weight2.contiguous() - bias2 = bias2.contiguous() if bias2 is not None else None - if process_group is not None and sequence_parallel: - handle_x.wait() - batch_shape, n = total_x.shape[:-1], total_x.shape[-1] - batch_dim = batch_shape.numel() - # https://github.com/pytorch/pytorch/blob/5b51849b48a7dbccd297286cc0110def4706f9e7/aten/src/ATen/native/cuda/Blas.cpp#L174 - if min(batch_dim, n, *weight1.shape, *weight2.shape) > 65535 * 32: - raise RuntimeError("fused_dense only supports matrix dims <= 2M") - if heuristic == -1: - pre_act = F.linear(total_x, weight1, bias1) - activation_fn = ( - partial(F.gelu, approximate="tanh") - if activation == "gelu_approx" - else (sqrelu_fwd if activation == "sqrelu" else F.relu) - ) - with torch.jit.fuser("fuser2"): - output1 = activation_fn(pre_act) - # This is before adding bias1 - # pre_act = F.linear(total_x.reshape(batch_dim, n), weight1) - # with torch.jit.fuser('fuser2'): - # output1 = bias_gelu(pre_act, bias1) - else: - is_gelu = activation == "gelu_approx" - output1, *rest = fused_dense_cuda.linear_act_forward( - total_x.reshape(batch_dim, n), weight1, bias1, is_gelu, save_pre_act, heuristic - ) - if save_pre_act: - pre_act = rest[0] - output2 = F.linear(output1, weight2, bias2) - if checkpoint_lvl == 0 or (checkpoint_lvl == 1 and activation == "relu"): - # For RELU the pre_act is very small (just a bit-mask) so we just save it - ctx.save_for_backward(x, weight1, weight2, pre_act, output1) - elif checkpoint_lvl == 1: - ctx.save_for_backward(x, weight1, weight2, pre_act) - elif checkpoint_lvl == 2: - ctx.save_for_backward(x, weight1, weight2, bias1) - output2 = output2.reshape(*batch_shape, output2.shape[-1]) - return output2 if not return_residual else (output2, x) - - @staticmethod - @custom_bwd - def backward(ctx, grad_output, *args): - grad_output = grad_output.contiguous() - checkpoint_lvl = ctx.checkpoint_lvl - activation = ctx.activation - activation_fn = ( - partial(F.gelu, approximate="tanh") - if activation == "gelu_approx" - else (sqrelu_fwd if activation == "sqrelu" else F.relu) - ) - if ctx.return_residual: - (grad_input,) = args - grad_input = grad_input.contiguous() - process_group = ctx.process_group - sequence_parallel = ctx.sequence_parallel - x, weight1, weight2, *rest = ctx.saved_tensors - if process_group is None or not sequence_parallel: - total_x = x - batch_shape = grad_output.shape[:-1] - batch_dim = batch_shape.numel() - if checkpoint_lvl in [0, 1]: - if process_group is not None and sequence_parallel: - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - if checkpoint_lvl == 0 or (checkpoint_lvl == 1 and activation == "relu"): - pre_act, output1 = rest - elif checkpoint_lvl == 1: - (pre_act,) = rest - with torch.jit.fuser("fuser2"): - output1 = activation_fn(pre_act) - elif checkpoint_lvl == 2: - (bias1,) = rest - if process_group is not None and sequence_parallel: - total_x, _ = all_gather_raw(x, process_group) - if ctx.heuristic == -1: - pre_act = F.linear(total_x, weight1, bias1) - with torch.jit.fuser("fuser2"): - output1 = activation_fn(pre_act) - else: - output1, pre_act = fused_dense_cuda.linear_act_forward( - total_x.reshape(batch_dim, total_x.shape[-1]), - weight1, - bias1, - activation == "gelu_approx", - True, - ctx.heuristic, - ) - - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - output1 = output1.reshape(batch_dim, output1.shape[-1]) - pre_act = pre_act.reshape(batch_dim, pre_act.shape[-1]) - if ctx.needs_input_grad[3]: - grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad( - output1, grad_output, ctx.needs_input_grad[4] - ) - else: - grad_weight2 = None - grad_bias2 = grad_output if ctx.needs_input_grad[4] else None - if ctx.heuristic == -1: - # grad_pre_act = matmul_dgelu(grad_output, weight2, pre_act) - grad_output1 = F.linear(grad_output, weight2.t()) - activation_grad_fn = ( - gelu_bwd - if activation == "gelu_approx" - else (sqrelu_bwd if activation == "sqrelu" else relu_bwd) - ) - with torch.jit.fuser("fuser2"): - grad_pre_act = activation_grad_fn(grad_output1, pre_act) - else: - # The cublasLt epilogue has to compute both gelu/relu grad and bias grad, we can't - # just compute gelu/relu grad - grad_pre_act, grad_bias1 = fused_dense_cuda.bias_act_linear_dgrad_bgrad( - weight2, grad_output, pre_act, activation == "gelu_approx", ctx.heuristic - ) - if not ctx.needs_input_grad[2]: - grad_bias1 = None - if ctx.needs_input_grad[0]: - if not ctx.return_residual: - grad_input = F.linear(grad_pre_act, weight1.t()) - else: - grad_input = torch.addmm( - grad_input.reshape(batch_dim, grad_input.shape[-1]), grad_pre_act, weight1 - ) - grad_input = grad_input.reshape(*batch_shape, grad_input.shape[-1]) - if process_group is not None: - reduce_fn = reduce_scatter_raw if sequence_parallel else all_reduce_raw - grad_input, handle_grad_input = reduce_fn(grad_input, process_group, async_op=True) - else: - grad_input = None - if ctx.heuristic == -1: - if ctx.needs_input_grad[1]: - if process_group is not None and sequence_parallel and checkpoint_lvl != 2: - handle_x.wait() - grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_wgrad( - total_x.reshape(batch_dim, total_x.shape[-1]), - grad_pre_act, - ctx.needs_input_grad[2], - ) - else: - grad_weight1 = None - grad_bias1 = grad_pre_act if ctx.needs_input_grad[2] else None - else: - if ctx.needs_input_grad[1]: - if process_group is not None and sequence_parallel and checkpoint_lvl != 2: - handle_x.wait() - grad_weight1 = F.linear( - grad_pre_act.t(), total_x.reshape(batch_dim, total_x.shape[-1]).t() - ) - else: - grad_weight1 = None - if process_group is not None and ctx.needs_input_grad[0]: - handle_grad_input.wait() - return ( - grad_input, - grad_weight1, - grad_bias1, - grad_weight2, - grad_bias2, - None, - None, - None, - None, - None, - None, - None, - ) - - -def fused_mlp_func( - x: Tensor, - weight1: Tensor, - weight2: Tensor, - bias1: Optional[Tensor] = None, - bias2: Optional[Tensor] = None, - activation: str = "gelu_approx", - save_pre_act: bool = True, - return_residual: bool = False, - checkpoint_lvl: int = 0, - heuristic: int = 0, - process_group: Optional[ProcessGroup] = None, - sequence_parallel: bool = True, -): - assert activation in ["gelu_approx", "relu", "sqrelu"] - dtype_eligible = x.dtype in [torch.float16, torch.bfloat16] or ( - x.dtype == torch.float32 and torch.is_autocast_enabled() - ) - # If we save pre-activation, dimension must be divisible by 128 (relu) or 8 (gelu) - dim_eligible = not save_pre_act or (x.shape[-1] % (128 if activation == "relu" else 8) == 0) - if ( - x.is_cuda - and weight1.is_cuda - and weight2.is_cuda - and (bias1 is None or bias1.is_cuda) - and (bias2 is None or bias2.is_cuda) - and dtype_eligible - and dim_eligible - ): - return FusedMLPFunc.apply( - x, - weight1, - bias1, - weight2, - bias2, - activation, - save_pre_act, - return_residual, - checkpoint_lvl, - heuristic, - process_group, - sequence_parallel, - ) - else: - assert process_group is None - pre_act = F.linear(x, weight1, bias1) - activation_fn = ( - partial(F.gelu, approximate="tanh") - if activation == "gelu_approx" - else partial(F.relu, inplace=True) - ) - output1 = activation_fn(pre_act) - output2 = F.linear(output1, weight2, bias2) - return output2 if not return_residual else (output2, x) - - -class FusedMLP(nn.Module): - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - bias1=True, - bias2=True, - activation="gelu_approx", - return_residual=False, - checkpoint_lvl=0, - heuristic="auto", - device=None, - dtype=None, - ): - """ - If process_group is not None, we're doing Tensor Parallel with sequence parallelism: - we do an all_gather of x before doing the matmul, gelu, then matmul. - Finally we do a reduce_scatter of the output. - - checkpoint_lvl (increasing lvl means slower but more memory saving): - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute pre_act and gelu_out in the bwd - heuristic: - -1: don't fuse gemm + gelu (separate kernel) - 0..4: use this heuristic for the algo section in the fused gemm + gelu - 'auto': heuristic will be picked automatically: - For CUDA >= 11.8, we set heuristic=0 for both fp16 and bf16 for best perf. - For CUDA <= 11.7, we set heuristic=1 for fp16 and heuristic=-1 for bf16. - For H100, we set heuristic=-1 for both fp16 and bf16 as the fused cuBlasLt implementation - is slower than the unfused version. - return_residual: whether to return the input x along with the output. This is for - performance reason: for post-norm architecture, returning the input allows us - to fuse the backward of nn.Linear with the residual connection. - """ - assert checkpoint_lvl in [0, 1, 2] - assert activation in ["gelu_approx", "relu", "sqrelu"] - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features * 4 - self.activation = activation - self.return_residual = return_residual - self.checkpoint_lvl = checkpoint_lvl - self.heuristic = heuristic if activation != "sqrelu" else -1 - self.fc1 = nn.Linear(in_features, hidden_features, bias=bias1, **factory_kwargs) - self.fc2 = nn.Linear(hidden_features, out_features, bias=bias2, **factory_kwargs) - - def forward(self, x, process_group=None): - dtype = x.dtype if not torch.is_autocast_enabled() else torch.get_autocast_gpu_dtype() - if self.heuristic == "auto": - if self.activation == "gelu_approx": - if torch.cuda.get_device_capability("cuda") == (9, 0): - heuristic = -1 - else: - cuda_ver = tuple(map(int, torch.version.cuda.split("."))) - heuristic = 0 if cuda_ver >= (11, 8) else (1 if dtype == torch.float16 else -1) - else: - heuristic = 0 - else: - heuristic = self.heuristic - out = fused_mlp_func( - x, - self.fc1.weight, - self.fc2.weight, - self.fc1.bias, - self.fc2.bias, - activation=self.activation, - save_pre_act=self.training, - return_residual=self.return_residual, - checkpoint_lvl=self.checkpoint_lvl, - heuristic=heuristic, - process_group=process_group, - ) - if self.return_residual: - out, x = out - if process_group is not None: - out = reduce_scatter(out, process_group) - return out if not self.return_residual else (out, x) - - -class ParallelFusedMLP(nn.Module): - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - activation="gelu_approx", - process_group: ProcessGroup = None, - bias1=True, - bias2=True, - sequence_parallel=True, - checkpoint_lvl=0, - heuristic="auto", - device=None, - dtype=None, - ): - """ - process_group is required. We're doing Tensor Parallel with sequence parallelism: - we do an all_gather of x before doing the matmul, gelu, then matmul. - Finally we do a reduce_scatter of the output. - - checkpoint_lvl (increasing lvl means slower but more memory saving): - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute pre_act and gelu_out in the bwd - heuristic: - -1: don't fuse gemm + gelu (separate kernel) - 0..4: use this heuristic for the algo section in the fused gemm + gelu - 'auto': heuristic will be picked automatically: - For CUDA >= 11.8, we set heuristic=0 for both fp16 and bf16 for best perf. - For CUDA <= 11.7, we set heuristic=1 for fp16 and heuristic=-1 for bf16. - """ - assert checkpoint_lvl in [0, 1, 2] - assert activation in ["gelu_approx", "relu", "sqrelu"] - assert process_group is not None - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features * 4 - self.activation = activation - self.process_group = process_group - self.sequence_parallel = sequence_parallel - self.checkpoint_lvl = checkpoint_lvl - self.heuristic = heuristic if activation != "sqrelu" else -1 - self.fc1 = ColumnParallelLinear( - in_features, hidden_features, process_group, bias=bias1, **factory_kwargs - ) - self.fc2 = RowParallelLinear( - hidden_features, out_features, process_group, bias=bias2, **factory_kwargs - ) - - def forward(self, x): - dtype = x.dtype if not torch.is_autocast_enabled() else torch.get_autocast_gpu_dtype() - if self.heuristic == "auto": - if self.activation == "gelu_approx": - cuda_ver = tuple(map(int, torch.version.cuda.split("."))) - heuristic = 0 if cuda_ver >= (11, 8) else (1 if dtype == torch.float16 else -1) - else: - heuristic = 0 - else: - heuristic = self.heuristic - out = fused_mlp_func( - x, - self.fc1.weight, - self.fc2.weight, - self.fc1.bias, - self.fc2.bias, - activation=self.activation, - save_pre_act=self.training, - checkpoint_lvl=self.checkpoint_lvl, - heuristic=heuristic, - process_group=self.process_group, - sequence_parallel=self.sequence_parallel, - ) - reduce_fn = reduce_scatter if self.sequence_parallel else all_reduce - return reduce_fn(out, self.process_group) diff --git a/build/torch29-cxx11-cu129-x86_64-linux/ops/layer_norm.py b/build/torch29-cxx11-cu129-x86_64-linux/ops/layer_norm.py deleted file mode 100644 index 4b6cd798fd02844ef9cd3897f8ab95e490e638bf..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-x86_64-linux/ops/layer_norm.py +++ /dev/null @@ -1,800 +0,0 @@ -# Copyright (c) 2022, Tri Dao. -# Adapted from https://github.com/NVIDIA/apex/blob/master/apex/contrib/layer_norm/layer_norm.py - -import dropout_layer_norm -import torch -from torch.nn import init - - -def maybe_align(x, alignment_in_bytes=16): - """Assume that x already has last dim divisible by alignment_in_bytes""" - # TD [2023-07-04] I'm not 100% sure that clone will align the memory - # https://discuss.pytorch.org/t/how-to-ensure-that-tensor-data-ptr-is-aligned-to-16-bytes/183440 - return x if x.data_ptr() % alignment_in_bytes == 0 else x.clone() - - -def _dropout_add_layer_norm_forward( - x0, - residual, - gamma, - beta, - rowscale, - colscale, - dropout_p, - epsilon, - residual_in_fp32=False, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes""" - hidden_size = gamma.numel() - x0mat = x0.view((-1, hidden_size)) - residualmat = residual.view((-1, hidden_size)) if residual is not None else None - rowscale = rowscale.view(-1) if rowscale is not None else None - zmat, xmat, dmask, mu, rsigma = dropout_layer_norm.dropout_add_ln_fwd( - x0mat, - residualmat, - gamma, - beta, - rowscale, - colscale, - None, - None, - dropout_p, - epsilon, - 1.0, - 0, - None, - residual_in_fp32, - is_rms_norm, - ) - # dmask is None if dropout_p == 0.0 - # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype - return zmat, xmat if xmat is not None else x0mat, dmask, mu, rsigma - - -def _dropout_add_layer_norm_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - rowscale, - colscale, - dropout_p, - has_residual, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes - dx == None means that it was a post-norm architecture - (x = drop(x0) + residual was not returned in the fwd). - x0 must not be None if we have colscale. - """ - hidden_size = gamma.numel() - xmat = x.view((-1, hidden_size)) - dzmat = dz.view(xmat.shape) - dxmat = dx.view(xmat.shape) if dx is not None else None - x0mat = x0.view((-1, hidden_size)) if x0 is not None else None - rowscale = rowscale.view(-1) if rowscale is not None else None - if colscale is not None: - assert x0 is not None, "x0 is required to compute the gradient of colscale" - dx0mat, dresidualmat, dgamma, dbeta, _, _, *rest = dropout_layer_norm.dropout_add_ln_bwd( - dzmat, - dxmat, - xmat, - x0mat, - dmask, - mu, - rsigma, - gamma, - rowscale, - colscale, - None, - None, - dropout_p, - 1.0, - 0, - has_residual, - is_rms_norm, - ) - # dresidualmat is None if not has_residual - if colscale is None: - return dx0mat, dresidualmat, dgamma, dbeta - else: - dcolscale = rest[0] - return dx0mat, dresidualmat, dgamma, dbeta, dcolscale - - -def _dropout_add_layer_norm_subset_forward( - x0, - residual, - gamma, - beta, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32=False, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes""" - hidden_size = gamma.numel() - x0mat = x0.view((-1, hidden_size)) - residualmat = residual.view((-1, hidden_size)) if residual is not None else None - x0_subset = x0_subset.view(-1) if x0_subset is not None else None - out_subset = out_subset.view(-1) if out_subset is not None else None - zmat, xmat, dmask, mu, rsigma = dropout_layer_norm.dropout_add_ln_fwd( - x0mat, - residualmat, - gamma, - beta, - None, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - None, - residual_in_fp32, - is_rms_norm, - ) - # dmask is None if dropout_p == 0.0 - # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype - return zmat, xmat if xmat is not None else x0mat, dmask, mu, rsigma - - -def _dropout_add_layer_norm_subset_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - colscale, - x0_subset, - out_subset, - dropout_p, - rowscale_const, - x0_numrows, - has_residual, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes - dx == None means that it was a post-norm architecture - (x = drop(x0) + residual was not returned in the fwd). - x0 must not be None if we have colscale. - """ - hidden_size = gamma.numel() - xmat = x.view((-1, hidden_size)) - dzmat = dz.view(-1, hidden_size) - dxmat = dx.view(xmat.shape) if dx is not None else None - x0mat = x0.view((-1, hidden_size)) if x0 is not None else None - x0_subset = x0_subset.view(-1) if x0_subset is not None else None - out_subset = out_subset.view(-1) if out_subset is not None else None - if colscale is not None: - assert x0 is not None, "x0 is required to compute the gradient of colscale" - dx0mat, dresidualmat, dgamma, dbeta, _, _, *rest = dropout_layer_norm.dropout_add_ln_bwd( - dzmat, - dxmat, - xmat, - x0mat, - dmask, - mu, - rsigma, - gamma, - None, - colscale, - x0_subset, - out_subset, - dropout_p, - rowscale_const, - x0_numrows, - has_residual, - is_rms_norm, - ) - # dresidualmat is None if not has_residual - if colscale is None: - return dx0mat, dresidualmat, dgamma, dbeta - else: - dcolscale = rest[0] - return dx0mat, dresidualmat, dgamma, dbeta, dcolscale - - -def _dropout_add_layer_norm_parallel_residual_forward( - x0, - x1, - residual, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - residual_in_fp32=False, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes""" - hidden_size = gamma0.numel() - x0mat = x0.view((-1, hidden_size)) - x1mat = x1.view((-1, hidden_size)) if x1 is not None else None - residualmat = residual.view((-1, hidden_size)) if residual is not None else None - ( - z0mat, - z1mat, - xmat, - dmask0, - dmask1, - mu, - rsigma, - ) = dropout_layer_norm.dropout_add_ln_parallel_residual_fwd( - x0mat, - x1mat, - residualmat, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - None, - residual_in_fp32, - is_rms_norm, - ) - # dmask0 and dmask1 are None if dropout_p == 0.0 - # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype - return z0mat, z1mat, xmat if xmat is not None else x0mat, dmask0, dmask1, mu, rsigma - - -def _dropout_add_layer_norm_parallel_residual_backward( - dz0, - dz1, - dx, - x, - dmask0, - dmask1, - mu, - rsigma, - gamma0, - gamma1, - dropout_p, - has_x1, - has_residual, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes - dx == None means that it was a post-norm architecture - (x = drop(x0) + residual was not returned in the fwd). - """ - hidden_size = gamma0.numel() - xmat = x.view((-1, hidden_size)) - dz0mat = dz0.view(xmat.shape) - dz1mat = dz1.view(xmat.shape) if dz1 is not None else None - dxmat = dx.view(xmat.shape) if dx is not None else None - ( - dx0mat, - dx1mat, - dresidualmat, - dgamma0, - dbeta0, - dgamma1, - dbeta1, - *rest, - ) = dropout_layer_norm.dropout_add_ln_parallel_residual_bwd( - dz0mat, - dz1mat, - dxmat, - xmat, - dmask0, - dmask1, - mu, - rsigma, - gamma0, - gamma1, - dropout_p, - has_x1, - has_residual, - is_rms_norm, - ) - # dresidualmat is None if not has_residual - return dx0mat, dx1mat, dresidualmat, dgamma0, dbeta0, dgamma1, dbeta1 - - -class DropoutAddLayerNormFn(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x0, - residual, - gamma, - beta, - rowscale, - colscale, - dropout_p, - epsilon, - residual_in_fp32=False, - prenorm=False, - is_rms_norm=False, - return_dmask=False, - ): - x0 = maybe_align(x0.contiguous(), 16) - residual = maybe_align(residual.contiguous(), 16) if residual is not None else None - gamma = maybe_align(gamma.contiguous(), 16) - beta = maybe_align(beta.contiguous(), 16) if beta is not None else None - rowscale = maybe_align(rowscale.contiguous(), 16) if rowscale is not None else None - colscale = maybe_align(colscale.contiguous(), 16) if colscale is not None else None - zmat, xmat, dmask, mu, rsigma = _dropout_add_layer_norm_forward( - x0, - residual, - gamma, - beta, - rowscale, - colscale, - dropout_p, - epsilon, - residual_in_fp32, - is_rms_norm, - ) - # Only need to save x0 if we need to compute gradient wrt colscale - x0_saved = x0 if colscale is not None else None - ctx.save_for_backward( - xmat.view(x0.shape), x0_saved, dmask, gamma, mu, rsigma, rowscale, colscale - ) - ctx.prenorm = prenorm - ctx.dropout_p = dropout_p - ctx.has_residual = residual is not None - ctx.is_rms_norm = is_rms_norm - ctx.has_beta = beta is not None - if not return_dmask: - return ( - zmat.view(x0.shape) if not prenorm else (zmat.view(x0.shape), xmat.view(x0.shape)) - ) - else: - dmask = ( - dmask.view(x0.shape) - if dropout_p > 0.0 - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - ctx.mark_non_differentiable(dmask) - return ( - (zmat.view(x0.shape), dmask) - if not prenorm - else (zmat.view(x0.shape), xmat.view(x0.shape), dmask) - ) - - @staticmethod - def backward(ctx, dz, *args): - # assert dz.is_contiguous() - dz = maybe_align(dz.contiguous(), 16) # this happens! - dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None - x, x0, dmask, gamma, mu, rsigma, rowscale, colscale = ctx.saved_tensors - # x0 is None if colscale is None - dropout_p = ctx.dropout_p - has_residual = ctx.has_residual - dx0mat, dresidualmat, dgamma, dbeta, *rest = _dropout_add_layer_norm_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - rowscale, - colscale, - dropout_p, - has_residual, - ctx.is_rms_norm, - ) - dx0 = dx0mat.view(x.shape) - dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None - dcolscale = rest[0] if colscale is not None else None - return ( - dx0, - dresidual, - dgamma, - dbeta if ctx.has_beta else None, - None, - dcolscale, - None, - None, - None, - None, - None, - None, - ) - - -class DropoutAddLayerNormSubsetFn(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x0, - residual, - gamma, - beta, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32=False, - prenorm=False, - is_rms_norm=False, - return_dmask=False, - ): - x0 = maybe_align(x0.contiguous(), 16) - residual = maybe_align(residual.contiguous(), 16) if residual is not None else None - gamma = maybe_align(gamma.contiguous(), 16) - beta = maybe_align(beta.contiguous(), 16) if beta is not None else None - colscale = maybe_align(colscale.contiguous(), 16) if colscale is not None else None - zmat, xmat, dmask, mu, rsigma = _dropout_add_layer_norm_subset_forward( - x0, - residual, - gamma, - beta, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32, - is_rms_norm, - ) - # Only need to save x0 if we need to compute gradient wrt colscale - x0_saved = x0 if colscale is not None else None - x_shape = (-1, *x0.shape[1:]) - ctx.save_for_backward( - xmat.view(x_shape), x0_saved, dmask, gamma, mu, rsigma, colscale, x0_subset, out_subset - ) - ctx.prenorm = prenorm - ctx.dropout_p = dropout_p - ctx.rowscale_const = rowscale_const - ctx.x0_numrows = x0.shape[:-1].numel() - ctx.has_residual = residual is not None - ctx.is_rms_norm = is_rms_norm - ctx.has_beta = beta is not None - z_shape = (-1, *x0.shape[1:]) - if not return_dmask: - return zmat.view(z_shape) if not prenorm else (zmat.view(z_shape), xmat.view(x0.shape)) - else: - z = zmat.view(z_shape) - dmask = ( - dmask.view(x0.shape) - if dropout_p > 0.0 - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - ctx.mark_non_differentiable(dmask) - return (z, dmask) if not prenorm else (z, xmat.view(x_shape), dmask) - - @staticmethod - def backward(ctx, dz, *args): - # assert dz.is_contiguous() - dz = maybe_align(dz.contiguous(), 16) # this happens! - dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None - x, x0, dmask, gamma, mu, rsigma, colscale, x0_subset, out_subset = ctx.saved_tensors - # x0 is None if colscale is None - dropout_p = ctx.dropout_p - has_residual = ctx.has_residual - dx0mat, dresidualmat, dgamma, dbeta, *rest = _dropout_add_layer_norm_subset_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - colscale, - x0_subset, - out_subset, - dropout_p, - ctx.rowscale_const, - ctx.x0_numrows, - has_residual, - ctx.is_rms_norm, - ) - dx0 = dx0mat.view(-1, *x.shape[1:]) - dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None - dcolscale = rest[0] if colscale is not None else None - return ( - dx0, - dresidual, - dgamma, - dbeta if ctx.has_beta else None, - dcolscale, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) - - -class DropoutAddLayerNormParallelResidualFn(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x0, - x1, - residual, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - residual_in_fp32=False, - prenorm=False, - is_rms_norm=False, - return_dmask=False, - ): - x0 = maybe_align(x0.contiguous(), 16) - x1 = maybe_align(x1.contiguous(), 16) if x1 is not None else None - residual = maybe_align(residual.contiguous(), 16) if residual is not None else None - gamma0 = maybe_align(gamma0.contiguous(), 16) - beta0 = maybe_align(beta0.contiguous(), 16) if beta0 is not None else None - gamma1 = maybe_align(gamma1.contiguous(), 16) if gamma1 is not None else None - beta1 = maybe_align(beta1.contiguous(), 16) if beta1 is not None else None - ( - z0mat, - z1mat, - xmat, - dmask0, - dmask1, - mu, - rsigma, - ) = _dropout_add_layer_norm_parallel_residual_forward( - x0, - x1, - residual, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - residual_in_fp32, - is_rms_norm, - ) - ctx.save_for_backward(xmat.view(x0.shape), dmask0, dmask1, gamma0, gamma1, mu, rsigma) - ctx.prenorm = prenorm - ctx.dropout_p = dropout_p - ctx.has_x1 = x1 is not None - ctx.has_residual = residual is not None - ctx.is_rms_norm = is_rms_norm - ctx.has_beta = beta0 is not None - z = (z0mat.view(x0.shape), z1mat.view(x0.shape) if z1mat is not None else None) - if not return_dmask: - return z if not prenorm else (*z, xmat.view(x0.shape)) - else: - dmask0 = ( - dmask0.view(x0.shape) - if dropout_p > 0.0 - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - dmask1 = ( - dmask1.view(x0.shape) - if dropout_p > 0.0 and x1 is not None - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - ctx.mark_non_differentiable(dmask0) - ctx.mark_non_differentiable(dmask1) - return ( - (*z, dmask0, dmask1) if not prenorm else (*z, xmat.view(x0.shape), dmask0, dmask1) - ) - - @staticmethod - def backward(ctx, dz0, dz1, *args): - dz0 = maybe_align(dz0.contiguous(), 16) # this happens! - dz1 = maybe_align(dz1.contiguous(), 16) if dz1 is not None else None - dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None - x, dmask0, dmask1, gamma0, gamma1, mu, rsigma = ctx.saved_tensors - dropout_p = ctx.dropout_p - has_x1 = ctx.has_x1 - has_residual = ctx.has_residual - ( - dx0mat, - dx1mat, - dresidualmat, - dgamma0, - dbeta0, - dgamma1, - dbeta1, - ) = _dropout_add_layer_norm_parallel_residual_backward( - dz0, - dz1, - dx, - x, - dmask0, - dmask1, - mu, - rsigma, - gamma0, - gamma1, - dropout_p, - has_x1, - has_residual, - ctx.is_rms_norm, - ) - dx0 = dx0mat.view(x.shape) - dx1 = dx1mat.view(x.shape) if dx1mat is not None else None - dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None - return ( - dx0, - dx1, - dresidual, - dgamma0, - dbeta0 if ctx.has_beta else None, - dgamma1, - dbeta1 if ctx.has_beta else None, - None, - None, - None, - None, - None, - None, - ) - - -def layer_norm(x, weight, bias, epsilon): - return DropoutAddLayerNormFn.apply(x, None, weight, bias, None, None, 0.0, epsilon, False) - - -def dropout_add_layer_norm( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - rowscale=None, - layerscale=None, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormFn.apply( - x0, - residual, - weight, - bias, - rowscale, - layerscale, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - False, - return_dropout_mask, - ) - - -def dropout_add_layer_norm_subset( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - layerscale=None, - x0_subset=None, - out_subset=None, - rowscale_const=1.0, - out_numrows=0, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormSubsetFn.apply( - x0, - residual, - weight, - bias, - layerscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32, - prenorm, - False, - return_dropout_mask, - ) - - -def dropout_add_layer_norm_parallel_residual( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormParallelResidualFn.apply( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - False, - return_dropout_mask, - ) - - -class DropoutAddLayerNorm(torch.nn.Module): - def __init__( - self, - hidden_size, - prenorm=False, - p=0.0, - eps=1e-5, - residual_in_fp32=False, - device=None, - dtype=None, - ): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.prenorm = prenorm - self.p = p - self.eps = eps - self.residual_in_fp32 = residual_in_fp32 - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.bias = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.reset_parameters() - - def reset_parameters(self): - init.ones_(self.weight) - init.zeros_(self.bias) - - def forward(self, x0, residual=None): - return dropout_add_layer_norm( - x0, - residual, - self.weight, - self.bias, - self.p if self.training else 0.0, - self.eps, - prenorm=self.prenorm, - residual_in_fp32=self.residual_in_fp32, - ) diff --git a/build/torch29-cxx11-cu129-x86_64-linux/ops/rms_norm.py b/build/torch29-cxx11-cu129-x86_64-linux/ops/rms_norm.py deleted file mode 100644 index 068348d61290e3839dd082b540d898578ba1e8e2..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-x86_64-linux/ops/rms_norm.py +++ /dev/null @@ -1,174 +0,0 @@ -# Copyright (c) 2022, Tri Dao. -# Adapted from https://github.com/NVIDIA/apex/blob/master/apex/contrib/layer_norm/layer_norm.py - -import torch -from torch.nn import init - -from flash_attn.ops.layer_norm import ( - DropoutAddLayerNormFn, - DropoutAddLayerNormParallelResidualFn, - DropoutAddLayerNormSubsetFn, -) - - -def rms_norm(x, weight, epsilon): - return DropoutAddLayerNormFn.apply( - x, None, weight, None, None, None, 0.0, epsilon, False, False, True - ) - - -def dropout_add_rms_norm( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - rowscale=None, - layerscale=None, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormFn.apply( - x0, - residual, - weight, - bias, - rowscale, - layerscale, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - True, - return_dropout_mask, - ) - - -def dropout_add_rms_norm_subset( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - layerscale=None, - x0_subset=None, - out_subset=None, - rowscale_const=1.0, - out_numrows=0, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormSubsetFn.apply( - x0, - residual, - weight, - bias, - layerscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32, - prenorm, - True, - return_dropout_mask, - ) - - -def dropout_add_rms_norm_parallel_residual( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormParallelResidualFn.apply( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - True, - return_dropout_mask, - ) - - -class RMSNorm(torch.nn.Module): - def __init__(self, hidden_size, eps=1e-5, device=None, dtype=None): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.eps = eps - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.register_parameter("bias", None) - self.reset_parameters() - - def reset_parameters(self): - init.ones_(self.weight) - - def forward(self, x): - return rms_norm(x, self.weight, self.eps) - - -class DropoutAddRMSNorm(torch.nn.Module): - def __init__( - self, - hidden_size, - prenorm=False, - p=0.0, - eps=1e-5, - residual_in_fp32=False, - device=None, - dtype=None, - ): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.prenorm = prenorm - self.p = p - self.eps = eps - self.residual_in_fp32 = residual_in_fp32 - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.register_parameter("bias", None) - self.reset_parameters() - - def reset_parameters(self): - init.ones_(self.weight) - - def forward(self, x0, residual=None): - return dropout_add_rms_norm( - x0, - residual, - self.weight, - None, - self.p if self.training else 0.0, - self.eps, - prenorm=self.prenorm, - residual_in_fp32=self.residual_in_fp32, - ) diff --git a/build/torch29-cxx11-cu129-x86_64-linux/ops/triton/__init__.py b/build/torch29-cxx11-cu129-x86_64-linux/ops/triton/__init__.py deleted file mode 100644 index 8b137891791fe96927ad78e64b0aad7bded08bdc..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-x86_64-linux/ops/triton/__init__.py +++ /dev/null @@ -1 +0,0 @@ - diff --git a/build/torch29-cxx11-cu129-x86_64-linux/ops/triton/cross_entropy.py b/build/torch29-cxx11-cu129-x86_64-linux/ops/triton/cross_entropy.py deleted file mode 100644 index 1b5a415b73f236f3e05fb14b9141959559e18526..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-x86_64-linux/ops/triton/cross_entropy.py +++ /dev/null @@ -1,330 +0,0 @@ -# Copyright (c) 2023, Tri Dao. - -from typing import Tuple, Optional, Union - -import torch -import torch.nn.functional as F - -import triton -import triton.language as tl - -# `all_gather_into_tensor` and `reduce_scatter_tensor` are new placeholders for -# `_all_gather_base` and `_reduce_scatter_base`. They require the most recent -# version of PyTorch. The following 2 lines are for backward compatibility with -# older PyTorch. -if "all_gather_into_tensor" not in dir(torch.distributed): - torch.distributed.all_gather_into_tensor = torch.distributed._all_gather_base - - -@triton.heuristics( - { - "HAS_SMOOTHING": lambda args: args["smoothing"] > 0.0, - } -) -@triton.jit -def cross_entropy_fwd_kernel( - loss_ptr, # data ptrs - lse_ptr, - z_loss_ptr, - logits_ptr, - labels_ptr, - smoothing, - logit_scale, - lse_square_scale, - ignore_index, - total_classes, - class_start_idx, # Useful for tensor parallel when each rank only has a subset of classes - n_cols, # shapes - logits_row_stride, # strides - BLOCK_SIZE: tl.constexpr, - HAS_SMOOTHING: tl.constexpr, - # if SPLIT (e.g. tensor parallel), don't include the LSE in the loss since it's not the final LSE - SPLIT: tl.constexpr, - PRECOMPUTED_LSE: tl.constexpr, # If LSE is already computed (also no smoothing and logit_scale == 1.0) -): - row_idx = tl.program_id(0) - logits_ptr = logits_ptr + row_idx * logits_row_stride.to(tl.int64) - sum_logits = 0.0 # For smoothing - if not PRECOMPUTED_LSE: - # Statistics for online softmax - m_i = -float("inf") - l_i = 0.0 - for col_offset in range(0, n_cols, BLOCK_SIZE): - cols = col_offset + tl.arange(0, BLOCK_SIZE) - logits = tl.load(logits_ptr + cols, mask=cols < n_cols, other=-float("inf")).to( - tl.float32 - ) * logit_scale - if HAS_SMOOTHING: - sum_logits += tl.sum(tl.where(cols < n_cols, logits, 0.0)) - m_i_new = tl.maximum(m_i, tl.max(logits)) - l_i = tl.exp(m_i - m_i_new) * l_i + tl.sum(tl.exp(logits - m_i_new)) - m_i = m_i_new - lse = tl.log(l_i) + m_i - tl.store(lse_ptr + row_idx, lse) - else: - lse = tl.load(lse_ptr + row_idx) - label_idx = tl.load(labels_ptr + row_idx) - if label_idx == ignore_index: - loss = 0.0 - z_loss = 0.0 - else: - label_idx -= class_start_idx - if label_idx >= 0 and label_idx < n_cols: - logits_label = tl.load(logits_ptr + label_idx) * logit_scale - if HAS_SMOOTHING: - loss = ( - (lse if not SPLIT else 0.0) - - smoothing * sum_logits / total_classes - - (1 - smoothing) * logits_label - ) - else: - loss = (lse if not SPLIT else 0.0) - logits_label - else: - # If label is out of bounds, we set the CE loss to 0.0. But we still want the smoothing loss - if HAS_SMOOTHING: - loss = smoothing * ((lse if not SPLIT else 0.0) - sum_logits / total_classes) - else: - loss = 0.0 - if not SPLIT: - z_loss = lse_square_scale * lse * lse - loss += z_loss - else: - z_loss = 0.0 - tl.store(loss_ptr + row_idx, loss) - if not SPLIT: - tl.store(z_loss_ptr + row_idx, z_loss) - - -@triton.heuristics( - { - "HAS_SMOOTHING": lambda args: args["smoothing"] > 0.0, - } -) -@triton.jit -def cross_entropy_bwd_kernel( - dlogits_ptr, # data ptrs - dloss_ptr, - logits_ptr, - lse_ptr, - labels_ptr, - smoothing, - logit_scale, - lse_square_scale, - ignore_index, - total_classes, - class_start_idx, # Useful for tensor parallel when each rank only has a subset of classes - n_cols, # shapes - logits_row_stride, # strides - dlogits_row_stride, - dloss_row_stride, - BLOCK_SIZE: tl.constexpr, - HAS_SMOOTHING: tl.constexpr, -): - row_idx = tl.program_id(0) - col_block_idx = tl.program_id(1) - logits_ptr = logits_ptr + row_idx * logits_row_stride.to(tl.int64) - dlogits_ptr = dlogits_ptr + row_idx * dlogits_row_stride.to(tl.int64) - col_offsets = col_block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - label_idx = tl.load(labels_ptr + row_idx) - if label_idx != ignore_index: - dloss = tl.load(dloss_ptr + row_idx * dloss_row_stride) - else: - dloss = 0.0 - logits = tl.load(logits_ptr + col_offsets, mask=col_offsets < n_cols, other=-float("inf")).to( - tl.float32 - ) * logit_scale - lse = tl.load(lse_ptr + row_idx) - probs = tl.exp(logits - lse) - probs += 2.0 * lse_square_scale * lse * probs - label_idx -= class_start_idx - if HAS_SMOOTHING: - smooth_positive = 1.0 - smoothing - smooth_negative = smoothing / total_classes - probs = tl.where(col_offsets == label_idx, probs - smooth_positive, probs) - smooth_negative - else: - probs = tl.where(col_offsets == label_idx, probs - 1.0, probs) - tl.store(dlogits_ptr + col_offsets, (dloss * logit_scale) * probs, mask=col_offsets < n_cols) - - -class CrossEntropyLoss(torch.autograd.Function): - - @staticmethod - def forward( - ctx, - logits, - labels, - precomputed_lse=None, - smoothing=0.0, - logit_scale=1.0, - lse_square_scale=0.0, - ignore_index=-100, - inplace_backward=False, - process_group=None, - ): - # For some reason Triton generates wrong code when labels has dtype long and its address - # is not aligned to 16 bytes. The ld.global.b64 seems to load the wrong label index. - if labels.dtype == torch.long and labels.data_ptr() % 16 != 0: - labels = F.pad(labels, (0, 1))[..., :-1] - assert labels.data_ptr() % 16 == 0 - assert logit_scale > 0.0 - n_rows, n_cols = logits.shape - assert labels.shape == (n_rows,) - world_size = 1 if process_group is None else torch.distributed.get_world_size(process_group) - total_classes = world_size * n_cols - rank = 0 if process_group is None else torch.distributed.get_rank(process_group) - class_start_idx = rank * n_cols - use_precomputed_lse = precomputed_lse is not None and logit_scale == 1.0 and smoothing == 0.0 - - if logits.stride(-1) != 1: - logits = logits.contiguous() - MAX_BLOCK_SIZE = 16 * 1024 - BLOCK_SIZE = min(triton.next_power_of_2(n_cols), MAX_BLOCK_SIZE) - num_warps = ( - 4 - if BLOCK_SIZE < 2048 - else (8 if BLOCK_SIZE < 8192 else (16 if BLOCK_SIZE < 128 * 1024 else 32)) - ) - losses = torch.empty(n_rows, dtype=torch.float, device=logits.device) - if use_precomputed_lse: - assert precomputed_lse.shape == (n_rows,) - lse = precomputed_lse.contiguous() - else: - lse = torch.empty(n_rows, dtype=torch.float, device=logits.device) - z_losses = torch.empty(n_rows, dtype=torch.float, device=logits.device) - # Need this, otherwise Triton tries to launch from cuda:0 and we get - # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?) - with torch.cuda.device(logits.device.index): - cross_entropy_fwd_kernel[(n_rows,)]( - losses, # data ptrs - lse, - z_losses, - logits, - labels, - smoothing, - logit_scale, - lse_square_scale, - ignore_index, - total_classes, - class_start_idx, - n_cols, # shapes - logits.stride(0), # strides - BLOCK_SIZE=BLOCK_SIZE, # constants - SPLIT=world_size > 1, - PRECOMPUTED_LSE=use_precomputed_lse, - num_warps=num_warps, - ) - - if world_size > 1: - # If there's no smoothing, if labels are in the vocab of this partition, losses contains - # - predicted logit, and 0 otherwise. - # If there's smoothing=0.1, for labels in the vocab of this partition, losses contains - # -0.9 * predicted logit - 0.1 * sum logit / total_classes. - # For labels not in the vocab of this partition, losses contains - # -0.1 * sum logit / total_classes. - if world_size > 1: - lse_allgather = torch.empty(world_size, n_rows, dtype=lse.dtype, device=lse.device) - torch.distributed.all_gather_into_tensor(lse_allgather, lse, group=process_group) - handle_losses = torch.distributed.all_reduce( - losses, op=torch.distributed.ReduceOp.SUM, group=process_group, async_op=True - ) - lse = torch.logsumexp(lse_allgather, dim=0) - handle_losses.wait() - # After the allreduce, if there's no smoothing, the total losses are - predicted_logit, - # we just have to add the (global) lse. - # If there's smoothing=0.1, the total losses are - # -0.9 * predicted_logit - 0.1 * sum logit / total_classes. - # Again, we just have to add the (global) lse. - losses += lse - if lse_square_scale != 0.0: - z_losses = lse_square_scale * lse.square() - z_losses.masked_fill_(labels == ignore_index, 0.0) - losses += z_losses - else: - z_losses = torch.zeros_like(losses) - losses.masked_fill_(labels == ignore_index, 0.0) - - ctx.save_for_backward(logits, lse, labels) - ctx.mark_non_differentiable(z_losses) - ctx.smoothing = smoothing - ctx.logit_scale = logit_scale - ctx.lse_square_scale = lse_square_scale - ctx.ignore_index = ignore_index - ctx.total_classes = total_classes - ctx.class_start_idx = class_start_idx - ctx.inplace_backward = inplace_backward - return losses, z_losses - - @staticmethod - def backward(ctx, grad_losses, grad_z_losses): - del grad_z_losses # z_losses are only for logging. - - logits, lse, labels = ctx.saved_tensors - dlogits = logits if ctx.inplace_backward else torch.empty_like(logits) - n_rows, n_cols = logits.shape - BLOCK_SIZE = min(triton.next_power_of_2(n_cols), 4 * 1024) - num_warps = 4 if BLOCK_SIZE < 2048 else (8 if BLOCK_SIZE < 8192 else 16) - grid = lambda META: (n_rows, triton.cdiv(n_cols, META["BLOCK_SIZE"])) # noqa - # Need this, otherwise Triton tries to launch from cuda:0 and we get - # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?) - with torch.cuda.device(logits.device.index): - cross_entropy_bwd_kernel[grid]( - dlogits, # data ptrs - grad_losses, - logits, - lse, - labels, - ctx.smoothing, - ctx.logit_scale, - ctx.lse_square_scale, - ctx.ignore_index, - ctx.total_classes, - ctx.class_start_idx, - n_cols, # shapes - logits.stride(0), # strides - dlogits.stride(0), - grad_losses.stride(0), - BLOCK_SIZE=BLOCK_SIZE, # constants - num_warps=num_warps, - ) - return dlogits, None, None, None, None, None, None, None, None, None - - -def cross_entropy_loss( - logits: torch.Tensor, - labels: torch.Tensor, - precomputed_lse: Optional[torch.Tensor] = None, - label_smoothing: float = 0.0, - logit_scale: float = 1.0, - lse_square_scale: float = 0.0, - ignore_index=-100, - inplace_backward: bool = False, - process_group=None, -) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Arguments: - logits: (batch, vocab_size) - labels: (batch,) - label_smoothing: float - logit_scale: float. Multiply logits by this scale before calculating the loss. - lse_square_scale: float. If > 0, we add lse_square_scale * lse(logits) ^ 2 to the loss. - This is also referred to as "z-loss". - ignore_index: int. If labels == ignore_index, the loss is set to 0.0. - inplace_backward: bool. If True, we do the backward pass in-place by modifying the logits. - This saves memory. - process_group: if not None, we're doing Tensor Parallel: each process is responsible for - one part of the vocab. The loss will be aggregated across processes. - Returns: - losses: (batch,), float - z_losses: (batch,), float - """ - return CrossEntropyLoss.apply( - logits, - labels, - precomputed_lse, - label_smoothing, - logit_scale, - lse_square_scale, - ignore_index, - inplace_backward, - process_group, - ) diff --git a/build/torch29-cxx11-cu129-x86_64-linux/ops/triton/k_activations.py b/build/torch29-cxx11-cu129-x86_64-linux/ops/triton/k_activations.py deleted file mode 100644 index efb83c358eb4a85d069ee340a3c83f418f9a805b..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-x86_64-linux/ops/triton/k_activations.py +++ /dev/null @@ -1,162 +0,0 @@ -# Adapted from https://github.com/facebookresearch/xformers/blob/main/xformers/triton/k_activations.py -# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved. -# -# This source code is licensed under the BSD license found in the -# LICENSE file in the root directory of this source tree. - -import math -from enum import Enum -from typing import Optional - -import triton -import triton.language as tl - -_sqrt2pi = math.sqrt(2.0 / math.pi) -_sqrt1_2 = math.sqrt(1.0 / 2) -_gaussian_pdf_normalization = 1.0 / math.sqrt(2 * math.pi) - - -class Activation(str, Enum): - SquaredReLU = "squared_relu" - GeLU = "gelu" - GeLUApprox = "gelu_approx" - LeakyReLU = "leaky_relu" - ReLU = "relu" - - -def get_triton_activation_kernel(activation: Optional[Activation]): - return ( - { - Activation.ReLU: relu, - Activation.LeakyReLU: leaky_relu, - Activation.GeLU: gelu, - Activation.GeLUApprox: gelu_approx, - Activation.SquaredReLU: squared_relu, - }[activation] - if activation - else None - ) - - -def get_triton_activation_bwd_kernel(activation: Optional[Activation]): - return ( - { - Activation.ReLU: relu_grad, - Activation.LeakyReLU: leaky_relu_grad, - Activation.GeLU: gelu_grad, - Activation.GeLUApprox: gelu_approx_grad, - Activation.SquaredReLU: squared_relu_grad, - }[activation] - if activation - else None - ) - - -@triton.jit -def tanh(x): - # Tanh is just a scaled sigmoid - return 2 * tl.sigmoid(2 * x) - 1 - - -@triton.jit -def cosh(x): - exp_x = tl.exp(x) - return (exp_x + 1.0 / exp_x) * 0.5 - - -# a Triton implementation of the most used activations -# See for instance http://arxiv.org/abs/1606.08415 for an overview - -# ReLU -@triton.jit -def relu(x): - """ - ReLU_ activation function - - .. _ReLU: https://pytorch.org/docs/stable/generated/torch.nn.ReLU.html - """ - zero = 0.0 - return tl.where(x >= 0, x, zero.to(x.dtype)) - - -@triton.jit -def relu_grad(x): - # ReLU is different from other activations - # in that it does not require the input to retrospectively compute its gradient - # here the input is the downstream gradient, and we return the upstream gradient directly - zero = 0.0 - one = 1.0 - return tl.where(x >= 0, one.to(x.dtype), zero.to(x.dtype)) - - -@triton.jit -def squared_relu(x): - """ - Squared ReLU activation, as proposed in the Primer_ paper. - - .. _Primer: https://arxiv.org/abs/2109.08668 - """ - x_ = relu(x) - return (x_ * x_).to(x.dtype) - - -@triton.jit -def squared_relu_grad(x): - return tl.where(x >= 0, 2.0 * x, 0.0) - - -# Leaky ReLU -@triton.jit -def leaky_relu(x): - """ - LeakyReLU_ activation - - .. _LeakyReLU: https://pytorch.org/docs/stable/generated/torch.nn.LeakyReLU.html - """ - scale = 0.01 + 0.0 - scale = scale.to(x.dtype) - return tl.where(x >= 0, x, scale * x) - - -@triton.jit -def leaky_relu_grad(x): - min_grad = 0.01 - max_grad = 1 - - min_grad = min_grad.to(x.dtype) - max_grad = max_grad.to(x.dtype) - - return tl.where(x >= 0, max_grad, min_grad) - - -@triton.jit -def gelu(x): - """Gaussian Error Linear Unit (GELU)""" - return x * 0.5 * (1.0 + tl.libdevice.erf(x * _sqrt1_2)) - - -@triton.jit -def gelu_grad(x): - cdf = 0.5 * (1.0 + tl.libdevice.erf(x * _sqrt1_2)) - pdf = tl.exp(-0.5 * x * x) * _gaussian_pdf_normalization - return cdf + x * pdf - - -@triton.jit -def gelu_approx(x): - """ - GeLU_ activation - Gaussian error linear unit, with tanh approximation - - .. _GeLU: https://arxiv.org/pdf/1606.08415.pdf - """ - return 0.5 * x * (1.0 + tanh(_sqrt2pi * x * (1.0 + 0.044715 * x * x))) - - -@triton.jit -def gelu_approx_grad(x): - # CREDITS: Fast implementation proposed in - # https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/fused_bias_gelu.py#L30 - tanh_out = tanh(0.79788456 * x * (1 + 0.044715 * x * x)) - return 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( - 1 + tanh_out - ) diff --git a/build/torch29-cxx11-cu129-x86_64-linux/ops/triton/linear.py b/build/torch29-cxx11-cu129-x86_64-linux/ops/triton/linear.py deleted file mode 100644 index a8966dbc345ab0e593df0124451ee7be3dae131a..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-x86_64-linux/ops/triton/linear.py +++ /dev/null @@ -1,594 +0,0 @@ -# Adapted from https://github.com/ELS-RD/kernl/blob/main/src/kernl/implementations/linear_layer.py -# and https://github.com/openai/triton/blob/master/python/triton/ops/matmul.py -from typing import Optional - -import torch -import triton -import triton.language as tl -from triton.ops.matmul_perf_model import early_config_prune, estimate_matmul_time - -from flash_attn.ops.triton.k_activations import ( - gelu, - gelu_approx, - gelu_approx_grad, - gelu_grad, - squared_relu, - squared_relu_grad, -) - -# CREDITS: Initially inspired by the Triton tutorial on matrix multiplications - - -def init_to_zero(name): - return lambda nargs: nargs[name].zero_() - - -def get_configs_io_bound(): - configs = [] - for num_stages in [2, 3, 4, 5, 6]: - for block_m in [16, 32]: - for block_k in [32, 64]: - for block_n in [32, 64, 128, 256]: - num_warps = 2 if block_n <= 64 else 4 - configs.append( - triton.Config( - { - "BLOCK_M": block_m, - "BLOCK_N": block_n, - "BLOCK_K": block_k, - "SPLIT_K": 1, - }, - num_stages=num_stages, - num_warps=num_warps, - ) - ) - # split_k not used - # for split_k in [2, 4, 8, 16]: - # configs.append(triton.Config( - # {'BLOCK_M': block_m, 'BLOCK_N': block_n, 'BLOCK_K': block_k, 'SPLIT_K': split_k}, - # num_stages=num_stages, num_warps=num_warps, pre_hook=init_to_zero('C'))) - return configs - - -@triton.autotune( - configs=[ - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - # good for int8 - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=4, - num_warps=4, - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - ] - + get_configs_io_bound(), - key=["CACHE_KEY_M", "CACHE_KEY_N", "CACHE_KEY_K"], - prune_configs_by={ - "early_config_prune": early_config_prune, - "perf_model": estimate_matmul_time, - "top_k": 10, - }, -) -@triton.heuristics( - { - "EVEN_K": lambda args: args["K"] % (args["BLOCK_K"] * args["SPLIT_K"]) == 0, - } -) -@triton.jit -def kernel_fwd( - C, # Pointers to matrices - ACT_INPUT, - A, - B, - bias, - # Matrix dimensions - M, - N, - K, - CACHE_KEY_M, - CACHE_KEY_N, - CACHE_KEY_K, - # The stride variables represent how much to increase the ptr by when moving by 1 - # element in a particular dimension. E.g. stride_am is how much to increase a_ptr - # by to get the element one row down (A has M rows) - stride_cm, - # stride_cn, # Assume that stride_cn == 1 - stride_am, - stride_ak, - stride_bn, - stride_bk, - # Meta-parameters - BLOCK_M: tl.constexpr, - GROUP_M: tl.constexpr, - BLOCK_N: tl.constexpr, - BLOCK_K: tl.constexpr, - # split k not used, not performant with activation, kept because early_config_prune is expecting it - SPLIT_K: tl.constexpr, - EVEN_K: tl.constexpr, - A_ROWMAJOR: tl.constexpr, - B_COLMAJOR: tl.constexpr, - BIAS: tl.constexpr, - SAVE_ACT_INPUT: tl.constexpr, - ACTIVATION: tl.constexpr, -): - - """ - Kernel for computing Out = activation(A x W + C) - - Input has shape (M, K) - - Weight has shape (K, N) - - Bias has shape (N,) - - Output has shape (M, N) - - ActInputs (optional) has shape (M, N) - 'ActInputs' optionally saves the A x W + C intermediate for backward computations - This kernel will consolidate over K - """ - - pid = tl.program_id(axis=0) - - grid_m = (M + BLOCK_M - 1) // BLOCK_M - grid_n = (N + BLOCK_N - 1) // BLOCK_N - # re-order program ID for better L2 performance - width = GROUP_M * grid_n - group_id = pid // width - group_size = min(grid_m - group_id * GROUP_M, GROUP_M) - pid_m = group_id * GROUP_M + (pid % group_size) - pid_n = (pid % width) // (group_size) - - # now compute the block that each program will go through - # rm (resp. rn) denotes a range of indices - # for rows (resp. col) of C - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - # trick to avoid masking on M and N axis - ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M) - rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N) - rk = tl.arange(0, BLOCK_K) - - if A_ROWMAJOR: - A = A + (ram[:, None] * stride_am + rk[None, :]) - else: - A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak) - if B_COLMAJOR: - B = B + (rk[:, None] + rbn[None, :] * stride_bn) - else: - B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn) - - acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) - - for k in range(K, 0, -BLOCK_K): - if EVEN_K: - a = tl.load(A) - b = tl.load(B) - else: - a = tl.load(A, mask=rk[None, :] < k, other=0.0) - b = tl.load(B, mask=rk[:, None] < k, other=0.0) - acc += tl.dot(a, b) - - if A_ROWMAJOR: - A += BLOCK_K - else: - A += BLOCK_K * stride_ak - if B_COLMAJOR: - B += BLOCK_K - else: - B += BLOCK_K * stride_bk - - # Putting bias after the matmul (instead of before) is faster, idk why - if BIAS: - bias = tl.load(bias + rn, mask=rn < N, other=0.0).to(tl.float32) - acc += bias[None, :] - - # optional: save the activation inputs - if SAVE_ACT_INPUT: - # act_in_ptrs = ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :] * stride_cn - act_in_ptrs = ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :] - tl.store(act_in_ptrs, acc) - - # optional: fused activation (while the data is in shared memory) - if ACTIVATION == "gelu": - acc = gelu(acc) - elif ACTIVATION == "gelu_approx": - acc = gelu_approx(acc) - elif ACTIVATION == "squared_relu": - acc = squared_relu(acc) - # rematerialize rm and rn to save registers - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - - # write back result - # C = C + rm[:, None] * stride_cm + rn[None, :] * stride_cn - C = C + rm[:, None] * stride_cm + rn[None, :] - mask = (rm < M)[:, None] & (rn < N)[None, :] - tl.store(C, acc) - - -def triton_linear_act( - x: torch.Tensor, - weight: torch.Tensor, - bias: Optional[torch.Tensor] = None, - activation: str = "id", - save_act_input: bool = False, -) -> torch.Tensor: - """ - Compute e = activation(x @ weight.T + bias). - This wrapper kicks the `kernel_fwd` Triton kernel - :param x: input tensor - :param weight: weight matrix - :param bias: an optional bias tensor - :param activation: Activation name. Needs to be a Triton kernel. - :param act_input: an optional tensor to save the activation inputs (for backward) - :return: result tensor - """ - # if torch.is_autocast_enabled(): - # dtype = torch.get_autocast_gpu_dtype() - # x, weight, bias = [a.to(dtype=dtype) for a in [x, weight, bias]] - - assert activation in ["id", "gelu", "gelu_approx", "squared_relu"] - - batch_shape, n = x.shape[:-1], x.shape[-1] - batch_dim = batch_shape.numel() - x_reshaped = x.reshape(batch_dim, n) - - if x_reshaped.stride(0) > 1 and x_reshaped.stride(1) > 1: - x_reshaped = x_reshaped.contiguous() - if weight.stride(0) > 1 and weight.stride(1) > 1: - weight = weight.contiguous() - bias = bias.contiguous() if bias is not None else None - - assert ( - x.dtype == weight.dtype - ), f"Input and weight must have the same dtype, got {x.dtype} and {weight.dtype}" - if bias is not None: - assert ( - x.dtype == bias.dtype - ), f"Input and bias must have the same dtype, got {x.dtype} and {bias.dtype}" - assert ( - x_reshaped.shape[1] == weight.shape[1] - ), f"Incompatible dimensions: {x_reshaped.shape} - {weight.shape}" - - assert ( - bias is None or bias.shape[0] == weight.shape[0] - ), "Incompatible dimensions in between weight and bias" - - M, K = x_reshaped.shape - N, K = weight.shape - - output = torch.empty((M, N), device=x.device, dtype=x.dtype) - act_input = torch.empty_like(output) if save_act_input else None - - # 1D launch kernel where each block gets its own program. - grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),) # noqa - - kernel_fwd[grid]( - output, - act_input, - x_reshaped, - weight, # data ptrs - bias if bias is not None else x, # auto skip bias if not present - M, # shapes - N, - K, - M // 32, # key for triton cache (limit number of compilations) - N // 32, - K // 32, - stride_cm=output.stride(0), # strides - # stride_cn=output.stride(1), - stride_am=x_reshaped.stride(0), - stride_ak=x_reshaped.stride(1), - stride_bk=weight.stride(1), - stride_bn=weight.stride(0), - BIAS=bias is not None, # optional fused bias - SAVE_ACT_INPUT=save_act_input, # optional save activation inputs - ACTIVATION=activation, # optional fused activation - A_ROWMAJOR=x_reshaped.stride(1) == 1, - B_COLMAJOR=weight.stride(1) == 1, - GROUP_M=8, # speed optimization: group the programs - ) - - if not save_act_input: - return output.reshape(*batch_shape, output.shape[-1]) - else: - return ( - output.reshape(*batch_shape, output.shape[-1]), - act_input.reshape(*batch_shape, act_input.shape[-1]), - ) - - -@triton.autotune( - configs=[ - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - # good for int8 - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=4, - num_warps=4, - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - ] - + get_configs_io_bound(), - key=["CACHE_KEY_M", "CACHE_KEY_N", "CACHE_KEY_K"], - prune_configs_by={ - "early_config_prune": early_config_prune, - "perf_model": estimate_matmul_time, - "top_k": 10, - }, -) -@triton.heuristics( - { - "EVEN_K": lambda args: args["K"] % (args["BLOCK_K"] * args["SPLIT_K"]) == 0, - } -) -@triton.jit -def kernel_bwd( - C, # Pointers to matrices - ACT_INPUT, - A, - B, - # Matrix dimensions - M, - N, - K, - CACHE_KEY_M, - CACHE_KEY_N, - CACHE_KEY_K, - # The stride variables represent how much to increase the ptr by when moving by 1 - # element in a particular dimension. E.g. stride_am is how much to increase a_ptr - # by to get the element one row down (A has M rows) - stride_cm, - # stride_cn, # Assume that stride_cn == 1 - stride_am, - stride_ak, - stride_bk, - stride_bn, - # Meta-parameters - BLOCK_M: tl.constexpr, - GROUP_M: tl.constexpr, - BLOCK_N: tl.constexpr, - BLOCK_K: tl.constexpr, - # split k not used, not performant with activation, kept because early_config_prune is expecting it - SPLIT_K: tl.constexpr, - EVEN_K: tl.constexpr, - ACTIVATION: tl.constexpr, -): - - """ - Kernel for computing Out = activation(A x W + C) - - Input has shape (M, K) - - Weight has shape (K, N) - - Output has shape (M, N) - - ActInputs (optional) has shape (M, N) - 'ActInputs' optionally saves the A x W + C intermediate for backward computations - This kernel will consolidate over K - """ - - pid = tl.program_id(axis=0) - - grid_m = (M + BLOCK_M - 1) // BLOCK_M - grid_n = (N + BLOCK_N - 1) // BLOCK_N - # re-order program ID for better L2 performance - width = GROUP_M * grid_n - group_id = pid // width - group_size = min(grid_m - group_id * GROUP_M, GROUP_M) - pid_m = group_id * GROUP_M + (pid % group_size) - pid_n = (pid % width) // (group_size) - - # now compute the block that each program will go through - # rm (resp. rn) denotes a range of indices - # for rows (resp. col) of C - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - # trick to avoid masking on M and N axis - ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M) - rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N) - rk = tl.arange(0, BLOCK_K) - - A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak) - B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn) - - acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) - - for k in range(K, 0, -BLOCK_K): - if EVEN_K: - a = tl.load(A) - b = tl.load(B) - else: - a = tl.load(A, mask=rk[None, :] < k, other=0.0) - b = tl.load(B, mask=rk[:, None] < k, other=0.0) - acc += tl.dot(a, b) - - A += BLOCK_K * stride_ak - B += BLOCK_K * stride_bk - - # optional: fused activation (while the data is in shared memory) - if ACTIVATION != "id": - act_in_ptrs = ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :] - act_input = tl.load(act_in_ptrs).to(acc.dtype) - if ACTIVATION == "gelu": - acc *= gelu_grad(act_input) - elif ACTIVATION == "gelu_approx": - acc *= gelu_approx_grad(act_input) - elif ACTIVATION == "squared_relu": - acc *= squared_relu_grad(act_input) - - # rematerialize rm and rn to save registers - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - - # write back result - C = C + rm[:, None] * stride_cm + rn[None, :] - mask = (rm < M)[:, None] & (rn < N)[None, :] - tl.store(C, acc, mask=mask) - - -def triton_dgrad_act( - grad_output: torch.Tensor, - weight: torch.Tensor, - activation: str = "id", - act_input: Optional[torch.Tensor] = None, -) -> torch.Tensor: - """ - Compute e = activation(grad_output @ weight + bias). - This wrapper kicks the `kernel_fwd` Triton kernel - :param grad_output: input tensor - :param weight: weight matrix - :param activation: Activation name. Needs to be a Triton kernel. - :param act_input: an optional tensor to save the activation inputs (for backward) - :return: result tensor - """ - assert activation in ["id", "gelu", "gelu_approx", "squared_relu"] - - batch_shape, n = grad_output.shape[:-1], grad_output.shape[-1] - batch_dim = batch_shape.numel() - grad_output_reshaped = grad_output.reshape(batch_dim, n) - - if grad_output_reshaped.stride(0) > 1 and grad_output_reshaped.stride(1) > 1: - grad_output_reshaped = grad_output_reshaped.contiguous() - if weight.stride(0) > 1 and weight.stride(1) > 1: - weight = weight.contiguous() - - assert ( - grad_output.dtype == weight.dtype - ), f"grad_output and weight must have the same dtype, got {grad_output.dtype} and {weight.dtype}" - assert ( - grad_output_reshaped.shape[1] == weight.shape[0] - ), f"Incompatible dimensions: {grad_output_reshaped.shape} - {weight.shape}" - if activation != "id": - assert act_input is not None, f"act_input is required for activation {activation}" - - # M, N, K in bwd are different from M, N, K in fwd - M, K = grad_output_reshaped.shape - K, N = weight.shape - - grad_input = torch.empty((M, N), device=grad_output.device, dtype=grad_output.dtype) - - # 1D launch kernel where each block gets its own program. - grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),) # noqa - - kernel_bwd[grid]( - grad_input, - act_input, - grad_output_reshaped, - weight, # data ptrs - M, # shapes - N, - K, - M // 32, # key for triton cache (limit number of compilations) - N // 32, - K // 32, - stride_cm=grad_input.stride(0), # strides - # stride_cn=grad_input.stride(1), - stride_am=grad_output_reshaped.stride(0), - stride_ak=grad_output_reshaped.stride(1), - stride_bk=weight.stride(0), - stride_bn=weight.stride(1), - ACTIVATION=activation, # optional fused activation - GROUP_M=8, # speed optimization: group the programs - ) - - return grad_input.reshape(*batch_shape, grad_input.shape[-1]) diff --git a/build/torch29-cxx11-cu129-x86_64-linux/ops/triton/mlp.py b/build/torch29-cxx11-cu129-x86_64-linux/ops/triton/mlp.py deleted file mode 100644 index 059f4f8a5e174c1f4824e43d313fca18eaa799b8..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-x86_64-linux/ops/triton/mlp.py +++ /dev/null @@ -1,149 +0,0 @@ -# The triton fused matmul + sqrelu is faster for fp16 but slower for bf16, compared -# to naive implementation. -import fused_dense_lib as fused_dense_cuda -import torch -import torch.nn as nn -import torch.nn.functional as F - -from flash_attn.utils.torch import custom_fwd, custom_bwd -from flash_attn.ops.activations import sqrelu_bwd, sqrelu_fwd -from flash_attn.ops.triton.linear import triton_dgrad_act, triton_linear_act - - -class FusedDenseSqreluDenseFunc(torch.autograd.Function): - @staticmethod - @custom_fwd - def forward(ctx, x, weight1, bias1, weight2, bias2, checkpoint_lvl=0): - """checkpoint_lvl: - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute act_input and gelu_out in the bwd - """ - if torch.is_autocast_enabled(): - dtype = torch.get_autocast_gpu_dtype() - x, weight1, bias1, weight2, bias2 = [ - a.to(dtype=dtype) for a in [x, weight1, bias1, weight2, bias2] - ] - is_bf16 = x.dtype == torch.bfloat16 - assert checkpoint_lvl in [0, 1, 2] - x = x.contiguous() - weight1 = weight1.contiguous() - bias1 = bias1.contiguous() - weight2 = weight2.contiguous() - bias2 = bias2.contiguous() - batch_shape, n = x.shape[:-1], x.shape[-1] - batch_dim = batch_shape.numel() - if is_bf16: - act_input = fused_dense_cuda.linear_bias_forward( - x.reshape(batch_dim, n), weight1, bias1 - ) - output1 = sqrelu_fwd(act_input) - else: - save_act_input = checkpoint_lvl != 2 - result = triton_linear_act( - x.reshape(batch_dim, n), - weight1, - bias1, - activation="squared_relu", - save_act_input=save_act_input, - ) - if save_act_input: - output1, act_input = result - else: - output1 = result - output2 = fused_dense_cuda.linear_bias_forward(output1, weight2, bias2) - ctx.checkpoint_lvl = checkpoint_lvl - if checkpoint_lvl == 0: - ctx.save_for_backward(x, weight1, bias1, weight2, act_input, output1) - elif checkpoint_lvl == 1: - ctx.save_for_backward(x, weight1, bias1, weight2, act_input) - elif checkpoint_lvl == 2: - ctx.save_for_backward(x, weight1, bias1, weight2) - return output2.reshape(*batch_shape, output2.shape[-1]) - - @staticmethod - @custom_bwd - def backward(ctx, grad_output): - grad_output = grad_output.contiguous() - checkpoint_lvl = ctx.checkpoint_lvl - x, weight1, bias1, weight2, *rest = ctx.saved_tensors - batch_shape, n = x.shape[:-1], x.shape[-1] - batch_dim = batch_shape.numel() - is_bf16 = x.dtype == torch.bfloat16 - if checkpoint_lvl == 0: - act_input, output1 = rest - elif checkpoint_lvl == 1: - (act_input,) = rest - output1 = sqrelu_fwd(act_input) - elif checkpoint_lvl == 2: - if is_bf16: - act_input = fused_dense_cuda.linear_bias_forward( - x.reshape(batch_dim, n), weight1, bias1 - ) - output1 = sqrelu_fwd(act_input) - else: - output1, act_input = triton_linear_act( - x.reshape(batch_dim, n), - weight1, - bias1, - activation="squared_relu", - save_act_input=True, - ) - - if is_bf16: - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad(output1, grad_output) - grad_output1 = grad_output @ weight2 - grad_act_input = sqrelu_bwd(grad_output1, act_input) - grad_input, grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_backward( - x.reshape(batch_dim, n), weight1, grad_act_input - ) - else: - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad(output1, grad_output) - grad_act_input = triton_dgrad_act( - grad_output, weight2, activation="squared_relu", act_input=act_input - ) - grad_input, grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_backward( - x.reshape(batch_dim, n), weight1, grad_act_input - ) - return grad_input.reshape_as(x), grad_weight1, grad_bias1, grad_weight2, grad_bias2, None - - -fused_dense_sqrelu_dense_function = FusedDenseSqreluDenseFunc.apply - - -class FusedDenseSqreluDense(nn.Module): - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - bias1=True, - bias2=True, - checkpoint_lvl=0, - device=None, - dtype=None, - ): - """ - checkpoint_lvl (increasing lvl means slower but more memory saving): - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute gelu_in and gelu_out in the bwd - """ - assert checkpoint_lvl in [0, 1, 2] - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features * 4 - assert bias1 == True, "DenseSqreluDense module without bias is currently not supported" - assert bias2 == True, "DenseSqreluDense module without bias is currently not supported" - self.checkpoint_lvl = checkpoint_lvl - self.fc1 = nn.Linear(in_features, hidden_features, bias=bias1, **factory_kwargs) - self.fc2 = nn.Linear(hidden_features, out_features, bias=bias2, **factory_kwargs) - - def forward(self, x): - assert x.is_cuda - return fused_dense_sqrelu_dense_function( - x, self.fc1.weight, self.fc1.bias, self.fc2.weight, self.fc2.bias, self.checkpoint_lvl - ) diff --git a/build/torch29-cxx11-cu129-x86_64-linux/ops/triton/rotary.py b/build/torch29-cxx11-cu129-x86_64-linux/ops/triton/rotary.py deleted file mode 100644 index 6ad3ab412a084270120167ad32b67ad1720f2c37..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-x86_64-linux/ops/triton/rotary.py +++ /dev/null @@ -1,186 +0,0 @@ -# Copyright (c) 2025, Tri Dao. -# As of 2025-04-23, we require triton >= 3.0 - -from typing import Optional, Union - -import torch - -import triton -import triton.language as tl - - -@triton.jit -def rotary_kernel( - OUT, # Pointers to matrices - X, - COS, - SIN, - CU_SEQLENS, - SEQLEN_OFFSETS, # this could be int or a pointer - # Matrix dimensions - seqlen, - nheads, - seqlen_ro, - # strides - stride_out_batch, - stride_out_seqlen, - stride_out_nheads, - stride_out_headdim, - stride_x_batch, - stride_x_seqlen, - stride_x_nheads, - stride_x_headdim, - # Meta-parameters - # We want ROTARY_DIM to be constexpr, otherwise the triton compiler doesn't know that - # the mask is constant every 8 elements, and it will generate LDG.16 instead of LDG.128 - ROTARY_DIM: tl.constexpr, - IS_SEQLEN_OFFSETS_TENSOR: tl.constexpr, - IS_VARLEN: tl.constexpr, - INTERLEAVED: tl.constexpr, - CONJUGATE: tl.constexpr, - BLOCK_H: tl.constexpr, - BLOCK_M: tl.constexpr, -): - BLOCK_K: tl.constexpr = triton.next_power_of_2(ROTARY_DIM) - ROTARY_DIM_HALF = ROTARY_DIM // 2 - pid_head = tl.program_id(axis=0) - pid_m = tl.program_id(axis=1) - pid_batch = tl.program_id(axis=2) - - if not IS_VARLEN: - X = X + pid_batch * stride_x_batch - OUT = OUT + pid_batch * stride_out_batch - else: - start_idx = tl.load(CU_SEQLENS + pid_batch) - seqlen = tl.load(CU_SEQLENS + pid_batch + 1) - start_idx - X = X + start_idx * stride_x_seqlen - OUT = OUT + start_idx * stride_out_seqlen - - if pid_m * BLOCK_M >= seqlen: - return - - rh = pid_head * BLOCK_H + tl.arange(0, BLOCK_H) - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - if not IS_SEQLEN_OFFSETS_TENSOR: - rm_cs = rm + SEQLEN_OFFSETS - else: - rm_cs = rm + tl.load(SEQLEN_OFFSETS + pid_batch) - - rk_half = tl.arange(0, BLOCK_K // 2) - COS = COS + (rm_cs[:, None] * ROTARY_DIM_HALF + rk_half[None, :]) - SIN = SIN + (rm_cs[:, None] * ROTARY_DIM_HALF + rk_half[None, :]) - mask_cs = (rm_cs[:, None] < seqlen_ro) & (rk_half[None, :] < ROTARY_DIM_HALF) - cos = tl.load(COS, mask=mask_cs, other=1.0).to(tl.float32) - sin = tl.load(SIN, mask=mask_cs, other=0.0).to(tl.float32) - if CONJUGATE: - sin = -sin - - if not INTERLEAVED: - # Load the 1st and 2nd halves of X, do calculation, then store to 1st and 2nd halves of OUT - X = X + (rh[:, None, None] * stride_x_nheads + rm[None, :, None] * stride_x_seqlen + rk_half[None, None, :] * stride_x_headdim) - OUT = OUT + (rh[:, None, None] * stride_out_nheads + rm[None, :, None] * stride_out_seqlen + rk_half[None, None, :] * stride_out_headdim) - mask = (rh[:, None, None] < nheads) & (rm[None, :, None] < seqlen) & (rk_half[None, None, :] < ROTARY_DIM_HALF) - x0 = tl.load(X, mask=mask, other=0.0).to(tl.float32) - x1 = tl.load(X + ROTARY_DIM_HALF * stride_x_headdim, mask=mask, other=0.0,).to(tl.float32) - o0 = x0 * cos - x1 * sin - o1 = x0 * sin + x1 * cos - tl.store(OUT, o0, mask=mask) - tl.store(OUT + ROTARY_DIM_HALF * stride_out_headdim, o1, mask=mask) - else: - rk = tl.arange(0, BLOCK_K) - X = X + (rh[:, None, None] * stride_x_nheads + rm[None, :, None] * stride_x_seqlen + rk[None, None, :] * stride_x_headdim) - OUT = OUT + (rh[:, None, None] * stride_out_nheads + rm[None, :, None] * stride_out_seqlen + rk[None, None, :] * stride_out_headdim) - mask = (rh[:, None, None] < nheads) & (rm[None, :, None] < seqlen) & (rk[None, None, :] < ROTARY_DIM) - x = tl.load(X, mask=mask, other=0.0).to(tl.float32) - x0, x1 = tl.split(tl.reshape(x, [BLOCK_H, BLOCK_M, BLOCK_K // 2, 2])) - o0 = x0 * cos - x1 * sin - o1 = x0 * sin + x1 * cos - o = tl.reshape(tl.join(o0, o1), [BLOCK_H, BLOCK_M, BLOCK_K]) - tl.store(OUT, o, mask=mask) - - -def apply_rotary( - x: torch.Tensor, - cos: torch.Tensor, - sin: torch.Tensor, - seqlen_offsets: Union[int, torch.Tensor] = 0, - cu_seqlens: Optional[torch.Tensor] = None, - max_seqlen: Optional[int] = None, - interleaved=False, - inplace=False, - conjugate=False, -) -> torch.Tensor: - """ - Arguments: - x: (batch, seqlen, nheads, headdim) if cu_seqlens is None - else (total_seqlen, nheads, headdim). - cos: (seqlen_ro, rotary_dim / 2) - sin: (seqlen_ro, rotary_dim / 2) - seqlen_offsets: integer or integer tensor of size (batch,) - cu_seqlens: (batch + 1,) or None - max_seqlen: int - Returns: - y: (batch, seqlen, nheads, headdim) - """ - is_varlen = cu_seqlens is not None - if not is_varlen: - batch, seqlen, nheads, headdim = x.shape - else: - assert max_seqlen is not None, "If cu_seqlens is passed in, then max_seqlen must be passed" - total_seqlen, nheads, headdim = x.shape - batch_p_1 = cu_seqlens.shape[0] - batch = batch_p_1 - 1 - seqlen = max_seqlen - seqlen_ro, rotary_dim = cos.shape - assert sin.shape == cos.shape - rotary_dim *= 2 - assert rotary_dim <= headdim, "rotary_dim must be <= headdim" - assert headdim <= 256, "Only support headdim <= 256" - assert seqlen_ro >= seqlen, "seqlen_ro must be >= seqlen" - - cos, sin = cos.contiguous(), sin.contiguous() - if isinstance(seqlen_offsets, torch.Tensor): - assert seqlen_offsets.shape == (batch,) - assert seqlen_offsets.dtype in [torch.int32, torch.int64] - seqlen_offsets = seqlen_offsets.contiguous() - else: - assert seqlen_offsets + seqlen <= seqlen_ro - - output = torch.empty_like(x) if not inplace else x - if rotary_dim < headdim and not inplace: - output[..., rotary_dim:].copy_(x[..., rotary_dim:]) - - grid = lambda META: (triton.cdiv(nheads, META["BLOCK_H"]), triton.cdiv(seqlen, META["BLOCK_M"]), batch) # noqa - BLOCK_M = 8 if rotary_dim <= 128 else 4 - - # Need this, otherwise Triton tries to launch from cuda:0 and we get - # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?) - device_ctx = torch.cuda.device(x.device.index) if x.device.type == 'cuda' else torch.xpu.device(x.device.index) - with device_ctx: - torch.library.wrap_triton(rotary_kernel)[grid]( - output, # data ptrs - x, - cos, - sin, - cu_seqlens, - seqlen_offsets, - seqlen, # shapes - nheads, - seqlen_ro, - output.stride(0) if not is_varlen else 0, # batch_strides if not varlen else 0 - output.stride(-3), # seqlen_stride or total_seqlen_stride - output.stride(-2), # nheads_stride - output.stride(-1), # headdim_stride - x.stride(0) if not is_varlen else 0, # batch_strides if not varlen else 0 - x.stride(-3), # seqlen stride or total_seqlen_stride - x.stride(-2), # nheads stride - x.stride(-1), # headdim stride - rotary_dim, - isinstance(seqlen_offsets, torch.Tensor), - is_varlen, - interleaved, - conjugate, - BLOCK_M=BLOCK_M, - BLOCK_H=2, - ) - return output diff --git a/build/torch29-cxx11-cu130-aarch64-linux/__init__.py b/build/torch29-cxx11-cu130-aarch64-linux/__init__.py deleted file mode 100644 index ecc2f9d896b6c93f90b0a1499856dc0612177422..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-aarch64-linux/__init__.py +++ /dev/null @@ -1,393 +0,0 @@ -from typing import Optional, List -import torch -from ._ops import ops as flash_attn_ops -from .flash_attn_interface import ( - flash_attn_func, - flash_attn_kvpacked_func, - flash_attn_qkvpacked_func, - flash_attn_varlen_func, - flash_attn_varlen_kvpacked_func, - flash_attn_varlen_qkvpacked_func, - flash_attn_with_kvcache, -) - - -def fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - return_softmax: bool = False, - gen: Optional[torch.Generator] = None, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention. - - Args: - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Optional output tensor, same shape as q - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - return_softmax: Whether to return softmax weights - gen: Optional random number generator - - Returns: - List of tensors: [output, softmax_lse, (softmax if return_softmax)] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.fwd( - q, - k, - v, - out, - alibi_slopes, - p_dropout, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - gen, - ) - - -def varlen_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - out: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - block_table: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - max_seqlen_q: int = 0, - max_seqlen_k: int = 0, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - zero_tensors: bool = False, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - return_softmax: bool = False, - gen: Optional[torch.Generator] = None, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention with variable sequence lengths. - - Args: - q: Query tensor of shape [total_q, num_heads, head_size] - k: Key tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - v: Value tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1] - cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1] - out: Optional output tensor of shape [total_q, num_heads, head_size] - seqused_k: Optional tensor specifying how many keys to use per batch element [batch_size] - leftpad_k: Optional left padding for keys of shape [batch_size] - block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq] - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - max_seqlen_q: Maximum sequence length for queries - max_seqlen_k: Maximum sequence length for keys - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - zero_tensors: Whether to zero tensors before computation - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - return_softmax: Whether to return softmax weights - gen: Optional random number generator - - Returns: - List of tensors: [output, softmax_lse, (softmax if return_softmax)] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.varlen_fwd( - q, - k, - v, - out, - cu_seqlens_q, - cu_seqlens_k, - seqused_k, - leftpad_k, - block_table, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - p_dropout, - softmax_scale, - zero_tensors, - is_causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - gen, - ) - - -def bwd( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - p_dropout: float = 0.0, - 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, - gen: Optional[torch.Generator] = None, - rng_state: Optional[torch.Tensor] = None, -) -> List[torch.Tensor]: - """ - Backward pass for multi-head attention. - - Args: - dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size] - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size] - softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q] - dq: Optional gradient tensor for queries, same shape as q - dk: Optional gradient tensor for keys, same shape as k - dv: Optional gradient tensor for values, same shape as v - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - deterministic: Whether to use deterministic algorithms - gen: Optional random number generator - rng_state: Optional RNG state from forward pass - - Returns: - List of tensors: [dq, dk, dv] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - alibi_slopes, - p_dropout, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - deterministic, - gen, - rng_state, - ) - - -def varlen_bwd( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - max_seqlen_q: int = 0, - max_seqlen_k: int = 0, - p_dropout: float = 0.0, - softmax_scale: Optional[float] = None, - zero_tensors: bool = False, - is_causal: bool = False, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - deterministic: bool = False, - gen: Optional[torch.Generator] = None, - rng_state: Optional[torch.Tensor] = None, -) -> List[torch.Tensor]: - """ - Backward pass for multi-head attention with variable sequence lengths. - - Args: - dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size] - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size] - out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size] - softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q] - cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1] - cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1] - dq: Optional gradient tensor for queries, same shape as q - dk: Optional gradient tensor for keys, same shape as k - dv: Optional gradient tensor for values, same shape as v - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - max_seqlen_q: Maximum sequence length for queries - max_seqlen_k: Maximum sequence length for keys - p_dropout: Dropout probability - softmax_scale: Scale factor for softmax - zero_tensors: Whether to zero tensors before computation - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - deterministic: Whether to use deterministic algorithms - gen: Optional random number generator - rng_state: Optional RNG state from forward pass - - Returns: - List of tensors: [dq, dk, dv] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.varlen_bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - p_dropout, - softmax_scale, - zero_tensors, - is_causal, - window_size_left, - window_size_right, - softcap, - deterministic, - gen, - rng_state, - ) - - -def fwd_kvcache( - q: torch.Tensor, - kcache: torch.Tensor, - vcache: torch.Tensor, - k: Optional[torch.Tensor] = None, - v: Optional[torch.Tensor] = None, - seqlens_k: Optional[torch.Tensor] = None, - rotary_cos: Optional[torch.Tensor] = None, - rotary_sin: Optional[torch.Tensor] = None, - cache_batch_idx: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - block_table: Optional[torch.Tensor] = None, - alibi_slopes: Optional[torch.Tensor] = None, - out: 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, - is_rotary_interleaved: bool = False, - num_splits: int = 1, -) -> List[torch.Tensor]: - """ - Forward pass for multi-head attention with KV cache. - - Args: - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - kcache: Key cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - vcache: Value cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size] - k: Optional new keys tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size] - v: Optional new values tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size] - seqlens_k: Optional sequence lengths for keys of shape [batch_size] - rotary_cos: Optional rotary cosine tensor of shape [seqlen_ro, rotary_dim/2] - rotary_sin: Optional rotary sine tensor of shape [seqlen_ro, rotary_dim/2] - cache_batch_idx: Optional indices to index into the KV cache - leftpad_k: Optional left padding for keys of shape [batch_size] - block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq] - alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads] - out: Optional output tensor, same shape as q - softmax_scale: Scale factor for softmax - is_causal: Whether to use causal attention - window_size_left: Window size for left context (-1 for unlimited) - window_size_right: Window size for right context (-1 for unlimited) - softcap: Soft cap for attention weights - is_rotary_interleaved: Whether rotary embeddings are interleaved - num_splits: Number of splits for computation - - Returns: - List of tensors: [output, softmax_lse] - """ - if softmax_scale is None: - attention_head_dim = q.shape[-1] - softmax_scale = 1.0 / (attention_head_dim**0.5) - - return flash_attn_ops.fwd_kvcache( - q, - kcache, - vcache, - k, - v, - seqlens_k, - rotary_cos, - rotary_sin, - cache_batch_idx, - leftpad_k, - block_table, - alibi_slopes, - out, - softmax_scale, - is_causal, - window_size_left, - window_size_right, - softcap, - is_rotary_interleaved, - num_splits, - ) diff --git a/build/torch29-cxx11-cu130-aarch64-linux/_flash_attn2_cuda_abda3a0.abi3.so b/build/torch29-cxx11-cu130-aarch64-linux/_flash_attn2_cuda_abda3a0.abi3.so deleted file mode 100644 index 088f0fe7212b7ba57b138c290f9721d0770c1e8c..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-aarch64-linux/_flash_attn2_cuda_abda3a0.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:f050d21698e38a8353748b4e5bcb4ee962b059d32cedc40313eafd637f23d7b2 -size 1008640480 diff --git a/build/torch29-cxx11-cu130-aarch64-linux/_ops.py b/build/torch29-cxx11-cu130-aarch64-linux/_ops.py deleted file mode 100644 index fb32434e09e651a658beb38be8c101e5db89374f..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-aarch64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _flash_attn2_cuda_abda3a0 -ops = torch.ops._flash_attn2_cuda_abda3a0 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_flash_attn2_cuda_abda3a0::{op_name}" diff --git a/build/torch29-cxx11-cu130-aarch64-linux/bert_padding.py b/build/torch29-cxx11-cu130-aarch64-linux/bert_padding.py deleted file mode 100644 index 3c2d35159a014a9d03aabead9e52e009168696ea..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-aarch64-linux/bert_padding.py +++ /dev/null @@ -1,218 +0,0 @@ -# Adapted from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/padding.py - -import torch -import torch.nn.functional as F -from einops import rearrange, repeat - - -class IndexFirstAxis(torch.autograd.Function): - @staticmethod - def forward(ctx, input, indices): - ctx.save_for_backward(indices) - assert input.ndim >= 2 - ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] - second_dim = other_shape.numel() - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - # return input[indices] - return torch.gather( - rearrange(input, "b ... -> b (...)"), 0, repeat(indices, "z -> z d", d=second_dim) - ).reshape(-1, *other_shape) - - @staticmethod - def backward(ctx, grad_output): - (indices,) = ctx.saved_tensors - assert grad_output.ndim >= 2 - other_shape = grad_output.shape[1:] - grad_output = rearrange(grad_output, "b ... -> b (...)") - grad_input = torch.zeros( - [ctx.first_axis_dim, grad_output.shape[1]], - device=grad_output.device, - dtype=grad_output.dtype, - ) - # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. - # grad_input[indices] = grad_output - grad_input.scatter_(0, repeat(indices, "z -> z d", d=grad_output.shape[1]), grad_output) - return grad_input.reshape(ctx.first_axis_dim, *other_shape), None - - -index_first_axis = IndexFirstAxis.apply - - -class IndexPutFirstAxis(torch.autograd.Function): - @staticmethod - def forward(ctx, values, indices, first_axis_dim): - ctx.save_for_backward(indices) - assert indices.ndim == 1 - assert values.ndim >= 2 - output = torch.zeros( - first_axis_dim, *values.shape[1:], device=values.device, dtype=values.dtype - ) - # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. - output[indices] = values - # output.scatter_(0, repeat(indices, 'z -> z d', d=values.shape[1]), values) - return output - - @staticmethod - def backward(ctx, grad_output): - (indices,) = ctx.saved_tensors - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - grad_values = grad_output[indices] - # grad_values = torch.gather(grad_output, 0, repeat(indices, 'z -> z d', d=grad_output.shape[1])) - return grad_values, None, None - - -index_put_first_axis = IndexPutFirstAxis.apply - - -class IndexFirstAxisResidual(torch.autograd.Function): - @staticmethod - def forward(ctx, input, indices): - ctx.save_for_backward(indices) - assert input.ndim >= 2 - ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] - second_dim = other_shape.numel() - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - output = input[indices] - # We don't want to reshape input (b ... -> b (...)) since it could change the channel_last - # memory format to channel_first. In other words, input might not be contiguous. - # If we don't detach, Pytorch complains about output being a view and is being modified inplace - return output, input.detach() - - @staticmethod - def backward(ctx, grad_output, grad_residual): - (indices,) = ctx.saved_tensors - assert grad_output.ndim >= 2 - other_shape = grad_output.shape[1:] - assert grad_residual.shape[1:] == other_shape - grad_input = grad_residual - # grad_input[indices] += grad_output - indices = indices.reshape(indices.shape[0], *((1,) * (grad_output.ndim - 1))) - indices = indices.expand_as(grad_output) - grad_input.scatter_add_(0, indices, grad_output) - return grad_input.reshape(ctx.first_axis_dim, *other_shape), None - - -index_first_axis_residual = IndexFirstAxisResidual.apply - - -def unpad_input(hidden_states, attention_mask, unused_mask=None): - """ - Arguments: - hidden_states: (batch, seqlen, ...) - attention_mask: (batch, seqlen), bool / int, 1 means valid and 0 means not valid. - unused_mask: (batch, seqlen), bool / int, 1 means the element is allocated but unused. - Return: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens selected in attention_mask + unused_mask. - indices: (total_nnz), the indices of masked tokens from the flattened input sequence. - cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. - max_seqlen_in_batch: int - seqused: (batch), returns the number of tokens selected in attention_mask + unused_mask. - """ - all_masks = (attention_mask + unused_mask) if unused_mask is not None else attention_mask - seqlens_in_batch = all_masks.sum(dim=-1, dtype=torch.int32) - used_seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) - indices = torch.nonzero(all_masks.flatten(), as_tuple=False).flatten() - max_seqlen_in_batch = seqlens_in_batch.max().item() - cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) - # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the - # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim - # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to - # index with integer indices. Moreover, torch's index is a bit slower than it needs to be, - # so we write custom forward and backward to make it a bit faster. - return ( - index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices), - indices, - cu_seqlens, - max_seqlen_in_batch, - used_seqlens_in_batch, - ) - - -def unpad_input_for_concatenated_sequences(hidden_states, attention_mask_in_length): - """ - Supports concatenating short samples in one sequence. The attention_mask_in_length is utilized to mask other short samples. It helps efficient training of variant lengths-based samples (e.g., the supervised fine-tuning task in large language model). - The motivation for this function is explained [here](https://github.com/Dao-AILab/flash-attention/issues/432#issuecomment-1668822286). - - For example, if batch = 3 and seqlen = 6, the attention_mask_in_length is: - ``` - [ - [2, 3, 0, 0, 0, 0], - [3, 2, 0, 0, 0, 0], - [6, 0, 0, 0, 0, 0] - ] - ``` - , which refers to the 3D-attention mask: - ``` - [ - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [0, 0, 1, 0, 0, 0], - [0, 0, 1, 1, 0, 0], - [0, 0, 1, 1, 1, 0], - [0, 0, 0, 0, 0, 1] - ], - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [0, 0, 0, 1, 0, 0], - [0, 0, 0, 1, 1, 0], - [0, 0, 0, 0, 0, 1] - ], - [ - [1, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [1, 1, 1, 1, 0, 0], - [1, 1, 1, 1, 1, 0], - [1, 1, 1, 1, 1, 1] - ] - ] - ```. - - Arguments: - hidden_states: (batch, seqlen, ...) - attention_mask_in_length: (batch, seqlen), int, a nonzero number (e.g., 1, 2, 3, etc.) means length of concatenated sequence in b-th batch, and 0 means none. - Return: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. - indices: (total_nnz), the indices of non-masked tokens from the flattened input sequence. - cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. - max_seqlen_in_batch: int - """ - length = attention_mask_in_length.sum(dim=-1) - seqlen = attention_mask_in_length.size(-1) - attention_mask_2d = torch.arange(seqlen, device=length.device, dtype=length.dtype).expand(len(length), seqlen) < length.unsqueeze(1) - real_indices_idx = torch.nonzero(attention_mask_in_length.flatten(), as_tuple=False).flatten() - seqlens_in_batch = attention_mask_in_length.flatten()[real_indices_idx] - indices = torch.nonzero(attention_mask_2d.flatten(), as_tuple=False).flatten() - max_seqlen_in_batch = seqlens_in_batch.max().item() - cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) - # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the - # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim - # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to - # index with integer indices. Moreover, torch's index is a bit slower than it needs to be, - # so we write custom forward and backward to make it a bit faster. - return ( - index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices), - indices, - cu_seqlens, - max_seqlen_in_batch, - ) - - -def pad_input(hidden_states, indices, batch, seqlen): - """ - Arguments: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. - indices: (total_nnz), the indices that represent the non-masked tokens of the original padded input sequence. - batch: int, batch size for the padded sequence. - seqlen: int, maximum sequence length for the padded sequence. - Return: - hidden_states: (batch, seqlen, ...) - """ - dim = hidden_states.shape[-1] - # output = torch.zeros((batch * seqlen), dim, device=hidden_states.device, dtype=hidden_states.dtype) - # output[indices] = hidden_states - output = index_put_first_axis(hidden_states, indices, batch * seqlen) - return rearrange(output, "(b s) ... -> b s ...", b=batch) diff --git a/build/torch29-cxx11-cu130-aarch64-linux/flash_attn2/__init__.py b/build/torch29-cxx11-cu130-aarch64-linux/flash_attn2/__init__.py deleted file mode 100644 index 03dbc1afe1cf156661a2b1b22003cd5f599a0309..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-aarch64-linux/flash_attn2/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -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-aarch64-linux/flash_attn_interface.py b/build/torch29-cxx11-cu130-aarch64-linux/flash_attn_interface.py deleted file mode 100644 index 6a0157f25449eb90c6c688bc8ef3f4fcf84d46dc..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-aarch64-linux/flash_attn_interface.py +++ /dev/null @@ -1,1622 +0,0 @@ -# Copyright (c) 2023, Tri Dao. - -from typing import Optional, Sequence, Tuple, Union - -import torch -import torch.nn as nn -import os - -# # isort: off -# # We need to import the CUDA kernels after importing torch -# USE_TRITON_ROCM = os.getenv("FLASH_ATTENTION_TRITON_AMD_ENABLE", "FALSE") == "TRUE" -# if USE_TRITON_ROCM: -# from .flash_attn_triton_amd import interface_fa as flash_attn -# else: -# import flash_attn_2_cuda as flash_attn - - -from ._ops import ops as flash_attn - -# # isort: on - -def maybe_contiguous(x): - return x.contiguous() if x is not None and x.stride(-1) != 1 else x - - -def _get_device(): - if torch.xpu.is_available(): - return "xpu" - elif torch.cuda.is_available(): - return "cuda" - else: - return "cpu" - - -def _get_block_size_n(device, head_dim, is_dropout, is_causal): - assert head_dim <= 256 - - if device.type == "xpu": - return 64 - - # This should match the block sizes in the CUDA kernel - major, minor = torch.cuda.get_device_capability(device) - is_sm8x = major == 8 and minor > 0 # Only include sm86 and sm89, exclude sm80 (A100) - is_sm80 = major == 8 and minor == 0 - is_sm90 = major == 9 and minor == 0 - if head_dim <= 32: - return 128 - if head_dim <= 64: - return 128 if not is_dropout else 64 - elif head_dim <= 96: - return 64 - elif head_dim <= 128: - if is_sm8x: - return 64 if (not is_dropout and is_causal) else 32 - else: - return 64 if not is_dropout else 32 - elif head_dim <= 192: - return 64 - elif head_dim <= 224: - return 64 - elif head_dim <= 256: - return 64 - - -def round_multiple(x, m): - return (x + m - 1) // m * m - - -# torch.compile() support is only enabled for pytorch >= 2.4 -# The reason for this is that we are using the new custom_op and register_fake -# APIs, which support inplace modification of inputs in the function itself -if torch.__version__ >= "2.4.0": - _torch_custom_op_wrapper = torch.library.custom_op - _torch_register_fake_wrapper = torch.library.register_fake -else: - def noop_custom_op_wrapper(name, fn=None, /, *, mutates_args, device_types=None, schema=None): - def wrap(func): - return func - if fn is None: - return wrap - return fn - def noop_register_fake_wrapper(op, fn=None, /, *, lib=None, _stacklevel=1): - def wrap(func): - return func - if fn is None: - return wrap - return fn - _torch_custom_op_wrapper = noop_custom_op_wrapper - _torch_register_fake_wrapper = noop_register_fake_wrapper - - -@_torch_custom_op_wrapper("flash_attn::_flash_attn_forward", mutates_args=(), device_types=_get_device()) -def _flash_attn_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - return_softmax: bool -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - out, softmax_lse, S_dmask, rng_state = flash_attn.fwd( - q, - k, - v, - None, - alibi_slopes, - dropout_p, - softmax_scale, - causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - None, - ) - return out, softmax_lse, S_dmask, rng_state - - -@_torch_register_fake_wrapper("flash_attn::_flash_attn_forward") -def _flash_attn_forward_fake( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - return_softmax: bool -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - batch_size, seqlen_q, num_heads, head_size = q.shape - seqlen_k = k.shape[1] - out = torch.empty_like(q) - softmax_lse = torch.empty((batch_size, num_heads, seqlen_q), dtype=torch.float32, device=q.device, layout=q.layout) - p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) - if return_softmax: - p = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128), round_multiple(seqlen_k, 128)), dtype=q.dtype, device=q.device, layout=q.layout) - rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) - - return out, softmax_lse, p, rng_state - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_forward = torch.ops.flash_attn._flash_attn_forward -else: - _wrapped_flash_attn_forward = _flash_attn_forward - - -@_torch_custom_op_wrapper("flash_attn::_flash_attn_varlen_forward", mutates_args=(), device_types=_get_device()) -def _flash_attn_varlen_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - alibi_slopes: Optional[torch.Tensor] = None, - return_softmax: bool = False, - block_table: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - out, softmax_lse, S_dmask, rng_state = flash_attn.varlen_fwd( - q, - k, - v, - None, - cu_seqlens_q, - cu_seqlens_k, - seqused_k, - leftpad_k, - block_table, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - zero_tensors, - causal, - window_size_left, - window_size_right, - softcap, - return_softmax, - None, - ) - # if out.isnan().any() or softmax_lse.isnan().any(): - # breakpoint() - return out, softmax_lse, S_dmask, rng_state - - -@_torch_register_fake_wrapper("flash_attn::_flash_attn_varlen_forward") -def _flash_attn_varlen_forward_fake( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int = -1, - window_size_right: int = -1, - softcap: float = 0.0, - alibi_slopes: Optional[torch.Tensor] = None, - return_softmax: bool = False, - block_table: Optional[torch.Tensor] = None, - leftpad_k: Optional[torch.Tensor] = None, - seqused_k: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - paged_kv = block_table is not None - batch_size = cu_seqlens_q.numel() - 1 - total_q, num_heads, _ = q.shape - - out = torch.empty_like(q) - softmax_lse = torch.empty((num_heads, total_q), dtype=torch.float32, device=q.device, layout=q.layout) - p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) - seqlen_q_rounded = round_multiple(max_seqlen_q, 128) - seqlen_k_rounded = round_multiple(max_seqlen_k, 128) - if return_softmax: - p = torch.empty((batch_size, num_heads, seqlen_q_rounded, seqlen_k_rounded), dtype=q.dtype, device=q.device, layout=q.layout) - rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) - return out, softmax_lse, p, rng_state - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_varlen_forward = torch.ops.flash_attn._flash_attn_varlen_forward -else: - _wrapped_flash_attn_varlen_forward = _flash_attn_varlen_forward - - -@_torch_custom_op_wrapper("flash_attn::_flash_attn_backward", mutates_args=("dq", "dk", "dv"), device_types=_get_device()) -def _flash_attn_backward( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, -) -> 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)] - ( - dq, - dk, - dv, - softmax_d, - ) = flash_attn.bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - alibi_slopes, - dropout_p, - softmax_scale, - causal, - window_size_left, - window_size_right, - softcap, - deterministic, - None, - rng_state, - ) - return softmax_d - - -@_torch_register_fake_wrapper("flash_attn::_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, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, -) -> torch.Tensor: - dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] - if dq is None: - dq = torch.empty_like(q) - if dk is None: - dk = torch.empty_like(k) - if dv is None: - dv = torch.empty_like(v) - batch_size, seqlen_q, num_heads, _ = q.shape - softmax_d = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128)), device=q.device, dtype=torch.float32) - - return softmax_d - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_backward = torch.ops.flash_attn._flash_attn_backward -else: - _wrapped_flash_attn_backward = _flash_attn_backward - - -@_torch_custom_op_wrapper("flash_attn::_flash_attn_varlen_backward", mutates_args=("dq", "dk", "dv"), device_types=_get_device()) -def _flash_attn_varlen_backward( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> 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)] - ( - dq, - dk, - dv, - softmax_d, - ) = flash_attn.varlen_bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - zero_tensors, - causal, - window_size_left, - window_size_right, - softcap, - deterministic, - None, - rng_state, - ) - # if dk.isnan().any() or dk.isnan().any() or dv.isnan().any() or softmax_d.isnan().any(): - # breakpoint() - return softmax_d - - -@_torch_register_fake_wrapper("flash_attn::_flash_attn_varlen_backward") -def _flash_attn_varlen_backward_fake( - dout: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - softmax_lse: torch.Tensor, - dq: Optional[torch.Tensor], - dk: Optional[torch.Tensor], - dv: Optional[torch.Tensor], - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - max_seqlen_k: int, - dropout_p: float, - softmax_scale: float, - causal: bool, - window_size_left: int, - window_size_right: int, - softcap: float, - alibi_slopes: Optional[torch.Tensor], - deterministic: bool, - rng_state: Optional[torch.Tensor] = None, - zero_tensors: bool = False, -) -> torch.Tensor: - dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] - batch_size = cu_seqlens_q.numel() - 1 - total_q, num_heads, _ = q.shape - - if dq is None: - dq = torch.empty_like(q) - if dk is None: - dk = torch.empty_like(k) - if dv is None: - dv = torch.empty_like(v) - softmax_d = torch.empty((num_heads, total_q + 128 * batch_size), device=q.device, dtype=torch.float32) - - return softmax_d - - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_varlen_backward = torch.ops.flash_attn._flash_attn_varlen_backward -else: - _wrapped_flash_attn_varlen_backward = _flash_attn_varlen_backward - - -class FlashAttnQKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and qkv.requires_grad - if softmax_scale is None: - softmax_scale = qkv.shape[-1] ** (-0.5) - q, k, v = qkv[:, :, 0].detach(), qkv[:, :, 1].detach(), qkv[:, :, 2].detach() - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) - dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dqkv[:, :, 0], - dqkv[:, :, 1], - dqkv[:, :, 2], - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dqkv = dqkv[..., : dout.shape[-1]] # We could have padded the head dimension - return dqkv, None, None, None, None, None, None, None, None, None - - -class FlashAttnVarlenQKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - cu_seqlens, - max_seqlen, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and qkv.requires_grad - if softmax_scale is None: - softmax_scale = qkv.shape[-1] ** (-0.5) - q, k, v = qkv[:, 0].detach(), qkv[:, 1].detach(), qkv[:, 2].detach() - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens, - cu_seqlens, - max_seqlen, - max_seqlen, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=None, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens, rng_state) - ctx.dropout_p = dropout_p - ctx.max_seqlen = max_seqlen - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens, rng_state = ctx.saved_tensors - qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) - dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dqkv[:, 0], - dqkv[:, 1], - dqkv[:, 2], - cu_seqlens, - cu_seqlens, - ctx.max_seqlen, - ctx.max_seqlen, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - 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 FlashAttnKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - kv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, kv] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - k, v = kv[:, :, 0].detach(), kv[:, :, 1].detach() - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - dq = torch.empty_like(q) - kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) - dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dkv[:, :, 0], - dkv[:, :, 1], - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dkv = dkv[..., : dout.shape[-1]] - return dq, dkv, None, None, None, None, None, None, None, None, None - - -class FlashAttnVarlenKVPackedFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, kv] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - k, v = kv[:, 0].detach(), kv[:, 1].detach() - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=None, - ) - if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state - ) - ctx.dropout_p = dropout_p - 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.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors - dq = torch.empty_like(q) - kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) - dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dkv[:, 0], - dkv[:, 1], - cu_seqlens_q, - cu_seqlens_k, - ctx.max_seqlen_q, - ctx.max_seqlen_k, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dkv = dkv[..., : dout.shape[-1]] - return dq, dkv, None, None, None, None, None, None, None, None, None, None, None, None, None - - -class FlashAttnFunc(torch.autograd.Function): - @staticmethod - def forward( - ctx, - q, - k, - v, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, k, v] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - head_size_og = q.size(3) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( - q, - k, - v, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - ) - if is_grad: - ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) - ctx.dropout_p = dropout_p - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors - dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) - head_size_og = dout.size(3) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dk = dk[..., : dout.shape[-1]] - dv = dv[..., : dout.shape[-1]] - return dq, dk, dv, 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, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_softmax, - block_table, - is_grad_enabled, - ): - is_grad = is_grad_enabled and any( - x.requires_grad for x in [q, k, v] - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - head_size_og = q.size(2) - if head_size_og % 8 != 0: - q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) - k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) - v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) - out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal=causal, - window_size_left=window_size[0], - window_size_right=window_size[1], - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=return_softmax and dropout_p > 0, - block_table=block_table, - ) - if is_grad: - ctx.save_for_backward( - q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state - ) - ctx.dropout_p = dropout_p - 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.softcap = softcap - ctx.alibi_slopes = alibi_slopes - ctx.deterministic = deterministic - - out = out_padded[..., :head_size_og] - return out if not return_softmax else (out, softmax_lse, S_dmask) - - @staticmethod - def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors - dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) - head_size_og = dout.size(2) - dout_padded = dout - if head_size_og % 8 != 0: - dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) - _wrapped_flash_attn_varlen_backward( - dout_padded, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - ctx.max_seqlen_q, - ctx.max_seqlen_k, - ctx.dropout_p, - ctx.softmax_scale, - ctx.causal, - ctx.window_size[0], - ctx.window_size[1], - ctx.softcap, - ctx.alibi_slopes, - ctx.deterministic, - rng_state=rng_state, - ) - dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension - dk = dk[..., : dout.shape[-1]] - dv = dv[..., : dout.shape[-1]] - return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None - - -def flash_attn_qkvpacked_func( - qkv, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # <=0.0 means deactivate - alibi_slopes=None, - deterministic=False, - 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, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_kvpacked_func( - q, - kv, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, -): - """dropout_p should be set to 0.0 during evaluation - If 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 K, V. - 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) - kv: (batch_size, seqlen, 2, 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. - softcap: float. Anything > 0 activates softcapping 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). - 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 FlashAttnKVPackedFunc.apply( - q, - kv, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_func( - q, - k, - v, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - 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). - 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 FlashAttnFunc.apply( - q, - k, - v, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_qkvpacked_func( - qkv, - cu_seqlens, - max_seqlen, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - 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_varlen_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_varlen_kvpacked_func and flash_attn_varlen_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: (total, 3, nheads, headdim), where total = total number of tokens in the batch. - cu_seqlens: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into qkv. - max_seqlen: int. Maximum sequence length in the batch. - 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenQKVPackedFunc.apply( - qkv, - cu_seqlens, - max_seqlen, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_kvpacked_func( - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, -): - """dropout_p should be set to 0.0 during evaluation - If 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 K, V. - 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: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch. - kv: (total_k, 2, nheads_k, headdim), where total_k = total number of key tokens in the batch. - cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into q. - cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into kv. - max_seqlen_q: int. Maximum query sequence length in the batch. - max_seqlen_k: int. Maximum key sequence length in the batch. - 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 + 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenKVPackedFunc.apply( - q, - kv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - torch.is_grad_enabled(), - ) - - -def flash_attn_varlen_func( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p=0.0, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, - block_table=None, -): - """dropout_p should be set to 0.0 during evaluation - Supports multi-query and grouped-query attention (MQA/GQA) by passing in K, V 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: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch. - k: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. - v: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. - cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into q. - cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into kv. - max_seqlen_q: int. Maximum query sequence length in the batch. - max_seqlen_k: int. Maximum key sequence length in the batch. - 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 + 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: (total, nheads, headdim). - softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_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 FlashAttnVarlenFunc.apply( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - causal, - window_size, - softcap, - alibi_slopes, - deterministic, - return_attn_probs, - block_table, - False if q.device.type == "cpu" else torch.is_grad_enabled(), - ) - - -def flash_attn_with_kvcache( - q, - k_cache, - v_cache, - k=None, - v=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, - block_table: Optional[torch.Tensor] = None, - softmax_scale=None, - causal=False, - window_size=(-1, -1), # -1 means infinite context window - softcap=0.0, # 0.0 means deactivated - rotary_interleaved=True, - alibi_slopes=None, - num_splits=0, - 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 block_table, - or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_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) if there's no block_table, - or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_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). Similar to k. - 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. - block_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). - 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. - 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" - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-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) - cache_batch_idx = maybe_contiguous(cache_batch_idx) - block_table = maybe_contiguous(block_table) - out, softmax_lse = flash_attn.fwd_kvcache( - q, - k_cache, - v_cache, - k, - v, - cache_seqlens, - rotary_cos, - rotary_sin, - cache_batch_idx, - cache_leftpad, - block_table, - alibi_slopes, - None, - softmax_scale, - causal, - window_size[0], - window_size[1], - softcap, - rotary_interleaved, - num_splits, - ) - return (out, softmax_lse) if return_softmax_lse else out diff --git a/build/torch29-cxx11-cu130-aarch64-linux/layers/__init__.py b/build/torch29-cxx11-cu130-aarch64-linux/layers/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/build/torch29-cxx11-cu130-aarch64-linux/layers/patch_embed.py b/build/torch29-cxx11-cu130-aarch64-linux/layers/patch_embed.py deleted file mode 100644 index 05562f8e8bcdb58e947c6f402a49eacd2d031871..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-aarch64-linux/layers/patch_embed.py +++ /dev/null @@ -1,67 +0,0 @@ -# We use the same API as https://github.com/rwightman/pytorch-image-models/blob/v0.6.11/timm/models/layers/patch_embed.py -# But we use nn.Linear instead of Conv2d and it's about 8x faster. - -from functools import partial - -import torch.nn as nn -from einops import rearrange -from torch import _assert -from torch.nn.modules.utils import _pair - -try: - from flash_attn.ops.fused_dense import FusedDense -except ImportError: - FusedDense = None - - -class PatchEmbed(nn.Module): - """2D Image to Patch Embedding""" - - def __init__( - self, - img_size=224, - patch_size=16, - in_chans=3, - embed_dim=768, - norm_layer=None, - flatten=True, - bias=True, - fused_bias_fc=False, - ): - super().__init__() - img_size = _pair(img_size) - patch_size = _pair(patch_size) - self.img_size = img_size - self.patch_size = patch_size - self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1]) - self.num_patches = self.grid_size[0] * self.grid_size[1] - self.flatten = flatten - if fused_bias_fc and FusedDense is None: - raise ImportError("fused_dense is not installed") - - linear_cls = nn.Linear if not fused_bias_fc or not bias else FusedDense - self.proj = linear_cls(in_chans * patch_size[0] * patch_size[1], embed_dim, bias=bias) - self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() - - def forward(self, x): - _, _, H, W = x.shape - _assert( - H == self.img_size[0], - f"Input image height ({H}) doesn't match model ({self.img_size[0]}).", - ) - _assert( - W == self.img_size[1], - f"Input image width ({W}) doesn't match model ({self.img_size[1]}).", - ) - x = self.proj( - rearrange( - x, - "b c (h p1) (w p2) -> b h w (c p1 p2)", - p1=self.patch_size[0], - p2=self.patch_size[1], - ) - ) - if self.flatten: - x = rearrange(x, "b h w c -> b (h w) c") - x = self.norm(x) - return x diff --git a/build/torch29-cxx11-cu130-aarch64-linux/layers/rotary.py b/build/torch29-cxx11-cu130-aarch64-linux/layers/rotary.py deleted file mode 100644 index d1bfc21fc7de1dd287e8f382847b194a48075981..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-aarch64-linux/layers/rotary.py +++ /dev/null @@ -1,483 +0,0 @@ -# Copyright (c) 2025, Tri Dao - -import math -from functools import partial -from typing import Optional, Tuple, Union - -import torch -from torch import Tensor - -from einops import rearrange, repeat -# from flash_attn.ops.triton.rotary import apply_rotary -from ..ops.triton.rotary import apply_rotary - - -def rotate_half(x, interleaved=False): - if not interleaved: - x1, x2 = x.chunk(2, dim=-1) - return torch.cat((-x2, x1), dim=-1) - else: - x1, x2 = x[..., ::2], x[..., 1::2] - return rearrange(torch.stack((-x2, x1), dim=-1), "... d two -> ... (d two)", two=2) - - -def apply_rotary_emb_torch(x, cos, sin, interleaved=False): - """ - x: (batch_size, seqlen, nheads, headdim) - cos, sin: (seqlen, rotary_dim / 2) or (batch_size, seqlen, rotary_dim / 2) - """ - ro_dim = cos.shape[-1] * 2 - assert ro_dim <= x.shape[-1] - cos = repeat(cos, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)") - sin = repeat(sin, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)") - return torch.cat( - [x[..., :ro_dim] * cos + rotate_half(x[..., :ro_dim], interleaved) * sin, x[..., ro_dim:]], - dim=-1, - ) - - -class ApplyRotaryEmb(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x, - cos, - sin, - interleaved=False, - inplace=False, - seqlen_offsets: Union[int, Tensor] = 0, - cu_seqlens: Optional[Tensor] = None, - max_seqlen: Optional[int] = None, - ): - out = apply_rotary( - x, - cos, - sin, - seqlen_offsets=seqlen_offsets, - cu_seqlens=cu_seqlens, - max_seqlen=max_seqlen, - interleaved=interleaved, - inplace=inplace, - ) - if isinstance(seqlen_offsets, int): - ctx.save_for_backward(cos, sin, cu_seqlens) # Can't save int with save_for_backward - ctx.seqlen_offsets = seqlen_offsets - else: - ctx.save_for_backward(cos, sin, cu_seqlens, seqlen_offsets) - ctx.seqlen_offsets = None - ctx.interleaved = interleaved - ctx.inplace = inplace - ctx.max_seqlen = max_seqlen - return out if not inplace else x - - @staticmethod - def backward(ctx, do): - seqlen_offsets = ctx.seqlen_offsets - if seqlen_offsets is None: - cos, sin, cu_seqlens, seqlen_offsets = ctx.saved_tensors - else: - cos, sin, cu_seqlens = ctx.saved_tensors - dx = apply_rotary( - do, - cos, - sin, - seqlen_offsets=seqlen_offsets, - cu_seqlens=cu_seqlens, - max_seqlen=ctx.max_seqlen, - interleaved=ctx.interleaved, - inplace=ctx.inplace, - conjugate=True, - ) - return dx, None, None, None, None, None, None, None - - -def apply_rotary_emb( - x, - cos, - sin, - interleaved=False, - inplace=False, - seqlen_offsets: Union[int, Tensor] = 0, - cu_seqlens: Optional[Tensor] = None, - max_seqlen: Optional[int] = None, -): - """ - Arguments: - x: (batch_size, seqlen, nheads, headdim) if cu_seqlens is None - else (total_seqlen, nheads, headdim) - cos, sin: (seqlen_rotary, rotary_dim / 2) - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead - of 1st half and 2nd half (GPT-NeoX style). - inplace: if True, apply rotary embedding in-place. - seqlen_offsets: (batch_size,) or int. Each sequence in x is shifted by this amount. - Most commonly used in inference when we have KV cache. - cu_seqlens: (batch + 1,) or None - max_seqlen: int - Return: - out: (batch_size, seqlen, nheads, headdim) if cu_seqlens is None - else (total_seqlen, nheads, headdim) - rotary_dim must be <= headdim - Apply rotary embedding to the first rotary_dim of x. - """ - return ApplyRotaryEmb.apply( - x, cos, sin, interleaved, inplace, seqlen_offsets, cu_seqlens, max_seqlen - ) - - -# For backward compatibility -apply_rotary_emb_func = apply_rotary_emb - - -def _apply_rotary_emb_qkv( - qkv, - cos, - sin, - cos_k=None, - sin_k=None, - interleaved=False, - inplace=False, - conjugate=False, - seqlen_offsets: Union[int, Tensor] = 0, - num_heads_q: Optional[int] = None, -): - apply_rotary_fn = partial( - apply_rotary, - interleaved=interleaved, - inplace=inplace, - conjugate=conjugate, - seqlen_offsets=seqlen_offsets - ) - if cos_k is None and sin_k is None and qkv.is_contiguous(): - # Call 1 kernel instead of 2 kernels - # We need qkv to be contiguous so that when we reshape to combine (3, nheads) - # dimensions, we get the same tensor - if qkv.dim() == 5: - batch, seqlen, three, nheads, headdim = qkv.shape - assert three == 3 - # qk = rearrange(qkv[:, :, :2], "b s t h d -> b s (t h) d") - qk = qkv[:, :, :2].reshape(batch, seqlen, -1, headdim) - qk = apply_rotary_fn(qk, cos, sin) - else: - assert qkv.dim() == 4 - assert num_heads_q is not None - num_heads_k = (qkv.shape[2] - num_heads_q) // 2 - assert qkv.shape[2] == num_heads_q + 2 * num_heads_k - qk = qkv[:, :, :num_heads_q + num_heads_k] - qk = apply_rotary_fn(qk, cos, sin) - if not inplace: - if qkv.dim() == 5: - qkv = torch.cat([rearrange(qk, "b s (t h) d -> b s t h d", t=2), qkv[:, :, 2:]], dim=2) - else: - qkv = torch.cat([qk, qkv[:, :, num_heads_q + num_heads_k :]], dim=2) - else: - cos_k = cos if cos_k is None else cos_k - sin_k = sin if sin_k is None else sin_k - if qkv.dim() == 5: - batch, seqlen, three, nheads, headdim = qkv.shape - assert three == 3 - q, k = qkv[:, :, 0], qkv[:, :, 1] - else: - assert qkv.dim() == 4 - assert num_heads_q is not None - num_heads_k = (qkv.shape[2] - num_heads_q) // 2 - assert qkv.shape[2] == num_heads_q + 2 * num_heads_k - q, k = qkv[:, :, :num_heads_q], qkv[:, :, num_heads_q : num_heads_q + num_heads_k] - q = apply_rotary_fn(q, cos, sin) - k = apply_rotary_fn(k, cos_k, sin_k) - if not inplace: - if qkv.dim() == 5: - qkv = torch.stack([q, k, qkv[:, :, 2]], dim=2) - else: - qkv = torch.cat([q, k, qkv[:, :, num_heads_q + num_heads_k:]], dim=2) - return qkv - - -class ApplyRotaryEmbQKV_(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qkv, - cos, - sin, - cos_k=None, - sin_k=None, - interleaved=False, - seqlen_offsets: Union[int, torch.Tensor] = 0, - num_heads_q: Optional[int] = None, - ): - # apply_rotary_emb_qkv_inplace( - qkv = _apply_rotary_emb_qkv( - qkv, cos, sin, cos_k, sin_k, interleaved=interleaved, inplace=True, - seqlen_offsets=seqlen_offsets, num_heads_q=num_heads_q, - ) - if isinstance(seqlen_offsets, int): - ctx.save_for_backward(cos, sin, cos_k, sin_k) - ctx.seqlen_offsets = seqlen_offsets - else: - ctx.save_for_backward(cos, sin, cos_k, sin_k, seqlen_offsets) - ctx.seqlen_offsets = None - ctx.interleaved = interleaved - ctx.num_heads_q = num_heads_q - return qkv - - @staticmethod - def backward(ctx, dqkv): - seqlen_offsets = ctx.seqlen_offsets - if seqlen_offsets is None: - cos, sin, cos_k, sin_k, seqlen_offsets = ctx.saved_tensors - else: - cos, sin, cos_k, sin_k = ctx.saved_tensors - dqkv = _apply_rotary_emb_qkv( - dqkv, cos, sin, cos_k, sin_k, interleaved=ctx.interleaved, inplace=True, - seqlen_offsets=seqlen_offsets, num_heads_q=ctx.num_heads_q, conjugate=True, - ) - return dqkv, None, None, None, None, None, None, None - - -def apply_rotary_emb_qkv_( - qkv, - cos, - sin, - cos_k=None, - sin_k=None, - interleaved=False, - seqlen_offsets: Union[int, torch.Tensor] = 0, - num_heads_q: Optional[int] = None, -): - """ - Arguments: - qkv: (batch_size, seqlen, 3, nheads, headdim) or (batch_size, seqlen, num_heads_q + 2 * num_heads_k, headdim). - If qkv has shape (batch_size, seqlen, num_heads_q + 2 * num_heads_k, headdim) (e.g. MQA / GQA), - then num_heads_q must be provided. - cos, sin: (seqlen, rotary_dim / 2) - cos_k, sin_k: (seqlen, rotary_dim / 2), optional - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead of - 1st half and 2nd half (GPT-NeoX style). - seqlen_offsets: (batch_size,) or int. Each sequence in Q and K is shifted by this amount. - Most commonly used in inference when we have KV cache. - Return: - qkv: (batch_size, seqlen, 3, nheads, headdim) or (batch_size, seqlen, num_heads_q + 2 * num_heads_k, headdim) - rotary_dim must be <= headdim - Apply rotary embedding *inplace* to the first rotary_dim of Q and K. - """ - return ApplyRotaryEmbQKV_.apply( - qkv, cos, sin, cos_k, sin_k, interleaved, seqlen_offsets, num_heads_q - ) - - -class ApplyRotaryEmbKV_(torch.autograd.Function): - - @staticmethod - def forward(ctx, kv, cos, sin, interleaved=False, seqlen_offsets: Union[int, torch.Tensor] = 0): - batch, seqlen, two, nheads, headdim = kv.shape - assert two == 2 - k = kv[:, :, 0] - apply_rotary( - k, cos, sin, seqlen_offsets=seqlen_offsets, interleaved=interleaved, inplace=True - ) - if isinstance(seqlen_offsets, int): - ctx.save_for_backward(cos, sin) # Can't save int with save_for_backward - ctx.seqlen_offsets = seqlen_offsets - else: - ctx.save_for_backward(cos, sin, seqlen_offsets) - ctx.seqlen_offsets = None - ctx.interleaved = interleaved - return kv - - @staticmethod - def backward(ctx, dkv): - seqlen_offsets = ctx.seqlen_offsets - if seqlen_offsets is None: - cos, sin, seqlen_offsets = ctx.saved_tensors - else: - cos, sin = ctx.saved_tensors - apply_rotary( - dkv[:, :, 0], - cos, - sin, - seqlen_offsets=seqlen_offsets, - interleaved=ctx.interleaved, - inplace=True, - conjugate=True, - ) - return dkv, None, None, None, None - - -apply_rotary_emb_kv_ = ApplyRotaryEmbKV_.apply - - -def apply_rotary_emb_kv_( - kv, - cos, - sin, - interleaved=False, - seqlen_offsets: Union[int, torch.Tensor] = 0, -): - """ - Arguments: - kv: (batch_size, seqlen, 2, nheads, headdim) - cos, sin: (seqlen, rotary_dim / 2) - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead of - 1st half and 2nd half (GPT-NeoX style). - seqlen_offsets: (batch_size,) or int. Each sequence in Q and K is shifted by this amount. - Most commonly used in inference when we have KV cache. - Return: - kv: (batch_size, seqlen, 2, nheads, headdim) - rotary_dim must be <= headdim - Apply rotary embedding *inplace* to the first rotary_dim of K. - """ - return ApplyRotaryEmbKV_.apply(kv, cos, sin, interleaved, seqlen_offsets) - - -class RotaryEmbedding(torch.nn.Module): - """ - The rotary position embeddings from RoFormer_ (Su et. al). - A crucial insight from the method is that the query and keys are - transformed by rotation matrices which depend on the relative positions. - - Other implementations are available in the Rotary Transformer repo_ and in - GPT-NeoX_, GPT-NeoX was an inspiration - - .. _RoFormer: https://arxiv.org/abs/2104.09864 - .. _repo: https://github.com/ZhuiyiTechnology/roformer - .. _GPT-NeoX: https://github.com/EleutherAI/gpt-neox - - If scale_base is not None, this implements XPos (Sun et al., https://arxiv.org/abs/2212.10554). - A recommended value for scale_base is 512: https://github.com/HazyResearch/flash-attention/issues/96 - Reference: https://github.com/sunyt32/torchscale/blob/main/torchscale/component/xpos_relative_position.py - """ - - def __init__( - self, - dim: int, - base=10000.0, - interleaved=False, - scale_base=None, - device=None, - ): - """ - interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead - of 1st half and 2nd half (GPT-NeoX style). - """ - super().__init__() - self.dim = dim - self.base = float(base) - # Generate and save the inverse frequency buffer (non trainable) - inv_freq = self._compute_inv_freq(device) - self.register_buffer("inv_freq", inv_freq, persistent=False) - self.interleaved = interleaved - self.scale_base = scale_base - scale = ( - (torch.arange(0, dim, 2, device=device, dtype=torch.float32) + 0.4 * dim) / (1.4 * dim) - if scale_base is not None - else None - ) - self.register_buffer("scale", scale, persistent=False) - - self._seq_len_cached = 0 - self._cos_cached = None - self._sin_cached = None - self._cos_k_cached = None - self._sin_k_cached = None - - def _compute_inv_freq(self, device=None): - return 1.0 / ( - self.base - ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim) - ) - - def _update_cos_sin_cache(self, seqlen, device=None, dtype=None): - # Reset the tables if the sequence length has changed, - # if we're on a new device (possibly due to tracing for instance), - # or if we're switching from inference mode to training - if ( - seqlen > self._seq_len_cached - or self._cos_cached is None - or self._cos_cached.device != device - or self._cos_cached.dtype != dtype - or (self.training and self._cos_cached.is_inference()) - ): - self._seq_len_cached = seqlen - # We want fp32 here, not self.inv_freq.dtype, since the model could be loaded in bf16 - # And the output of arange can be quite large, so bf16 would lose a lot of precision. - t = torch.arange(seqlen, device=device, dtype=torch.float32) - # We want fp32 here as well since inv_freq will be multiplied with t, and the output - # will be large. Having it in bf16 will lose a lot of precision and cause the - # cos & sin output to change significantly. - # We want to recompute self.inv_freq if it was not loaded in fp32 - if self.inv_freq.dtype != torch.float32: - inv_freq = self._compute_inv_freq(device=device) - else: - inv_freq = self.inv_freq - # Don't do einsum, it converts fp32 to bf16 under AMP - # freqs = torch.einsum("i,j->ij", t, self.inv_freq) - freqs = torch.outer(t, inv_freq) - if self.scale is None: - self._cos_cached = torch.cos(freqs).to(dtype) - self._sin_cached = torch.sin(freqs).to(dtype) - else: - power = ( - torch.arange(seqlen, dtype=self.scale.dtype, device=self.scale.device) - - seqlen // 2 - ) / self.scale_base - scale = self.scale.to(device=power.device) ** rearrange(power, "s -> s 1") - # We want the multiplication by scale to happen in fp32 - self._cos_cached = (torch.cos(freqs) * scale).to(dtype) - self._sin_cached = (torch.sin(freqs) * scale).to(dtype) - self._cos_k_cached = (torch.cos(freqs) / scale).to(dtype) - self._sin_k_cached = (torch.sin(freqs) / scale).to(dtype) - - def forward( - self, - qkv: torch.Tensor, - kv: Optional[torch.Tensor] = None, - seqlen_offset: Union[int, torch.Tensor] = 0, - max_seqlen: Optional[int] = None, - num_heads_q: Optional[int] = None, - ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: - """ - qkv: (batch, seqlen, 3, nheads, headdim) or (batch, seqlen, num_heads_q + 2 * num_heads_k, headdim) - if kv is none, else it's just q of shape (batch, seqlen, nheads, headdim). - If qkv has shape (batch, seqlen, num_heads_q + 2 * num_heads_k, headdim) (e.g. MQA / GQA), - then num_heads_q must be provided. - kv: (batch, seqlen, 2, nheads, headdim) - seqlen_offset: (batch_size,) or int. Each sequence in x is shifted by this amount. - Most commonly used in inference when we have KV cache. - If it's a tensor of shape (batch_size,), then to update the cos / sin cache, one - should pass in max_seqlen, which will update the cos / sin cache up to that length. - Apply rotary embedding *inplace* to qkv and / or kv. - """ - seqlen = qkv.shape[1] - if max_seqlen is not None: - self._update_cos_sin_cache(max_seqlen, device=qkv.device, dtype=qkv.dtype) - elif isinstance(seqlen_offset, int): - self._update_cos_sin_cache(seqlen + seqlen_offset, device=qkv.device, dtype=qkv.dtype) - if kv is None: - return apply_rotary_emb_qkv_( - qkv, - self._cos_cached, - self._sin_cached, - self._cos_k_cached if self.scale is not None else None, - self._sin_k_cached if self.scale is not None else None, - interleaved=self.interleaved, - seqlen_offsets=seqlen_offset, - num_heads_q=num_heads_q, - ) - else: - q = qkv - q = apply_rotary_emb_func( - q, - self._cos_cached, - self._sin_cached, - interleaved=self.interleaved, - inplace=True, - seqlen_offsets=seqlen_offset, - ) - kv = apply_rotary_emb_kv_( - kv, - self._cos_cached if self.scale is None else self._cos_k_cached, - self._sin_cached if self.scale is None else self._sin_k_cached, - interleaved=self.interleaved, - seqlen_offsets=seqlen_offset, - ) - return q, kv diff --git a/build/torch29-cxx11-cu130-aarch64-linux/metadata.json b/build/torch29-cxx11-cu130-aarch64-linux/metadata.json deleted file mode 100644 index 8b00d7fb7685121edf7a97bb007415122444f1d5..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-aarch64-linux/metadata.json +++ /dev/null @@ -1,14 +0,0 @@ -{ - "version": 1, - "license": "BSD-3-Clause", - "python-depends": [], - "backend": { - "type": "cuda", - "archs": [ - "10.0", - "12.0", - "8.0", - "9.0" - ] - } -} diff --git a/build/torch29-cxx11-cu130-aarch64-linux/ops/__init__.py b/build/torch29-cxx11-cu130-aarch64-linux/ops/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/build/torch29-cxx11-cu130-aarch64-linux/ops/activations.py b/build/torch29-cxx11-cu130-aarch64-linux/ops/activations.py deleted file mode 100644 index 7c09649fc41e12d5a360c5672825d8380bc7ec80..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-aarch64-linux/ops/activations.py +++ /dev/null @@ -1,135 +0,0 @@ -# Copied from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/model/layers/activations.py -import math - -import torch -import torch.nn as nn -import torch.nn.functional as F - -# 1/sqrt(2*pi)-> 0.3989423 -# 1/sqrt(2) -> 0.70710678 -# sqrt(2/pi) -> 0.79788456 - -# this function is tanh approximation of gelu -# actual gelu is: -# x * 0.5 * (1.0 + torch.erf(x * 0.70710678)) -@torch.jit.script -def bias_gelu(y, bias): - x = bias + y - return (x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))).to(dtype=y.dtype) - - -# gradient of tanh approximation of gelu -# gradient of actual gelu is: -# 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x) -@torch.jit.script -def bias_gelu_back(g, y, bias): - """Assume that y has shape (B, D) and bias has shape (D)""" - x = bias + y - tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)) - # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243 - ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( - 1 + tanh_out - ) - grad_y = ff * g - return grad_y.to(dtype=y.dtype), grad_y.sum(dim=(0), dtype=bias.dtype) - - -class GeLUFunction(torch.autograd.Function): - @staticmethod - # bias is an optional argument - def forward(ctx, input, bias): - ctx.save_for_backward(input, bias) - return bias_gelu(input, bias) - - @staticmethod - def backward(ctx, grad_output): - input, bias = ctx.saved_tensors - tmp = bias_gelu_back(grad_output, input, bias) - return tmp, tmp - - -bias_gelu_impl = GeLUFunction.apply - -# this function is tanh approximation of gelu -# actual gelu is: -# x * 0.5 * (1.0 + torch.erf(x * 0.70710678)) -@torch.jit.script -def gelu_fwd(x): - return (x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))).to(dtype=x.dtype) - - -# gradient of tanh approximation of gelu -# gradient of actual gelu is: -# 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x) -@torch.jit.script -def gelu_bwd(g, x): - tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)) - # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243 - ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( - 1 + tanh_out - ) - return (ff * g).to(dtype=x.dtype) - - -class FastGeLUFunction(torch.autograd.Function): - @staticmethod - # bias is an optional argument - def forward(ctx, input): - ctx.save_for_backward(input) - return gelu_fwd(input) - - @staticmethod - def backward(ctx, grad_output): - (input,) = ctx.saved_tensors - tmp = gelu_bwd(grad_output, input) - return tmp - - -fast_gelu_impl = FastGeLUFunction.apply - - -@torch.jit.script -def relu_bwd(g, x): - return torch.where(x >= 0, g, 0.0).to(dtype=x.dtype) - - -@torch.jit.script -def sqrelu_fwd(x): - r = F.relu(x) - return (r * r).to(dtype=x.dtype) - - -@torch.jit.script -def sqrelu_bwd(g, x): - return (2.0 * g * F.relu(x)).to(dtype=x.dtype) - - -swiglu_fwd_codestring = """ -template T swiglu_fwd(T x, T y) { - return float(x) * float(y) / (1.0f + ::exp(-float(x))); -} -""" -swiglu_bwd_codestring = """ -template void swiglu_bwd(T x, T y, T g, T& dx, T& dy) { - float x_sigmoid = 1.0f / (1.0f + ::exp(-float(x))); - dx = x_sigmoid * (1 + float(x) * (1.0f - x_sigmoid)) * float(g) * float(y); - dy = float(x) * x_sigmoid * float(g); -} -""" -swiglu_fwd = torch.cuda.jiterator._create_jit_fn(swiglu_fwd_codestring) -swiglu_bwd = torch.cuda.jiterator._create_multi_output_jit_fn(swiglu_bwd_codestring, num_outputs=2) - - -class SwiGLUFunction(torch.autograd.Function): - - @staticmethod - def forward(ctx, x, y): - ctx.save_for_backward(x, y) - return swiglu_fwd(x, y) - - @staticmethod - def backward(ctx, dout): - x, y = ctx.saved_tensors - return swiglu_bwd(x, y, dout) - -swiglu = SwiGLUFunction.apply diff --git a/build/torch29-cxx11-cu130-aarch64-linux/ops/fused_dense.py b/build/torch29-cxx11-cu130-aarch64-linux/ops/fused_dense.py deleted file mode 100644 index 6b4033d134e4093fe278f7b3f8c7d3128ce9f36d..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-aarch64-linux/ops/fused_dense.py +++ /dev/null @@ -1,688 +0,0 @@ -# Copyright (c) 2023, Tri Dao. -# Inspired by https://github.com/NVIDIA/apex/blob/master/apex/fused_dense/fused_dense.py -# We make it work with pytorch amp and with bfloat16. -# The TensorParallel linear modules are inspired by https://github.com/NVIDIA/apex/blob/master/apex/transformer/tensor_parallel/layers.py -from functools import partial -from typing import Optional - -# import fused_dense_cuda # from apex -import fused_dense_lib as fused_dense_cuda -import torch -import torch.nn as nn -import torch.nn.functional as F -from torch import Tensor -from torch.distributed import ProcessGroup - -from flash_attn.utils.torch import custom_fwd, custom_bwd -from flash_attn.ops.activations import gelu_bwd, relu_bwd, sqrelu_bwd, sqrelu_fwd -from flash_attn.utils.distributed import ( - all_gather_raw, - all_reduce, - all_reduce_raw, - reduce_scatter, - reduce_scatter_raw, -) - - -class FusedDenseFunc(torch.autograd.Function): - @staticmethod - @custom_fwd - def forward( - ctx, x, weight, bias, return_residual=False, process_group=None, sequence_parallel=True - ): - """ - If process_group is not None and sequence_parallel=True, we're doing Tensor Parallel - with sequence parallelism: we do an all_gather_raw of x before doing the matmul. - """ - ctx.compute_weight_gradient = weight.requires_grad - ctx.return_residual = return_residual - ctx.process_group = process_group - ctx.sequence_parallel = sequence_parallel - - if torch.is_autocast_enabled(): - x = x.to(dtype=torch.get_autocast_gpu_dtype()) - x = x.contiguous() - if process_group is not None and sequence_parallel: - # We want to kick off the all_gather early, before weight dtype conversion - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - else: - total_x = x - - if torch.is_autocast_enabled(): - weight = weight.to(dtype=torch.get_autocast_gpu_dtype()) - bias = bias.to(dtype=torch.get_autocast_gpu_dtype()) if bias is not None else None - weight = weight.contiguous() - if process_group is not None and sequence_parallel: - handle_x.wait() - batch_shape, n = total_x.shape[:-1], total_x.shape[-1] - batch_dim = batch_shape.numel() - # https://github.com/pytorch/pytorch/blob/5b51849b48a7dbccd297286cc0110def4706f9e7/aten/src/ATen/native/cuda/Blas.cpp#L174 - if min(batch_dim, n, *weight.shape) > 65535 * 32: - raise RuntimeError("fused_dense only supports matrix dims <= 2M") - output = F.linear(total_x, weight, bias) - if ctx.compute_weight_gradient: - ctx.save_for_backward(x, weight) - else: - ctx.save_for_backward(weight) - return output if not return_residual else (output, x) - - @staticmethod - @custom_bwd - def backward(ctx, grad_output, *args): - grad_output = grad_output.contiguous() - if ctx.return_residual: - (grad_input,) = args - grad_input = grad_input.contiguous() - process_group = ctx.process_group - sequence_parallel = ctx.sequence_parallel - if ctx.compute_weight_gradient: - x, weight = ctx.saved_tensors - if process_group is not None and sequence_parallel: - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - else: - total_x = x - else: - (weight,) = ctx.saved_tensors - total_x = None - batch_shape = grad_output.shape[:-1] - batch_dim = batch_shape.numel() - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - if ctx.needs_input_grad[0]: - if not ctx.return_residual: - grad_input = F.linear(grad_output, weight.t()) - else: - grad_input = torch.addmm( - grad_input.reshape(batch_dim, grad_input.shape[-1]), grad_output, weight - ) - grad_input = grad_input.reshape(*batch_shape, grad_input.shape[-1]) - if process_group is not None: - reduce_fn = reduce_scatter_raw if sequence_parallel else all_reduce_raw - grad_input, handle_grad_input = reduce_fn(grad_input, process_group, async_op=True) - else: - grad_input = None - if ctx.needs_input_grad[1]: - assert ctx.compute_weight_gradient - if process_group is not None and sequence_parallel: - handle_x.wait() - grad_weight, grad_bias = fused_dense_cuda.linear_bias_wgrad( - total_x.reshape(batch_dim, total_x.shape[-1]), grad_output, ctx.needs_input_grad[2] - ) - else: - grad_weight = None - grad_bias = grad_output if ctx.needs_input_grad[2] else None - if process_group is not None and ctx.needs_input_grad[0]: - handle_grad_input.wait() - return grad_input, grad_weight, grad_bias, None, None, None - - -def fused_dense_func( - x: Tensor, - weight: Tensor, - bias: Optional[Tensor] = None, - return_residual: bool = False, - process_group: Optional[ProcessGroup] = None, - sequence_parallel: bool = True, -): - dtype_eligible = x.dtype in [torch.float16, torch.bfloat16] or ( - x.dtype == torch.float32 and torch.is_autocast_enabled() - ) - if x.is_cuda and weight.is_cuda and (bias is None or bias.is_cuda) and dtype_eligible: - return FusedDenseFunc.apply( - x, weight, bias, return_residual, process_group, sequence_parallel - ) - else: - assert process_group is None - out = F.linear(x, weight, bias) - return out if not return_residual else (out, x) - - -class FusedDense(nn.Linear): - def __init__( - self, - in_features: int, - out_features: int, - bias: bool = True, - return_residual: bool = False, - device=None, - dtype=None, - ) -> None: - super().__init__(in_features, out_features, bias=bias, device=device, dtype=dtype) - self.return_residual = return_residual - - def forward(self, x, process_group=None): - """ - If process_group is not None, we're doing Tensor Parallel with sequence parallelism: - we do an all_gather of x before doing the matmul. - """ - return fused_dense_func( - x, - self.weight, - self.bias, - return_residual=self.return_residual, - process_group=process_group, - ) - - -class ColumnParallelLinear(nn.Linear): - def __init__( - self, - in_features: int, - out_features: int, - process_group: ProcessGroup, - bias: bool = True, - sequence_parallel=True, - multiple_of=1, - device=None, - dtype=None, - ) -> None: - world_size = torch.distributed.get_world_size(process_group) - if out_features % multiple_of: - raise ValueError(f"out_features ({out_features}) must be a multiple of {multiple_of}") - multiple = out_features // multiple_of - # We want to split @multiple across world_size, but it could be an uneven split - div = multiple // world_size - mod = multiple % world_size - # The first @mod ranks get @div + 1 copies, the rest get @div copies - local_multiple = div + int(torch.distributed.get_rank(process_group) < mod) - super().__init__( - in_features, local_multiple * multiple_of, bias=bias, device=device, dtype=dtype - ) - self.process_group = process_group - self.sequence_parallel = sequence_parallel - - def forward(self, x): - # If self.sequence_parallel is True, we're doing Tensor Parallel with sequence parallelism: - # we do an all_gather of x before doing the matmul. - # If not, then the input is already gathered. - return fused_dense_func( - x, - self.weight, - self.bias, - process_group=self.process_group, - sequence_parallel=self.sequence_parallel, - ) - - -class RowParallelLinear(nn.Linear): - def __init__( - self, - in_features: int, - out_features: int, - process_group: ProcessGroup, - bias: bool = True, - sequence_parallel=True, - multiple_of=1, - device=None, - dtype=None, - ) -> None: - world_size = torch.distributed.get_world_size(process_group) - rank = torch.distributed.get_rank(process_group) - if in_features % multiple_of: - raise ValueError(f"in_features ({in_features}) must be a multiple of {multiple_of}") - multiple = in_features // multiple_of - # We want to split @multiple across world_size, but it could be an uneven split - div = multiple // world_size - mod = multiple % world_size - # The first @mod ranks get @div + 1 copies, the rest get @div copies - local_multiple = div + int(torch.distributed.get_rank(process_group) < mod) - # Only rank 0 will have bias - super().__init__( - local_multiple * multiple_of, - out_features, - bias=bias and rank == 0, - device=device, - dtype=dtype, - ) - self.process_group = process_group - self.sequence_parallel = sequence_parallel - - def forward(self, x): - """ - We're doing Tensor Parallel with sequence parallelism: we do the matmul and then - a reduce_scatter of the result. - """ - out = fused_dense_func(x, self.weight, self.bias) - reduce_fn = reduce_scatter if self.sequence_parallel else all_reduce - return reduce_fn(out, self.process_group) - - -class FusedMLPFunc(torch.autograd.Function): - @staticmethod - @custom_fwd - def forward( - ctx, - x, - weight1, - bias1, - weight2, - bias2, - activation="gelu_approx", - save_pre_act=True, - return_residual=False, - checkpoint_lvl=0, - heuristic=0, - process_group=None, - sequence_parallel=True, - ): - """ - If process_group is not None and sequence_parallel=True, we're doing Tensor Parallel - with sequence parallelism: we do an all_gather of x before doing the matmul. - If sequence_parallel=False, then the input is already gathered. - - checkpoint_lvl: - 0: no recomputation in the bwd - 1: recompute gelu_out / relu_out in the bwd - 2: recompute pre_act and gelu_out / relu_out in the bwd - """ - assert -1 <= heuristic <= 4 - assert activation in ["gelu_approx", "relu", "sqrelu"] - if activation == "sqrelu": - assert heuristic == -1 - if not save_pre_act: - checkpoint_lvl = 2 - assert checkpoint_lvl in [0, 1, 2] - ctx.return_residual = return_residual - ctx.process_group = process_group - ctx.sequence_parallel = sequence_parallel - ctx.checkpoint_lvl = checkpoint_lvl - ctx.activation = activation - ctx.heuristic = heuristic - - if torch.is_autocast_enabled(): - x = x.to(dtype=torch.get_autocast_gpu_dtype()) - x = x.contiguous() - if process_group is not None and sequence_parallel: - # We want to kick off the all_gather early, before weight dtype conversion - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - else: - total_x = x - - if torch.is_autocast_enabled(): - dtype = torch.get_autocast_gpu_dtype() - weight1, weight2 = [a.to(dtype=dtype) for a in [weight1, weight2]] - bias1 = bias1.to(dtype=dtype) if bias1 is not None else None - bias2 = bias2.to(dtype=dtype) if bias2 is not None else None - weight1 = weight1.contiguous() - bias1 = bias1.contiguous() if bias1 is not None else None - weight2 = weight2.contiguous() - bias2 = bias2.contiguous() if bias2 is not None else None - if process_group is not None and sequence_parallel: - handle_x.wait() - batch_shape, n = total_x.shape[:-1], total_x.shape[-1] - batch_dim = batch_shape.numel() - # https://github.com/pytorch/pytorch/blob/5b51849b48a7dbccd297286cc0110def4706f9e7/aten/src/ATen/native/cuda/Blas.cpp#L174 - if min(batch_dim, n, *weight1.shape, *weight2.shape) > 65535 * 32: - raise RuntimeError("fused_dense only supports matrix dims <= 2M") - if heuristic == -1: - pre_act = F.linear(total_x, weight1, bias1) - activation_fn = ( - partial(F.gelu, approximate="tanh") - if activation == "gelu_approx" - else (sqrelu_fwd if activation == "sqrelu" else F.relu) - ) - with torch.jit.fuser("fuser2"): - output1 = activation_fn(pre_act) - # This is before adding bias1 - # pre_act = F.linear(total_x.reshape(batch_dim, n), weight1) - # with torch.jit.fuser('fuser2'): - # output1 = bias_gelu(pre_act, bias1) - else: - is_gelu = activation == "gelu_approx" - output1, *rest = fused_dense_cuda.linear_act_forward( - total_x.reshape(batch_dim, n), weight1, bias1, is_gelu, save_pre_act, heuristic - ) - if save_pre_act: - pre_act = rest[0] - output2 = F.linear(output1, weight2, bias2) - if checkpoint_lvl == 0 or (checkpoint_lvl == 1 and activation == "relu"): - # For RELU the pre_act is very small (just a bit-mask) so we just save it - ctx.save_for_backward(x, weight1, weight2, pre_act, output1) - elif checkpoint_lvl == 1: - ctx.save_for_backward(x, weight1, weight2, pre_act) - elif checkpoint_lvl == 2: - ctx.save_for_backward(x, weight1, weight2, bias1) - output2 = output2.reshape(*batch_shape, output2.shape[-1]) - return output2 if not return_residual else (output2, x) - - @staticmethod - @custom_bwd - def backward(ctx, grad_output, *args): - grad_output = grad_output.contiguous() - checkpoint_lvl = ctx.checkpoint_lvl - activation = ctx.activation - activation_fn = ( - partial(F.gelu, approximate="tanh") - if activation == "gelu_approx" - else (sqrelu_fwd if activation == "sqrelu" else F.relu) - ) - if ctx.return_residual: - (grad_input,) = args - grad_input = grad_input.contiguous() - process_group = ctx.process_group - sequence_parallel = ctx.sequence_parallel - x, weight1, weight2, *rest = ctx.saved_tensors - if process_group is None or not sequence_parallel: - total_x = x - batch_shape = grad_output.shape[:-1] - batch_dim = batch_shape.numel() - if checkpoint_lvl in [0, 1]: - if process_group is not None and sequence_parallel: - total_x, handle_x = all_gather_raw(x, process_group, async_op=True) - if checkpoint_lvl == 0 or (checkpoint_lvl == 1 and activation == "relu"): - pre_act, output1 = rest - elif checkpoint_lvl == 1: - (pre_act,) = rest - with torch.jit.fuser("fuser2"): - output1 = activation_fn(pre_act) - elif checkpoint_lvl == 2: - (bias1,) = rest - if process_group is not None and sequence_parallel: - total_x, _ = all_gather_raw(x, process_group) - if ctx.heuristic == -1: - pre_act = F.linear(total_x, weight1, bias1) - with torch.jit.fuser("fuser2"): - output1 = activation_fn(pre_act) - else: - output1, pre_act = fused_dense_cuda.linear_act_forward( - total_x.reshape(batch_dim, total_x.shape[-1]), - weight1, - bias1, - activation == "gelu_approx", - True, - ctx.heuristic, - ) - - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - output1 = output1.reshape(batch_dim, output1.shape[-1]) - pre_act = pre_act.reshape(batch_dim, pre_act.shape[-1]) - if ctx.needs_input_grad[3]: - grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad( - output1, grad_output, ctx.needs_input_grad[4] - ) - else: - grad_weight2 = None - grad_bias2 = grad_output if ctx.needs_input_grad[4] else None - if ctx.heuristic == -1: - # grad_pre_act = matmul_dgelu(grad_output, weight2, pre_act) - grad_output1 = F.linear(grad_output, weight2.t()) - activation_grad_fn = ( - gelu_bwd - if activation == "gelu_approx" - else (sqrelu_bwd if activation == "sqrelu" else relu_bwd) - ) - with torch.jit.fuser("fuser2"): - grad_pre_act = activation_grad_fn(grad_output1, pre_act) - else: - # The cublasLt epilogue has to compute both gelu/relu grad and bias grad, we can't - # just compute gelu/relu grad - grad_pre_act, grad_bias1 = fused_dense_cuda.bias_act_linear_dgrad_bgrad( - weight2, grad_output, pre_act, activation == "gelu_approx", ctx.heuristic - ) - if not ctx.needs_input_grad[2]: - grad_bias1 = None - if ctx.needs_input_grad[0]: - if not ctx.return_residual: - grad_input = F.linear(grad_pre_act, weight1.t()) - else: - grad_input = torch.addmm( - grad_input.reshape(batch_dim, grad_input.shape[-1]), grad_pre_act, weight1 - ) - grad_input = grad_input.reshape(*batch_shape, grad_input.shape[-1]) - if process_group is not None: - reduce_fn = reduce_scatter_raw if sequence_parallel else all_reduce_raw - grad_input, handle_grad_input = reduce_fn(grad_input, process_group, async_op=True) - else: - grad_input = None - if ctx.heuristic == -1: - if ctx.needs_input_grad[1]: - if process_group is not None and sequence_parallel and checkpoint_lvl != 2: - handle_x.wait() - grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_wgrad( - total_x.reshape(batch_dim, total_x.shape[-1]), - grad_pre_act, - ctx.needs_input_grad[2], - ) - else: - grad_weight1 = None - grad_bias1 = grad_pre_act if ctx.needs_input_grad[2] else None - else: - if ctx.needs_input_grad[1]: - if process_group is not None and sequence_parallel and checkpoint_lvl != 2: - handle_x.wait() - grad_weight1 = F.linear( - grad_pre_act.t(), total_x.reshape(batch_dim, total_x.shape[-1]).t() - ) - else: - grad_weight1 = None - if process_group is not None and ctx.needs_input_grad[0]: - handle_grad_input.wait() - return ( - grad_input, - grad_weight1, - grad_bias1, - grad_weight2, - grad_bias2, - None, - None, - None, - None, - None, - None, - None, - ) - - -def fused_mlp_func( - x: Tensor, - weight1: Tensor, - weight2: Tensor, - bias1: Optional[Tensor] = None, - bias2: Optional[Tensor] = None, - activation: str = "gelu_approx", - save_pre_act: bool = True, - return_residual: bool = False, - checkpoint_lvl: int = 0, - heuristic: int = 0, - process_group: Optional[ProcessGroup] = None, - sequence_parallel: bool = True, -): - assert activation in ["gelu_approx", "relu", "sqrelu"] - dtype_eligible = x.dtype in [torch.float16, torch.bfloat16] or ( - x.dtype == torch.float32 and torch.is_autocast_enabled() - ) - # If we save pre-activation, dimension must be divisible by 128 (relu) or 8 (gelu) - dim_eligible = not save_pre_act or (x.shape[-1] % (128 if activation == "relu" else 8) == 0) - if ( - x.is_cuda - and weight1.is_cuda - and weight2.is_cuda - and (bias1 is None or bias1.is_cuda) - and (bias2 is None or bias2.is_cuda) - and dtype_eligible - and dim_eligible - ): - return FusedMLPFunc.apply( - x, - weight1, - bias1, - weight2, - bias2, - activation, - save_pre_act, - return_residual, - checkpoint_lvl, - heuristic, - process_group, - sequence_parallel, - ) - else: - assert process_group is None - pre_act = F.linear(x, weight1, bias1) - activation_fn = ( - partial(F.gelu, approximate="tanh") - if activation == "gelu_approx" - else partial(F.relu, inplace=True) - ) - output1 = activation_fn(pre_act) - output2 = F.linear(output1, weight2, bias2) - return output2 if not return_residual else (output2, x) - - -class FusedMLP(nn.Module): - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - bias1=True, - bias2=True, - activation="gelu_approx", - return_residual=False, - checkpoint_lvl=0, - heuristic="auto", - device=None, - dtype=None, - ): - """ - If process_group is not None, we're doing Tensor Parallel with sequence parallelism: - we do an all_gather of x before doing the matmul, gelu, then matmul. - Finally we do a reduce_scatter of the output. - - checkpoint_lvl (increasing lvl means slower but more memory saving): - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute pre_act and gelu_out in the bwd - heuristic: - -1: don't fuse gemm + gelu (separate kernel) - 0..4: use this heuristic for the algo section in the fused gemm + gelu - 'auto': heuristic will be picked automatically: - For CUDA >= 11.8, we set heuristic=0 for both fp16 and bf16 for best perf. - For CUDA <= 11.7, we set heuristic=1 for fp16 and heuristic=-1 for bf16. - For H100, we set heuristic=-1 for both fp16 and bf16 as the fused cuBlasLt implementation - is slower than the unfused version. - return_residual: whether to return the input x along with the output. This is for - performance reason: for post-norm architecture, returning the input allows us - to fuse the backward of nn.Linear with the residual connection. - """ - assert checkpoint_lvl in [0, 1, 2] - assert activation in ["gelu_approx", "relu", "sqrelu"] - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features * 4 - self.activation = activation - self.return_residual = return_residual - self.checkpoint_lvl = checkpoint_lvl - self.heuristic = heuristic if activation != "sqrelu" else -1 - self.fc1 = nn.Linear(in_features, hidden_features, bias=bias1, **factory_kwargs) - self.fc2 = nn.Linear(hidden_features, out_features, bias=bias2, **factory_kwargs) - - def forward(self, x, process_group=None): - dtype = x.dtype if not torch.is_autocast_enabled() else torch.get_autocast_gpu_dtype() - if self.heuristic == "auto": - if self.activation == "gelu_approx": - if torch.cuda.get_device_capability("cuda") == (9, 0): - heuristic = -1 - else: - cuda_ver = tuple(map(int, torch.version.cuda.split("."))) - heuristic = 0 if cuda_ver >= (11, 8) else (1 if dtype == torch.float16 else -1) - else: - heuristic = 0 - else: - heuristic = self.heuristic - out = fused_mlp_func( - x, - self.fc1.weight, - self.fc2.weight, - self.fc1.bias, - self.fc2.bias, - activation=self.activation, - save_pre_act=self.training, - return_residual=self.return_residual, - checkpoint_lvl=self.checkpoint_lvl, - heuristic=heuristic, - process_group=process_group, - ) - if self.return_residual: - out, x = out - if process_group is not None: - out = reduce_scatter(out, process_group) - return out if not self.return_residual else (out, x) - - -class ParallelFusedMLP(nn.Module): - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - activation="gelu_approx", - process_group: ProcessGroup = None, - bias1=True, - bias2=True, - sequence_parallel=True, - checkpoint_lvl=0, - heuristic="auto", - device=None, - dtype=None, - ): - """ - process_group is required. We're doing Tensor Parallel with sequence parallelism: - we do an all_gather of x before doing the matmul, gelu, then matmul. - Finally we do a reduce_scatter of the output. - - checkpoint_lvl (increasing lvl means slower but more memory saving): - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute pre_act and gelu_out in the bwd - heuristic: - -1: don't fuse gemm + gelu (separate kernel) - 0..4: use this heuristic for the algo section in the fused gemm + gelu - 'auto': heuristic will be picked automatically: - For CUDA >= 11.8, we set heuristic=0 for both fp16 and bf16 for best perf. - For CUDA <= 11.7, we set heuristic=1 for fp16 and heuristic=-1 for bf16. - """ - assert checkpoint_lvl in [0, 1, 2] - assert activation in ["gelu_approx", "relu", "sqrelu"] - assert process_group is not None - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features * 4 - self.activation = activation - self.process_group = process_group - self.sequence_parallel = sequence_parallel - self.checkpoint_lvl = checkpoint_lvl - self.heuristic = heuristic if activation != "sqrelu" else -1 - self.fc1 = ColumnParallelLinear( - in_features, hidden_features, process_group, bias=bias1, **factory_kwargs - ) - self.fc2 = RowParallelLinear( - hidden_features, out_features, process_group, bias=bias2, **factory_kwargs - ) - - def forward(self, x): - dtype = x.dtype if not torch.is_autocast_enabled() else torch.get_autocast_gpu_dtype() - if self.heuristic == "auto": - if self.activation == "gelu_approx": - cuda_ver = tuple(map(int, torch.version.cuda.split("."))) - heuristic = 0 if cuda_ver >= (11, 8) else (1 if dtype == torch.float16 else -1) - else: - heuristic = 0 - else: - heuristic = self.heuristic - out = fused_mlp_func( - x, - self.fc1.weight, - self.fc2.weight, - self.fc1.bias, - self.fc2.bias, - activation=self.activation, - save_pre_act=self.training, - checkpoint_lvl=self.checkpoint_lvl, - heuristic=heuristic, - process_group=self.process_group, - sequence_parallel=self.sequence_parallel, - ) - reduce_fn = reduce_scatter if self.sequence_parallel else all_reduce - return reduce_fn(out, self.process_group) diff --git a/build/torch29-cxx11-cu130-aarch64-linux/ops/layer_norm.py b/build/torch29-cxx11-cu130-aarch64-linux/ops/layer_norm.py deleted file mode 100644 index 4b6cd798fd02844ef9cd3897f8ab95e490e638bf..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-aarch64-linux/ops/layer_norm.py +++ /dev/null @@ -1,800 +0,0 @@ -# Copyright (c) 2022, Tri Dao. -# Adapted from https://github.com/NVIDIA/apex/blob/master/apex/contrib/layer_norm/layer_norm.py - -import dropout_layer_norm -import torch -from torch.nn import init - - -def maybe_align(x, alignment_in_bytes=16): - """Assume that x already has last dim divisible by alignment_in_bytes""" - # TD [2023-07-04] I'm not 100% sure that clone will align the memory - # https://discuss.pytorch.org/t/how-to-ensure-that-tensor-data-ptr-is-aligned-to-16-bytes/183440 - return x if x.data_ptr() % alignment_in_bytes == 0 else x.clone() - - -def _dropout_add_layer_norm_forward( - x0, - residual, - gamma, - beta, - rowscale, - colscale, - dropout_p, - epsilon, - residual_in_fp32=False, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes""" - hidden_size = gamma.numel() - x0mat = x0.view((-1, hidden_size)) - residualmat = residual.view((-1, hidden_size)) if residual is not None else None - rowscale = rowscale.view(-1) if rowscale is not None else None - zmat, xmat, dmask, mu, rsigma = dropout_layer_norm.dropout_add_ln_fwd( - x0mat, - residualmat, - gamma, - beta, - rowscale, - colscale, - None, - None, - dropout_p, - epsilon, - 1.0, - 0, - None, - residual_in_fp32, - is_rms_norm, - ) - # dmask is None if dropout_p == 0.0 - # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype - return zmat, xmat if xmat is not None else x0mat, dmask, mu, rsigma - - -def _dropout_add_layer_norm_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - rowscale, - colscale, - dropout_p, - has_residual, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes - dx == None means that it was a post-norm architecture - (x = drop(x0) + residual was not returned in the fwd). - x0 must not be None if we have colscale. - """ - hidden_size = gamma.numel() - xmat = x.view((-1, hidden_size)) - dzmat = dz.view(xmat.shape) - dxmat = dx.view(xmat.shape) if dx is not None else None - x0mat = x0.view((-1, hidden_size)) if x0 is not None else None - rowscale = rowscale.view(-1) if rowscale is not None else None - if colscale is not None: - assert x0 is not None, "x0 is required to compute the gradient of colscale" - dx0mat, dresidualmat, dgamma, dbeta, _, _, *rest = dropout_layer_norm.dropout_add_ln_bwd( - dzmat, - dxmat, - xmat, - x0mat, - dmask, - mu, - rsigma, - gamma, - rowscale, - colscale, - None, - None, - dropout_p, - 1.0, - 0, - has_residual, - is_rms_norm, - ) - # dresidualmat is None if not has_residual - if colscale is None: - return dx0mat, dresidualmat, dgamma, dbeta - else: - dcolscale = rest[0] - return dx0mat, dresidualmat, dgamma, dbeta, dcolscale - - -def _dropout_add_layer_norm_subset_forward( - x0, - residual, - gamma, - beta, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32=False, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes""" - hidden_size = gamma.numel() - x0mat = x0.view((-1, hidden_size)) - residualmat = residual.view((-1, hidden_size)) if residual is not None else None - x0_subset = x0_subset.view(-1) if x0_subset is not None else None - out_subset = out_subset.view(-1) if out_subset is not None else None - zmat, xmat, dmask, mu, rsigma = dropout_layer_norm.dropout_add_ln_fwd( - x0mat, - residualmat, - gamma, - beta, - None, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - None, - residual_in_fp32, - is_rms_norm, - ) - # dmask is None if dropout_p == 0.0 - # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype - return zmat, xmat if xmat is not None else x0mat, dmask, mu, rsigma - - -def _dropout_add_layer_norm_subset_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - colscale, - x0_subset, - out_subset, - dropout_p, - rowscale_const, - x0_numrows, - has_residual, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes - dx == None means that it was a post-norm architecture - (x = drop(x0) + residual was not returned in the fwd). - x0 must not be None if we have colscale. - """ - hidden_size = gamma.numel() - xmat = x.view((-1, hidden_size)) - dzmat = dz.view(-1, hidden_size) - dxmat = dx.view(xmat.shape) if dx is not None else None - x0mat = x0.view((-1, hidden_size)) if x0 is not None else None - x0_subset = x0_subset.view(-1) if x0_subset is not None else None - out_subset = out_subset.view(-1) if out_subset is not None else None - if colscale is not None: - assert x0 is not None, "x0 is required to compute the gradient of colscale" - dx0mat, dresidualmat, dgamma, dbeta, _, _, *rest = dropout_layer_norm.dropout_add_ln_bwd( - dzmat, - dxmat, - xmat, - x0mat, - dmask, - mu, - rsigma, - gamma, - None, - colscale, - x0_subset, - out_subset, - dropout_p, - rowscale_const, - x0_numrows, - has_residual, - is_rms_norm, - ) - # dresidualmat is None if not has_residual - if colscale is None: - return dx0mat, dresidualmat, dgamma, dbeta - else: - dcolscale = rest[0] - return dx0mat, dresidualmat, dgamma, dbeta, dcolscale - - -def _dropout_add_layer_norm_parallel_residual_forward( - x0, - x1, - residual, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - residual_in_fp32=False, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes""" - hidden_size = gamma0.numel() - x0mat = x0.view((-1, hidden_size)) - x1mat = x1.view((-1, hidden_size)) if x1 is not None else None - residualmat = residual.view((-1, hidden_size)) if residual is not None else None - ( - z0mat, - z1mat, - xmat, - dmask0, - dmask1, - mu, - rsigma, - ) = dropout_layer_norm.dropout_add_ln_parallel_residual_fwd( - x0mat, - x1mat, - residualmat, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - None, - residual_in_fp32, - is_rms_norm, - ) - # dmask0 and dmask1 are None if dropout_p == 0.0 - # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype - return z0mat, z1mat, xmat if xmat is not None else x0mat, dmask0, dmask1, mu, rsigma - - -def _dropout_add_layer_norm_parallel_residual_backward( - dz0, - dz1, - dx, - x, - dmask0, - dmask1, - mu, - rsigma, - gamma0, - gamma1, - dropout_p, - has_x1, - has_residual, - is_rms_norm=False, -): - """Assume that arguments are contiguous and aligned to 16 bytes - dx == None means that it was a post-norm architecture - (x = drop(x0) + residual was not returned in the fwd). - """ - hidden_size = gamma0.numel() - xmat = x.view((-1, hidden_size)) - dz0mat = dz0.view(xmat.shape) - dz1mat = dz1.view(xmat.shape) if dz1 is not None else None - dxmat = dx.view(xmat.shape) if dx is not None else None - ( - dx0mat, - dx1mat, - dresidualmat, - dgamma0, - dbeta0, - dgamma1, - dbeta1, - *rest, - ) = dropout_layer_norm.dropout_add_ln_parallel_residual_bwd( - dz0mat, - dz1mat, - dxmat, - xmat, - dmask0, - dmask1, - mu, - rsigma, - gamma0, - gamma1, - dropout_p, - has_x1, - has_residual, - is_rms_norm, - ) - # dresidualmat is None if not has_residual - return dx0mat, dx1mat, dresidualmat, dgamma0, dbeta0, dgamma1, dbeta1 - - -class DropoutAddLayerNormFn(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x0, - residual, - gamma, - beta, - rowscale, - colscale, - dropout_p, - epsilon, - residual_in_fp32=False, - prenorm=False, - is_rms_norm=False, - return_dmask=False, - ): - x0 = maybe_align(x0.contiguous(), 16) - residual = maybe_align(residual.contiguous(), 16) if residual is not None else None - gamma = maybe_align(gamma.contiguous(), 16) - beta = maybe_align(beta.contiguous(), 16) if beta is not None else None - rowscale = maybe_align(rowscale.contiguous(), 16) if rowscale is not None else None - colscale = maybe_align(colscale.contiguous(), 16) if colscale is not None else None - zmat, xmat, dmask, mu, rsigma = _dropout_add_layer_norm_forward( - x0, - residual, - gamma, - beta, - rowscale, - colscale, - dropout_p, - epsilon, - residual_in_fp32, - is_rms_norm, - ) - # Only need to save x0 if we need to compute gradient wrt colscale - x0_saved = x0 if colscale is not None else None - ctx.save_for_backward( - xmat.view(x0.shape), x0_saved, dmask, gamma, mu, rsigma, rowscale, colscale - ) - ctx.prenorm = prenorm - ctx.dropout_p = dropout_p - ctx.has_residual = residual is not None - ctx.is_rms_norm = is_rms_norm - ctx.has_beta = beta is not None - if not return_dmask: - return ( - zmat.view(x0.shape) if not prenorm else (zmat.view(x0.shape), xmat.view(x0.shape)) - ) - else: - dmask = ( - dmask.view(x0.shape) - if dropout_p > 0.0 - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - ctx.mark_non_differentiable(dmask) - return ( - (zmat.view(x0.shape), dmask) - if not prenorm - else (zmat.view(x0.shape), xmat.view(x0.shape), dmask) - ) - - @staticmethod - def backward(ctx, dz, *args): - # assert dz.is_contiguous() - dz = maybe_align(dz.contiguous(), 16) # this happens! - dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None - x, x0, dmask, gamma, mu, rsigma, rowscale, colscale = ctx.saved_tensors - # x0 is None if colscale is None - dropout_p = ctx.dropout_p - has_residual = ctx.has_residual - dx0mat, dresidualmat, dgamma, dbeta, *rest = _dropout_add_layer_norm_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - rowscale, - colscale, - dropout_p, - has_residual, - ctx.is_rms_norm, - ) - dx0 = dx0mat.view(x.shape) - dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None - dcolscale = rest[0] if colscale is not None else None - return ( - dx0, - dresidual, - dgamma, - dbeta if ctx.has_beta else None, - None, - dcolscale, - None, - None, - None, - None, - None, - None, - ) - - -class DropoutAddLayerNormSubsetFn(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x0, - residual, - gamma, - beta, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32=False, - prenorm=False, - is_rms_norm=False, - return_dmask=False, - ): - x0 = maybe_align(x0.contiguous(), 16) - residual = maybe_align(residual.contiguous(), 16) if residual is not None else None - gamma = maybe_align(gamma.contiguous(), 16) - beta = maybe_align(beta.contiguous(), 16) if beta is not None else None - colscale = maybe_align(colscale.contiguous(), 16) if colscale is not None else None - zmat, xmat, dmask, mu, rsigma = _dropout_add_layer_norm_subset_forward( - x0, - residual, - gamma, - beta, - colscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32, - is_rms_norm, - ) - # Only need to save x0 if we need to compute gradient wrt colscale - x0_saved = x0 if colscale is not None else None - x_shape = (-1, *x0.shape[1:]) - ctx.save_for_backward( - xmat.view(x_shape), x0_saved, dmask, gamma, mu, rsigma, colscale, x0_subset, out_subset - ) - ctx.prenorm = prenorm - ctx.dropout_p = dropout_p - ctx.rowscale_const = rowscale_const - ctx.x0_numrows = x0.shape[:-1].numel() - ctx.has_residual = residual is not None - ctx.is_rms_norm = is_rms_norm - ctx.has_beta = beta is not None - z_shape = (-1, *x0.shape[1:]) - if not return_dmask: - return zmat.view(z_shape) if not prenorm else (zmat.view(z_shape), xmat.view(x0.shape)) - else: - z = zmat.view(z_shape) - dmask = ( - dmask.view(x0.shape) - if dropout_p > 0.0 - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - ctx.mark_non_differentiable(dmask) - return (z, dmask) if not prenorm else (z, xmat.view(x_shape), dmask) - - @staticmethod - def backward(ctx, dz, *args): - # assert dz.is_contiguous() - dz = maybe_align(dz.contiguous(), 16) # this happens! - dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None - x, x0, dmask, gamma, mu, rsigma, colscale, x0_subset, out_subset = ctx.saved_tensors - # x0 is None if colscale is None - dropout_p = ctx.dropout_p - has_residual = ctx.has_residual - dx0mat, dresidualmat, dgamma, dbeta, *rest = _dropout_add_layer_norm_subset_backward( - dz, - dx, - x, - x0, - dmask, - mu, - rsigma, - gamma, - colscale, - x0_subset, - out_subset, - dropout_p, - ctx.rowscale_const, - ctx.x0_numrows, - has_residual, - ctx.is_rms_norm, - ) - dx0 = dx0mat.view(-1, *x.shape[1:]) - dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None - dcolscale = rest[0] if colscale is not None else None - return ( - dx0, - dresidual, - dgamma, - dbeta if ctx.has_beta else None, - dcolscale, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) - - -class DropoutAddLayerNormParallelResidualFn(torch.autograd.Function): - @staticmethod - def forward( - ctx, - x0, - x1, - residual, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - residual_in_fp32=False, - prenorm=False, - is_rms_norm=False, - return_dmask=False, - ): - x0 = maybe_align(x0.contiguous(), 16) - x1 = maybe_align(x1.contiguous(), 16) if x1 is not None else None - residual = maybe_align(residual.contiguous(), 16) if residual is not None else None - gamma0 = maybe_align(gamma0.contiguous(), 16) - beta0 = maybe_align(beta0.contiguous(), 16) if beta0 is not None else None - gamma1 = maybe_align(gamma1.contiguous(), 16) if gamma1 is not None else None - beta1 = maybe_align(beta1.contiguous(), 16) if beta1 is not None else None - ( - z0mat, - z1mat, - xmat, - dmask0, - dmask1, - mu, - rsigma, - ) = _dropout_add_layer_norm_parallel_residual_forward( - x0, - x1, - residual, - gamma0, - beta0, - gamma1, - beta1, - dropout_p, - epsilon, - residual_in_fp32, - is_rms_norm, - ) - ctx.save_for_backward(xmat.view(x0.shape), dmask0, dmask1, gamma0, gamma1, mu, rsigma) - ctx.prenorm = prenorm - ctx.dropout_p = dropout_p - ctx.has_x1 = x1 is not None - ctx.has_residual = residual is not None - ctx.is_rms_norm = is_rms_norm - ctx.has_beta = beta0 is not None - z = (z0mat.view(x0.shape), z1mat.view(x0.shape) if z1mat is not None else None) - if not return_dmask: - return z if not prenorm else (*z, xmat.view(x0.shape)) - else: - dmask0 = ( - dmask0.view(x0.shape) - if dropout_p > 0.0 - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - dmask1 = ( - dmask1.view(x0.shape) - if dropout_p > 0.0 and x1 is not None - else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) - ) - ctx.mark_non_differentiable(dmask0) - ctx.mark_non_differentiable(dmask1) - return ( - (*z, dmask0, dmask1) if not prenorm else (*z, xmat.view(x0.shape), dmask0, dmask1) - ) - - @staticmethod - def backward(ctx, dz0, dz1, *args): - dz0 = maybe_align(dz0.contiguous(), 16) # this happens! - dz1 = maybe_align(dz1.contiguous(), 16) if dz1 is not None else None - dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None - x, dmask0, dmask1, gamma0, gamma1, mu, rsigma = ctx.saved_tensors - dropout_p = ctx.dropout_p - has_x1 = ctx.has_x1 - has_residual = ctx.has_residual - ( - dx0mat, - dx1mat, - dresidualmat, - dgamma0, - dbeta0, - dgamma1, - dbeta1, - ) = _dropout_add_layer_norm_parallel_residual_backward( - dz0, - dz1, - dx, - x, - dmask0, - dmask1, - mu, - rsigma, - gamma0, - gamma1, - dropout_p, - has_x1, - has_residual, - ctx.is_rms_norm, - ) - dx0 = dx0mat.view(x.shape) - dx1 = dx1mat.view(x.shape) if dx1mat is not None else None - dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None - return ( - dx0, - dx1, - dresidual, - dgamma0, - dbeta0 if ctx.has_beta else None, - dgamma1, - dbeta1 if ctx.has_beta else None, - None, - None, - None, - None, - None, - None, - ) - - -def layer_norm(x, weight, bias, epsilon): - return DropoutAddLayerNormFn.apply(x, None, weight, bias, None, None, 0.0, epsilon, False) - - -def dropout_add_layer_norm( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - rowscale=None, - layerscale=None, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormFn.apply( - x0, - residual, - weight, - bias, - rowscale, - layerscale, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - False, - return_dropout_mask, - ) - - -def dropout_add_layer_norm_subset( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - layerscale=None, - x0_subset=None, - out_subset=None, - rowscale_const=1.0, - out_numrows=0, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormSubsetFn.apply( - x0, - residual, - weight, - bias, - layerscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32, - prenorm, - False, - return_dropout_mask, - ) - - -def dropout_add_layer_norm_parallel_residual( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormParallelResidualFn.apply( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - False, - return_dropout_mask, - ) - - -class DropoutAddLayerNorm(torch.nn.Module): - def __init__( - self, - hidden_size, - prenorm=False, - p=0.0, - eps=1e-5, - residual_in_fp32=False, - device=None, - dtype=None, - ): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.prenorm = prenorm - self.p = p - self.eps = eps - self.residual_in_fp32 = residual_in_fp32 - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.bias = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.reset_parameters() - - def reset_parameters(self): - init.ones_(self.weight) - init.zeros_(self.bias) - - def forward(self, x0, residual=None): - return dropout_add_layer_norm( - x0, - residual, - self.weight, - self.bias, - self.p if self.training else 0.0, - self.eps, - prenorm=self.prenorm, - residual_in_fp32=self.residual_in_fp32, - ) diff --git a/build/torch29-cxx11-cu130-aarch64-linux/ops/rms_norm.py b/build/torch29-cxx11-cu130-aarch64-linux/ops/rms_norm.py deleted file mode 100644 index 068348d61290e3839dd082b540d898578ba1e8e2..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-aarch64-linux/ops/rms_norm.py +++ /dev/null @@ -1,174 +0,0 @@ -# Copyright (c) 2022, Tri Dao. -# Adapted from https://github.com/NVIDIA/apex/blob/master/apex/contrib/layer_norm/layer_norm.py - -import torch -from torch.nn import init - -from flash_attn.ops.layer_norm import ( - DropoutAddLayerNormFn, - DropoutAddLayerNormParallelResidualFn, - DropoutAddLayerNormSubsetFn, -) - - -def rms_norm(x, weight, epsilon): - return DropoutAddLayerNormFn.apply( - x, None, weight, None, None, None, 0.0, epsilon, False, False, True - ) - - -def dropout_add_rms_norm( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - rowscale=None, - layerscale=None, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormFn.apply( - x0, - residual, - weight, - bias, - rowscale, - layerscale, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - True, - return_dropout_mask, - ) - - -def dropout_add_rms_norm_subset( - x0, - residual, - weight, - bias, - dropout_p, - epsilon, - layerscale=None, - x0_subset=None, - out_subset=None, - rowscale_const=1.0, - out_numrows=0, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormSubsetFn.apply( - x0, - residual, - weight, - bias, - layerscale, - x0_subset, - out_subset, - dropout_p, - epsilon, - rowscale_const, - out_numrows, - residual_in_fp32, - prenorm, - True, - return_dropout_mask, - ) - - -def dropout_add_rms_norm_parallel_residual( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - prenorm=False, - residual_in_fp32=False, - return_dropout_mask=False, -): - """residual_in_fp32 only has an effect if residual is None. - Otherwise residual dtype is residual.dtype. - """ - return DropoutAddLayerNormParallelResidualFn.apply( - x0, - x1, - residual, - weight0, - bias0, - weight1, - bias1, - dropout_p, - epsilon, - residual_in_fp32, - prenorm, - True, - return_dropout_mask, - ) - - -class RMSNorm(torch.nn.Module): - def __init__(self, hidden_size, eps=1e-5, device=None, dtype=None): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.eps = eps - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.register_parameter("bias", None) - self.reset_parameters() - - def reset_parameters(self): - init.ones_(self.weight) - - def forward(self, x): - return rms_norm(x, self.weight, self.eps) - - -class DropoutAddRMSNorm(torch.nn.Module): - def __init__( - self, - hidden_size, - prenorm=False, - p=0.0, - eps=1e-5, - residual_in_fp32=False, - device=None, - dtype=None, - ): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.prenorm = prenorm - self.p = p - self.eps = eps - self.residual_in_fp32 = residual_in_fp32 - self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.register_parameter("bias", None) - self.reset_parameters() - - def reset_parameters(self): - init.ones_(self.weight) - - def forward(self, x0, residual=None): - return dropout_add_rms_norm( - x0, - residual, - self.weight, - None, - self.p if self.training else 0.0, - self.eps, - prenorm=self.prenorm, - residual_in_fp32=self.residual_in_fp32, - ) diff --git a/build/torch29-cxx11-cu130-aarch64-linux/ops/triton/__init__.py b/build/torch29-cxx11-cu130-aarch64-linux/ops/triton/__init__.py deleted file mode 100644 index 8b137891791fe96927ad78e64b0aad7bded08bdc..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-aarch64-linux/ops/triton/__init__.py +++ /dev/null @@ -1 +0,0 @@ - diff --git a/build/torch29-cxx11-cu130-aarch64-linux/ops/triton/cross_entropy.py b/build/torch29-cxx11-cu130-aarch64-linux/ops/triton/cross_entropy.py deleted file mode 100644 index 1b5a415b73f236f3e05fb14b9141959559e18526..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-aarch64-linux/ops/triton/cross_entropy.py +++ /dev/null @@ -1,330 +0,0 @@ -# Copyright (c) 2023, Tri Dao. - -from typing import Tuple, Optional, Union - -import torch -import torch.nn.functional as F - -import triton -import triton.language as tl - -# `all_gather_into_tensor` and `reduce_scatter_tensor` are new placeholders for -# `_all_gather_base` and `_reduce_scatter_base`. They require the most recent -# version of PyTorch. The following 2 lines are for backward compatibility with -# older PyTorch. -if "all_gather_into_tensor" not in dir(torch.distributed): - torch.distributed.all_gather_into_tensor = torch.distributed._all_gather_base - - -@triton.heuristics( - { - "HAS_SMOOTHING": lambda args: args["smoothing"] > 0.0, - } -) -@triton.jit -def cross_entropy_fwd_kernel( - loss_ptr, # data ptrs - lse_ptr, - z_loss_ptr, - logits_ptr, - labels_ptr, - smoothing, - logit_scale, - lse_square_scale, - ignore_index, - total_classes, - class_start_idx, # Useful for tensor parallel when each rank only has a subset of classes - n_cols, # shapes - logits_row_stride, # strides - BLOCK_SIZE: tl.constexpr, - HAS_SMOOTHING: tl.constexpr, - # if SPLIT (e.g. tensor parallel), don't include the LSE in the loss since it's not the final LSE - SPLIT: tl.constexpr, - PRECOMPUTED_LSE: tl.constexpr, # If LSE is already computed (also no smoothing and logit_scale == 1.0) -): - row_idx = tl.program_id(0) - logits_ptr = logits_ptr + row_idx * logits_row_stride.to(tl.int64) - sum_logits = 0.0 # For smoothing - if not PRECOMPUTED_LSE: - # Statistics for online softmax - m_i = -float("inf") - l_i = 0.0 - for col_offset in range(0, n_cols, BLOCK_SIZE): - cols = col_offset + tl.arange(0, BLOCK_SIZE) - logits = tl.load(logits_ptr + cols, mask=cols < n_cols, other=-float("inf")).to( - tl.float32 - ) * logit_scale - if HAS_SMOOTHING: - sum_logits += tl.sum(tl.where(cols < n_cols, logits, 0.0)) - m_i_new = tl.maximum(m_i, tl.max(logits)) - l_i = tl.exp(m_i - m_i_new) * l_i + tl.sum(tl.exp(logits - m_i_new)) - m_i = m_i_new - lse = tl.log(l_i) + m_i - tl.store(lse_ptr + row_idx, lse) - else: - lse = tl.load(lse_ptr + row_idx) - label_idx = tl.load(labels_ptr + row_idx) - if label_idx == ignore_index: - loss = 0.0 - z_loss = 0.0 - else: - label_idx -= class_start_idx - if label_idx >= 0 and label_idx < n_cols: - logits_label = tl.load(logits_ptr + label_idx) * logit_scale - if HAS_SMOOTHING: - loss = ( - (lse if not SPLIT else 0.0) - - smoothing * sum_logits / total_classes - - (1 - smoothing) * logits_label - ) - else: - loss = (lse if not SPLIT else 0.0) - logits_label - else: - # If label is out of bounds, we set the CE loss to 0.0. But we still want the smoothing loss - if HAS_SMOOTHING: - loss = smoothing * ((lse if not SPLIT else 0.0) - sum_logits / total_classes) - else: - loss = 0.0 - if not SPLIT: - z_loss = lse_square_scale * lse * lse - loss += z_loss - else: - z_loss = 0.0 - tl.store(loss_ptr + row_idx, loss) - if not SPLIT: - tl.store(z_loss_ptr + row_idx, z_loss) - - -@triton.heuristics( - { - "HAS_SMOOTHING": lambda args: args["smoothing"] > 0.0, - } -) -@triton.jit -def cross_entropy_bwd_kernel( - dlogits_ptr, # data ptrs - dloss_ptr, - logits_ptr, - lse_ptr, - labels_ptr, - smoothing, - logit_scale, - lse_square_scale, - ignore_index, - total_classes, - class_start_idx, # Useful for tensor parallel when each rank only has a subset of classes - n_cols, # shapes - logits_row_stride, # strides - dlogits_row_stride, - dloss_row_stride, - BLOCK_SIZE: tl.constexpr, - HAS_SMOOTHING: tl.constexpr, -): - row_idx = tl.program_id(0) - col_block_idx = tl.program_id(1) - logits_ptr = logits_ptr + row_idx * logits_row_stride.to(tl.int64) - dlogits_ptr = dlogits_ptr + row_idx * dlogits_row_stride.to(tl.int64) - col_offsets = col_block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - label_idx = tl.load(labels_ptr + row_idx) - if label_idx != ignore_index: - dloss = tl.load(dloss_ptr + row_idx * dloss_row_stride) - else: - dloss = 0.0 - logits = tl.load(logits_ptr + col_offsets, mask=col_offsets < n_cols, other=-float("inf")).to( - tl.float32 - ) * logit_scale - lse = tl.load(lse_ptr + row_idx) - probs = tl.exp(logits - lse) - probs += 2.0 * lse_square_scale * lse * probs - label_idx -= class_start_idx - if HAS_SMOOTHING: - smooth_positive = 1.0 - smoothing - smooth_negative = smoothing / total_classes - probs = tl.where(col_offsets == label_idx, probs - smooth_positive, probs) - smooth_negative - else: - probs = tl.where(col_offsets == label_idx, probs - 1.0, probs) - tl.store(dlogits_ptr + col_offsets, (dloss * logit_scale) * probs, mask=col_offsets < n_cols) - - -class CrossEntropyLoss(torch.autograd.Function): - - @staticmethod - def forward( - ctx, - logits, - labels, - precomputed_lse=None, - smoothing=0.0, - logit_scale=1.0, - lse_square_scale=0.0, - ignore_index=-100, - inplace_backward=False, - process_group=None, - ): - # For some reason Triton generates wrong code when labels has dtype long and its address - # is not aligned to 16 bytes. The ld.global.b64 seems to load the wrong label index. - if labels.dtype == torch.long and labels.data_ptr() % 16 != 0: - labels = F.pad(labels, (0, 1))[..., :-1] - assert labels.data_ptr() % 16 == 0 - assert logit_scale > 0.0 - n_rows, n_cols = logits.shape - assert labels.shape == (n_rows,) - world_size = 1 if process_group is None else torch.distributed.get_world_size(process_group) - total_classes = world_size * n_cols - rank = 0 if process_group is None else torch.distributed.get_rank(process_group) - class_start_idx = rank * n_cols - use_precomputed_lse = precomputed_lse is not None and logit_scale == 1.0 and smoothing == 0.0 - - if logits.stride(-1) != 1: - logits = logits.contiguous() - MAX_BLOCK_SIZE = 16 * 1024 - BLOCK_SIZE = min(triton.next_power_of_2(n_cols), MAX_BLOCK_SIZE) - num_warps = ( - 4 - if BLOCK_SIZE < 2048 - else (8 if BLOCK_SIZE < 8192 else (16 if BLOCK_SIZE < 128 * 1024 else 32)) - ) - losses = torch.empty(n_rows, dtype=torch.float, device=logits.device) - if use_precomputed_lse: - assert precomputed_lse.shape == (n_rows,) - lse = precomputed_lse.contiguous() - else: - lse = torch.empty(n_rows, dtype=torch.float, device=logits.device) - z_losses = torch.empty(n_rows, dtype=torch.float, device=logits.device) - # Need this, otherwise Triton tries to launch from cuda:0 and we get - # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?) - with torch.cuda.device(logits.device.index): - cross_entropy_fwd_kernel[(n_rows,)]( - losses, # data ptrs - lse, - z_losses, - logits, - labels, - smoothing, - logit_scale, - lse_square_scale, - ignore_index, - total_classes, - class_start_idx, - n_cols, # shapes - logits.stride(0), # strides - BLOCK_SIZE=BLOCK_SIZE, # constants - SPLIT=world_size > 1, - PRECOMPUTED_LSE=use_precomputed_lse, - num_warps=num_warps, - ) - - if world_size > 1: - # If there's no smoothing, if labels are in the vocab of this partition, losses contains - # - predicted logit, and 0 otherwise. - # If there's smoothing=0.1, for labels in the vocab of this partition, losses contains - # -0.9 * predicted logit - 0.1 * sum logit / total_classes. - # For labels not in the vocab of this partition, losses contains - # -0.1 * sum logit / total_classes. - if world_size > 1: - lse_allgather = torch.empty(world_size, n_rows, dtype=lse.dtype, device=lse.device) - torch.distributed.all_gather_into_tensor(lse_allgather, lse, group=process_group) - handle_losses = torch.distributed.all_reduce( - losses, op=torch.distributed.ReduceOp.SUM, group=process_group, async_op=True - ) - lse = torch.logsumexp(lse_allgather, dim=0) - handle_losses.wait() - # After the allreduce, if there's no smoothing, the total losses are - predicted_logit, - # we just have to add the (global) lse. - # If there's smoothing=0.1, the total losses are - # -0.9 * predicted_logit - 0.1 * sum logit / total_classes. - # Again, we just have to add the (global) lse. - losses += lse - if lse_square_scale != 0.0: - z_losses = lse_square_scale * lse.square() - z_losses.masked_fill_(labels == ignore_index, 0.0) - losses += z_losses - else: - z_losses = torch.zeros_like(losses) - losses.masked_fill_(labels == ignore_index, 0.0) - - ctx.save_for_backward(logits, lse, labels) - ctx.mark_non_differentiable(z_losses) - ctx.smoothing = smoothing - ctx.logit_scale = logit_scale - ctx.lse_square_scale = lse_square_scale - ctx.ignore_index = ignore_index - ctx.total_classes = total_classes - ctx.class_start_idx = class_start_idx - ctx.inplace_backward = inplace_backward - return losses, z_losses - - @staticmethod - def backward(ctx, grad_losses, grad_z_losses): - del grad_z_losses # z_losses are only for logging. - - logits, lse, labels = ctx.saved_tensors - dlogits = logits if ctx.inplace_backward else torch.empty_like(logits) - n_rows, n_cols = logits.shape - BLOCK_SIZE = min(triton.next_power_of_2(n_cols), 4 * 1024) - num_warps = 4 if BLOCK_SIZE < 2048 else (8 if BLOCK_SIZE < 8192 else 16) - grid = lambda META: (n_rows, triton.cdiv(n_cols, META["BLOCK_SIZE"])) # noqa - # Need this, otherwise Triton tries to launch from cuda:0 and we get - # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?) - with torch.cuda.device(logits.device.index): - cross_entropy_bwd_kernel[grid]( - dlogits, # data ptrs - grad_losses, - logits, - lse, - labels, - ctx.smoothing, - ctx.logit_scale, - ctx.lse_square_scale, - ctx.ignore_index, - ctx.total_classes, - ctx.class_start_idx, - n_cols, # shapes - logits.stride(0), # strides - dlogits.stride(0), - grad_losses.stride(0), - BLOCK_SIZE=BLOCK_SIZE, # constants - num_warps=num_warps, - ) - return dlogits, None, None, None, None, None, None, None, None, None - - -def cross_entropy_loss( - logits: torch.Tensor, - labels: torch.Tensor, - precomputed_lse: Optional[torch.Tensor] = None, - label_smoothing: float = 0.0, - logit_scale: float = 1.0, - lse_square_scale: float = 0.0, - ignore_index=-100, - inplace_backward: bool = False, - process_group=None, -) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Arguments: - logits: (batch, vocab_size) - labels: (batch,) - label_smoothing: float - logit_scale: float. Multiply logits by this scale before calculating the loss. - lse_square_scale: float. If > 0, we add lse_square_scale * lse(logits) ^ 2 to the loss. - This is also referred to as "z-loss". - ignore_index: int. If labels == ignore_index, the loss is set to 0.0. - inplace_backward: bool. If True, we do the backward pass in-place by modifying the logits. - This saves memory. - process_group: if not None, we're doing Tensor Parallel: each process is responsible for - one part of the vocab. The loss will be aggregated across processes. - Returns: - losses: (batch,), float - z_losses: (batch,), float - """ - return CrossEntropyLoss.apply( - logits, - labels, - precomputed_lse, - label_smoothing, - logit_scale, - lse_square_scale, - ignore_index, - inplace_backward, - process_group, - ) diff --git a/build/torch29-cxx11-cu130-aarch64-linux/ops/triton/k_activations.py b/build/torch29-cxx11-cu130-aarch64-linux/ops/triton/k_activations.py deleted file mode 100644 index efb83c358eb4a85d069ee340a3c83f418f9a805b..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-aarch64-linux/ops/triton/k_activations.py +++ /dev/null @@ -1,162 +0,0 @@ -# Adapted from https://github.com/facebookresearch/xformers/blob/main/xformers/triton/k_activations.py -# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved. -# -# This source code is licensed under the BSD license found in the -# LICENSE file in the root directory of this source tree. - -import math -from enum import Enum -from typing import Optional - -import triton -import triton.language as tl - -_sqrt2pi = math.sqrt(2.0 / math.pi) -_sqrt1_2 = math.sqrt(1.0 / 2) -_gaussian_pdf_normalization = 1.0 / math.sqrt(2 * math.pi) - - -class Activation(str, Enum): - SquaredReLU = "squared_relu" - GeLU = "gelu" - GeLUApprox = "gelu_approx" - LeakyReLU = "leaky_relu" - ReLU = "relu" - - -def get_triton_activation_kernel(activation: Optional[Activation]): - return ( - { - Activation.ReLU: relu, - Activation.LeakyReLU: leaky_relu, - Activation.GeLU: gelu, - Activation.GeLUApprox: gelu_approx, - Activation.SquaredReLU: squared_relu, - }[activation] - if activation - else None - ) - - -def get_triton_activation_bwd_kernel(activation: Optional[Activation]): - return ( - { - Activation.ReLU: relu_grad, - Activation.LeakyReLU: leaky_relu_grad, - Activation.GeLU: gelu_grad, - Activation.GeLUApprox: gelu_approx_grad, - Activation.SquaredReLU: squared_relu_grad, - }[activation] - if activation - else None - ) - - -@triton.jit -def tanh(x): - # Tanh is just a scaled sigmoid - return 2 * tl.sigmoid(2 * x) - 1 - - -@triton.jit -def cosh(x): - exp_x = tl.exp(x) - return (exp_x + 1.0 / exp_x) * 0.5 - - -# a Triton implementation of the most used activations -# See for instance http://arxiv.org/abs/1606.08415 for an overview - -# ReLU -@triton.jit -def relu(x): - """ - ReLU_ activation function - - .. _ReLU: https://pytorch.org/docs/stable/generated/torch.nn.ReLU.html - """ - zero = 0.0 - return tl.where(x >= 0, x, zero.to(x.dtype)) - - -@triton.jit -def relu_grad(x): - # ReLU is different from other activations - # in that it does not require the input to retrospectively compute its gradient - # here the input is the downstream gradient, and we return the upstream gradient directly - zero = 0.0 - one = 1.0 - return tl.where(x >= 0, one.to(x.dtype), zero.to(x.dtype)) - - -@triton.jit -def squared_relu(x): - """ - Squared ReLU activation, as proposed in the Primer_ paper. - - .. _Primer: https://arxiv.org/abs/2109.08668 - """ - x_ = relu(x) - return (x_ * x_).to(x.dtype) - - -@triton.jit -def squared_relu_grad(x): - return tl.where(x >= 0, 2.0 * x, 0.0) - - -# Leaky ReLU -@triton.jit -def leaky_relu(x): - """ - LeakyReLU_ activation - - .. _LeakyReLU: https://pytorch.org/docs/stable/generated/torch.nn.LeakyReLU.html - """ - scale = 0.01 + 0.0 - scale = scale.to(x.dtype) - return tl.where(x >= 0, x, scale * x) - - -@triton.jit -def leaky_relu_grad(x): - min_grad = 0.01 - max_grad = 1 - - min_grad = min_grad.to(x.dtype) - max_grad = max_grad.to(x.dtype) - - return tl.where(x >= 0, max_grad, min_grad) - - -@triton.jit -def gelu(x): - """Gaussian Error Linear Unit (GELU)""" - return x * 0.5 * (1.0 + tl.libdevice.erf(x * _sqrt1_2)) - - -@triton.jit -def gelu_grad(x): - cdf = 0.5 * (1.0 + tl.libdevice.erf(x * _sqrt1_2)) - pdf = tl.exp(-0.5 * x * x) * _gaussian_pdf_normalization - return cdf + x * pdf - - -@triton.jit -def gelu_approx(x): - """ - GeLU_ activation - Gaussian error linear unit, with tanh approximation - - .. _GeLU: https://arxiv.org/pdf/1606.08415.pdf - """ - return 0.5 * x * (1.0 + tanh(_sqrt2pi * x * (1.0 + 0.044715 * x * x))) - - -@triton.jit -def gelu_approx_grad(x): - # CREDITS: Fast implementation proposed in - # https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/fused_bias_gelu.py#L30 - tanh_out = tanh(0.79788456 * x * (1 + 0.044715 * x * x)) - return 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( - 1 + tanh_out - ) diff --git a/build/torch29-cxx11-cu130-aarch64-linux/ops/triton/linear.py b/build/torch29-cxx11-cu130-aarch64-linux/ops/triton/linear.py deleted file mode 100644 index a8966dbc345ab0e593df0124451ee7be3dae131a..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-aarch64-linux/ops/triton/linear.py +++ /dev/null @@ -1,594 +0,0 @@ -# Adapted from https://github.com/ELS-RD/kernl/blob/main/src/kernl/implementations/linear_layer.py -# and https://github.com/openai/triton/blob/master/python/triton/ops/matmul.py -from typing import Optional - -import torch -import triton -import triton.language as tl -from triton.ops.matmul_perf_model import early_config_prune, estimate_matmul_time - -from flash_attn.ops.triton.k_activations import ( - gelu, - gelu_approx, - gelu_approx_grad, - gelu_grad, - squared_relu, - squared_relu_grad, -) - -# CREDITS: Initially inspired by the Triton tutorial on matrix multiplications - - -def init_to_zero(name): - return lambda nargs: nargs[name].zero_() - - -def get_configs_io_bound(): - configs = [] - for num_stages in [2, 3, 4, 5, 6]: - for block_m in [16, 32]: - for block_k in [32, 64]: - for block_n in [32, 64, 128, 256]: - num_warps = 2 if block_n <= 64 else 4 - configs.append( - triton.Config( - { - "BLOCK_M": block_m, - "BLOCK_N": block_n, - "BLOCK_K": block_k, - "SPLIT_K": 1, - }, - num_stages=num_stages, - num_warps=num_warps, - ) - ) - # split_k not used - # for split_k in [2, 4, 8, 16]: - # configs.append(triton.Config( - # {'BLOCK_M': block_m, 'BLOCK_N': block_n, 'BLOCK_K': block_k, 'SPLIT_K': split_k}, - # num_stages=num_stages, num_warps=num_warps, pre_hook=init_to_zero('C'))) - return configs - - -@triton.autotune( - configs=[ - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - # good for int8 - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=4, - num_warps=4, - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - ] - + get_configs_io_bound(), - key=["CACHE_KEY_M", "CACHE_KEY_N", "CACHE_KEY_K"], - prune_configs_by={ - "early_config_prune": early_config_prune, - "perf_model": estimate_matmul_time, - "top_k": 10, - }, -) -@triton.heuristics( - { - "EVEN_K": lambda args: args["K"] % (args["BLOCK_K"] * args["SPLIT_K"]) == 0, - } -) -@triton.jit -def kernel_fwd( - C, # Pointers to matrices - ACT_INPUT, - A, - B, - bias, - # Matrix dimensions - M, - N, - K, - CACHE_KEY_M, - CACHE_KEY_N, - CACHE_KEY_K, - # The stride variables represent how much to increase the ptr by when moving by 1 - # element in a particular dimension. E.g. stride_am is how much to increase a_ptr - # by to get the element one row down (A has M rows) - stride_cm, - # stride_cn, # Assume that stride_cn == 1 - stride_am, - stride_ak, - stride_bn, - stride_bk, - # Meta-parameters - BLOCK_M: tl.constexpr, - GROUP_M: tl.constexpr, - BLOCK_N: tl.constexpr, - BLOCK_K: tl.constexpr, - # split k not used, not performant with activation, kept because early_config_prune is expecting it - SPLIT_K: tl.constexpr, - EVEN_K: tl.constexpr, - A_ROWMAJOR: tl.constexpr, - B_COLMAJOR: tl.constexpr, - BIAS: tl.constexpr, - SAVE_ACT_INPUT: tl.constexpr, - ACTIVATION: tl.constexpr, -): - - """ - Kernel for computing Out = activation(A x W + C) - - Input has shape (M, K) - - Weight has shape (K, N) - - Bias has shape (N,) - - Output has shape (M, N) - - ActInputs (optional) has shape (M, N) - 'ActInputs' optionally saves the A x W + C intermediate for backward computations - This kernel will consolidate over K - """ - - pid = tl.program_id(axis=0) - - grid_m = (M + BLOCK_M - 1) // BLOCK_M - grid_n = (N + BLOCK_N - 1) // BLOCK_N - # re-order program ID for better L2 performance - width = GROUP_M * grid_n - group_id = pid // width - group_size = min(grid_m - group_id * GROUP_M, GROUP_M) - pid_m = group_id * GROUP_M + (pid % group_size) - pid_n = (pid % width) // (group_size) - - # now compute the block that each program will go through - # rm (resp. rn) denotes a range of indices - # for rows (resp. col) of C - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - # trick to avoid masking on M and N axis - ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M) - rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N) - rk = tl.arange(0, BLOCK_K) - - if A_ROWMAJOR: - A = A + (ram[:, None] * stride_am + rk[None, :]) - else: - A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak) - if B_COLMAJOR: - B = B + (rk[:, None] + rbn[None, :] * stride_bn) - else: - B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn) - - acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) - - for k in range(K, 0, -BLOCK_K): - if EVEN_K: - a = tl.load(A) - b = tl.load(B) - else: - a = tl.load(A, mask=rk[None, :] < k, other=0.0) - b = tl.load(B, mask=rk[:, None] < k, other=0.0) - acc += tl.dot(a, b) - - if A_ROWMAJOR: - A += BLOCK_K - else: - A += BLOCK_K * stride_ak - if B_COLMAJOR: - B += BLOCK_K - else: - B += BLOCK_K * stride_bk - - # Putting bias after the matmul (instead of before) is faster, idk why - if BIAS: - bias = tl.load(bias + rn, mask=rn < N, other=0.0).to(tl.float32) - acc += bias[None, :] - - # optional: save the activation inputs - if SAVE_ACT_INPUT: - # act_in_ptrs = ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :] * stride_cn - act_in_ptrs = ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :] - tl.store(act_in_ptrs, acc) - - # optional: fused activation (while the data is in shared memory) - if ACTIVATION == "gelu": - acc = gelu(acc) - elif ACTIVATION == "gelu_approx": - acc = gelu_approx(acc) - elif ACTIVATION == "squared_relu": - acc = squared_relu(acc) - # rematerialize rm and rn to save registers - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - - # write back result - # C = C + rm[:, None] * stride_cm + rn[None, :] * stride_cn - C = C + rm[:, None] * stride_cm + rn[None, :] - mask = (rm < M)[:, None] & (rn < N)[None, :] - tl.store(C, acc) - - -def triton_linear_act( - x: torch.Tensor, - weight: torch.Tensor, - bias: Optional[torch.Tensor] = None, - activation: str = "id", - save_act_input: bool = False, -) -> torch.Tensor: - """ - Compute e = activation(x @ weight.T + bias). - This wrapper kicks the `kernel_fwd` Triton kernel - :param x: input tensor - :param weight: weight matrix - :param bias: an optional bias tensor - :param activation: Activation name. Needs to be a Triton kernel. - :param act_input: an optional tensor to save the activation inputs (for backward) - :return: result tensor - """ - # if torch.is_autocast_enabled(): - # dtype = torch.get_autocast_gpu_dtype() - # x, weight, bias = [a.to(dtype=dtype) for a in [x, weight, bias]] - - assert activation in ["id", "gelu", "gelu_approx", "squared_relu"] - - batch_shape, n = x.shape[:-1], x.shape[-1] - batch_dim = batch_shape.numel() - x_reshaped = x.reshape(batch_dim, n) - - if x_reshaped.stride(0) > 1 and x_reshaped.stride(1) > 1: - x_reshaped = x_reshaped.contiguous() - if weight.stride(0) > 1 and weight.stride(1) > 1: - weight = weight.contiguous() - bias = bias.contiguous() if bias is not None else None - - assert ( - x.dtype == weight.dtype - ), f"Input and weight must have the same dtype, got {x.dtype} and {weight.dtype}" - if bias is not None: - assert ( - x.dtype == bias.dtype - ), f"Input and bias must have the same dtype, got {x.dtype} and {bias.dtype}" - assert ( - x_reshaped.shape[1] == weight.shape[1] - ), f"Incompatible dimensions: {x_reshaped.shape} - {weight.shape}" - - assert ( - bias is None or bias.shape[0] == weight.shape[0] - ), "Incompatible dimensions in between weight and bias" - - M, K = x_reshaped.shape - N, K = weight.shape - - output = torch.empty((M, N), device=x.device, dtype=x.dtype) - act_input = torch.empty_like(output) if save_act_input else None - - # 1D launch kernel where each block gets its own program. - grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),) # noqa - - kernel_fwd[grid]( - output, - act_input, - x_reshaped, - weight, # data ptrs - bias if bias is not None else x, # auto skip bias if not present - M, # shapes - N, - K, - M // 32, # key for triton cache (limit number of compilations) - N // 32, - K // 32, - stride_cm=output.stride(0), # strides - # stride_cn=output.stride(1), - stride_am=x_reshaped.stride(0), - stride_ak=x_reshaped.stride(1), - stride_bk=weight.stride(1), - stride_bn=weight.stride(0), - BIAS=bias is not None, # optional fused bias - SAVE_ACT_INPUT=save_act_input, # optional save activation inputs - ACTIVATION=activation, # optional fused activation - A_ROWMAJOR=x_reshaped.stride(1) == 1, - B_COLMAJOR=weight.stride(1) == 1, - GROUP_M=8, # speed optimization: group the programs - ) - - if not save_act_input: - return output.reshape(*batch_shape, output.shape[-1]) - else: - return ( - output.reshape(*batch_shape, output.shape[-1]), - act_input.reshape(*batch_shape, act_input.shape[-1]), - ) - - -@triton.autotune( - configs=[ - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - # good for int8 - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, - num_stages=4, - num_warps=4, - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 - ), - triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=5, num_warps=2 - ), - ] - + get_configs_io_bound(), - key=["CACHE_KEY_M", "CACHE_KEY_N", "CACHE_KEY_K"], - prune_configs_by={ - "early_config_prune": early_config_prune, - "perf_model": estimate_matmul_time, - "top_k": 10, - }, -) -@triton.heuristics( - { - "EVEN_K": lambda args: args["K"] % (args["BLOCK_K"] * args["SPLIT_K"]) == 0, - } -) -@triton.jit -def kernel_bwd( - C, # Pointers to matrices - ACT_INPUT, - A, - B, - # Matrix dimensions - M, - N, - K, - CACHE_KEY_M, - CACHE_KEY_N, - CACHE_KEY_K, - # The stride variables represent how much to increase the ptr by when moving by 1 - # element in a particular dimension. E.g. stride_am is how much to increase a_ptr - # by to get the element one row down (A has M rows) - stride_cm, - # stride_cn, # Assume that stride_cn == 1 - stride_am, - stride_ak, - stride_bk, - stride_bn, - # Meta-parameters - BLOCK_M: tl.constexpr, - GROUP_M: tl.constexpr, - BLOCK_N: tl.constexpr, - BLOCK_K: tl.constexpr, - # split k not used, not performant with activation, kept because early_config_prune is expecting it - SPLIT_K: tl.constexpr, - EVEN_K: tl.constexpr, - ACTIVATION: tl.constexpr, -): - - """ - Kernel for computing Out = activation(A x W + C) - - Input has shape (M, K) - - Weight has shape (K, N) - - Output has shape (M, N) - - ActInputs (optional) has shape (M, N) - 'ActInputs' optionally saves the A x W + C intermediate for backward computations - This kernel will consolidate over K - """ - - pid = tl.program_id(axis=0) - - grid_m = (M + BLOCK_M - 1) // BLOCK_M - grid_n = (N + BLOCK_N - 1) // BLOCK_N - # re-order program ID for better L2 performance - width = GROUP_M * grid_n - group_id = pid // width - group_size = min(grid_m - group_id * GROUP_M, GROUP_M) - pid_m = group_id * GROUP_M + (pid % group_size) - pid_n = (pid % width) // (group_size) - - # now compute the block that each program will go through - # rm (resp. rn) denotes a range of indices - # for rows (resp. col) of C - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - # trick to avoid masking on M and N axis - ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M) - rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N) - rk = tl.arange(0, BLOCK_K) - - A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak) - B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn) - - acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) - - for k in range(K, 0, -BLOCK_K): - if EVEN_K: - a = tl.load(A) - b = tl.load(B) - else: - a = tl.load(A, mask=rk[None, :] < k, other=0.0) - b = tl.load(B, mask=rk[:, None] < k, other=0.0) - acc += tl.dot(a, b) - - A += BLOCK_K * stride_ak - B += BLOCK_K * stride_bk - - # optional: fused activation (while the data is in shared memory) - if ACTIVATION != "id": - act_in_ptrs = ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :] - act_input = tl.load(act_in_ptrs).to(acc.dtype) - if ACTIVATION == "gelu": - acc *= gelu_grad(act_input) - elif ACTIVATION == "gelu_approx": - acc *= gelu_approx_grad(act_input) - elif ACTIVATION == "squared_relu": - acc *= squared_relu_grad(act_input) - - # rematerialize rm and rn to save registers - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - - # write back result - C = C + rm[:, None] * stride_cm + rn[None, :] - mask = (rm < M)[:, None] & (rn < N)[None, :] - tl.store(C, acc, mask=mask) - - -def triton_dgrad_act( - grad_output: torch.Tensor, - weight: torch.Tensor, - activation: str = "id", - act_input: Optional[torch.Tensor] = None, -) -> torch.Tensor: - """ - Compute e = activation(grad_output @ weight + bias). - This wrapper kicks the `kernel_fwd` Triton kernel - :param grad_output: input tensor - :param weight: weight matrix - :param activation: Activation name. Needs to be a Triton kernel. - :param act_input: an optional tensor to save the activation inputs (for backward) - :return: result tensor - """ - assert activation in ["id", "gelu", "gelu_approx", "squared_relu"] - - batch_shape, n = grad_output.shape[:-1], grad_output.shape[-1] - batch_dim = batch_shape.numel() - grad_output_reshaped = grad_output.reshape(batch_dim, n) - - if grad_output_reshaped.stride(0) > 1 and grad_output_reshaped.stride(1) > 1: - grad_output_reshaped = grad_output_reshaped.contiguous() - if weight.stride(0) > 1 and weight.stride(1) > 1: - weight = weight.contiguous() - - assert ( - grad_output.dtype == weight.dtype - ), f"grad_output and weight must have the same dtype, got {grad_output.dtype} and {weight.dtype}" - assert ( - grad_output_reshaped.shape[1] == weight.shape[0] - ), f"Incompatible dimensions: {grad_output_reshaped.shape} - {weight.shape}" - if activation != "id": - assert act_input is not None, f"act_input is required for activation {activation}" - - # M, N, K in bwd are different from M, N, K in fwd - M, K = grad_output_reshaped.shape - K, N = weight.shape - - grad_input = torch.empty((M, N), device=grad_output.device, dtype=grad_output.dtype) - - # 1D launch kernel where each block gets its own program. - grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),) # noqa - - kernel_bwd[grid]( - grad_input, - act_input, - grad_output_reshaped, - weight, # data ptrs - M, # shapes - N, - K, - M // 32, # key for triton cache (limit number of compilations) - N // 32, - K // 32, - stride_cm=grad_input.stride(0), # strides - # stride_cn=grad_input.stride(1), - stride_am=grad_output_reshaped.stride(0), - stride_ak=grad_output_reshaped.stride(1), - stride_bk=weight.stride(0), - stride_bn=weight.stride(1), - ACTIVATION=activation, # optional fused activation - GROUP_M=8, # speed optimization: group the programs - ) - - return grad_input.reshape(*batch_shape, grad_input.shape[-1]) diff --git a/build/torch29-cxx11-cu130-aarch64-linux/ops/triton/mlp.py b/build/torch29-cxx11-cu130-aarch64-linux/ops/triton/mlp.py deleted file mode 100644 index 059f4f8a5e174c1f4824e43d313fca18eaa799b8..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-aarch64-linux/ops/triton/mlp.py +++ /dev/null @@ -1,149 +0,0 @@ -# The triton fused matmul + sqrelu is faster for fp16 but slower for bf16, compared -# to naive implementation. -import fused_dense_lib as fused_dense_cuda -import torch -import torch.nn as nn -import torch.nn.functional as F - -from flash_attn.utils.torch import custom_fwd, custom_bwd -from flash_attn.ops.activations import sqrelu_bwd, sqrelu_fwd -from flash_attn.ops.triton.linear import triton_dgrad_act, triton_linear_act - - -class FusedDenseSqreluDenseFunc(torch.autograd.Function): - @staticmethod - @custom_fwd - def forward(ctx, x, weight1, bias1, weight2, bias2, checkpoint_lvl=0): - """checkpoint_lvl: - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute act_input and gelu_out in the bwd - """ - if torch.is_autocast_enabled(): - dtype = torch.get_autocast_gpu_dtype() - x, weight1, bias1, weight2, bias2 = [ - a.to(dtype=dtype) for a in [x, weight1, bias1, weight2, bias2] - ] - is_bf16 = x.dtype == torch.bfloat16 - assert checkpoint_lvl in [0, 1, 2] - x = x.contiguous() - weight1 = weight1.contiguous() - bias1 = bias1.contiguous() - weight2 = weight2.contiguous() - bias2 = bias2.contiguous() - batch_shape, n = x.shape[:-1], x.shape[-1] - batch_dim = batch_shape.numel() - if is_bf16: - act_input = fused_dense_cuda.linear_bias_forward( - x.reshape(batch_dim, n), weight1, bias1 - ) - output1 = sqrelu_fwd(act_input) - else: - save_act_input = checkpoint_lvl != 2 - result = triton_linear_act( - x.reshape(batch_dim, n), - weight1, - bias1, - activation="squared_relu", - save_act_input=save_act_input, - ) - if save_act_input: - output1, act_input = result - else: - output1 = result - output2 = fused_dense_cuda.linear_bias_forward(output1, weight2, bias2) - ctx.checkpoint_lvl = checkpoint_lvl - if checkpoint_lvl == 0: - ctx.save_for_backward(x, weight1, bias1, weight2, act_input, output1) - elif checkpoint_lvl == 1: - ctx.save_for_backward(x, weight1, bias1, weight2, act_input) - elif checkpoint_lvl == 2: - ctx.save_for_backward(x, weight1, bias1, weight2) - return output2.reshape(*batch_shape, output2.shape[-1]) - - @staticmethod - @custom_bwd - def backward(ctx, grad_output): - grad_output = grad_output.contiguous() - checkpoint_lvl = ctx.checkpoint_lvl - x, weight1, bias1, weight2, *rest = ctx.saved_tensors - batch_shape, n = x.shape[:-1], x.shape[-1] - batch_dim = batch_shape.numel() - is_bf16 = x.dtype == torch.bfloat16 - if checkpoint_lvl == 0: - act_input, output1 = rest - elif checkpoint_lvl == 1: - (act_input,) = rest - output1 = sqrelu_fwd(act_input) - elif checkpoint_lvl == 2: - if is_bf16: - act_input = fused_dense_cuda.linear_bias_forward( - x.reshape(batch_dim, n), weight1, bias1 - ) - output1 = sqrelu_fwd(act_input) - else: - output1, act_input = triton_linear_act( - x.reshape(batch_dim, n), - weight1, - bias1, - activation="squared_relu", - save_act_input=True, - ) - - if is_bf16: - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad(output1, grad_output) - grad_output1 = grad_output @ weight2 - grad_act_input = sqrelu_bwd(grad_output1, act_input) - grad_input, grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_backward( - x.reshape(batch_dim, n), weight1, grad_act_input - ) - else: - grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad(output1, grad_output) - grad_act_input = triton_dgrad_act( - grad_output, weight2, activation="squared_relu", act_input=act_input - ) - grad_input, grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_backward( - x.reshape(batch_dim, n), weight1, grad_act_input - ) - return grad_input.reshape_as(x), grad_weight1, grad_bias1, grad_weight2, grad_bias2, None - - -fused_dense_sqrelu_dense_function = FusedDenseSqreluDenseFunc.apply - - -class FusedDenseSqreluDense(nn.Module): - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - bias1=True, - bias2=True, - checkpoint_lvl=0, - device=None, - dtype=None, - ): - """ - checkpoint_lvl (increasing lvl means slower but more memory saving): - 0: no recomputation in the bwd - 1: recompute gelu_out in the bwd - 2: recompute gelu_in and gelu_out in the bwd - """ - assert checkpoint_lvl in [0, 1, 2] - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features * 4 - assert bias1 == True, "DenseSqreluDense module without bias is currently not supported" - assert bias2 == True, "DenseSqreluDense module without bias is currently not supported" - self.checkpoint_lvl = checkpoint_lvl - self.fc1 = nn.Linear(in_features, hidden_features, bias=bias1, **factory_kwargs) - self.fc2 = nn.Linear(hidden_features, out_features, bias=bias2, **factory_kwargs) - - def forward(self, x): - assert x.is_cuda - return fused_dense_sqrelu_dense_function( - x, self.fc1.weight, self.fc1.bias, self.fc2.weight, self.fc2.bias, self.checkpoint_lvl - ) diff --git a/build/torch29-cxx11-cu130-aarch64-linux/ops/triton/rotary.py b/build/torch29-cxx11-cu130-aarch64-linux/ops/triton/rotary.py deleted file mode 100644 index 6ad3ab412a084270120167ad32b67ad1720f2c37..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-aarch64-linux/ops/triton/rotary.py +++ /dev/null @@ -1,186 +0,0 @@ -# Copyright (c) 2025, Tri Dao. -# As of 2025-04-23, we require triton >= 3.0 - -from typing import Optional, Union - -import torch - -import triton -import triton.language as tl - - -@triton.jit -def rotary_kernel( - OUT, # Pointers to matrices - X, - COS, - SIN, - CU_SEQLENS, - SEQLEN_OFFSETS, # this could be int or a pointer - # Matrix dimensions - seqlen, - nheads, - seqlen_ro, - # strides - stride_out_batch, - stride_out_seqlen, - stride_out_nheads, - stride_out_headdim, - stride_x_batch, - stride_x_seqlen, - stride_x_nheads, - stride_x_headdim, - # Meta-parameters - # We want ROTARY_DIM to be constexpr, otherwise the triton compiler doesn't know that - # the mask is constant every 8 elements, and it will generate LDG.16 instead of LDG.128 - ROTARY_DIM: tl.constexpr, - IS_SEQLEN_OFFSETS_TENSOR: tl.constexpr, - IS_VARLEN: tl.constexpr, - INTERLEAVED: tl.constexpr, - CONJUGATE: tl.constexpr, - BLOCK_H: tl.constexpr, - BLOCK_M: tl.constexpr, -): - BLOCK_K: tl.constexpr = triton.next_power_of_2(ROTARY_DIM) - ROTARY_DIM_HALF = ROTARY_DIM // 2 - pid_head = tl.program_id(axis=0) - pid_m = tl.program_id(axis=1) - pid_batch = tl.program_id(axis=2) - - if not IS_VARLEN: - X = X + pid_batch * stride_x_batch - OUT = OUT + pid_batch * stride_out_batch - else: - start_idx = tl.load(CU_SEQLENS + pid_batch) - seqlen = tl.load(CU_SEQLENS + pid_batch + 1) - start_idx - X = X + start_idx * stride_x_seqlen - OUT = OUT + start_idx * stride_out_seqlen - - if pid_m * BLOCK_M >= seqlen: - return - - rh = pid_head * BLOCK_H + tl.arange(0, BLOCK_H) - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - if not IS_SEQLEN_OFFSETS_TENSOR: - rm_cs = rm + SEQLEN_OFFSETS - else: - rm_cs = rm + tl.load(SEQLEN_OFFSETS + pid_batch) - - rk_half = tl.arange(0, BLOCK_K // 2) - COS = COS + (rm_cs[:, None] * ROTARY_DIM_HALF + rk_half[None, :]) - SIN = SIN + (rm_cs[:, None] * ROTARY_DIM_HALF + rk_half[None, :]) - mask_cs = (rm_cs[:, None] < seqlen_ro) & (rk_half[None, :] < ROTARY_DIM_HALF) - cos = tl.load(COS, mask=mask_cs, other=1.0).to(tl.float32) - sin = tl.load(SIN, mask=mask_cs, other=0.0).to(tl.float32) - if CONJUGATE: - sin = -sin - - if not INTERLEAVED: - # Load the 1st and 2nd halves of X, do calculation, then store to 1st and 2nd halves of OUT - X = X + (rh[:, None, None] * stride_x_nheads + rm[None, :, None] * stride_x_seqlen + rk_half[None, None, :] * stride_x_headdim) - OUT = OUT + (rh[:, None, None] * stride_out_nheads + rm[None, :, None] * stride_out_seqlen + rk_half[None, None, :] * stride_out_headdim) - mask = (rh[:, None, None] < nheads) & (rm[None, :, None] < seqlen) & (rk_half[None, None, :] < ROTARY_DIM_HALF) - x0 = tl.load(X, mask=mask, other=0.0).to(tl.float32) - x1 = tl.load(X + ROTARY_DIM_HALF * stride_x_headdim, mask=mask, other=0.0,).to(tl.float32) - o0 = x0 * cos - x1 * sin - o1 = x0 * sin + x1 * cos - tl.store(OUT, o0, mask=mask) - tl.store(OUT + ROTARY_DIM_HALF * stride_out_headdim, o1, mask=mask) - else: - rk = tl.arange(0, BLOCK_K) - X = X + (rh[:, None, None] * stride_x_nheads + rm[None, :, None] * stride_x_seqlen + rk[None, None, :] * stride_x_headdim) - OUT = OUT + (rh[:, None, None] * stride_out_nheads + rm[None, :, None] * stride_out_seqlen + rk[None, None, :] * stride_out_headdim) - mask = (rh[:, None, None] < nheads) & (rm[None, :, None] < seqlen) & (rk[None, None, :] < ROTARY_DIM) - x = tl.load(X, mask=mask, other=0.0).to(tl.float32) - x0, x1 = tl.split(tl.reshape(x, [BLOCK_H, BLOCK_M, BLOCK_K // 2, 2])) - o0 = x0 * cos - x1 * sin - o1 = x0 * sin + x1 * cos - o = tl.reshape(tl.join(o0, o1), [BLOCK_H, BLOCK_M, BLOCK_K]) - tl.store(OUT, o, mask=mask) - - -def apply_rotary( - x: torch.Tensor, - cos: torch.Tensor, - sin: torch.Tensor, - seqlen_offsets: Union[int, torch.Tensor] = 0, - cu_seqlens: Optional[torch.Tensor] = None, - max_seqlen: Optional[int] = None, - interleaved=False, - inplace=False, - conjugate=False, -) -> torch.Tensor: - """ - Arguments: - x: (batch, seqlen, nheads, headdim) if cu_seqlens is None - else (total_seqlen, nheads, headdim). - cos: (seqlen_ro, rotary_dim / 2) - sin: (seqlen_ro, rotary_dim / 2) - seqlen_offsets: integer or integer tensor of size (batch,) - cu_seqlens: (batch + 1,) or None - max_seqlen: int - Returns: - y: (batch, seqlen, nheads, headdim) - """ - is_varlen = cu_seqlens is not None - if not is_varlen: - batch, seqlen, nheads, headdim = x.shape - else: - assert max_seqlen is not None, "If cu_seqlens is passed in, then max_seqlen must be passed" - total_seqlen, nheads, headdim = x.shape - batch_p_1 = cu_seqlens.shape[0] - batch = batch_p_1 - 1 - seqlen = max_seqlen - seqlen_ro, rotary_dim = cos.shape - assert sin.shape == cos.shape - rotary_dim *= 2 - assert rotary_dim <= headdim, "rotary_dim must be <= headdim" - assert headdim <= 256, "Only support headdim <= 256" - assert seqlen_ro >= seqlen, "seqlen_ro must be >= seqlen" - - cos, sin = cos.contiguous(), sin.contiguous() - if isinstance(seqlen_offsets, torch.Tensor): - assert seqlen_offsets.shape == (batch,) - assert seqlen_offsets.dtype in [torch.int32, torch.int64] - seqlen_offsets = seqlen_offsets.contiguous() - else: - assert seqlen_offsets + seqlen <= seqlen_ro - - output = torch.empty_like(x) if not inplace else x - if rotary_dim < headdim and not inplace: - output[..., rotary_dim:].copy_(x[..., rotary_dim:]) - - grid = lambda META: (triton.cdiv(nheads, META["BLOCK_H"]), triton.cdiv(seqlen, META["BLOCK_M"]), batch) # noqa - BLOCK_M = 8 if rotary_dim <= 128 else 4 - - # Need this, otherwise Triton tries to launch from cuda:0 and we get - # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?) - device_ctx = torch.cuda.device(x.device.index) if x.device.type == 'cuda' else torch.xpu.device(x.device.index) - with device_ctx: - torch.library.wrap_triton(rotary_kernel)[grid]( - output, # data ptrs - x, - cos, - sin, - cu_seqlens, - seqlen_offsets, - seqlen, # shapes - nheads, - seqlen_ro, - output.stride(0) if not is_varlen else 0, # batch_strides if not varlen else 0 - output.stride(-3), # seqlen_stride or total_seqlen_stride - output.stride(-2), # nheads_stride - output.stride(-1), # headdim_stride - x.stride(0) if not is_varlen else 0, # batch_strides if not varlen else 0 - x.stride(-3), # seqlen stride or total_seqlen_stride - x.stride(-2), # nheads stride - x.stride(-1), # headdim stride - rotary_dim, - isinstance(seqlen_offsets, torch.Tensor), - is_varlen, - interleaved, - conjugate, - BLOCK_M=BLOCK_M, - BLOCK_H=2, - ) - return output diff --git a/build/torch29-cxx11-cu130-x86_64-linux/_flash_attn2_588b404.abi3.so b/build/torch29-cxx11-cu130-x86_64-linux/_flash_attn2_588b404.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..a95f730a8a5b3af46b3a862e7a94b6acb0b2da35 --- /dev/null +++ b/build/torch29-cxx11-cu130-x86_64-linux/_flash_attn2_588b404.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:568eb670747b578b865649894b7674d9053a2ba660ba2e491030c788e3d5936a +size 1009019168 diff --git a/build/torch29-cxx11-cu130-x86_64-linux/_flash_attn2_cuda_abda3a0.abi3.so b/build/torch29-cxx11-cu130-x86_64-linux/_flash_attn2_cuda_abda3a0.abi3.so deleted file mode 100644 index 0b9816e0cea7ab5bfb85afbe9f8893e8b4ce58d8..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-x86_64-linux/_flash_attn2_cuda_abda3a0.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:45ed1be1f8bf3a148e9c6ebf07796d20bda6b889c52192deb4bef381f9772543 -size 1009019192 diff --git a/build/torch29-cxx11-cu130-x86_64-linux/_ops.py b/build/torch29-cxx11-cu130-x86_64-linux/_ops.py index fb32434e09e651a658beb38be8c101e5db89374f..03506d8a18d17ca7296a5f1b7a18ec7dc4ad73bb 100644 --- a/build/torch29-cxx11-cu130-x86_64-linux/_ops.py +++ b/build/torch29-cxx11-cu130-x86_64-linux/_ops.py @@ -1,9 +1,9 @@ import torch -from . import _flash_attn2_cuda_abda3a0 -ops = torch.ops._flash_attn2_cuda_abda3a0 +from . import _flash_attn2_588b404 +ops = torch.ops._flash_attn2_588b404 def add_op_namespace_prefix(op_name: str): """ Prefix op by namespace. """ - return f"_flash_attn2_cuda_abda3a0::{op_name}" + return f"_flash_attn2_588b404::{op_name}" \ No newline at end of file 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 index 6a0157f25449eb90c6c688bc8ef3f4fcf84d46dc..66e8d3665834b141d703998496b3da6fe1f6fe06 100644 --- a/build/torch29-cxx11-cu130-x86_64-linux/flash_attn_interface.py +++ b/build/torch29-cxx11-cu130-x86_64-linux/flash_attn_interface.py @@ -31,14 +31,12 @@ def _get_device(): else: return "cpu" +_XPU_AVAILABLE = torch.xpu.is_available() if hasattr(torch, "xpu") else False # TODO remove hasattr check when bwd is supported on XPU -def _get_block_size_n(device, head_dim, is_dropout, is_causal): - assert head_dim <= 256 - - if device.type == "xpu": - return 64 +def _get_block_size_n(device, head_dim, is_dropout, is_causal): # This should match the block sizes in the CUDA kernel + assert head_dim <= 256 major, minor = torch.cuda.get_device_capability(device) is_sm8x = major == 8 and minor > 0 # Only include sm86 and sm89, exclude sm80 (A100) is_sm80 = major == 8 and minor == 0 @@ -1066,7 +1064,7 @@ def flash_attn_qkvpacked_func( alibi_slopes, deterministic, return_attn_probs, - torch.is_grad_enabled(), + False if _XPU_AVAILABLE else torch.is_grad_enabled(), ) @@ -1144,7 +1142,7 @@ def flash_attn_kvpacked_func( alibi_slopes, deterministic, return_attn_probs, - torch.is_grad_enabled(), + False if _XPU_AVAILABLE else torch.is_grad_enabled(), ) @@ -1221,7 +1219,7 @@ def flash_attn_func( alibi_slopes, deterministic, return_attn_probs, - torch.is_grad_enabled(), + False if _XPU_AVAILABLE else torch.is_grad_enabled(), ) @@ -1287,7 +1285,7 @@ def flash_attn_varlen_qkvpacked_func( alibi_slopes, deterministic, return_attn_probs, - torch.is_grad_enabled(), + False if _XPU_AVAILABLE else torch.is_grad_enabled(), ) @@ -1379,7 +1377,7 @@ def flash_attn_varlen_kvpacked_func( alibi_slopes, deterministic, return_attn_probs, - torch.is_grad_enabled(), + False if _XPU_AVAILABLE else torch.is_grad_enabled(), ) @@ -1473,7 +1471,7 @@ def flash_attn_varlen_func( deterministic, return_attn_probs, block_table, - False if q.device.type == "cpu" else torch.is_grad_enabled(), + False if _XPU_AVAILABLE or q.device.type == "cpu" else torch.is_grad_enabled(), ) diff --git a/build/torch29-cxx11-cu130-x86_64-linux/metadata.json b/build/torch29-cxx11-cu130-x86_64-linux/metadata.json index 8b00d7fb7685121edf7a97bb007415122444f1d5..9cf5deed9898dce769f4cc73913d3530b92a0bd8 100644 --- a/build/torch29-cxx11-cu130-x86_64-linux/metadata.json +++ b/build/torch29-cxx11-cu130-x86_64-linux/metadata.json @@ -1,14 +1,4 @@ { "version": 1, - "license": "BSD-3-Clause", - "python-depends": [], - "backend": { - "type": "cuda", - "archs": [ - "10.0", - "12.0", - "8.0", - "9.0" - ] - } -} + "python-depends": [] +} \ No newline at end of file diff --git a/build/torch29-cxx11-xpu20252-x86_64-linux/_flash_attn2_588b404.abi3.so b/build/torch29-cxx11-xpu20252-x86_64-linux/_flash_attn2_588b404.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..c082acd4b4ec32cb86c1526b8007bfbf37a2735f --- /dev/null +++ b/build/torch29-cxx11-xpu20252-x86_64-linux/_flash_attn2_588b404.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c30735ef658c0500deaa387ce38b1fd0380be8db2e4d02758cbae454cb7f9161 +size 9382080 diff --git a/build/torch29-cxx11-xpu20252-x86_64-linux/_flash_attn2_xpu_8641a05.abi3.so b/build/torch29-cxx11-xpu20252-x86_64-linux/_flash_attn2_xpu_8641a05.abi3.so deleted file mode 100644 index 686af1199589603561e9c57342c2cd02f9862740..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-xpu20252-x86_64-linux/_flash_attn2_xpu_8641a05.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:e5aa9382ee56f3cf74d15bf2de14b4ebe5831f3fd74f4129f9da8f3fb104bd24 -size 14205096 diff --git a/build/torch29-cxx11-xpu20252-x86_64-linux/_ops.py b/build/torch29-cxx11-xpu20252-x86_64-linux/_ops.py index 66ae708ea2db048a7252cd2b816268ff3dda2170..03506d8a18d17ca7296a5f1b7a18ec7dc4ad73bb 100644 --- a/build/torch29-cxx11-xpu20252-x86_64-linux/_ops.py +++ b/build/torch29-cxx11-xpu20252-x86_64-linux/_ops.py @@ -1,9 +1,9 @@ import torch -from . import _flash_attn2_xpu_8641a05 -ops = torch.ops._flash_attn2_xpu_8641a05 +from . import _flash_attn2_588b404 +ops = torch.ops._flash_attn2_588b404 def add_op_namespace_prefix(op_name: str): """ Prefix op by namespace. """ - return f"_flash_attn2_xpu_8641a05::{op_name}" + return f"_flash_attn2_588b404::{op_name}" \ No newline at end of file diff --git a/build/torch29-cxx11-xpu20252-x86_64-linux/flash_attn2/__init__.py b/build/torch29-cxx11-xpu20252-x86_64-linux/flash_attn2/__init__.py index a9b2672c1cd85b74c1b3ded0fc0b2100e1aeac23..03dbc1afe1cf156661a2b1b22003cd5f599a0309 100644 --- a/build/torch29-cxx11-xpu20252-x86_64-linux/flash_attn2/__init__.py +++ b/build/torch29-cxx11-xpu20252-x86_64-linux/flash_attn2/__init__.py @@ -1,10 +1,10 @@ import ctypes -import importlib.util 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 diff --git a/build/torch29-cxx11-xpu20252-x86_64-linux/flash_attn_interface.py b/build/torch29-cxx11-xpu20252-x86_64-linux/flash_attn_interface.py index 6a0157f25449eb90c6c688bc8ef3f4fcf84d46dc..66e8d3665834b141d703998496b3da6fe1f6fe06 100644 --- a/build/torch29-cxx11-xpu20252-x86_64-linux/flash_attn_interface.py +++ b/build/torch29-cxx11-xpu20252-x86_64-linux/flash_attn_interface.py @@ -31,14 +31,12 @@ def _get_device(): else: return "cpu" +_XPU_AVAILABLE = torch.xpu.is_available() if hasattr(torch, "xpu") else False # TODO remove hasattr check when bwd is supported on XPU -def _get_block_size_n(device, head_dim, is_dropout, is_causal): - assert head_dim <= 256 - - if device.type == "xpu": - return 64 +def _get_block_size_n(device, head_dim, is_dropout, is_causal): # This should match the block sizes in the CUDA kernel + assert head_dim <= 256 major, minor = torch.cuda.get_device_capability(device) is_sm8x = major == 8 and minor > 0 # Only include sm86 and sm89, exclude sm80 (A100) is_sm80 = major == 8 and minor == 0 @@ -1066,7 +1064,7 @@ def flash_attn_qkvpacked_func( alibi_slopes, deterministic, return_attn_probs, - torch.is_grad_enabled(), + False if _XPU_AVAILABLE else torch.is_grad_enabled(), ) @@ -1144,7 +1142,7 @@ def flash_attn_kvpacked_func( alibi_slopes, deterministic, return_attn_probs, - torch.is_grad_enabled(), + False if _XPU_AVAILABLE else torch.is_grad_enabled(), ) @@ -1221,7 +1219,7 @@ def flash_attn_func( alibi_slopes, deterministic, return_attn_probs, - torch.is_grad_enabled(), + False if _XPU_AVAILABLE else torch.is_grad_enabled(), ) @@ -1287,7 +1285,7 @@ def flash_attn_varlen_qkvpacked_func( alibi_slopes, deterministic, return_attn_probs, - torch.is_grad_enabled(), + False if _XPU_AVAILABLE else torch.is_grad_enabled(), ) @@ -1379,7 +1377,7 @@ def flash_attn_varlen_kvpacked_func( alibi_slopes, deterministic, return_attn_probs, - torch.is_grad_enabled(), + False if _XPU_AVAILABLE else torch.is_grad_enabled(), ) @@ -1473,7 +1471,7 @@ def flash_attn_varlen_func( deterministic, return_attn_probs, block_table, - False if q.device.type == "cpu" else torch.is_grad_enabled(), + False if _XPU_AVAILABLE or q.device.type == "cpu" else torch.is_grad_enabled(), ) diff --git a/build/torch29-cxx11-xpu20252-x86_64-linux/metadata.json b/build/torch29-cxx11-xpu20252-x86_64-linux/metadata.json index 8f032899cf61212add2325c22107252842bd1588..9cf5deed9898dce769f4cc73913d3530b92a0bd8 100644 --- a/build/torch29-cxx11-xpu20252-x86_64-linux/metadata.json +++ b/build/torch29-cxx11-xpu20252-x86_64-linux/metadata.json @@ -1,8 +1,4 @@ { "version": 1, - "license": "BSD-3-Clause", - "python-depends": [], - "backend": { - "type": "xpu" - } + "python-depends": [] } \ No newline at end of file diff --git a/media/benches_dark_animation.svg b/media/benches_dark_animation.svg deleted file mode 100644 index 3cbe6f193be051d7ad4f30114d14eb6333ca339a..0000000000000000000000000000000000000000 --- a/media/benches_dark_animation.svg +++ /dev/null @@ -1,96 +0,0 @@ - -kernels-community/flash-attn2 vs Torch - Relative Speed -PyTorch 2.11.0+cu130 · CPU - -FlashAttn.large -19.16x - - - - - - - -FlashAttn.medium -17.94x - - - - - - - -FlashAttn.small -2.39x - - - - - - - -FlashAttnCausal.large -31.85x - - - - - - - -FlashAttnCausal.medium -23.83x - - - - - - - -FlashAttnCausal.small -2.83x - - - - - - - -FlashAttnVarlen.large -15.44x - - - - - - - -FlashAttnVarlen.medium -15.91x - - - - - - - -FlashAttnVarlen.small -9.58x - - - - - - - -Kernel - -Torch (ref) - - - - - - - - \ No newline at end of file diff --git a/media/benches_dark_latency.svg b/media/benches_dark_latency.svg deleted file mode 100644 index b1a19697665aa99c3d1a55914ccf65744bf329b7..0000000000000000000000000000000000000000 --- a/media/benches_dark_latency.svg +++ /dev/null @@ -1,2608 +0,0 @@ - - - - - - - - 2026-04-20T20:36:43.330821 - image/svg+xml - - - Matplotlib v3.10.8, https://matplotlib.org/ - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - diff --git a/media/benches_dark_throughput.svg b/media/benches_dark_throughput.svg deleted file mode 100644 index 3d66c7187f9cadc1d15f525bed88bec5a55aae74..0000000000000000000000000000000000000000 --- a/media/benches_dark_throughput.svg +++ /dev/null @@ -1,2858 +0,0 @@ - - - - - - - - 2026-04-20T20:36:43.610631 - image/svg+xml - - - Matplotlib v3.10.8, https://matplotlib.org/ - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - diff --git a/media/benches_light_animation.svg b/media/benches_light_animation.svg deleted file mode 100644 index c4384caf1667c158ac8d6572b030752064939975..0000000000000000000000000000000000000000 --- a/media/benches_light_animation.svg +++ /dev/null @@ -1,96 +0,0 @@ - -kernels-community/flash-attn2 vs Torch - Relative Speed -PyTorch 2.11.0+cu130 · CPU - -FlashAttn.large -19.16x - - - - - - - -FlashAttn.medium -17.94x - - - - - - - -FlashAttn.small -2.39x - - - - - - - -FlashAttnCausal.large -31.85x - - - - - - - -FlashAttnCausal.medium -23.83x - - - - - - - -FlashAttnCausal.small -2.83x - - - - - - - -FlashAttnVarlen.large -15.44x - - - - - - - -FlashAttnVarlen.medium -15.91x - - - - - - - -FlashAttnVarlen.small -9.58x - - - - - - - -Kernel - -Torch (ref) - - - - - - - - \ No newline at end of file diff --git a/media/benches_light_latency.svg b/media/benches_light_latency.svg deleted file mode 100644 index c01bfd6981be3171e2b2daebe1e3e318799f4005..0000000000000000000000000000000000000000 --- a/media/benches_light_latency.svg +++ /dev/null @@ -1,2608 +0,0 @@ - - - - - - - - 2026-04-20T20:36:42.034782 - image/svg+xml - - - Matplotlib v3.10.8, https://matplotlib.org/ - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - diff --git a/media/benches_light_throughput.svg b/media/benches_light_throughput.svg deleted file mode 100644 index b6565a130cc6268925983fa94e965b1ca23839a1..0000000000000000000000000000000000000000 --- a/media/benches_light_throughput.svg +++ /dev/null @@ -1,2858 +0,0 @@ - - - - - - - - 2026-04-20T20:36:42.598908 - image/svg+xml - - - Matplotlib v3.10.8, https://matplotlib.org/ - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -